Repository: ariga/atlas
Pull request #1052 (2022-08-15 14:05:57)

sql/migrate: add support for replaying goose migration files


Patch:
diff --git a/sql/migrate/dir.go b/sql/migrate/dir.go
index e63520c0..da6921c5 100644
--- a/sql/migrate/dir.go
+++ b/sql/migrate/dir.go
@@ -181,7 +181,6 @@ type TemplateFormatter struct {
 //		template.Must(template.New("").Parse("{{now.Unix}}{{.Name}}.sql")),                 // name template
 //		template.Must(template.New("").Parse("{{range .Changes}}{{println .Cmd}}{{end}}")), // content template
 //	)
-//
 func NewTemplateFormatter(templates ...*template.Template) (*TemplateFormatter, error) {
 	if n := len(templates); n == 0 || n%2 == 1 {
 		return nil, fmt.Errorf("zero or odd number of templates given")
@@ -363,7 +362,6 @@ var reDirective = regexp.MustCompile(`^([ -~]*)atlas:(\w+)(?: +([ -~]*))*`)
 //	directive(c, "delimiter", "-- ")	// '-- atlas:delimiter.*'
 //	directive(c, "sum", "")				// 'atlas:sum.*'
 //	directive(c, "sum")					// '.*atlas:sum'
-//
 func directive(content, name string, prefix ...string) (string, bool) {
 	m := reDirective.FindStringSubmatch(content)
 	// In case the prefix was provided ensures it is matched.
diff --git a/sql/sqltool/tool.go b/sql/sqltool/tool.go
index b3dc61c4..158ed25d 100644
--- a/sql/sqltool/tool.go
+++ b/sql/sqltool/tool.go
@@ -5,8 +5,11 @@
 package sqltool
 
 import (
+	"bufio"
+	"bytes"
 	"fmt"
 	"io/fs"
+	"regexp"
 	"sort"
 	"strings"
 	"text/template"
@@ -61,8 +64,7 @@ var (
 	)
 )
 
-// GolangMigrateDir wraps migrate.LocalDir and provides
-// implementation compatible with golang-migrate/migrate.
+// GolangMigrateDir wraps migrate.LocalDir and provides implementation compatible with golang-migrate/migrate.
 type GolangMigrateDir struct{ *migrate.LocalDir }
 
 // NewGolangMigrateDir returns a new GolangMigrateDir.
@@ -74,7 +76,7 @@ func NewGolangMigrateDir(path string) (*GolangMigrateDir, error) {
 	return &GolangMigrateDir{dir}, nil
 }
 
-// Files implements Scanner.Files. It looks for all files with up.sql suffix and orders them by filename-
+// Files implements Scanner.Files. It looks for all files with up.sql suffix and orders them by filename.
 func (d *GolangMigrateDir) Files() ([]migrate.File, error) {
 	names, err := fs.Glob(d, "*.up.sql")
 	if err != nil {
@@ -95,8 +97,7 @@ func (d *GolangMigrateDir) Files() ([]migrate.File, error) {
 	return ret, nil
 }
 
-// GolangMigrateFile wraps migrate.LocalFile
-// with custom description function.
+// GolangMigrateFile wraps migrate.LocalFile with custom description function.
 type GolangMigrateFile struct {
 	*migrate.LocalFile
 }
@@ -106,6 +107,126 @@ func (f *GolangMigrateFile) Desc() string {
 	return strings.TrimSuffix(f.LocalFile.Desc(), ".up")
 }
 
+// GooseDir wraps migrate.LocalDir and provides implementation compatible with pressly/goose.
+type GooseDir struct{ *migrate.LocalDir }
+
+// NewGooseDir returns a new GooseDir.
+func NewGooseDir(path string) (*GooseDir, error) {
+	dir, err := migrate.NewLocalDir(path)
+	if err != nil {
+		return nil, err
+	}
+	return &GooseDir{dir}, nil
+}
+
+// GooseFile wraps migrate.LocalFile with custom statements function.
+type GooseFile struct {
+	*migrate.LocalFile
+}
+
+// Files implements Scanner.Files. It looks for all files with up.sql suffix and orders them by filename.
+func (d *GooseDir) Files() ([]migrate.File, error) {
+	files, err := d.LocalDir.Files()
+	if err != nil {
+		return nil, err
+	}
+	for i, f := range files {
+		files[i] = &GooseFile{f.(*migrate.LocalFile)}
+	}
+	return files, nil
+}
+
+// Stmts implements Scanner.Stmts. It understands the migration format used by pressly/goose sql migration files.
+func (f *GooseFile) Stmts() ([]string, error) {
+	var (
+		state, line int
+		stmts       []string
+		buf         strings.Builder
+		sc          = bufio.NewScanner(bytes.NewReader(f.Bytes()))
+	)
+	for sc.Scan() {
+		line++
+		s := sc.Text()
+		// Handle pragmas.
+		if strings.HasPrefix(s, pragma) {
+			switch strings.TrimSpace(strings.TrimPrefix(s, pragma)) {
+			case "Up":
+				switch state {
+				case none: // found the "up" part of the file
+					state = up
+				default:
+					return nil, unexpectedGoosePragmaErr(f, line, "Up")
+				}
+			case "Down":
+				switch state {
+				case up: // found the "down" part
+					if rest := strings.TrimSpace(buf.String()); len(rest) > 0 {
+						return nil, unexpectedGoosePragmaErr(f, line, "Down")
+					}
+					return stmts, nil
+				default:
+					return nil, unexpectedGoosePragmaErr(f, line, "Down")
+				}
+			case "StatementBegin":
+				switch state {
+				case up:
+					state = begin // begin of a statement
+				default:
+					return nil, unexpectedGoosePragmaErr(f, line, "StatementBegin")
+				}
+			case "StatementEnd":
+				switch state {
+				case begin:
+					state = end // end of a statement
+				default:
+					return nil, unexpectedGoosePragmaErr(f, line, "StatementEnd")
+				}
+			}
+		}
+		// Write the line of the statement.
+		if !rePragma.MatchString(s) {
+			if _, err := buf.WriteString(s + "\n"); err != nil {
+				return nil, err
+			}
+		}
+		switch state {
+		case up: // end of statement if line ends with semicolon
+			if s := strings.TrimSpace(s); strings.HasSuffix(s, ";") && !strings.HasPrefix(s, "--") {
+				stmts = append(stmts, strings.TrimSpace(buf.String()))
+				buf.Reset()
+			}
+		case end: // end of statement marked by pragma
+			stmts = append(stmts, strings.TrimSpace(buf.String()))
+			buf.Reset()
+			state = up // back in up block
+		}
+	}
+	if err := sc.Err(); err != nil {
+		return nil, fmt.Errorf("sql/migrate: scanning migration %q: %w", f.Name(), err)
+	}
+	if state == none {
+		return nil, fmt.Errorf("sql/migrate: empty migration %q", f.Name())
+	}
+	return stmts, nil
+}
+
+const (
+	none int = iota // state when parsing goose sql file
+	up
+	begin
+	end
+	pragma = "-- +goose"
+)
+
+var rePragma = regexp.MustCompile("-- \\+goose Up|Down|StatementBegin|StatementEnd")
+
+func unexpectedGoosePragmaErr(f *GooseFile, line int, pragma string) error {
+	return fmt.Errorf(
+		"sql/migrate: goose: %s:%d unexpected pragma '%s'",
+		f.Name(), line, pragma,
+	)
+}
+
 // funcs contains the template.FuncMap for the different formatters.
 var funcs = template.FuncMap{
 	"inc": func(x int) int { return x + 1 },


Test patch:
diff --git a/internal/integration/testdata/mysql/cli-migrate-diff-format.txt b/internal/integration/testdata/mysql/cli-migrate-diff-format.txt
new file mode 100644
index 00000000..cfbac315
--- /dev/null
+++ b/internal/integration/testdata/mysql/cli-migrate-diff-format.txt
@@ -0,0 +1,65 @@
+only maria103
+
+mkdir migrations
+
+atlas migrate diff --dev-url URL --to file://1.hcl --dir-format golang-migrate first
+cmpmig 0 golang-migrate/1.down.sql
+cmpmig 1 golang-migrate/1.up.sql
+
+atlas migrate diff --dev-url URL --to file://2.hcl --dir-format golang-migrate second
+cmpmig 2 golang-migrate/2.down.sql
+cmpmig 3 golang-migrate/2.up.sql
+
+-- golang-migrate/1.up.sql --
+-- create "users" table
+CREATE TABLE `users` (`id` bigint NOT NULL AUTO_INCREMENT, PRIMARY KEY (`id`)) CHARSET utf8mb4 COLLATE utf8mb4_general_ci;
+
+-- golang-migrate/1.down.sql --
+-- reverse: create "users" table
+DROP TABLE `users`;
+
+-- golang-migrate/2.up.sql --
+-- modify "users" table
+ALTER TABLE `users` ADD COLUMN `email` varchar(100) NULL;
+
+-- golang-migrate/2.down.sql --
+-- reverse: modify "users" table
+ALTER TABLE `users` DROP COLUMN `email`;
+
+-- 1.hcl --
+schema "script_cli_migrate_diff_format" {}
+
+table "users" {
+  schema = schema.script_cli_migrate_diff_format
+  column "id" {
+    null = false
+    type = bigint
+    auto_increment = true
+  }
+  primary_key  {
+    columns = [column.id]
+  }
+  charset = "utf8mb4"
+  collate = "utf8mb4_general_ci"
+}
+
+-- 2.hcl --
+schema "script_cli_migrate_diff_format" {}
+
+table "users" {
+  schema = schema.script_cli_migrate_diff_format
+  column "id" {
+    null = false
+    type = bigint
+    auto_increment = true
+  }
+  column "email" {
+    null = true
+    type = varchar(100)
+  }
+  primary_key  {
+    columns = [column.id]
+  }
+  charset = "utf8mb4"
+  collate = "utf8mb4_general_ci"
+}
\ No newline at end of file
diff --git a/internal/integration/testdata/mysql/cli-migrate-diff.txt b/internal/integration/testdata/mysql/cli-migrate-diff.txt
index b63ffa2c..05152585 100644
--- a/internal/integration/testdata/mysql/cli-migrate-diff.txt
+++ b/internal/integration/testdata/mysql/cli-migrate-diff.txt
@@ -1,4 +1,4 @@
-only maria
+only maria107 maria102
 
 exec mkdir migrations
 
@@ -74,10 +74,6 @@ ALTER TABLE `users` ADD COLUMN `create_time` timestamp(6) NOT NULL DEFAULT CURRE
 -- modify "users" table
 ALTER TABLE `users` ADD COLUMN `create_time` timestamp(6) NOT NULL DEFAULT (current_timestamp(6));
 
--- maria103/2.sql --
--- modify "users" table
-ALTER TABLE `users` ADD COLUMN `create_time` timestamp(6) NOT NULL DEFAULT (current_timestamp(6));
-
 -- maria107/2.sql --
 -- modify "users" table
 ALTER TABLE `users` ADD COLUMN `create_time` timestamp(6) NOT NULL DEFAULT (current_timestamp(6));
diff --git a/sql/sqltool/testdata/goose/1_initial.sql b/sql/sqltool/testdata/goose/1_initial.sql
new file mode 100644
index 00000000..111699b5
--- /dev/null
+++ b/sql/sqltool/testdata/goose/1_initial.sql
@@ -0,0 +1,19 @@
+-- +goose Up
+CREATE TABLE post
+(
+    id    int NOT NULL,
+    title text,
+    body  text,
+    PRIMARY KEY (id)
+);
+
+ALTER TABLE post ADD created_at TIMESTAMP NOT NULL;
+
+INSERT INTO post (title) VALUES (
+'This is
+my multiline
+
+value');
+
+-- +goose Down
+DROP TABLE post;
\ No newline at end of file
diff --git a/sql/sqltool/testdata/goose/2_second_migration.sql b/sql/sqltool/testdata/goose/2_second_migration.sql
new file mode 100644
index 00000000..416c9c5a
--- /dev/null
+++ b/sql/sqltool/testdata/goose/2_second_migration.sql
@@ -0,0 +1,30 @@
+
+
+
+-- +goose Up
+
+
+
+
+-- +goose StatementBegin
+CREATE
+OR REPLACE FUNCTION histories_partition_creation( DATE, DATE )
+returns void AS $$
+DECLARE
+create_query text;
+BEGIN
+FOR create_query IN
+SELECT 'CREATE TABLE IF NOT EXISTS histories_'
+           || TO_CHAR(d, 'YYYY_MM')
+           || ' ( CHECK( created_at >= timestamp '''
+           || TO_CHAR(d, 'YYYY-MM-DD 00:00:00')
+           || ''' AND created_at < timestamp '''
+           || TO_CHAR(d + INTERVAL '1 month', 'YYYY-MM-DD 00:00:00')
+           || ''' ) ) inherits ( histories );'
+FROM generate_series($1, $2, '1 month') AS d LOOP
+    EXECUTE create_query;
+END LOOP;  -- LOOP END
+END;         -- FUNCTION END
+$$
+language plpgsql;
+-- +goose StatementEnd
\ No newline at end of file
diff --git a/sql/sqltool/testdata/goose/atlas.sum b/sql/sqltool/testdata/goose/atlas.sum
new file mode 100644
index 00000000..3d2ec76e
--- /dev/null
+++ b/sql/sqltool/testdata/goose/atlas.sum
@@ -0,0 +1,3 @@
+h1:MFwmYzrSq09/I/jZgUuUoiRTBCtz+0UWOTotkcBdkIE=
+1_initial.sql h1:mOgAYIay1aUHP9BgbtNlwjSPsiw0yXJ5HIW+kOXyt1s=
+2_second_migration.sql h1:kOL5bbnrzs3M/MDRmLLvuwrO95p/pocjHAP9hdrR0e8=
diff --git a/sql/sqltool/tool_test.go b/sql/sqltool/tool_test.go
index 1c6a9b57..2cf16498 100644
--- a/sql/sqltool/tool_test.go
+++ b/sql/sqltool/tool_test.go
@@ -131,28 +131,61 @@ DROP TABLE t1 IF EXISTS;
 }
 
 func TestScanners(t *testing.T) {
-	d, err := sqltool.NewGolangMigrateDir("testdata/golang-migrate")
-	require.NoError(t, err)
-
-	files, err := d.Files()
-	require.NoError(t, err)
-	require.Len(t, files, 2)
-	require.Equal(t, "1_initial.up.sql", files[0].Name())
-	require.Equal(t, "2_second_migration.up.sql", files[1].Name())
-
-	first, second := files[0], files[1]
-
-	stmts, err := first.Stmts()
-	require.NoError(t, err)
-	require.Equal(t, []string{"CREATE TABLE tbl\n(\n    col INT\n);"}, stmts)
-	stmts, err = second.Stmts()
-	require.NoError(t, err)
-	require.Equal(t, []string{"CREATE TABLE tbl_2 (col INT);"}, stmts)
-
-	require.Equal(t, "1", first.Version())
-	require.Equal(t, "initial", first.Desc())
-	require.Equal(t, "2", second.Version())
-	require.Equal(t, "second_migration", second.Desc())
+	for _, tt := range []struct {
+		name                   string
+		dir                    migrate.Dir
+		versions, descriptions []string
+		stmts                  [][]string
+	}{
+		{
+			name: "golang-migrate",
+			dir: func() migrate.Dir {
+				d, err := sqltool.NewGolangMigrateDir("testdata/golang-migrate")
+				require.NoError(t, err)
+				return d
+			}(),
+			versions:     []string{"1", "2"},
+			descriptions: []string{"initial", "second_migration"},
+			stmts: [][]string{
+				{"CREATE TABLE tbl\n(\n    col INT\n);"},
+				{"CREATE TABLE tbl_2 (col INT);"},
+			},
+		},
+		{
+			name: "goose",
+			dir: func() migrate.Dir {
+				d, err := sqltool.NewGooseDir("testdata/goose")
+				require.NoError(t, err)
+				return d
+			}(),
+			versions:     []string{"1", "2"},
+			descriptions: []string{"initial", "second_migration"},
+			stmts: [][]string{
+				{
+					"CREATE TABLE post\n(\n    id    int NOT NULL,\n    title text,\n    body  text,\n    PRIMARY KEY (id)\n);",
+					"ALTER TABLE post ADD created_at TIMESTAMP NOT NULL;",
+					"INSERT INTO post (title) VALUES (\n'This is\nmy multiline\n\nvalue');",
+				},
+				{"CREATE\nOR REPLACE FUNCTION histories_partition_creation( DATE, DATE )\nreturns void AS $$\nDECLARE\ncreate_query text;\nBEGIN\nFOR create_query IN\nSELECT 'CREATE TABLE IF NOT EXISTS histories_'\n           || TO_CHAR(d, 'YYYY_MM')\n           || ' ( CHECK( created_at >= timestamp '''\n           || TO_CHAR(d, 'YYYY-MM-DD 00:00:00')\n           || ''' AND created_at < timestamp '''\n           || TO_CHAR(d + INTERVAL '1 month', 'YYYY-MM-DD 00:00:00')\n           || ''' ) ) inherits ( histories );'\nFROM generate_series($1, $2, '1 month') AS d LOOP\n    EXECUTE create_query;\nEND LOOP;  -- LOOP END\nEND;         -- FUNCTION END\n$$\nlanguage plpgsql;"},
+			},
+		},
+	} {
+		t.Run(tt.name, func(t *testing.T) {
+			files, err := tt.dir.Files()
+			require.NoError(t, err)
+			require.Len(t, files, len(tt.versions))
+			for i := range tt.versions {
+				require.Equal(t, tt.versions[i], files[i].Version())
+				require.Equal(t, tt.descriptions[i], files[i].Desc())
+				stmts, err := files[i].Stmts()
+				require.NoError(t, err)
+				require.Len(t, stmts, len(tt.stmts[i]))
+				for j, stmt := range stmts {
+					require.Equal(t, tt.stmts[i][j], stmt)
+				}
+			}
+		})
+	}
 }
 
 func dir(t *testing.T) migrate.Dir {