Add WithNoVersionSchemaForRawSQL option (#365)

Add a new `WithNoVersionSchemaForRawSQL` option to control whether or
not version schema should be created for raw SQL migrations.

With the option set, a raw SQL migration:
* Has no version schema created on migration start.
* Leaves the previous version schema in place on migration completion.

Author: andrew-farries

pkg/migrations/migrations.go

@@ -85,3 +85,13 @@ func (m *Migration) Validate(ctx context.Context, s *schema.Schema) error {
 
 	return nil
 }
+
+// ContainsRawSQLOperation returns true if the migration contains a raw SQL operation
+func (m *Migration) ContainsRawSQLOperation() bool {
+	for _, op := range m.Operations {
+		if _, ok := op.(*OpRawSQL); ok {
+			return true
+		}
+	}
+	return false
+}

pkg/roll/execute.go

@@ -90,7 +90,7 @@ func (m *Roll) StartDDLOperations(ctx context.Context, migration *migrations.Mig
 		}
 	}
 
-	if m.disableVersionSchemas {
+	if m.disableVersionSchemas || migration.ContainsRawSQLOperation() && m.noVersionSchemaForRawSQL {
 		// skip creating version schemas
 		return tablesToBackfill, nil
 	}
@@ -131,7 +131,7 @@ func (m *Roll) Complete(ctx context.Context) error {
 	}
 
 	// Drop the old schema
-	if !m.disableVersionSchemas {
+	if !m.disableVersionSchemas && (!migration.ContainsRawSQLOperation() || !m.noVersionSchemaForRawSQL) {
 		prevVersion, err := m.state.PreviousVersion(ctx, m.schema)
 		if err != nil {
 			return fmt.Errorf("unable to get name of previous version: %w", err)
@@ -172,7 +172,9 @@ func (m *Roll) Complete(ctx context.Context) error {
 		}
 
 		if _, ok := op.(migrations.RequiresSchemaRefreshOperation); ok {
-			refreshViews = true
+			if _, ok := op.(*migrations.OpRawSQL); !ok || !m.noVersionSchemaForRawSQL {
+				refreshViews = true
+			}
 		}
 	}
 

pkg/roll/execute_test.go

@@ -155,6 +155,54 @@ func TestPreviousVersionIsDroppedAfterMigrationCompletion(t *testing.T) {
 	})
 }
 
+func TestNoVersionSchemaForRawSQLMigrationsOptionIsRespected(t *testing.T) {
+	t.Parallel()
+
+	opts := []roll.Option{roll.WithNoVersionSchemaForRawSQL()}
+
+	testutils.WithMigratorAndStateAndConnectionToContainerWithOptions(t, opts, func(mig *roll.Roll, st *state.State, db *sql.DB) {
+		ctx := context.Background()
+
+		// Apply a create table migration
+		err := mig.Start(ctx, &migrations.Migration{Name: "01_create_table", Operations: migrations.Operations{createTableOp("table1")}})
+		require.NoError(t, err)
+		err = mig.Complete(ctx)
+		require.NoError(t, err)
+
+		// Apply a raw SQL migration - no version schema should be created for this version
+		err = mig.Start(ctx, &migrations.Migration{
+			Name: "02_create_table",
+			Operations: migrations.Operations{&migrations.OpRawSQL{
+				Up: "CREATE TABLE table2(a int)",
+			}},
+		})
+		require.NoError(t, err)
+		err = mig.Complete(ctx)
+		require.NoError(t, err)
+
+		// Start a third create table migration
+		err = mig.Start(ctx, &migrations.Migration{Name: "03_create_table", Operations: migrations.Operations{createTableOp("table3")}})
+		require.NoError(t, err)
+
+		// The previous version is migration 01 because there is no version schema
+		// for migration 02 due to the `WithNoVersionSchemaForRawSQL` option
+		prevVersion, err := st.PreviousVersion(ctx, "public")
+		require.NoError(t, err)
+		require.NotNil(t, prevVersion)
+		assert.Equal(t, "01_create_table", *prevVersion)
+
+		// Complete the third migration
+		err = mig.Complete(ctx)
+		require.NoError(t, err)
+
+		// The latest version version is migration 03
+		latestVersion, err := st.LatestVersion(ctx, "public")
+		require.NoError(t, err)
+		require.NotNil(t, latestVersion)
+		assert.Equal(t, "03_create_table", *latestVersion)
+	})
+}
+
 func TestSchemaIsDroppedAfterMigrationRollback(t *testing.T) {
 	t.Parallel()
 

pkg/roll/options.go

@@ -18,7 +18,11 @@ type options struct {
 
 	// disable pgroll version schemas creation and deletion
 	disableVersionSchemas bool
-	migrationHooks        MigrationHooks
+
+	// disable creation of version schema for raw SQL migrations
+	noVersionSchemaForRawSQL bool
+
+	migrationHooks MigrationHooks
 }
 
 // MigrationHooks defines hooks that can be set to be called at various points
@@ -58,6 +62,12 @@ func WithDisableViewsManagement() Option {
 	}
 }
 
+func WithNoVersionSchemaForRawSQL() Option {
+	return func(o *options) {
+		o.noVersionSchemaForRawSQL = true
+	}
+}
+
 // WithMigrationHooks sets the migration hooks for the Roll instance
 // Migration hooks are called at various points during the migration process
 // to allow for custom behavior to be injected

pkg/roll/roll.go

@@ -27,6 +27,9 @@ type Roll struct {
 	// disable pgroll version schemas creation and deletion
 	disableVersionSchemas bool
 
+	// disable creation of version schema for raw SQL migrations
+	noVersionSchemaForRawSQL bool
+
 	migrationHooks MigrationHooks
 	state          *state.State
 	pgVersion      PGVersion
@@ -58,13 +61,14 @@ func New(ctx context.Context, pgURL, schema string, state *state.State, opts ...
 	}
 
 	return &Roll{
-		pgConn:                &db.RDB{DB: conn},
-		schema:                schema,
-		state:                 state,
-		pgVersion:             PGVersion(pgMajorVersion),
-		disableVersionSchemas: rollOpts.disableVersionSchemas,
-		migrationHooks:        rollOpts.migrationHooks,
-		sqlTransformer:        sqlTransformer,
+		pgConn:                   &db.RDB{DB: conn},
+		schema:                   schema,
+		state:                    state,
+		pgVersion:                PGVersion(pgMajorVersion),
+		disableVersionSchemas:    rollOpts.disableVersionSchemas,
+		noVersionSchemaForRawSQL: rollOpts.noVersionSchemaForRawSQL,
+		migrationHooks:           rollOpts.migrationHooks,
+		sqlTransformer:           sqlTransformer,
 	}, nil
 }
 

pkg/state/state.go

@@ -70,22 +70,27 @@ LANGUAGE SQL
 STABLE;
 
 -- Get the name of the previous version of the schema, or NULL if there is none.
+-- This ignores previous versions for which no version schema exists, such as
+-- versions corresponding to inferred migrations.
 CREATE OR REPLACE FUNCTION %[1]s.previous_version(schemaname NAME) RETURNS text
 AS $$
-  WITH RECURSIVE find_ancestor AS (
-    SELECT schema, name, parent, migration_type FROM %[1]s.migrations
-      WHERE name = (SELECT %[1]s.latest_version(schemaname)) AND schema = schemaname
+  WITH RECURSIVE ancestors AS (
+  SELECT name, parent, migration_type, 0 AS depth FROM %[1]s.migrations
+    WHERE name = %[1]s.latest_version(schemaname)
 
-    UNION ALL
+  UNION ALL
 
-    SELECT m.schema, m.name, m.parent, m.migration_type FROM %[1]s.migrations m
-      INNER JOIN find_ancestor fa ON fa.parent = m.name AND fa.schema = m.schema
-      WHERE m.migration_type = 'inferred'
+  SELECT m.name, m.parent, m.migration_type, a.depth + 1
+    FROM %[1]s.migrations m
+    JOIN ancestors a ON m.name = a.parent
   )
-  SELECT a.parent
-  FROM find_ancestor AS a
-  JOIN %[1]s.migrations AS b ON a.parent = b.name AND a.schema = b.schema
-  WHERE b.migration_type = 'pgroll';
+  SELECT a.name FROM ancestors a
+    JOIN information_schema.schemata s
+    ON s.schema_name = schemaname || '_' || a.name
+    WHERE migration_type = 'pgroll'
+    AND a.depth > 0
+    ORDER by a.depth ASC
+    LIMIT 1;
 $$
 LANGUAGE SQL
 STABLE;

pkg/testutils/util.go

@@ -174,6 +174,51 @@ func WithMigratorInSchemaAndConnectionToContainerWithOptions(t *testing.T, schem
 	fn(mig, db)
 }
 
+func WithMigratorAndStateAndConnectionToContainerWithOptions(t *testing.T, opts []roll.Option, fn func(*roll.Roll, *state.State, *sql.DB)) {
+	t.Helper()
+	ctx := context.Background()
+
+	db, connStr, dbName := setupTestDatabase(t)
+
+	st, err := state.New(ctx, connStr, "pgroll")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	err = st.Init(ctx)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	mig, err := roll.New(ctx, connStr, "public", st, opts...)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	t.Cleanup(func() {
+		if err := mig.Close(); err != nil {
+			t.Fatalf("Failed to close migrator connection: %v", err)
+		}
+	})
+
+	_, err = db.ExecContext(ctx, fmt.Sprintf("CREATE SCHEMA IF NOT EXISTS %s", "public"))
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	_, err = db.ExecContext(ctx, fmt.Sprintf("GRANT ALL PRIVILEGES ON SCHEMA %s TO pgroll", "public"))
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	_, err = db.ExecContext(ctx, fmt.Sprintf("GRANT ALL PRIVILEGES ON DATABASE %s TO pgroll", dbName))
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	fn(mig, st, db)
+}
+
 func WithMigratorInSchemaAndConnectionToContainer(t *testing.T, schema string, fn func(mig *roll.Roll, db *sql.DB)) {
 	WithMigratorInSchemaAndConnectionToContainerWithOptions(t, schema, []roll.Option{roll.WithLockTimeoutMs(500)}, fn)
 }