Repository: ariga/atlas
Pull request #2358 (2023-12-10 21:58:26)

sql/migrate: support scanning compound statements


Patch:
diff --git a/sql/migrate/lex.go b/sql/migrate/lex.go
index 778935e4..7babf261 100644
--- a/sql/migrate/lex.go
+++ b/sql/migrate/lex.go
@@ -44,13 +44,47 @@ func (s *Stmt) Directive(name string) (ds []string) {
 
 // Stmts provides a generic implementation for extracting SQL statements from the given file contents.
 func Stmts(input string) ([]*Stmt, error) {
+	return (&StmtScanner{
+		ScannerOptions: ScannerOptions{
+			// Default options for backward compatibility.
+			MatchBegin:       false,
+			MatchBeginAtomic: true,
+			MatchDollarQuote: true,
+		},
+	}).Scan(input)
+}
+
+type (
+	// StmtScanner scanning SQL statements from migration and schema files.
+	StmtScanner struct {
+		ScannerOptions
+		// scanner state.
+		src, input string   // src and current input text
+		pos        int      // current phase position
+		total      int      // total bytes scanned so far
+		width      int      // size of latest rune
+		delim      string   // configured delimiter
+		comments   []string // collected comments
+	}
+	// ScannerOptions controls the behavior of the scanner.
+	ScannerOptions struct {
+		// MatchBegin enables matching for BEGIN ... END statements block.
+		MatchBegin bool
+		// MatchBeginAtomic enables matching for BEGIN ATOMIC ... END statements block.
+		MatchBeginAtomic bool
+		// MatchDollarQuote enables the PostgreSQL dollar-quoted string syntax.
+		MatchDollarQuote bool
+	}
+)
+
+// Scan scans the statement in the given input.
+func (s *StmtScanner) Scan(input string) ([]*Stmt, error) {
 	var stmts []*Stmt
-	l, err := newLex(input)
-	if err != nil {
+	if err := s.init(input); err != nil {
 		return nil, err
 	}
 	for {
-		s, err := l.stmt()
+		s, err := s.stmt()
 		if err == io.EOF {
 			return stmts, nil
 		}
@@ -61,277 +95,304 @@ func Stmts(input string) ([]*Stmt, error) {
 	}
 }
 
-type lex struct {
-	src, input string   // src and current input text
-	pos        int      // current phase position
-	total      int      // total bytes scanned so far
-	width      int      // size of latest rune
-	delim      string   // configured delimiter
-	comments   []string // collected comments
-}
-
-const (
-	eos          = -1
-	delimiter    = ";"
-	delimiterCmd = "delimiter"
-)
-
-func newLex(input string) (*lex, error) {
-	l := &lex{src: input, input: input, delim: delimiter}
+// init initializes the scanner state.
+func (s *StmtScanner) init(input string) error {
+	s.comments = nil
+	s.pos, s.total, s.width = 0, 0, 0
+	s.src, s.input, s.delim = input, input, delimiter
 	if d, ok := directive(input, directiveDelimiter, directivePrefixSQL); ok {
-		if err := l.setDelim(d); err != nil {
-			return nil, err
+		if err := s.setDelim(d); err != nil {
+			return err
 		}
 		parts := strings.SplitN(input, "\n", 2)
 		if len(parts) == 1 {
-			return nil, l.error(l.pos, "no input found after delimiter %q", d)
+			return s.error(s.pos, "no input found after delimiter %q", d)
 		}
-		l.input = parts[1]
+		s.input = parts[1]
 	}
-	return l, nil
+	return nil
 }
 
+const (
+	eos          = -1
+	delimiter    = ";"
+	delimiterCmd = "delimiter"
+)
+
 var (
 	// Dollar-quoted string as defined by the PostgreSQL scanner.
 	reDollarQuote = regexp.MustCompile(`^\$([A-Za-zÈ-ÿ_][\wÈ-ÿ]*)*\$`)
 	// The 'BEGIN ATOMIC' syntax as specified in the SQL 2003 standard.
 	reBeginAtomic = regexp.MustCompile(`(?i)^\s*BEGIN\s+ATOMIC\s+`)
+	reBegin       = regexp.MustCompile(`(?i)^\s*BEGIN\s+`)
 	reEnd         = regexp.MustCompile(`(?i)^\s*END\s*`)
 )
 
-func (l *lex) stmt() (*Stmt, error) {
+func (s *StmtScanner) stmt() (*Stmt, error) {
 	var (
 		depth, openingPos int
 		text              string
 	)
-	l.skipSpaces()
+	s.skipSpaces()
 Scan:
 	for {
-		switch r := l.next(); {
+		switch r := s.next(); {
 		case r == eos:
 			switch {
 			case depth > 0:
-				return nil, l.error(openingPos, "unclosed '('")
-			case l.pos > 0:
-				text = l.input
+				return nil, s.error(openingPos, "unclosed '('")
+			case s.pos > 0:
+				text = s.input
 				break Scan
 			default:
 				return nil, io.EOF
 			}
 		case r == '(':
 			if depth == 0 {
-				openingPos = l.pos
+				openingPos = s.pos
 			}
 			depth++
 		case r == ')':
 			if depth == 0 {
-				return nil, l.error(l.pos, "unexpected ')'")
+				return nil, s.error(s.pos, "unexpected ')'")
 			}
 			depth--
 		case r == '\'', r == '"', r == '`':
-			if err := l.skipQuote(r); err != nil {
+			if err := s.skipQuote(r); err != nil {
 				return nil, err
 			}
 		// Check if the start of the statement is the MySQL DELIMITER command.
 		// See https://dev.mysql.com/doc/refman/8.0/en/mysql-commands.html.
-		case l.pos == 1 && len(l.input) > len(delimiterCmd) && strings.EqualFold(l.input[:len(delimiterCmd)], delimiterCmd):
-			l.addPos(len(delimiterCmd) - 1)
-			if err := l.delimCmd(); err != nil {
+		case s.pos == 1 && len(s.input) > len(delimiterCmd) && strings.EqualFold(s.input[:len(delimiterCmd)], delimiterCmd):
+			s.addPos(len(delimiterCmd) - 1)
+			if err := s.delimCmd(); err != nil {
 				return nil, err
 			}
 		// Delimiters take precedence over comments.
-		case depth == 0 && strings.HasPrefix(l.input[l.pos-l.width:], l.delim):
-			l.addPos(len(l.delim) - l.width)
-			text = l.input[:l.pos]
+		case depth == 0 && strings.HasPrefix(s.input[s.pos-s.width:], s.delim):
+			s.addPos(len(s.delim) - s.width)
+			text = s.input[:s.pos]
 			break Scan
-		case r == '$' && reDollarQuote.MatchString(l.input[l.pos-1:]):
-			if err := l.skipDollarQuote(); err != nil {
+		case s.MatchDollarQuote && r == '$' && reDollarQuote.MatchString(s.input[s.pos-1:]):
+			if err := s.skipDollarQuote(); err != nil {
 				return nil, err
 			}
 		// Skip non-standard MySQL comments if they are inside
 		// expressions until we make the lexer driver-aware.
 		case depth == 0 && r == '#':
-			l.comment("#", "\n")
-		case r == '-' && l.next() == '-':
-			l.comment("--", "\n")
-		case r == '/' && l.next() == '*':
-			l.comment("/*", "*/")
-		case reBeginAtomic.MatchString(l.input[l.pos-1:]):
-			if err := l.skipBeginAtomic(); err != nil {
+			s.comment("#", "\n")
+		case r == '-' && s.next() == '-':
+			s.comment("--", "\n")
+		case r == '/' && s.next() == '*':
+			s.comment("/*", "*/")
+		case s.MatchBeginAtomic && reBeginAtomic.MatchString(s.input[s.pos-1:]):
+			if err := s.skipBeginAtomic(); err != nil {
 				return nil, err
 			}
-			text = l.input[:l.pos]
+			text = s.input[:s.pos]
+			break Scan
+		case s.MatchBegin && reBegin.MatchString(s.input[s.pos-1:]):
+			if err := s.skipBegin(); err != nil {
+				return nil, err
+			}
+			text = s.input[:s.pos]
 			break Scan
 		}
 	}
-	return l.emit(text), nil
+	return s.emit(text), nil
 }
 
-func (l *lex) next() rune {
-	if l.pos >= len(l.input) {
+func (s *StmtScanner) next() rune {
+	if s.pos >= len(s.input) {
 		return eos
 	}
-	r, w := utf8.DecodeRuneInString(l.input[l.pos:])
-	l.width = w
-	l.addPos(w)
+	r, w := utf8.DecodeRuneInString(s.input[s.pos:])
+	s.width = w
+	s.addPos(w)
 	return r
 }
 
-func (l *lex) pick() rune {
-	p, w := l.pos, l.width
-	r := l.next()
-	l.pos, l.width = p, w
+func (s *StmtScanner) pick() rune {
+	p, w := s.pos, s.width
+	r := s.next()
+	s.pos, s.width = p, w
 	return r
 }
 
-func (l *lex) addPos(p int) {
-	l.pos += p
-	l.total += p
+func (s *StmtScanner) addPos(p int) {
+	s.pos += p
+	s.total += p
 }
 
-func (l *lex) skipQuote(quote rune) error {
-	pos := l.pos
+func (s *StmtScanner) skipQuote(quote rune) error {
+	pos := s.pos
 	for {
-		switch r := l.next(); {
+		switch r := s.next(); {
 		case r == eos:
-			return l.error(pos, "unclosed quote %q", quote)
+			return s.error(pos, "unclosed quote %q", quote)
 		case r == '\\':
-			l.next()
+			s.next()
 		case r == quote:
 			return nil
 		}
 	}
 }
 
-func (l *lex) skipDollarQuote() error {
-	m := reDollarQuote.FindString(l.input[l.pos-1:])
+func (s *StmtScanner) skipDollarQuote() error {
+	m := reDollarQuote.FindString(s.input[s.pos-1:])
 	if m == "" {
-		return l.error(l.pos, "unexpected dollar quote")
+		return s.error(s.pos, "unexpected dollar quote")
 	}
-	l.addPos(len(m) - 1)
+	s.addPos(len(m) - 1)
 	for {
-		switch r := l.next(); {
+		switch r := s.next(); {
 		case r == eos:
 			// Fail only if a delimiter was not set.
-			if l.delim == "" {
-				return l.error(l.pos, "unclosed dollar-quoted string")
+			if s.delim == "" {
+				return s.error(s.pos, "unclosed dollar-quoted string")
 			}
 			return nil
-		case r == '$' && strings.HasPrefix(l.input[l.pos-1:], m):
-			l.addPos(len(m) - 1)
+		case r == '$' && strings.HasPrefix(s.input[s.pos-1:], m):
+			s.addPos(len(m) - 1)
 			return nil
 		}
 	}
 }
 
-func (l *lex) skipBeginAtomic() error {
-	m := reBeginAtomic.FindString(l.input[l.pos-1:])
+func (s *StmtScanner) skipBeginAtomic() error {
+	m := reBeginAtomic.FindString(s.input[s.pos-1:])
 	if m == "" {
-		return l.error(l.pos, "unexpected missing BEGIN ATOMIC block")
+		return s.error(s.pos, "unexpected missing BEGIN ATOMIC block")
 	}
-	l.addPos(len(m) - 1)
-	body, err := newLex(l.input[l.pos:])
-	if err != nil {
-		return l.error(l.pos, "create scanner for sql body: %v", err)
+	s.addPos(len(m) - 1)
+	body := &StmtScanner{ScannerOptions: s.ScannerOptions}
+	if err := body.init(s.input[s.pos:]); err != nil {
+		return err
 	}
 	for {
-		s, err := body.stmt()
+		stmt, err := body.stmt()
 		if err == io.EOF {
-			return l.error(l.pos, "unexpected eof when scanning sql body")
+			return s.error(s.pos, "unexpected eof when scanning sql body")
 		}
 		if err != nil {
-			return l.error(l.pos, "scan sql body: %v", err)
+			return s.error(s.pos, "scan sql body: %v", err)
 		}
-		if reEnd.MatchString(s.Text) {
+		if reEnd.MatchString(stmt.Text) {
 			break
 		}
 	}
-	l.addPos(body.total)
+	s.addPos(body.total)
+	return nil
+}
+
+func (s *StmtScanner) skipBegin() error {
+	m := reBegin.FindString(s.input[s.pos-1:])
+	if m == "" {
+		return s.error(s.pos, "unexpected missing BEGIN block")
+	}
+	s.addPos(len(m) - 1)
+	group := &StmtScanner{ScannerOptions: s.ScannerOptions}
+	if err := group.init(s.input[s.pos:]); err != nil {
+		return err
+	}
+	for depth := 1; depth > 0; {
+		switch stmt, err := group.stmt(); {
+		case err == io.EOF:
+			return s.error(s.pos, "unexpected eof when scanning compound statements")
+		case err != nil:
+			return s.error(s.pos, "scan compound statements: %v", err)
+		case reEnd.MatchString(stmt.Text):
+			if m := reEnd.FindString(stmt.Text); len(m) == len(stmt.Text) || strings.TrimPrefix(stmt.Text, m) == s.delim {
+				depth--
+			}
+		}
+	}
+	s.addPos(group.total)
 	return nil
 }
 
-func (l *lex) comment(left, right string) {
-	i := strings.Index(l.input[l.pos:], right)
+func (s *StmtScanner) comment(left, right string) {
+	i := strings.Index(s.input[s.pos:], right)
 	// Not a comment.
 	if i == -1 {
 		return
 	}
 	// If the comment reside inside a statement, collect it.
-	if l.pos != len(left) {
-		l.addPos(i + len(right))
+	if s.pos != len(left) {
+		s.addPos(i + len(right))
 		return
 	}
-	l.addPos(i + len(right))
+	s.addPos(i + len(right))
 	// If we did not scan any statement characters, it
 	// can be skipped and stored in the comments group.
-	l.comments = append(l.comments, l.input[:l.pos])
-	l.input = l.input[l.pos:]
-	l.pos = 0
+	s.comments = append(s.comments, s.input[:s.pos])
+	s.input = s.input[s.pos:]
+	s.pos = 0
 	// Double \n separate the comments group from the statement.
-	if strings.HasPrefix(l.input, "\n\n") || right == "\n" && strings.HasPrefix(l.input, "\n") {
-		l.comments = nil
+	if strings.HasPrefix(s.input, "\n\n") || right == "\n" && strings.HasPrefix(s.input, "\n") {
+		s.comments = nil
 	}
-	l.skipSpaces()
+	s.skipSpaces()
 }
 
-func (l *lex) skipSpaces() {
-	n := len(l.input)
-	l.input = strings.TrimLeftFunc(l.input, unicode.IsSpace)
-	l.total += n - len(l.input)
+func (s *StmtScanner) skipSpaces() {
+	n := len(s.input)
+	s.input = strings.TrimLeftFunc(s.input, unicode.IsSpace)
+	s.total += n - len(s.input)
 }
 
-func (l *lex) emit(text string) *Stmt {
-	s := &Stmt{Pos: l.total - len(text), Text: text, Comments: l.comments}
-	l.input = l.input[l.pos:]
-	l.pos = 0
-	l.comments = nil
+func (s *StmtScanner) emit(text string) *Stmt {
+	stmt := &Stmt{Pos: s.total - len(text), Text: text, Comments: s.comments}
+	s.input = s.input[s.pos:]
+	s.pos = 0
+	s.comments = nil
 	// Trim custom delimiter.
-	if l.delim != delimiter {
-		s.Text = strings.TrimSuffix(s.Text, l.delim)
+	if s.delim != delimiter {
+		stmt.Text = strings.TrimSuffix(stmt.Text, s.delim)
 	}
-	s.Text = strings.TrimSpace(s.Text)
-	return s
+	stmt.Text = strings.TrimSpace(stmt.Text)
+	return stmt
 }
 
 // delimCmd checks if the scanned "DELIMITER"
 // text represents an actual delimiter command.
-func (l *lex) delimCmd() error {
+func (s *StmtScanner) delimCmd() error {
 	// A space must come after the delimiter.
-	if l.pick() != ' ' {
+	if s.pick() != ' ' {
 		return nil
 	}
 	// Scan delimiter.
-	for r := l.pick(); r != eos && r != '\n'; r = l.next() {
+	for r := s.pick(); r != eos && r != '\n'; r = s.next() {
 	}
-	delim := strings.TrimSpace(l.input[len(delimiterCmd):l.pos])
+	delim := strings.TrimSpace(s.input[len(delimiterCmd):s.pos])
 	// MySQL client allows quoting delimiters.
 	if strings.HasPrefix(delim, "'") && strings.HasSuffix(delim, "'") {
 		delim = strings.ReplaceAll(delim[1:len(delim)-1], "''", "'")
 	}
-	if err := l.setDelim(delim); err != nil {
+	if err := s.setDelim(delim); err != nil {
 		return err
 	}
 	// Skip all we saw until now.
-	l.emit(l.input[:l.pos])
+	s.emit(s.input[:s.pos])
 	return nil
 }
 
-func (l *lex) setDelim(d string) error {
+func (s *StmtScanner) setDelim(d string) error {
 	if d == "" {
 		return errors.New("empty delimiter")
 	}
 	// Unescape delimiters. e.g. "\\n" => "\n".
-	l.delim = strings.NewReplacer(`\n`, "\n", `\r`, "\r", `\t`, "\t").Replace(d)
+	s.delim = strings.NewReplacer(`\n`, "\n", `\r`, "\r", `\t`, "\t").Replace(d)
 	return nil
 }
 
-func (l *lex) error(pos int, format string, args ...any) error {
+func (s *StmtScanner) error(pos int, format string, args ...any) error {
 	format = "%d:%d: " + format
 	var (
-		p    = len(l.src) - len(l.input) + pos
-		s    = l.src[:p]
-		col  = strings.LastIndex(s, "\n")
-		line = 1 + strings.Count(s, "\n")
+		p    = len(s.src) - len(s.input) + pos
+		src  = s.src[:p]
+		col  = strings.LastIndex(src, "\n")
+		line = 1 + strings.Count(src, "\n")
 	)
 	if line == 1 {
 		col = p


Test patch:
diff --git a/sql/migrate/lex_test.go b/sql/migrate/lex_test.go
index 799086f2..25e4c1d8 100644
--- a/sql/migrate/lex_test.go
+++ b/sql/migrate/lex_test.go
@@ -21,13 +21,34 @@ func TestLocalFile_Stmts(t *testing.T) {
 	require.NoError(t, err)
 	for _, f := range files {
 		stmts, err := f.Stmts()
-		require.NoError(t, err)
+		require.NoErrorf(t, err, "file: %s", f.Name())
 		buf, err := os.ReadFile(filepath.Join(path, f.Name()+".golden"))
 		require.NoError(t, err)
 		require.Equalf(t, string(buf), strings.Join(stmts, "\n-- end --\n"), "mismatched statements in file %q", f.Name())
 	}
 }
 
+func TestScanner_StmtsGroup(t *testing.T) {
+	scan := &StmtScanner{}
+	scan.MatchBegin = true
+	path := filepath.Join("testdata", "lexgroup")
+	dir, err := NewLocalDir(path)
+	require.NoError(t, err)
+	files, err := dir.Files()
+	require.NoError(t, err)
+	for _, f := range files {
+		stmts, err := scan.Scan(string(f.Bytes()))
+		require.NoErrorf(t, err, "file: %s", f.Name())
+		buf, err := os.ReadFile(filepath.Join(path, f.Name()+".golden"))
+		require.NoError(t, err)
+		got := make([]string, len(stmts))
+		for i, s := range stmts {
+			got[i] = s.Text
+		}
+		require.Equalf(t, string(buf), strings.Join(got, "\n-- end --\n"), "mismatched statements in file %q", f.Name())
+	}
+}
+
 func TestLocalFile_StmtDecls(t *testing.T) {
 	f := `cmd0;
 -- test
@@ -157,9 +178,7 @@ func TestLex_Errors(t *testing.T) {
 		},
 	} {
 		t.Run(tt.name, func(t *testing.T) {
-			l, err := newLex(tt.stmt)
-			require.NoError(t, err)
-			_, err = l.stmt()
+			_, err := Stmts(tt.stmt)
 			require.EqualError(t, err, tt.err)
 		})
 	}
diff --git a/sql/migrate/testdata/lexgroup/1_trigger.sql b/sql/migrate/testdata/lexgroup/1_trigger.sql
new file mode 100644
index 00000000..da741c33
--- /dev/null
+++ b/sql/migrate/testdata/lexgroup/1_trigger.sql
@@ -0,0 +1,72 @@
+-- SQLite test cases.
+
+CREATE TRIGGER before_tbl_insert BEFORE INSERT ON tbl BEGIN SELECT CASE
+    WHEN (new.a = 4) THEN RAISE(IGNORE) END;
+END;
+
+CREATE TRIGGER after_tbl_insert AFTER INSERT ON tbl BEGIN SELECT CASE
+    WHEN (new.a = 1) THEN RAISE(ABORT,    'Trigger abort')
+    WHEN (new.a = 2) THEN RAISE(FAIL,     'Trigger fail')
+    WHEN (new.a = 3) THEN RAISE(ROLLBACK, 'Trigger rollback') END;
+END;
+
+CREATE TRIGGER after_tbl2_insert AFTER INSERT ON tbl2 BEGIN
+    UPDATE tbl SET c = 10;
+    INSERT INTO tbl2 VALUES (new.a, new.b, new.c);
+END;
+
+CREATE TABLE t2(x,y,z);
+CREATE TRIGGER t2r1 AFTER INSERT ON t2 BEGIN SELECT 1; END;
+CREATE TRIGGER t2r2 BEFORE INSERT ON t2 BEGIN SELECT 1; END;
+CREATE TRIGGER t2r3 AFTER UPDATE ON t2 BEGIN SELECT 1; END;
+CREATE TRIGGER t2r4 BEFORE UPDATE ON t2 BEGIN SELECT 1; END;
+CREATE TRIGGER t2r5 AFTER DELETE ON t2 BEGIN SELECT 1; END;
+CREATE TRIGGER t2r6 BEFORE DELETE ON t2 BEGIN SELECT 1; END;
+CREATE TRIGGER t2r7 AFTER INSERT ON t2 BEGIN SELECT 1; END;
+CREATE TRIGGER t2r8 BEFORE INSERT ON t2 BEGIN SELECT 1; END;
+CREATE TRIGGER t2r9 AFTER UPDATE ON t2 BEGIN SELECT 1; END;
+CREATE TRIGGER t2r10 BEFORE UPDATE ON t2 BEGIN SELECT 1; END;
+CREATE TRIGGER t2r11 AFTER DELETE ON t2 BEGIN SELECT 1; END;
+CREATE TRIGGER t2r12 BEFORE DELETE ON t2 BEGIN SELECT 1; END;
+
+create trigger I_test instead of insert on test
+  begin
+    insert into test1 (id,a) values (NEW.id,NEW.a);
+    insert into test2 (id,b) values (NEW.id,NEW.b);
+  end;
+
+CREATE TRIGGER tr AFTER INSERT ON t3 BEGIN
+  INSERT INTO t3 SELECT new.c+1 WHERE new.c<5;
+  INSERT INTO t2 SELECT new.c*10000+xx.a*100+yy.a
+                   FROM t1 AS xx, t1 AS yy
+                  WHERE xx.a IN (1,2,3,4)
+                    AND yy.a IN (2,3,4,5);
+END;
+
+CREATE TABLE Item(
+   a integer PRIMARY KEY NOT NULL ,
+   b double NULL ,
+   c int NOT NULL DEFAULT 0
+);
+CREATE TABLE Undo(UndoAction TEXT);
+INSERT INTO Item VALUES (1,38205.60865,340);
+CREATE TRIGGER trigItem_UNDO_AD AFTER DELETE ON Item FOR EACH ROW
+BEGIN
+  INSERT INTO Undo SELECT 'INSERT INTO Item (a,b,c) VALUES ('
+   || coalesce(old.a,'NULL') || ',' || quote(old.b) || ',' || old.c || ');';
+END;
+DELETE FROM Item WHERE a = 1;
+SELECT * FROM Undo;
+
+CREATE TRIGGER transactionTrigger AFTER INSERT ON tbl10 BEGIN
+    SAVEPOINT sp1;
+    INSERT INTO tbl11 VALUES (new.f);
+    ROLLBACK TO sp1;
+END;
+
+CREATE TRIGGER errorHandlingTrigger AFTER INSERT ON tbl12 BEGIN
+    BEGIN
+        INSERT INTO tbl13 VALUES (new.g);
+        INSERT INTO tbl13 VALUES (new.g);
+    END;
+END;
diff --git a/sql/migrate/testdata/lexgroup/1_trigger.sql.golden b/sql/migrate/testdata/lexgroup/1_trigger.sql.golden
new file mode 100644
index 00000000..306a5643
--- /dev/null
+++ b/sql/migrate/testdata/lexgroup/1_trigger.sql.golden
@@ -0,0 +1,87 @@
+CREATE TRIGGER before_tbl_insert BEFORE INSERT ON tbl BEGIN SELECT CASE
+    WHEN (new.a = 4) THEN RAISE(IGNORE) END;
+END;
+-- end --
+CREATE TRIGGER after_tbl_insert AFTER INSERT ON tbl BEGIN SELECT CASE
+    WHEN (new.a = 1) THEN RAISE(ABORT,    'Trigger abort')
+    WHEN (new.a = 2) THEN RAISE(FAIL,     'Trigger fail')
+    WHEN (new.a = 3) THEN RAISE(ROLLBACK, 'Trigger rollback') END;
+END;
+-- end --
+CREATE TRIGGER after_tbl2_insert AFTER INSERT ON tbl2 BEGIN
+    UPDATE tbl SET c = 10;
+    INSERT INTO tbl2 VALUES (new.a, new.b, new.c);
+END;
+-- end --
+CREATE TABLE t2(x,y,z);
+-- end --
+CREATE TRIGGER t2r1 AFTER INSERT ON t2 BEGIN SELECT 1; END;
+-- end --
+CREATE TRIGGER t2r2 BEFORE INSERT ON t2 BEGIN SELECT 1; END;
+-- end --
+CREATE TRIGGER t2r3 AFTER UPDATE ON t2 BEGIN SELECT 1; END;
+-- end --
+CREATE TRIGGER t2r4 BEFORE UPDATE ON t2 BEGIN SELECT 1; END;
+-- end --
+CREATE TRIGGER t2r5 AFTER DELETE ON t2 BEGIN SELECT 1; END;
+-- end --
+CREATE TRIGGER t2r6 BEFORE DELETE ON t2 BEGIN SELECT 1; END;
+-- end --
+CREATE TRIGGER t2r7 AFTER INSERT ON t2 BEGIN SELECT 1; END;
+-- end --
+CREATE TRIGGER t2r8 BEFORE INSERT ON t2 BEGIN SELECT 1; END;
+-- end --
+CREATE TRIGGER t2r9 AFTER UPDATE ON t2 BEGIN SELECT 1; END;
+-- end --
+CREATE TRIGGER t2r10 BEFORE UPDATE ON t2 BEGIN SELECT 1; END;
+-- end --
+CREATE TRIGGER t2r11 AFTER DELETE ON t2 BEGIN SELECT 1; END;
+-- end --
+CREATE TRIGGER t2r12 BEFORE DELETE ON t2 BEGIN SELECT 1; END;
+-- end --
+create trigger I_test instead of insert on test
+  begin
+    insert into test1 (id,a) values (NEW.id,NEW.a);
+    insert into test2 (id,b) values (NEW.id,NEW.b);
+  end;
+-- end --
+CREATE TRIGGER tr AFTER INSERT ON t3 BEGIN
+  INSERT INTO t3 SELECT new.c+1 WHERE new.c<5;
+  INSERT INTO t2 SELECT new.c*10000+xx.a*100+yy.a
+                   FROM t1 AS xx, t1 AS yy
+                  WHERE xx.a IN (1,2,3,4)
+                    AND yy.a IN (2,3,4,5);
+END;
+-- end --
+CREATE TABLE Item(
+   a integer PRIMARY KEY NOT NULL ,
+   b double NULL ,
+   c int NOT NULL DEFAULT 0
+);
+-- end --
+CREATE TABLE Undo(UndoAction TEXT);
+-- end --
+INSERT INTO Item VALUES (1,38205.60865,340);
+-- end --
+CREATE TRIGGER trigItem_UNDO_AD AFTER DELETE ON Item FOR EACH ROW
+BEGIN
+  INSERT INTO Undo SELECT 'INSERT INTO Item (a,b,c) VALUES ('
+   || coalesce(old.a,'NULL') || ',' || quote(old.b) || ',' || old.c || ');';
+END;
+-- end --
+DELETE FROM Item WHERE a = 1;
+-- end --
+SELECT * FROM Undo;
+-- end --
+CREATE TRIGGER transactionTrigger AFTER INSERT ON tbl10 BEGIN
+    SAVEPOINT sp1;
+    INSERT INTO tbl11 VALUES (new.f);
+    ROLLBACK TO sp1;
+END;
+-- end --
+CREATE TRIGGER errorHandlingTrigger AFTER INSERT ON tbl12 BEGIN
+    BEGIN
+        INSERT INTO tbl13 VALUES (new.g);
+        INSERT INTO tbl13 VALUES (new.g);
+    END;
+END;
\ No newline at end of file
diff --git a/sql/migrate/testdata/lexgroup/2_function.sql b/sql/migrate/testdata/lexgroup/2_function.sql
new file mode 100644
index 00000000..dc8037d5
--- /dev/null
+++ b/sql/migrate/testdata/lexgroup/2_function.sql
@@ -0,0 +1,126 @@
+
+CREATE PROCEDURE CompProc(IN p INT)
+BEGIN
+    DECLARE v1 INT DEFAULT 0;
+    DECLARE v2 INT;
+    DECLARE d INT DEFAULT FALSE;
+    DECLARE c1 CURSOR FOR SELECT a FROM t1 WHERE b = p;
+    DECLARE CONTINUE HANDLER FOR NOT FOUND SET d = TRUE;
+
+    OPEN c1;
+    rl: LOOP
+        FETCH c1 INTO v2;
+        IF d THEN
+            LEAVE rl;
+        END IF;
+
+        BEGIN
+            IF v2 > 0 THEN
+                UPDATE t2 SET c = c + v2 WHERE d = p;
+            END IF;
+        END;
+    END LOOP;
+    CLOSE c1;
+
+    BEGIN
+        SELECT COUNT(*) INTO v1 FROM t2 WHERE d = p;
+        IF v1 > 100 THEN
+            CALL OtherProc(p);
+        END IF;
+    END;
+
+END;
+
+
+CREATE FUNCTION CompFunc(eID INT) RETURNS INT
+BEGIN
+    DECLARE ts INT;
+    DECLARE b INT;
+
+    BEGIN
+        SELECT SUM(s) INTO ts FROM sales WHERE e = eID;
+    END;
+
+    BEGIN
+        IF ts > 10000 THEN
+            SET b = 500;
+        ELSEIF ts BETWEEN 5000 AND 10000 THEN
+            SET b = 300;
+        ELSE
+            SET b = 0;
+        END IF;
+    END;
+
+    RETURN b;
+END;
+
+
+CREATE PROCEDURE CompLogicProc()
+BEGIN
+    DECLARE i INT DEFAULT 1;
+    DECLARE t INT DEFAULT 0;
+
+    sl: WHILE i <= 10 DO
+        BEGIN
+            IF i MOD 2 = 0 THEN
+                SET t = t + i;
+            ELSE
+                BEGIN
+                    IF i MOD 3 = 0 THEN
+                        SET t = t - i;
+                    END IF;
+                END;
+            END IF;
+            SET i = i + 1;
+        END;
+    END WHILE;
+
+    BEGIN
+        IF t < 0 THEN
+            SET t = 0;
+        END IF;
+    END;
+
+    SELECT t;
+END;
+
+
+CREATE PROCEDURE ff1()
+BEGIN
+    DECLARE v1, v2, v3, i, j INT DEFAULT 0;
+    DECLARE flag INT DEFAULT FALSE;
+    DECLARE cur CURSOR FOR SELECT col1 FROM tableX;
+    DECLARE CONTINUE HANDLER FOR NOT FOUND SET flag = TRUE;
+
+    OPEN cur;
+    main_loop: LOOP
+        FETCH cur INTO v1;
+        IF flag THEN
+            LEAVE main_loop;
+        END IF;
+
+        SET i = 1;
+        WHILE i <= 5 DO
+            SET j = 1;
+            REPEAT
+                SET v2 = (SELECT COUNT(*) FROM tableY WHERE col2 = v1 AND col3 = j);
+                IF v2 > 3 THEN
+                    UPDATE tableZ SET col4 = col4 + 1 WHERE col5 = i;
+                END IF;
+                SET j = j + 1;
+            UNTIL j > 3 END REPEAT;
+            SET i = i + 1;
+        END WHILE;
+
+        compound_logic: BEGIN
+            IF v1 < 10 THEN
+                INSERT INTO tableA (colA) VALUES (v1);
+            ELSEIF v1 BETWEEN 10 AND 20 THEN
+                UPDATE tableB SET colB = v1 WHERE colB < v1;
+            ELSE
+                DELETE FROM tableC WHERE colC = v1;
+            END IF;
+        END;
+    END LOOP;
+    CLOSE cur;
+END;
\ No newline at end of file
diff --git a/sql/migrate/testdata/lexgroup/2_function.sql.golden b/sql/migrate/testdata/lexgroup/2_function.sql.golden
new file mode 100644
index 00000000..61c319d6
--- /dev/null
+++ b/sql/migrate/testdata/lexgroup/2_function.sql.golden
@@ -0,0 +1,122 @@
+CREATE PROCEDURE CompProc(IN p INT)
+BEGIN
+    DECLARE v1 INT DEFAULT 0;
+    DECLARE v2 INT;
+    DECLARE d INT DEFAULT FALSE;
+    DECLARE c1 CURSOR FOR SELECT a FROM t1 WHERE b = p;
+    DECLARE CONTINUE HANDLER FOR NOT FOUND SET d = TRUE;
+
+    OPEN c1;
+    rl: LOOP
+        FETCH c1 INTO v2;
+        IF d THEN
+            LEAVE rl;
+        END IF;
+
+        BEGIN
+            IF v2 > 0 THEN
+                UPDATE t2 SET c = c + v2 WHERE d = p;
+            END IF;
+        END;
+    END LOOP;
+    CLOSE c1;
+
+    BEGIN
+        SELECT COUNT(*) INTO v1 FROM t2 WHERE d = p;
+        IF v1 > 100 THEN
+            CALL OtherProc(p);
+        END IF;
+    END;
+
+END;
+-- end --
+CREATE FUNCTION CompFunc(eID INT) RETURNS INT
+BEGIN
+    DECLARE ts INT;
+    DECLARE b INT;
+
+    BEGIN
+        SELECT SUM(s) INTO ts FROM sales WHERE e = eID;
+    END;
+
+    BEGIN
+        IF ts > 10000 THEN
+            SET b = 500;
+        ELSEIF ts BETWEEN 5000 AND 10000 THEN
+            SET b = 300;
+        ELSE
+            SET b = 0;
+        END IF;
+    END;
+
+    RETURN b;
+END;
+-- end --
+CREATE PROCEDURE CompLogicProc()
+BEGIN
+    DECLARE i INT DEFAULT 1;
+    DECLARE t INT DEFAULT 0;
+
+    sl: WHILE i <= 10 DO
+        BEGIN
+            IF i MOD 2 = 0 THEN
+                SET t = t + i;
+            ELSE
+                BEGIN
+                    IF i MOD 3 = 0 THEN
+                        SET t = t - i;
+                    END IF;
+                END;
+            END IF;
+            SET i = i + 1;
+        END;
+    END WHILE;
+
+    BEGIN
+        IF t < 0 THEN
+            SET t = 0;
+        END IF;
+    END;
+
+    SELECT t;
+END;
+-- end --
+CREATE PROCEDURE ff1()
+BEGIN
+    DECLARE v1, v2, v3, i, j INT DEFAULT 0;
+    DECLARE flag INT DEFAULT FALSE;
+    DECLARE cur CURSOR FOR SELECT col1 FROM tableX;
+    DECLARE CONTINUE HANDLER FOR NOT FOUND SET flag = TRUE;
+
+    OPEN cur;
+    main_loop: LOOP
+        FETCH cur INTO v1;
+        IF flag THEN
+            LEAVE main_loop;
+        END IF;
+
+        SET i = 1;
+        WHILE i <= 5 DO
+            SET j = 1;
+            REPEAT
+                SET v2 = (SELECT COUNT(*) FROM tableY WHERE col2 = v1 AND col3 = j);
+                IF v2 > 3 THEN
+                    UPDATE tableZ SET col4 = col4 + 1 WHERE col5 = i;
+                END IF;
+                SET j = j + 1;
+            UNTIL j > 3 END REPEAT;
+            SET i = i + 1;
+        END WHILE;
+
+        compound_logic: BEGIN
+            IF v1 < 10 THEN
+                INSERT INTO tableA (colA) VALUES (v1);
+            ELSEIF v1 BETWEEN 10 AND 20 THEN
+                UPDATE tableB SET colB = v1 WHERE colB < v1;
+            ELSE
+                DELETE FROM tableC WHERE colC = v1;
+            END IF;
+        END;
+    END LOOP;
+    CLOSE cur;
+END;
\ No newline at end of file