migrate: Add support CQL comment callbacks

This patch adds a new migration event type CallComment that it triggered by adding `-- CALL <name>;` comment in a CQL file.

Fixes #101

migrate/README.md

@@ -2,15 +2,13 @@
 
 Package `migrate` provides simple and flexible CQL migrations. 
 Migrations can be read from a flat directory containing cql files.
-There is no imposed naming schema, migration name is file name and the
-migrations are processed in lexicographical order. Caller provides a
-`gocql.Session`, the session must use a desired keyspace as migrate would try
-to create migrations table.
+There is no imposed naming schema, migration name is file name and the migrations are processed in lexicographical order.
+Caller provides a `gocqlx.Session`, the session must use a desired keyspace as migrate would try to create migrations table.
 
 ## Features
 
 * Each CQL statement will run once
-* Go code migrations using callbacks 
+* Go code migrations using callbacks
 
 ## Example
 
@@ -20,7 +18,7 @@ package main
 import (
     "context"
 
-    "github.com/scylladb/gocqlx/migrate"
+    "github.com/scylladb/gocqlx/v2/migrate"
 )
 
 const dir = "./cql" 

migrate/callback.go

@@ -17,12 +17,17 @@ type CallbackEvent uint8
 const (
 	BeforeMigration CallbackEvent = iota
 	AfterMigration
+	CallComment
 )
 
-// CallbackFunc enables interrupting the migration process and executing code
-// while migrating. If error is returned the migration is aborted.
+// CallbackFunc enables execution of arbitrary Go code during migration.
+// If error is returned the migration is aborted.
+// BeforeMigration and AfterMigration are triggered before and after processing
+// of each migration file respectively.
+// CallComment is triggered for each comment in a form `-- CALL <name>;` (note the semicolon).
 type CallbackFunc func(ctx context.Context, session gocqlx.Session, ev CallbackEvent, name string) error
 
-// Callback is called before and after each migration.
+// Callback is means of executing Go code during migrations.
+// Use this variable to register a global callback dispatching function.
 // See CallbackFunc for details.
 var Callback CallbackFunc

migrate/doc.go

@@ -4,8 +4,6 @@
 
 // Package migrate provides simple and flexible CLQ migrations.
 // Migrations can be read from a flat directory containing cql files.
-// There is no imposed naming schema, migration name is file name and the
-// migrations are processed in lexicographical order. Caller provides a
-// gocql.Session, the session must use a desired keyspace as migrate would try
-// to create migrations table.
+// There is no imposed naming schema, migration name is file name and the migrations are processed in lexicographical order.
+// Caller provides a gocqlx.Session, the session must use a desired keyspace as migrate would try to create migrations table.
 package migrate

migrate/export_test.go

@@ -0,0 +1,5 @@
+package migrate
+
+func IsCallback(stmt string) (name string) {
+	return isCallback(stmt)
+}

migrate/migrate.go

