| 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 |
| } |