Repository: ariga/atlas
Pull request #2401 (2023-12-23 16:44:47)

sql/migrate: improve begin..end block detection


Patch:
diff --git a/sql/migrate/lex.go b/sql/migrate/lex.go
index 8bf6ec30..02521a3d 100644
--- a/sql/migrate/lex.go
+++ b/sql/migrate/lex.go
@@ -200,6 +200,7 @@ Scan:
 			if err := s.delimCmd(); err != nil {
 				return nil, err
 			}
+			s.skipSpaces()
 		// Delimiters take precedence over comments.
 		case depth == 0 && strings.HasPrefix(s.input[s.pos-s.width:], s.delim):
 			s.addPos(len(s.delim) - s.width)
@@ -223,7 +224,9 @@ Scan:
 				break Scan
 			}
 			// Not a "BEGIN ATOMIC" block.
-		case s.delim == delimiter && s.MatchBegin && reBegin.MatchString(s.input[s.pos-1:]):
+		case s.delim == delimiter && s.MatchBegin &&
+			// Either the current scanned statement starts with BEGIN, or we inside a statement and expects at least one ~space before).
+			(s.pos == 1 && reBegin.MatchString(s.input[s.pos-1:]) || s.pos > 1 && reBegin.MatchString(s.input[s.pos-2:])):
 			if err := s.skipBegin(); err == nil {
 				text = s.input[:s.pos]
 				break Scan


Test patch:
diff --git a/sql/migrate/testdata/lexgroup/2_function.sql b/sql/migrate/testdata/lexgroup/2_function.sql
index cfbc6baa..d09abab6 100644
--- a/sql/migrate/testdata/lexgroup/2_function.sql
+++ b/sql/migrate/testdata/lexgroup/2_function.sql
@@ -142,6 +142,46 @@ begin
     set new.begin = new.end;
 end;
 
+/*
+ The "NEW.begin" one is not scanned as a "BEGIN" because it is
+ not a beginning of s statement and there is no \s before it.
+*/
+CREATE TRIGGER begin1 BEFORE INSERT ON t5
+FOR EACH ROW
+BEGIN
+    SELECT 1;
+    SELECT 1;
+    SET NEW.begin = NEW.end;
+END;
+
+/*
+ The `begin` column is quoted as an identifier as therefore is skipped.
+*/
+CREATE TRIGGER begin2 BEFORE INSERT ON t5
+    FOR EACH ROW
+BEGIN
+    SELECT 1;
+    SELECT 1;
+    SET `begin` = `end`;
+END;
+
+/*
+ An unquoted begin confuses the lexer, and requires using the DELIMITER command.
+*/
+-- Set a special statement delimiter.
+DELIMITER //;
+
+CREATE TRIGGER begin3 BEFORE INSERT ON t5
+    FOR EACH ROW
+BEGIN
+    SELECT 1;
+    SELECT 1;
+    SET begin = end;
+END //;
+
+-- Unset the special statement delimiter.
+DELIMITER ;
+
 -- issue 2397.
 CREATE OR REPLACE VIEW claims.current_member_view AS
     SELECT l.member_load_id, n.name_last_or_organization, g.name_given, n.name_middle, n.name_prefix, n.name_suffix, ps.external_plan_sponsor_id, i.member_identification_code, i.date_of_birth, cgop.group_or_policy_number, a.plan_sponsor_name, stat.is_subscriber, stat.relationship_to_subscriber, ed.employment_begin, cd.benefit_begin, cd.benefit_end
diff --git a/sql/migrate/testdata/lexgroup/2_function.sql.golden b/sql/migrate/testdata/lexgroup/2_function.sql.golden
index 6b13384e..b6665d90 100644
--- a/sql/migrate/testdata/lexgroup/2_function.sql.golden
+++ b/sql/migrate/testdata/lexgroup/2_function.sql.golden
@@ -144,6 +144,30 @@ begin
     set new.begin = new.end;
 end;
 -- end --
+CREATE TRIGGER begin1 BEFORE INSERT ON t5
+FOR EACH ROW
+BEGIN
+    SELECT 1;
+    SELECT 1;
+    SET NEW.begin = NEW.end;
+END;
+-- end --
+CREATE TRIGGER begin2 BEFORE INSERT ON t5
+    FOR EACH ROW
+BEGIN
+    SELECT 1;
+    SELECT 1;
+    SET `begin` = `end`;
+END;
+-- end --
+CREATE TRIGGER begin3 BEFORE INSERT ON t5
+    FOR EACH ROW
+BEGIN
+    SELECT 1;
+    SELECT 1;
+    SET begin = end;
+END
+-- end --
 CREATE OR REPLACE VIEW claims.current_member_view AS
     SELECT l.member_load_id, n.name_last_or_organization, g.name_given, n.name_middle, n.name_prefix, n.name_suffix, ps.external_plan_sponsor_id, i.member_identification_code, i.date_of_birth, cgop.group_or_policy_number, a.plan_sponsor_name, stat.is_subscriber, stat.relationship_to_subscriber, ed.employment_begin, cd.benefit_begin, cd.benefit_end
     FROM claims.member_coverage_load mcl