@@ -13,6 +13,7 @@ import (
 	"io/ioutil"
 	"os"
 	"path/filepath"
+	"regexp"
 	"sort"
 	"strings"
 	"time"
@@ -90,6 +91,8 @@ func ensureInfoTable(ctx context.Context, session gocqlx.Session) error {
 }
 
 // Migrate reads the cql files from a directory and applies required migrations.
+// It also supports code based migrations, see Callback and CallbackFunc.
+// Any comment in form `-- CALL <name>;` will trigger an CallComment callback.
 func Migrate(ctx context.Context, session gocqlx.Session, dir string) error {
 	// get database migrations
 	dbm, err := List(ctx, session)
@@ -215,10 +218,21 @@ func applyMigration(ctx context.Context, session gocqlx.Session, path string, do
 			}
 		}
 
-		// execute
-		q := session.ContextQuery(ctx, stmt, nil).RetryPolicy(nil)
-		if err := q.ExecRelease(); err != nil {
-			return fmt.Errorf("statement %d failed: %s", i, err)
+		// trim new lines and all whitespace characters
+		stmt = strings.TrimSpace(stmt)
+
+		if cb := isCallback(stmt); cb != "" {
+			if Callback == nil {
+				return fmt.Errorf("statement %d failed: missing callback handler while trying to call %s", i, cb)
+			}
+			if err := Callback(ctx, session, CallComment, cb); err != nil {
+				return fmt.Errorf("callback %s failed: %s", cb, err)
+			}
+		} else {
+			q := session.ContextQuery(ctx, stmt, nil).RetryPolicy(nil)
+			if err := q.ExecRelease(); err != nil {
+				return fmt.Errorf("statement %d failed: %s", i, err)
+			}
 		}
 
 		// update info
@@ -240,3 +254,13 @@ func applyMigration(ctx context.Context, session gocqlx.Session, path string, do
 
 	return nil
 }
+
+var cbRegexp = regexp.MustCompile("^-- *CALL +(.+);$")
+
+func isCallback(stmt string) (name string) {
+	s := cbRegexp.FindStringSubmatch(stmt)
+	if len(s) == 0 {
+		return ""
+	}
+	return s[1]
+}

migrate/migrate_test.go

@@ -90,7 +90,7 @@ func TestMigration(t *testing.T) {
 		dir := makeMigrationDir(t, 4)
 		defer os.Remove(dir)
 
-		temperFile(t, dir, "3.cql")
+		appendMigrationFile(t, dir, 3, "\nSELECT * FROM bla;\n")
 
 		if err := migrate.Migrate(ctx, session, dir); err == nil || !strings.Contains(err.Error(), "tempered") {
 			t.Fatal("expected error")
@@ -129,10 +129,62 @@ func TestMigrationNoSemicolon(t *testing.T) {
 	}
 }
 
+func TestIsCallback(t *testing.T) {
+	table := []struct {
+		Name string
+		Stmt string
+		Cb   string
+	}{
+		{
+			Name: "CQL statement",
+			Stmt: "SELECT * from X;",
+		},
+		{
+			Name: "CQL comment",
+			Stmt: "-- Item",
+		},
+		{
+			Name: "CALL without space",
+			Stmt: "--CALL Foo;",
+			Cb:   "Foo",
+		},
+		{
+			Name: "CALL with space",
+			Stmt: "-- CALL Foo;",
+			Cb:   "Foo",
+		},
+		{
+			Name: "CALL with many spaces",
+			Stmt: "--   CALL Foo;",
+			Cb:   "Foo",
+		},
+		{
+			Name: "CALL with many spaces 2",
+			Stmt: "--   CALL   Foo;",
+			Cb:   "Foo",
+		},
+		{
+			Name: "CALL with unicode",
+			Stmt: "-- CALL α;",
+			Cb:   "α",
+		},
+	}
+
+	for i := range table {
+		test := table[i]
+		t.Run(test.Name, func(t *testing.T) {
+			if migrate.IsCallback(test.Stmt) != test.Cb {
+				t.Errorf("IsCallback(%s)=%s, expected %s", test.Stmt, migrate.IsCallback(test.Stmt), test.Cb)
+			}
+		})
+	}
+}
+
 func TestMigrationCallback(t *testing.T) {
 	var (
 		beforeCalled int
 		afterCalled  int
+		inCalled     int
 	)
 	migrate.Callback = func(ctx context.Context, session gocqlx.Session, ev migrate.CallbackEvent, name string) error {
 		switch ev {
@@ -140,6 +192,8 @@ func TestMigrationCallback(t *testing.T) {
 			beforeCalled += 1
 		case migrate.AfterMigration:
 			afterCalled += 1
+		case migrate.CallComment:
+			inCalled += 1
 		}
 		return nil
 	}
@@ -151,14 +205,18 @@ func TestMigrationCallback(t *testing.T) {
 	reset := func() {
 		beforeCalled = 0
 		afterCalled = 0
+		inCalled = 0
 	}
 
-	assertCallbacks := func(t *testing.T, b, a int) {
-		if beforeCalled != b {
-			t.Fatalf("expected %d before calls got %d", b, beforeCalled)
+	assertCallbacks := func(t *testing.T, before, afer, in int) {
+		if beforeCalled != before {
+			t.Fatalf("expected %d before calls got %d", before, beforeCalled)
 		}
-		if afterCalled != b {
-			t.Fatalf("expected %d after calls got %d", a, afterCalled)
+		if afterCalled != afer {
+			t.Fatalf("expected %d after calls got %d", afer, afterCalled)
+		}
+		if inCalled != in {
+			t.Fatalf("expected %d in calls got %d", in, inCalled)
 		}
 	}
 
@@ -180,7 +238,7 @@ func TestMigrationCallback(t *testing.T) {
 		if err := migrate.Migrate(ctx, session, dir); err != nil {
 			t.Fatal(err)
 		}
-		assertCallbacks(t, 2, 2)
+		assertCallbacks(t, 2, 2, 0)
 	})
 
 	t.Run("no duplicate calls", func(t *testing.T) {
@@ -191,7 +249,21 @@ func TestMigrationCallback(t *testing.T) {
 		if err := migrate.Migrate(ctx, session, dir); err != nil {
 			t.Fatal(err)
 		}
-		assertCallbacks(t, 2, 2)
+		assertCallbacks(t, 2, 2, 0)
+	})
+
+	t.Run("in calls", func(t *testing.T) {
+		dir := makeMigrationDir(t, 4)
+		defer os.Remove(dir)
+		reset()
+
+		appendMigrationFile(t, dir, 4, "\n-- CALL Foo;\n")
+		appendMigrationFile(t, dir, 5, "\n-- CALL Bar;\n")
+
+		if err := migrate.Migrate(ctx, session, dir); err != nil {
+			t.Fatal(err)
+		}
+		assertCallbacks(t, 2, 2, 2)
 	})
 }
 
@@ -204,7 +276,7 @@ func makeMigrationDir(tb testing.TB, n int) (dir string) {
 	}
 
 	for i := 0; i < n; i++ {
-		path := filepath.Join(dir, fmt.Sprint(i, ".cql"))
+		path := migrateFilePath(dir, i)
 		cql := []byte(fmt.Sprintf(insertMigrate, i) + ";")
 		if err := ioutil.WriteFile(path, cql, os.ModePerm); err != nil {
 			os.Remove(dir)
@@ -225,10 +297,17 @@ func countMigrations(tb testing.TB, session gocqlx.Session) int {
 	return v
 }
 
-func temperFile(tb testing.TB, dir, name string) {
-	tb.Helper()
-
-	if err := ioutil.WriteFile(filepath.Join(dir, name), []byte("SELECT * FROM bla;"), os.ModePerm); err != nil {
+func appendMigrationFile(tb testing.TB, dir string, i int, text string) {
+	path := migrateFilePath(dir, i)
+	f, err := os.OpenFile(path, os.O_WRONLY|os.O_CREATE|os.O_APPEND, os.ModePerm)
+	if err != nil {
+		tb.Fatal(err)
+	}
+	if _, err := f.WriteString(text); err != nil {
 		tb.Fatal(err)
 	}
 }
+
+func migrateFilePath(dir string, i int) string {
+	return filepath.Join(dir, fmt.Sprint(i, ".cql"))
+}