Repository: ariga/atlas
Pull request #3220 (2024-11-14 13:20:21)

sql/migrate: scan hash-style comments


Patch:
diff --git a/sql/migrate/lex.go b/sql/migrate/lex.go
index c8262199..d8be4903 100644
--- a/sql/migrate/lex.go
+++ b/sql/migrate/lex.go
@@ -112,6 +112,8 @@ type (
 		// EscapedStringExt enables the supported for PG extension for escaped strings and adopted by its flavors.
 		// See: https://www.postgresql.org/docs/current/sql-syntax-lexical.html#SQL-SYNTAX-STRINGS-ESCAPE.
 		EscapedStringExt bool
+		// HashComments enables MySQL/MariaDB hash-like (#) comments.
+		HashComments bool
 	}
 )
 
@@ -216,9 +218,7 @@ Scan:
 			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 == '#':
+		case r == '#' && s.HashComments:
 			s.comment("#", "\n")
 		case r == '-' && s.pick() == '-':
 			s.next()
diff --git a/sql/mysql/driver.go b/sql/mysql/driver.go
index 3360ea65..428ea91d 100644
--- a/sql/mysql/driver.go
+++ b/sql/mysql/driver.go
@@ -279,6 +279,7 @@ func (*Driver) ScanStmts(input string) ([]*migrate.Stmt, error) {
 		ScannerOptions: migrate.ScannerOptions{
 			MatchBegin:       true,
 			BackslashEscapes: true,
+			HashComments:     true,
 			// The following are not support by MySQL/MariaDB.
 			MatchBeginAtomic: false,
 			MatchDollarQuote: false,


Test patch:
diff --git a/sql/migrate/lex_test.go b/sql/migrate/lex_test.go
index efd21260..d6640a2d 100644
--- a/sql/migrate/lex_test.go
+++ b/sql/migrate/lex_test.go
@@ -20,10 +20,24 @@ func TestLocalFile_Stmts(t *testing.T) {
 	files, err := dir.Files()
 	require.NoError(t, err)
 	for _, f := range files {
-		stmts, err := f.Stmts()
+		sc := &Scanner{
+			ScannerOptions: ScannerOptions{
+				MatchBegin:       true,
+				MatchBeginAtomic: true,
+				MatchDollarQuote: true,
+				BackslashEscapes: true,
+				EscapedStringExt: true,
+				HashComments:     !strings.Contains(f.Name(), "_pg"),
+			},
+		}
+		decls, err := sc.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)
+		stmts := make([]string, len(decls))
+		for i, s := range decls {
+			stmts[i] = s.Text
+		}
 		require.Equalf(t, string(buf), strings.Join(stmts, "\n-- end --\n"), "mismatched statements in file %q", f.Name())
 	}
 }
@@ -125,7 +139,17 @@ cmd6;
 -- atlas:nolint
 cmd7;
 `
-	stmts, err := NewLocalFile("f", []byte(f)).StmtDecls()
+	sc := &Scanner{
+		ScannerOptions: ScannerOptions{
+			MatchBegin:       true,
+			MatchBeginAtomic: true,
+			MatchDollarQuote: true,
+			BackslashEscapes: true,
+			EscapedStringExt: true,
+			HashComments:     true,
+		},
+	}
+	stmts, err := sc.Scan(f)
 	require.NoError(t, err)
 	require.Len(t, stmts, 8)
 
diff --git a/sql/migrate/testdata/lex/18_pg_expr.sql b/sql/migrate/testdata/lex/18_pg_expr.sql
new file mode 100644
index 00000000..9c0f7f32
--- /dev/null
+++ b/sql/migrate/testdata/lex/18_pg_expr.sql
@@ -0,0 +1,66 @@
+-- Comment 1.
+CREATE INDEX "i" ON "s"."t" (((c #>> '{a,b,c}'::text[])));
+
+/*
+ Comment 2.
+ */
+SELECT name
+FROM company.employees
+WHERE info #>> '{department, name}' = 'Engineering';
+
+/*
+ SELECT name
+FROM company.employees
+WHERE info #>> '{contact, email}' = 'alice@company.com';
+ */
+SELECT name
+FROM company.employees
+WHERE info #>> '{contact, email}' = 'alice@company.com';
+
+-- Comment 3.
+CREATE INDEX "idx_emp_department" ON "company"."employees" (
+    (info #>> '{department, name}')
+);
+
+/*
+
+CREATE INDEX "idx_emp_contact" ON "company"."employees" (
+    LOWER(info #>> '{contact, email}')
+);
+ */
+CREATE INDEX "idx_emp_contact" ON "company"."employees" (
+    LOWER(info #>> '{contact, email}')
+);
+
+/**
+  SELECT
+    info #>> '{department, name}' AS department,
+    COUNT(*) AS emp_count
+FROM company.employees
+GROUP BY info #>> '{department, name}';
+ */
+SELECT
+    info #>> '{department, name}' AS department,
+    COUNT(*) AS emp_count
+FROM company.employees
+GROUP BY info #>> '{department, name}';
+
+/**
+  SELECT
+    info #>> '{department, name}' AS department,
+    COUNT(*) AS emp_count
+FROM company.employees
+GROUP BY info #>> '{department, name}';
+ */
+/**
+  SELECT
+    info #>> '{department, name}' AS department,
+    COUNT(*) AS emp_count
+FROM company.employees
+GROUP BY info #>> '{department, name}';
+ */
+SELECT name
+FROM company.employees
+WHERE (info #>> '{department, name}' = 'Engineering'
+       OR info #>> '{department, location}' = 'Building A')
+  AND info #>> '{contact, email}' LIKE '%@company.com';
\ No newline at end of file
diff --git a/sql/migrate/testdata/lex/18_pg_expr.sql.golden b/sql/migrate/testdata/lex/18_pg_expr.sql.golden
new file mode 100644
index 00000000..12c92409
--- /dev/null
+++ b/sql/migrate/testdata/lex/18_pg_expr.sql.golden
@@ -0,0 +1,29 @@
+CREATE INDEX "i" ON "s"."t" (((c #>> '{a,b,c}'::text[])));
+-- end --
+SELECT name
+FROM company.employees
+WHERE info #>> '{department, name}' = 'Engineering';
+-- end --
+SELECT name
+FROM company.employees
+WHERE info #>> '{contact, email}' = 'alice@company.com';
+-- end --
+CREATE INDEX "idx_emp_department" ON "company"."employees" (
+    (info #>> '{department, name}')
+);
+-- end --
+CREATE INDEX "idx_emp_contact" ON "company"."employees" (
+    LOWER(info #>> '{contact, email}')
+);
+-- end --
+SELECT
+    info #>> '{department, name}' AS department,
+    COUNT(*) AS emp_count
+FROM company.employees
+GROUP BY info #>> '{department, name}';
+-- end --
+SELECT name
+FROM company.employees
+WHERE (info #>> '{department, name}' = 'Engineering'
+       OR info #>> '{department, location}' = 'Building A')
+  AND info #>> '{contact, email}' LIKE '%@company.com';
\ No newline at end of file
diff --git a/sql/migrate/testdata/lex/6_skip_comment.sql b/sql/migrate/testdata/lex/6_skip_comment.sql
index d533424e..69882cc9 100644
--- a/sql/migrate/testdata/lex/6_skip_comment.sql
+++ b/sql/migrate/testdata/lex/6_skip_comment.sql
@@ -25,9 +25,9 @@ CREATE TABLE t5(
     -- comment
 ) ENGINE=InnoDB;
 
-# MySQL comment.
-CREATE INDEX "i" ON "s"."t" (((c #>> '{a,b,c}'::text[])));
-
 SELECT * FROM (
   SELECT * FROM t1 # comment
 );
+
+# This is a statement's comment.
+SELECT * FROM t2;
\ No newline at end of file
diff --git a/sql/migrate/testdata/lex/6_skip_comment.sql.golden b/sql/migrate/testdata/lex/6_skip_comment.sql.golden
index e9ee7629..ea5350ad 100644
--- a/sql/migrate/testdata/lex/6_skip_comment.sql.golden
+++ b/sql/migrate/testdata/lex/6_skip_comment.sql.golden
@@ -16,8 +16,8 @@ CREATE TABLE t5(
     -- comment
 ) ENGINE=InnoDB;
 -- end --
-CREATE INDEX "i" ON "s"."t" (((c #>> '{a,b,c}'::text[])));
--- end --
 SELECT * FROM (
   SELECT * FROM t1 # comment
-);
\ No newline at end of file
+);
+-- end --
+SELECT * FROM t2;
\ No newline at end of file