128 lines
2.9 KiB
Go
128 lines
2.9 KiB
Go
package sqlparse
|
|
|
|
import (
|
|
"bufio"
|
|
"bytes"
|
|
"errors"
|
|
"io"
|
|
|
|
"strings"
|
|
)
|
|
|
|
const sqlCmdPrefix = "-- +migrate "
|
|
|
|
// Checks the line to see if the line has a statement-ending semicolon
|
|
// or if the line contains a double-dash comment.
|
|
func endsWithSemicolon(line string) bool {
|
|
|
|
prev := ""
|
|
scanner := bufio.NewScanner(strings.NewReader(line))
|
|
scanner.Split(bufio.ScanWords)
|
|
|
|
for scanner.Scan() {
|
|
word := scanner.Text()
|
|
if strings.HasPrefix(word, "--") {
|
|
break
|
|
}
|
|
prev = word
|
|
}
|
|
|
|
return strings.HasSuffix(prev, ";")
|
|
}
|
|
|
|
// Split the given sql script into individual statements.
|
|
//
|
|
// The base case is to simply split on semicolons, as these
|
|
// naturally terminate a statement.
|
|
//
|
|
// However, more complex cases like pl/pgsql can have semicolons
|
|
// within a statement. For these cases, we provide the explicit annotations
|
|
// 'StatementBegin' and 'StatementEnd' to allow the script to
|
|
// tell us to ignore semicolons.
|
|
func SplitSQLStatements(r io.ReadSeeker, direction bool) ([]string, error) {
|
|
_, err := r.Seek(0, 0)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
var buf bytes.Buffer
|
|
scanner := bufio.NewScanner(r)
|
|
|
|
// track the count of each section
|
|
// so we can diagnose scripts with no annotations
|
|
upSections := 0
|
|
downSections := 0
|
|
|
|
statementEnded := false
|
|
ignoreSemicolons := false
|
|
directionIsActive := false
|
|
|
|
stmts := make([]string, 0)
|
|
|
|
for scanner.Scan() {
|
|
|
|
line := scanner.Text()
|
|
|
|
// handle any migrate-specific commands
|
|
if strings.HasPrefix(line, sqlCmdPrefix) {
|
|
cmd := strings.TrimSpace(line[len(sqlCmdPrefix):])
|
|
switch cmd {
|
|
case "Up":
|
|
directionIsActive = (direction == true)
|
|
upSections++
|
|
break
|
|
|
|
case "Down":
|
|
directionIsActive = (direction == false)
|
|
downSections++
|
|
break
|
|
|
|
case "StatementBegin":
|
|
if directionIsActive {
|
|
ignoreSemicolons = true
|
|
}
|
|
break
|
|
|
|
case "StatementEnd":
|
|
if directionIsActive {
|
|
statementEnded = (ignoreSemicolons == true)
|
|
ignoreSemicolons = false
|
|
}
|
|
break
|
|
}
|
|
}
|
|
|
|
if !directionIsActive {
|
|
continue
|
|
}
|
|
|
|
if _, err := buf.WriteString(line + "\n"); err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
// Wrap up the two supported cases: 1) basic with semicolon; 2) psql statement
|
|
// Lines that end with semicolon that are in a statement block
|
|
// do not conclude statement.
|
|
if (!ignoreSemicolons && endsWithSemicolon(line)) || statementEnded {
|
|
statementEnded = false
|
|
stmts = append(stmts, buf.String())
|
|
buf.Reset()
|
|
}
|
|
}
|
|
|
|
if err := scanner.Err(); err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
// diagnose likely migration script errors
|
|
if ignoreSemicolons {
|
|
return nil, errors.New("ERROR: saw '-- +migrate StatementBegin' with no matching '-- +migrate StatementEnd'")
|
|
}
|
|
|
|
if upSections == 0 && downSections == 0 {
|
|
return nil, errors.New(`ERROR: no Up/Down annotations found, so no statements were executed.
|
|
See https://github.com/rubenv/sql-migrate for details.`)
|
|
}
|
|
|
|
return stmts, nil
|
|
}
|