Repository: ariga/atlas
Pull request #3398 (2025-02-20 17:43:54)

sql/migrate: support BEGIN..END terminator


Patch:
diff --git a/sql/migrate/lex.go b/sql/migrate/lex.go
index 1b17810b..810dd7ba 100644
--- a/sql/migrate/lex.go
+++ b/sql/migrate/lex.go
@@ -97,6 +97,9 @@ type (
 		width      int      // size of latest rune
 		delim      string   // configured delimiter
 		comments   []string // collected comments
+		// internal option to indicate if the
+		// END word is parsed as a terminator.
+		endterm *regexp.Regexp
 	}
 
 	// ScannerOptions controls the behavior of the scanner.
@@ -119,6 +122,9 @@ type (
 		HashComments bool
 		// Enable the "GO" command as a delimiter.
 		GoCommand bool
+		// BeginEndTerminator is a T-SQL specific option that allows
+		// the scanner to terminate BEGIN/END blocks with a semicolon.
+		BeginEndTerminator bool
 	}
 )
 
@@ -247,6 +253,9 @@ Scan:
 		case r == '/' && s.pick() == '*':
 			s.next()
 			s.comment("/*", "*/")
+		case s.endterm != nil && s.endterm.MatchString(s.input[:s.pos]):
+			text = s.input[:s.pos]
+			break Scan
 		case s.delim == delimiter && s.MatchBeginAtomic && reBeginAtomic.MatchString(s.input[s.pos-1:]):
 			if err := s.skipBeginAtomic(); err == nil {
 				text = s.input[:s.pos]
@@ -389,6 +398,10 @@ func (s *Scanner) skipBeginTryCatch() error {
 	return nil
 }
 
+var (
+	reEndTerm = regexp.MustCompile(`(?i)\s*END\s*$`)
+)
+
 func (s *Scanner) skipBegin() error {
 	m := reBegin.FindString(s.input[s.pos-1:])
 	if m == "" {
@@ -396,10 +409,14 @@ func (s *Scanner) skipBegin() error {
 	}
 	s.addPos(len(m) - 1)
 	group := &Scanner{ScannerOptions: s.ScannerOptions}
+	if s.BeginEndTerminator {
+		group.endterm = reEndTerm
+	}
 	if err := group.init(s.input[s.pos:]); err != nil {
 		return err
 	}
-	for depth := 1; depth > 0; {
+Loop:
+	for {
 		switch stmt, err := group.stmt(); {
 		case err == io.EOF:
 			return s.error(s.pos, "unexpected eof when scanning compound statements")
@@ -407,8 +424,10 @@ func (s *Scanner) skipBegin() error {
 			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--
+				break Loop
 			}
+		case s.BeginEndTerminator && reEndTerm.MatchString(stmt.Text):
+			break Loop
 		}
 	}
 	s.addPos(group.total)


Test patch:
diff --git a/sql/migrate/lex_test.go b/sql/migrate/lex_test.go
index 04a3c618..b44e2cc9 100644
--- a/sql/migrate/lex_test.go
+++ b/sql/migrate/lex_test.go
@@ -130,6 +130,28 @@ func TestScanner_BeginTryCatch(t *testing.T) {
 	}
 }
 
+func TestScanner_SQLServer(t *testing.T) {
+	scan := &Scanner{}
+	scan.MatchBegin = true
+	scan.BeginEndTerminator = true
+	path := filepath.Join("testdata", "sqlserver")
+	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
diff --git a/sql/migrate/testdata/sqlserver/1_return_table.sql b/sql/migrate/testdata/sqlserver/1_return_table.sql
new file mode 100644
index 00000000..5dbabd3b
--- /dev/null
+++ b/sql/migrate/testdata/sqlserver/1_return_table.sql
@@ -0,0 +1,8 @@
+CREATE FUNCTION [f2](@a as INT, @b as INT = 1)
+RETURNS TABLE
+AS RETURN SELECT @a as [a], @b as [b], (@a+@b)*2 as [p], @a*@b as [s];
+CREATE FUNCTION [f3] (@a int, @b int = 1) RETURNS @t1 TABLE ([c1] int NOT NULL, [c2] nvarchar(255) COLLATE SQL_Latin1_General_CP1_CI_AS NOT NULL, [c3] nvarchar(255) COLLATE SQL_Latin1_General_CP1_CI_AS DEFAULT N'G' NULL, [c4] int NOT NULL, PRIMARY KEY CLUSTERED ([c1] ASC), INDEX [idx] NONCLUSTERED ([c2] ASC), UNIQUE NONCLUSTERED ([c2] ASC, [c3] DESC), UNIQUE NONCLUSTERED ([c3] DESC, [c4] ASC), CHECK ([c4]>(0))) AS BEGIN 
+  INSERT @t1
+  SELECT 1 AS [c1], 'A' AS [c2], NULL AS [c3], @a * @a + @b AS [c4];
+RETURN
+END
diff --git a/sql/migrate/testdata/sqlserver/1_return_table.sql.golden b/sql/migrate/testdata/sqlserver/1_return_table.sql.golden
new file mode 100644
index 00000000..81ce4e4d
--- /dev/null
+++ b/sql/migrate/testdata/sqlserver/1_return_table.sql.golden
@@ -0,0 +1,9 @@
+CREATE FUNCTION [f2](@a as INT, @b as INT = 1)
+RETURNS TABLE
+AS RETURN SELECT @a as [a], @b as [b], (@a+@b)*2 as [p], @a*@b as [s];
+-- end --
+CREATE FUNCTION [f3] (@a int, @b int = 1) RETURNS @t1 TABLE ([c1] int NOT NULL, [c2] nvarchar(255) COLLATE SQL_Latin1_General_CP1_CI_AS NOT NULL, [c3] nvarchar(255) COLLATE SQL_Latin1_General_CP1_CI_AS DEFAULT N'G' NULL, [c4] int NOT NULL, PRIMARY KEY CLUSTERED ([c1] ASC), INDEX [idx] NONCLUSTERED ([c2] ASC), UNIQUE NONCLUSTERED ([c2] ASC, [c3] DESC), UNIQUE NONCLUSTERED ([c3] DESC, [c4] ASC), CHECK ([c4]>(0))) AS BEGIN 
+  INSERT @t1
+  SELECT 1 AS [c1], 'A' AS [c2], NULL AS [c3], @a * @a + @b AS [c4];
+RETURN
+END
\ No newline at end of file
diff --git a/sql/migrate/testdata/sqlserver/2_function.sql b/sql/migrate/testdata/sqlserver/2_function.sql
new file mode 100644
index 00000000..63d8afe5
--- /dev/null
+++ b/sql/migrate/testdata/sqlserver/2_function.sql
@@ -0,0 +1,19 @@
+-- atlas:delimiter \nGO
+
+CREATE FUNCTION [f3] (@a int, @b int = 1) RETURNS @t1 TABLE ([c1] int NOT NULL, [c2] nvarchar(255) COLLATE SQL_Latin1_General_CP1_CI_AS NOT NULL, [c3] nvarchar(255) COLLATE SQL_Latin1_General_CP1_CI_AS DEFAULT N'G' NULL, [c4] int NOT NULL, PRIMARY KEY CLUSTERED ([c1] ASC), INDEX [idx] NONCLUSTERED ([c2] ASC), UNIQUE NONCLUSTERED ([c2] ASC, [c3] DESC), UNIQUE NONCLUSTERED ([c3] DESC, [c4] ASC), CHECK ([c4]>(0))) AS BEGIN
+  INSERT @t1
+  SELECT 1 AS [c1], 'A' AS [c2], NULL AS [c3], @a * @a + @b AS [c4];
+RETURN
+END
+GO
+CREATE FUNCTION [f3] (@a int, @b int = 1) RETURNS @t1 TABLE ([c1] int NOT NULL, [c2] nvarchar(255) COLLATE SQL_Latin1_General_CP1_CI_AS NOT NULL, [c3] nvarchar(255) COLLATE SQL_Latin1_General_CP1_CI_AS DEFAULT N'G' NULL, [c4] int NOT NULL, PRIMARY KEY CLUSTERED ([c1] ASC), INDEX [idx] NONCLUSTERED ([c2] ASC), UNIQUE NONCLUSTERED ([c2] ASC, [c3] DESC), UNIQUE NONCLUSTERED ([c3] DESC, [c4] ASC), CHECK ([c4]>(0))) AS BEGIN
+  INSERT @t1
+  SELECT 1 AS [c1], 'A' AS [c2], NULL AS [c3], @a * @a + @b AS [c4];
+RETURN
+END
+GO
+CREATE FUNCTION [f3] (@a int, @b int = 1) RETURNS @t1 TABLE ([c1] int NOT NULL, [c2] nvarchar(255) COLLATE SQL_Latin1_General_CP1_CI_AS NOT NULL, [c3] nvarchar(255) COLLATE SQL_Latin1_General_CP1_CI_AS DEFAULT N'G' NULL, [c4] int NOT NULL, PRIMARY KEY CLUSTERED ([c1] ASC), INDEX [idx] NONCLUSTERED ([c2] ASC), UNIQUE NONCLUSTERED ([c2] ASC, [c3] DESC), UNIQUE NONCLUSTERED ([c3] DESC, [c4] ASC), CHECK ([c4]>(0))) AS BEGIN
+  INSERT @t1
+  SELECT 1 AS [c1], 'A' AS [c2], NULL AS [c3], @a * @a + @b AS [c4];
+RETURN
+END
\ No newline at end of file
diff --git a/sql/migrate/testdata/sqlserver/2_function.sql.golden b/sql/migrate/testdata/sqlserver/2_function.sql.golden
new file mode 100644
index 00000000..57a57590
--- /dev/null
+++ b/sql/migrate/testdata/sqlserver/2_function.sql.golden
@@ -0,0 +1,17 @@
+CREATE FUNCTION [f3] (@a int, @b int = 1) RETURNS @t1 TABLE ([c1] int NOT NULL, [c2] nvarchar(255) COLLATE SQL_Latin1_General_CP1_CI_AS NOT NULL, [c3] nvarchar(255) COLLATE SQL_Latin1_General_CP1_CI_AS DEFAULT N'G' NULL, [c4] int NOT NULL, PRIMARY KEY CLUSTERED ([c1] ASC), INDEX [idx] NONCLUSTERED ([c2] ASC), UNIQUE NONCLUSTERED ([c2] ASC, [c3] DESC), UNIQUE NONCLUSTERED ([c3] DESC, [c4] ASC), CHECK ([c4]>(0))) AS BEGIN
+  INSERT @t1
+  SELECT 1 AS [c1], 'A' AS [c2], NULL AS [c3], @a * @a + @b AS [c4];
+RETURN
+END
+-- end --
+CREATE FUNCTION [f3] (@a int, @b int = 1) RETURNS @t1 TABLE ([c1] int NOT NULL, [c2] nvarchar(255) COLLATE SQL_Latin1_General_CP1_CI_AS NOT NULL, [c3] nvarchar(255) COLLATE SQL_Latin1_General_CP1_CI_AS DEFAULT N'G' NULL, [c4] int NOT NULL, PRIMARY KEY CLUSTERED ([c1] ASC), INDEX [idx] NONCLUSTERED ([c2] ASC), UNIQUE NONCLUSTERED ([c2] ASC, [c3] DESC), UNIQUE NONCLUSTERED ([c3] DESC, [c4] ASC), CHECK ([c4]>(0))) AS BEGIN
+  INSERT @t1
+  SELECT 1 AS [c1], 'A' AS [c2], NULL AS [c3], @a * @a + @b AS [c4];
+RETURN
+END
+-- end --
+CREATE FUNCTION [f3] (@a int, @b int = 1) RETURNS @t1 TABLE ([c1] int NOT NULL, [c2] nvarchar(255) COLLATE SQL_Latin1_General_CP1_CI_AS NOT NULL, [c3] nvarchar(255) COLLATE SQL_Latin1_General_CP1_CI_AS DEFAULT N'G' NULL, [c4] int NOT NULL, PRIMARY KEY CLUSTERED ([c1] ASC), INDEX [idx] NONCLUSTERED ([c2] ASC), UNIQUE NONCLUSTERED ([c2] ASC, [c3] DESC), UNIQUE NONCLUSTERED ([c3] DESC, [c4] ASC), CHECK ([c4]>(0))) AS BEGIN
+  INSERT @t1
+  SELECT 1 AS [c1], 'A' AS [c2], NULL AS [c3], @a * @a + @b AS [c4];
+RETURN
+END
\ No newline at end of file