Strengthen SQL transformation tests

pkg/migrations/op_common_test.go

@@ -18,12 +18,14 @@ import (
 )
 
 type TestCase struct {
-	name          string
-	migrations    []migrations.Migration
-	wantStartErr  error
-	afterStart    func(t *testing.T, db *sql.DB, schema string)
-	afterComplete func(t *testing.T, db *sql.DB, schema string)
-	afterRollback func(t *testing.T, db *sql.DB, schema string)
+	name            string
+	migrations      []migrations.Migration
+	wantStartErr    error
+	wantRollbackErr error
+	wantCompleteErr error
+	afterStart      func(t *testing.T, db *sql.DB, schema string)
+	afterComplete   func(t *testing.T, db *sql.DB, schema string)
+	afterRollback   func(t *testing.T, db *sql.DB, schema string)
 }
 
 type TestCases []TestCase
@@ -69,7 +71,14 @@ func ExecuteTests(t *testing.T, tests TestCases, opts ...roll.Option) {
 				}
 
 				// roll back the migration
-				if err := mig.Rollback(ctx); err != nil {
+				err = mig.Rollback(ctx)
+				if tt.wantRollbackErr != nil {
+					if !errors.Is(err, tt.wantRollbackErr) {
+						t.Fatalf("Expected error %q, got %q", tt.wantRollbackErr, err)
+					}
+					return
+				}
+				if err != nil {
 					t.Fatalf("Failed to roll back migration: %v", err)
 				}
 
@@ -84,7 +93,14 @@ func ExecuteTests(t *testing.T, tests TestCases, opts ...roll.Option) {
 				}
 
 				// complete the last migration
-				if err := mig.Complete(ctx); err != nil {
+				err = mig.Complete(ctx)
+				if tt.wantCompleteErr != nil {
+					if !errors.Is(err, tt.wantCompleteErr) {
+						t.Fatalf("Expected error %q, got %q", tt.wantCompleteErr, err)
+					}
+					return
+				}
+				if err != nil {
 					t.Fatalf("Failed to complete migration: %v", err)
 				}
 

pkg/migrations/op_raw_sql_test.go

@@ -4,11 +4,11 @@ package migrations_test
 
 import (
 	"database/sql"
-	"fmt"
 	"testing"
 
 	"github.com/xataio/pgroll/pkg/migrations"
 	"github.com/xataio/pgroll/pkg/roll"
+	"github.com/xataio/pgroll/pkg/testutils"
 )
 
 func TestRawSQL(t *testing.T) {
@@ -190,71 +190,102 @@ func TestRawSQL(t *testing.T) {
 func TestRawSQLTransformation(t *testing.T) {
 	t.Parallel()
 
-	t.Run("for normal raw SQL operations with up and down SQL", func(t *testing.T) {
-		ExecuteTests(t, TestCases{
-			{
-				name: "SQL transformer rewrites up and down SQL",
-				migrations: []migrations.Migration{
-					{
-						Name: "01_create_table",
-						Operations: migrations.Operations{
-							&migrations.OpRawSQL{
-								Up:   "CREATE TABLE apples(id int)",
-								Down: "CREATE TABLE bananas(id int)",
-							},
+	sqlTransformer := testutils.NewMockSQLTransformer(map[string]string{
+		"CREATE TABLE people(id int)":     "CREATE TABLE users(id int)",
+		"DROP TABLE people":               "DROP TABLE users",
+		"CREATE TABLE restricted(id int)": testutils.MockSQLTransformerError,
+	})
+
+	ExecuteTests(t, TestCases{
+		{
+			name: "SQL transformer rewrites up and down SQL",
+			migrations: []migrations.Migration{
+				{
+					Name: "01_create_table",
+					Operations: migrations.Operations{
+						&migrations.OpRawSQL{
+							Up:   "CREATE TABLE people(id int)",
+							Down: "DROP TABLE people",
 						},
 					},
 				},
-				afterStart: func(t *testing.T, db *sql.DB, schema string) {
-					// The transformed `up` SQL was used in place of the original SQL
-					TableMustExist(t, db, schema, "table_1")
-					TableMustNotExist(t, db, schema, "apples")
-				},
-				afterRollback: func(t *testing.T, db *sql.DB, schema string) {
-					// The transformed `down` SQL was used in place of the original SQL
-					TableMustExist(t, db, schema, "table_2")
-					TableMustNotExist(t, db, schema, "bananas")
-				},
-				afterComplete: func(t *testing.T, db *sql.DB, schema string) {
-				},
 			},
-		}, roll.WithSQLTransformer(&simpleSQLTransformer{}))
-	})
-
-	t.Run("for raw SQL operations that run on complete", func(t *testing.T) {
-		ExecuteTests(t, TestCases{
-			{
-				name: "SQL transformer rewrites up SQL when up is run on completion",
-				migrations: []migrations.Migration{
-					{
-						Name: "01_create_table",
-						Operations: migrations.Operations{
-							&migrations.OpRawSQL{
-								Up:         "CREATE TABLE apples(id int)",
-								OnComplete: true,
-							},
+			afterStart: func(t *testing.T, db *sql.DB, schema string) {
+				// The transformed `up` SQL was used in place of the original SQL
+				TableMustExist(t, db, schema, "users")
+			},
+			afterRollback: func(t *testing.T, db *sql.DB, schema string) {
+				// The transformed `down` SQL was used in place of the original SQL
+				TableMustNotExist(t, db, schema, "users")
+			},
+			afterComplete: func(t *testing.T, db *sql.DB, schema string) {
+			},
+		},
+		{
+			name: "SQL transformer rewrites up SQL when up is run on completion",
+			migrations: []migrations.Migration{
+				{
+					Name: "01_create_table",
+					Operations: migrations.Operations{
+						&migrations.OpRawSQL{
+							Up:         "CREATE TABLE people(id int)",
+							OnComplete: true,
 						},
 					},
 				},
-				afterStart: func(t *testing.T, db *sql.DB, schema string) {
+			},
+			afterStart: func(t *testing.T, db *sql.DB, schema string) {
+			},
+			afterRollback: func(t *testing.T, db *sql.DB, schema string) {
+			},
+			afterComplete: func(t *testing.T, db *sql.DB, schema string) {
+				// The transformed `up` SQL was used in place of the original SQL
+				TableMustExist(t, db, schema, "users")
+			},
+		},
+		{
+			name: "raw SQL operation fails when SQL transformer returns an error on up SQL",
+			migrations: []migrations.Migration{
+				{
+					Name: "01_create_table",
+					Operations: migrations.Operations{
+						&migrations.OpRawSQL{
+							Up: "CREATE TABLE restricted(id int)",
+						},
+					},
 				},
-				afterRollback: func(t *testing.T, db *sql.DB, schema string) {
+			},
+			wantStartErr: testutils.ErrMockSQLTransformer,
+		},
+		{
+			name: "raw SQL operation fails when SQL transformer returns an error on down SQL",
+			migrations: []migrations.Migration{
+				{
+					Name: "01_create_table",
+					Operations: migrations.Operations{
+						&migrations.OpRawSQL{
+							Up:   "CREATE TABLE products(id int)",
+							Down: "CREATE TABLE restricted(id int)",
+						},
+					},
 				},
-				afterComplete: func(t *testing.T, db *sql.DB, schema string) {
-					// The transformed `up` SQL was used in place of the original SQL
-					TableMustExist(t, db, schema, "table_1")
-					TableMustNotExist(t, db, schema, "apples")
+			},
+			wantRollbackErr: testutils.ErrMockSQLTransformer,
+		},
+		{
+			name: "raw SQL onComplete operation fails when SQL transformer returns an error on up SQL",
+			migrations: []migrations.Migration{
+				{
+					Name: "01_create_table",
+					Operations: migrations.Operations{
+						&migrations.OpRawSQL{
+							Up:         "CREATE TABLE restricted(id int)",
+							OnComplete: true,
+						},
+					},
 				},
 			},
-		}, roll.WithSQLTransformer(&simpleSQLTransformer{}))
-	})
-}
-
-type simpleSQLTransformer struct {
-	counter int
-}
-
-func (s *simpleSQLTransformer) TransformSQL(sql string) (string, error) {
-	s.counter++
-	return fmt.Sprintf("CREATE TABLE table_%d(id int)", s.counter), nil
+			wantCompleteErr: testutils.ErrMockSQLTransformer,
+		},
+	}, roll.WithSQLTransformer(sqlTransformer))
 }

pkg/roll/execute_test.go

@@ -665,12 +665,12 @@ func TestSQLTransformerOptionIsUsedWhenCreatingTriggers(t *testing.T) {
 	t.Parallel()
 
 	t.Run("when the SQL transformer is used to rewrite SQL", func(t *testing.T) {
-		var transformer migrations.SQLTransformer = migrations.SQLTransformerFunc(
-			func(sql string) (string, error) {
-				return "'rewritten'", nil
-			},
-		)
-		opts := []roll.Option{roll.WithSQLTransformer(transformer)}
+		t.Parallel()
+
+		sqlTransformer := testutils.NewMockSQLTransformer(map[string]string{
+			"'apples'": "'rewritten'",
+		})
+		opts := []roll.Option{roll.WithSQLTransformer(sqlTransformer)}
 
 		testutils.WithMigratorAndConnectionToContainerWithOptions(t, opts, func(mig *roll.Roll, db *sql.DB) {
 			ctx := context.Background()
@@ -696,7 +696,7 @@ func TestSQLTransformerOptionIsUsedWhenCreatingTriggers(t *testing.T) {
 				Operations: migrations.Operations{
 					&migrations.OpAddColumn{
 						Table: "table1",
-						Up:    "apples",
+						Up:    "'apples'",
 						Column: migrations.Column{
 							Name:     "description",
 							Type:     "text",
@@ -721,14 +721,12 @@ func TestSQLTransformerOptionIsUsedWhenCreatingTriggers(t *testing.T) {
 	})
 
 	t.Run("when the SQL transformer returns an error", func(t *testing.T) {
-		transformerError := errors.New("oops")
+		t.Parallel()
 
-		var transformer migrations.SQLTransformer = migrations.SQLTransformerFunc(
-			func(sql string) (string, error) {
-				return "", transformerError
-			},
-		)
-		opts := []roll.Option{roll.WithSQLTransformer(transformer)}
+		sqlTransformer := testutils.NewMockSQLTransformer(map[string]string{
+			"'apples'": testutils.MockSQLTransformerError,
+		})
+		opts := []roll.Option{roll.WithSQLTransformer(sqlTransformer)}
 
 		testutils.WithMigratorAndConnectionToContainerWithOptions(t, opts, func(mig *roll.Roll, db *sql.DB) {
 			ctx := context.Background()
@@ -754,7 +752,7 @@ func TestSQLTransformerOptionIsUsedWhenCreatingTriggers(t *testing.T) {
 				Operations: migrations.Operations{
 					&migrations.OpAddColumn{
 						Table: "table1",
-						Up:    "apples",
+						Up:    "'apples'",
 						Column: migrations.Column{
 							Name:     "description",
 							Type:     "text",
@@ -764,7 +762,7 @@ func TestSQLTransformerOptionIsUsedWhenCreatingTriggers(t *testing.T) {
 				},
 			})
 			// Ensure that the start phase has failed with a SQL transformer error
-			require.ErrorIs(t, err, transformerError)
+			require.ErrorIs(t, err, testutils.ErrMockSQLTransformer)
 		})
 	})
 }

pkg/testutils/sql_transformer.go

@@ -0,0 +1,38 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package testutils
+
+import "errors"
+
+type MockSQLTransformer struct {
+	transformations map[string]string
+}
+
+const MockSQLTransformerError = "ERROR"
+
+var ErrMockSQLTransformer = errors.New("SQL transformer error")
+
+// NewMockSQLTransformer creates a MockSQLTransformer with the given transformations.
+// The transformations map is a map of input SQL to output SQL. If the output
+// SQL is "ERROR", the transformer will return an error on that input.
+func NewMockSQLTransformer(ts map[string]string) *MockSQLTransformer {
+	return &MockSQLTransformer{
+		transformations: ts,
+	}
+}
+
+// TransformSQL transforms the given SQL string according to the transformations
+// provided to the MockSQLTransformer. If the input SQL is not in the transformations
+// map, the input SQL is returned unchanged.
+func (s *MockSQLTransformer) TransformSQL(sql string) (string, error) {
+	out, found := s.transformations[sql]
+	if !found {
+		return sql, nil
+	}
+
+	if out == MockSQLTransformerError {
+		return "", ErrMockSQLTransformer
+	}
+
+	return out, nil
+}