mirror of https://github.com/dexidp/dex.git
You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
128 lines
2.9 KiB
128 lines
2.9 KiB
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 |
|
}
|
|
|