Allow to disable version schemas (#233)

There are situatiosn were creating schema verisons is not required. This internal flag will allow to disable creating & managing these schemas when needed. Note: I opted to not expose this flag to the CLI as I believe it would make create more harm & confusion than not having the feature there.
xataio · Jan 15, 2024 · 83b5d76 · 83b5d76
1 parent 73c2016
commit 83b5d76
Show file tree

Hide file tree

Showing 4 changed files with 98 additions and 76 deletions.
diff --git a/pkg/roll/execute.go b/pkg/roll/execute.go
@@ -59,6 +59,11 @@ func (m *Roll) Start(ctx context.Context, migration *migrations.Migration, cbs .
 		}
 	}
 
+	if m.disableVersionSchemas {
+		// skip creating version schemas
+		return nil
+	}
+
 	// create schema for the new version
 	versionSchema := VersionedSchemaName(m.schema, migration.Name)
 	_, err = m.pgConn.ExecContext(ctx, fmt.Sprintf("CREATE SCHEMA IF NOT EXISTS %s", pq.QuoteIdentifier(versionSchema)))
@@ -86,15 +91,17 @@ func (m *Roll) Complete(ctx context.Context) error {
 	}
 
 	// Drop the old schema
-	prevVersion, err := m.state.PreviousVersion(ctx, m.schema)
-	if err != nil {
-		return fmt.Errorf("unable to get name of previous version: %w", err)
-	}
-	if prevVersion != nil {
-		versionSchema := VersionedSchemaName(m.schema, *prevVersion)
-		_, err = m.pgConn.ExecContext(ctx, fmt.Sprintf("DROP SCHEMA IF EXISTS %s CASCADE", pq.QuoteIdentifier(versionSchema)))
+	if !m.disableVersionSchemas {
+		prevVersion, err := m.state.PreviousVersion(ctx, m.schema)
 		if err != nil {
-			return fmt.Errorf("unable to drop previous version: %w", err)
+			return fmt.Errorf("unable to get name of previous version: %w", err)
+		}
+		if prevVersion != nil {
+			versionSchema := VersionedSchemaName(m.schema, *prevVersion)
+			_, err = m.pgConn.ExecContext(ctx, fmt.Sprintf("DROP SCHEMA IF EXISTS %s CASCADE", pq.QuoteIdentifier(versionSchema)))
+			if err != nil {
+				return fmt.Errorf("unable to drop previous version: %w", err)
+			}
 		}
 	}
 
@@ -122,11 +129,13 @@ func (m *Roll) Rollback(ctx context.Context) error {
 		return fmt.Errorf("unable to get active migration: %w", err)
 	}
 
-	// delete the schema and view for the new version
-	versionSchema := VersionedSchemaName(m.schema, migration.Name)
-	_, err = m.pgConn.ExecContext(ctx, fmt.Sprintf("DROP SCHEMA IF EXISTS %s CASCADE", pq.QuoteIdentifier(versionSchema)))
-	if err != nil {
-		return err
+	if !m.disableVersionSchemas {
+		// delete the schema and view for the new version
+		versionSchema := VersionedSchemaName(m.schema, migration.Name)
+		_, err = m.pgConn.ExecContext(ctx, fmt.Sprintf("DROP SCHEMA IF EXISTS %s CASCADE", pq.QuoteIdentifier(versionSchema)))
+		if err != nil {
+			return err
+		}
 	}
 
 	// execute operations

diff --git a/pkg/roll/execute_test.go b/pkg/roll/execute_test.go
@@ -25,7 +25,7 @@ func TestMain(m *testing.M) {
 	testutils.SharedTestMain(m)
 }
 
-func TestSchemaIsCreatedfterMigrationStart(t *testing.T) {
+func TestSchemaIsCreatedAfterMigrationStart(t *testing.T) {
 	t.Parallel()
 
 	testutils.WithMigratorAndConnectionToContainer(t, func(mig *roll.Roll, db *sql.DB) {
@@ -39,23 +39,64 @@ func TestSchemaIsCreatedfterMigrationStart(t *testing.T) {
 		//
 		// Check that the schema exists
 		//
-		var exists bool
-		err := db.QueryRow(`
-    SELECT EXISTS(
-      SELECT 1
-      FROM pg_catalog.pg_namespace
-      WHERE nspname = $1
-    )`, roll.VersionedSchemaName(schema, version)).Scan(&exists)
-		if err != nil {
-			t.Fatal(err)
+		if !schemaExists(t, db, roll.VersionedSchemaName(schema, version)) {
+			t.Errorf("Expected schema %q to exist", version)
 		}
+	})
+}
 
-		if !exists {
-			t.Errorf("Expected schema %q to exist", version)
+func TestDisabledSchemaManagement(t *testing.T) {
+	t.Parallel()
+
+	testutils.WithMigratorInSchemaAndConnectionToContainerWithOptions(t, "public", []roll.Option{roll.WithDisableViewsManagement()}, func(mig *roll.Roll, db *sql.DB) {
+		ctx := context.Background()
+		version := "1_create_table"
+
+		if err := mig.Start(ctx, &migrations.Migration{Name: version, Operations: migrations.Operations{createTableOp("table1")}}); err != nil {
+			t.Fatalf("Failed to start migration: %v", err)
+		}
+
+		//
+		// Check that the schema doesn't get created
+		//
+		if schemaExists(t, db, roll.VersionedSchemaName(schema, version)) {
+			t.Errorf("Expected schema %q to not exist", version)
+		}
+
+		if err := mig.Rollback(ctx); err != nil {
+			t.Fatalf("Failed to rollback migration: %v", err)
+		}
+
+		if err := mig.Start(ctx, &migrations.Migration{Name: version, Operations: migrations.Operations{createTableOp("table1")}}); err != nil {
+			t.Fatalf("Failed to start migration again: %v", err)
+		}
+
+		// complete the migration, check that the schema still doesn't exist
+		if err := mig.Complete(ctx); err != nil {
+			t.Fatalf("Failed to complete migration: %v", err)
+		}
+
+		if schemaExists(t, db, roll.VersionedSchemaName(schema, version)) {
+			t.Errorf("Expected schema %q to not exist", version)
 		}
 	})
 }
 
+func schemaExists(t *testing.T, db *sql.DB, schema string) bool {
+	t.Helper()
+	var exists bool
+	err := db.QueryRow(`
+	SELECT EXISTS(
+		SELECT 1
+		FROM pg_catalog.pg_namespace
+		WHERE nspname = $1
+	)`, schema).Scan(&exists)
+	if err != nil {
+		t.Fatal(err)
+	}
+	return exists
+}
+
 func TestPreviousVersionIsDroppedAfterMigrationCompletion(t *testing.T) {
 	t.Parallel()
 
@@ -83,18 +124,7 @@ func TestPreviousVersionIsDroppedAfterMigrationCompletion(t *testing.T) {
 			//
 			// Check that the schema for the first version has been dropped
 			//
-			var exists bool
-			err := db.QueryRow(`
-    SELECT EXISTS(
-      SELECT 1
-      FROM pg_catalog.pg_namespace
-      WHERE nspname = $1
-    )`, roll.VersionedSchemaName(schema, firstVersion)).Scan(&exists)
-			if err != nil {
-				t.Fatal(err)
-			}
-
-			if exists {
+			if schemaExists(t, db, roll.VersionedSchemaName(schema, firstVersion)) {
 				t.Errorf("Expected schema %q to not exist", firstVersion)
 			}
 		})
@@ -133,18 +163,7 @@ func TestPreviousVersionIsDroppedAfterMigrationCompletion(t *testing.T) {
 			//
 			// Check that the schema for the first version has been dropped
 			//
-			var exists bool
-			err = db.QueryRow(`
-    SELECT EXISTS(
-      SELECT 1
-      FROM pg_catalog.pg_namespace
-      WHERE nspname = $1
-    )`, roll.VersionedSchemaName(schema, firstVersion)).Scan(&exists)
-			if err != nil {
-				t.Fatal(err)
-			}
-
-			if exists {
+			if schemaExists(t, db, roll.VersionedSchemaName(schema, firstVersion)) {
 				t.Errorf("Expected schema %q to not exist", firstVersion)
 			}
 		})
@@ -168,18 +187,7 @@ func TestSchemaIsDroppedAfterMigrationRollback(t *testing.T) {
 		//
 		// Check that the schema has been dropped
 		//
-		var exists bool
-		err := db.QueryRow(`
-    SELECT EXISTS(
-      SELECT 1
-      FROM pg_catalog.pg_namespace
-      WHERE nspname = $1
-    )`, roll.VersionedSchemaName(schema, version)).Scan(&exists)
-		if err != nil {
-			t.Fatal(err)
-		}
-
-		if exists {
+		if schemaExists(t, db, roll.VersionedSchemaName(schema, version)) {
 			t.Errorf("Expected schema %q to not exist", version)
 		}
 	})
@@ -234,17 +242,7 @@ func TestSchemaOptionIsRespected(t *testing.T) {
 		}
 
 		// Ensure that the versioned schema for the first migration has been dropped
-		err = db.QueryRow(`
-    SELECT EXISTS(
-      SELECT 1
-      FROM pg_catalog.pg_namespace
-      WHERE nspname = $1
-    )`, roll.VersionedSchemaName("schema1", version1)).Scan(&exists)
-		if err != nil {
-			t.Fatal(err)
-		}
-
-		if exists {
+		if schemaExists(t, db, roll.VersionedSchemaName("schema1", version1)) {
 			t.Errorf("Expected schema %q to not exist", version1)
 		}
 	})

diff --git a/pkg/roll/options.go b/pkg/roll/options.go
@@ -8,6 +8,9 @@ type options struct {
 
 	// optional role to set before executing migrations
 	role string
+
+	// disable pgroll version schemas creation and deletion
+	disableVersionSchemas bool
 }
 
 type Option func(*options)
@@ -25,3 +28,11 @@ func WithRole(role string) Option {
 		o.role = role
 	}
 }
+
+// WithDisableViewsManagement disables pgroll version schemas management
+// when passed, pgroll will not create or drop version schemas
+func WithDisableViewsManagement() Option {
+	return func(o *options) {
+		o.disableVersionSchemas = true
+	}
+}
diff --git a/pkg/roll/roll.go b/pkg/roll/roll.go
@@ -22,6 +22,9 @@ type Roll struct {
 	// schema we are acting on
 	schema string
 
+	// disable pgroll version schemas creation and deletion
+	disableVersionSchemas bool
+
 	state     *state.State
 	pgVersion PGVersion
 }
@@ -74,10 +77,11 @@ func New(ctx context.Context, pgURL, schema string, state *state.State, opts ...
 	}
 
 	return &Roll{
-		pgConn:    conn,
-		schema:    schema,
-		state:     state,
-		pgVersion: PGVersion(pgMajorVersion),
+		pgConn:                conn,
+		schema:                schema,
+		state:                 state,
+		pgVersion:             PGVersion(pgMajorVersion),
+		disableVersionSchemas: options.disableVersionSchemas,
 	}, nil
 }