Repository: ariga/atlas
Pull request #2165 (2023-10-09 15:59:34)

sql/migrate: scan sql-standard body in functions/procedures


Patch:
diff --git a/sql/migrate/lex.go b/sql/migrate/lex.go
index a6968112..778935e4 100644
--- a/sql/migrate/lex.go
+++ b/sql/migrate/lex.go
@@ -91,8 +91,13 @@ func newLex(input string) (*lex, error) {
 	return l, nil
 }
 
-// Dollar-quoted string as defined by the PostgreSQL scanner.
-var reDollarQuote = regexp.MustCompile(`^\$([A-Za-zÈ-ÿ_][\wÈ-ÿ]*)*\$`)
+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+`)
+	reEnd         = regexp.MustCompile(`(?i)^\s*END\s*`)
+)
 
 func (l *lex) stmt() (*Stmt, error) {
 	var (
@@ -151,6 +156,12 @@ Scan:
 			l.comment("--", "\n")
 		case r == '/' && l.next() == '*':
 			l.comment("/*", "*/")
+		case reBeginAtomic.MatchString(l.input[l.pos-1:]):
+			if err := l.skipBeginAtomic(); err != nil {
+				return nil, err
+			}
+			text = l.input[:l.pos]
+			break Scan
 		}
 	}
 	return l.emit(text), nil
@@ -213,6 +224,32 @@ func (l *lex) skipDollarQuote() error {
 	}
 }
 
+func (l *lex) skipBeginAtomic() error {
+	m := reBeginAtomic.FindString(l.input[l.pos-1:])
+	if m == "" {
+		return l.error(l.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)
+	}
+	for {
+		s, err := body.stmt()
+		if err == io.EOF {
+			return l.error(l.pos, "unexpected eof when scanning sql body")
+		}
+		if err != nil {
+			return l.error(l.pos, "scan sql body: %v", err)
+		}
+		if reEnd.MatchString(s.Text) {
+			break
+		}
+	}
+	l.addPos(body.total)
+	return nil
+}
+
 func (l *lex) comment(left, right string) {
 	i := strings.Index(l.input[l.pos:], right)
 	// Not a comment.


Test patch:
diff --git a/sql/migrate/testdata/lex/16_begin_atomic.sql b/sql/migrate/testdata/lex/16_begin_atomic.sql
new file mode 100644
index 00000000..dff21054
--- /dev/null
+++ b/sql/migrate/testdata/lex/16_begin_atomic.sql
@@ -0,0 +1,51 @@
+CREATE FUNCTION functest_S_1(a text, b date) RETURNS boolean
+    LANGUAGE SQL
+    RETURN a = 'abcd' AND b > '2001-01-01';
+CREATE FUNCTION functest_S_2(a text[]) RETURNS int
+    RETURN a[1]::int;
+CREATE FUNCTION functest_S_3() RETURNS boolean
+    RETURN false;
+CREATE FUNCTION functest_S_10(a text, b date) RETURNS boolean
+    LANGUAGE SQL
+BEGIN ATOMIC
+SELECT a = 'abcd' AND b > '2001-01-01';
+END;
+CREATE FUNCTION functest_S_13() RETURNS boolean
+BEGIN ATOMIC
+SELECT 1;
+SELECT false;
+END;
+CREATE FUNCTION functest_S_15(x int) RETURNS boolean
+    LANGUAGE SQL
+BEGIN ATOMIC
+select case when x % 2 = 0 then true else false end; -- tricky parsing
+END;
+CREATE FUNCTION functest_sri2() RETURNS SETOF int
+LANGUAGE SQL
+STABLE
+BEGIN ATOMIC
+SELECT * FROM functest3;
+END;
+CREATE TABLE functest1 (i int);
+CREATE FUNCTION functest_S_16(a int, b int) RETURNS void
+    LANGUAGE SQL
+BEGIN ATOMIC
+INSERT INTO functest1 SELECT a + $2;
+END;
+
+CREATE FUNCTION functest_S_15(x int) RETURNS boolean
+    LANGUAGE SQL
+BEGIN ATOMIC
+select case when x % 2 = 0 then true else false end;
+select case when x % 2 = 0 then true else false end;
+select case when x % 2 = 0 then true else false end;
+select case when x % 2 = 0 then true else false end;
+END;
+
+CREATE FUNCTION "functest_s_15" ("x" integer) RETURNS boolean LANGUAGE SQL BEGIN ATOMIC
+SELECT
+         CASE
+             WHEN ((x % 2) = 0) THEN true
+             ELSE false
+         END AS "case";
+END;
\ No newline at end of file
diff --git a/sql/migrate/testdata/lex/16_begin_atomic.sql.golden b/sql/migrate/testdata/lex/16_begin_atomic.sql.golden
new file mode 100644
index 00000000..eb30d4f8
--- /dev/null
+++ b/sql/migrate/testdata/lex/16_begin_atomic.sql.golden
@@ -0,0 +1,59 @@
+CREATE FUNCTION functest_S_1(a text, b date) RETURNS boolean
+    LANGUAGE SQL
+    RETURN a = 'abcd' AND b > '2001-01-01';
+-- end --
+CREATE FUNCTION functest_S_2(a text[]) RETURNS int
+    RETURN a[1]::int;
+-- end --
+CREATE FUNCTION functest_S_3() RETURNS boolean
+    RETURN false;
+-- end --
+CREATE FUNCTION functest_S_10(a text, b date) RETURNS boolean
+    LANGUAGE SQL
+BEGIN ATOMIC
+SELECT a = 'abcd' AND b > '2001-01-01';
+END;
+-- end --
+CREATE FUNCTION functest_S_13() RETURNS boolean
+BEGIN ATOMIC
+SELECT 1;
+SELECT false;
+END;
+-- end --
+CREATE FUNCTION functest_S_15(x int) RETURNS boolean
+    LANGUAGE SQL
+BEGIN ATOMIC
+select case when x % 2 = 0 then true else false end; -- tricky parsing
+END;
+-- end --
+CREATE FUNCTION functest_sri2() RETURNS SETOF int
+LANGUAGE SQL
+STABLE
+BEGIN ATOMIC
+SELECT * FROM functest3;
+END;
+-- end --
+CREATE TABLE functest1 (i int);
+-- end --
+CREATE FUNCTION functest_S_16(a int, b int) RETURNS void
+    LANGUAGE SQL
+BEGIN ATOMIC
+INSERT INTO functest1 SELECT a + $2;
+END;
+-- end --
+CREATE FUNCTION functest_S_15(x int) RETURNS boolean
+    LANGUAGE SQL
+BEGIN ATOMIC
+select case when x % 2 = 0 then true else false end;
+select case when x % 2 = 0 then true else false end;
+select case when x % 2 = 0 then true else false end;
+select case when x % 2 = 0 then true else false end;
+END;
+-- end --
+CREATE FUNCTION "functest_s_15" ("x" integer) RETURNS boolean LANGUAGE SQL BEGIN ATOMIC
+SELECT
+         CASE
+             WHEN ((x % 2) = 0) THEN true
+             ELSE false
+         END AS "case";
+END;
\ No newline at end of file