Change Complete method signature (#231)

Change the signature of the `Complete` method on the `Operation` interface to take a `*schema.Schema` argument like the `Start` and `Validate` methods already do. This allows for implementations of `Complete` to be aware of the entire database schema, supporting things like renaming of any temporary constraints created on columns affected the operation being completed. See #230, which depends on this PR.
xataio · Jan 15, 2024 · 4b124b0 · 4b124b0
1 parent 83b5d76
commit 4b124b0
Show file tree

Hide file tree

Showing 19 changed files with 26 additions and 20 deletions.
diff --git a/pkg/migrations/migrations.go b/pkg/migrations/migrations.go
@@ -22,7 +22,7 @@ type Operation interface {
 	// Complete will update the database schema to match the current version
 	// after calling Start.
 	// This method should be called once the previous version is no longer used
-	Complete(ctx context.Context, conn *sql.DB) error
+	Complete(ctx context.Context, conn *sql.DB, s *schema.Schema) error
 
 	// Rollback will revert the changes made by Start. It is not possible to
 	// rollback a completed migration.

diff --git a/pkg/migrations/op_add_column.go b/pkg/migrations/op_add_column.go
@@ -65,7 +65,7 @@ func (o *OpAddColumn) Start(ctx context.Context, conn *sql.DB, stateSchema strin
 	return nil
 }
 
-func (o *OpAddColumn) Complete(ctx context.Context, conn *sql.DB) error {
+func (o *OpAddColumn) Complete(ctx context.Context, conn *sql.DB, s *schema.Schema) error {
 	tempName := TemporaryName(o.Column.Name)
 
 	_, err := conn.ExecContext(ctx, fmt.Sprintf("ALTER TABLE IF EXISTS %s RENAME COLUMN %s TO %s",

diff --git a/pkg/migrations/op_alter_column.go b/pkg/migrations/op_alter_column.go
@@ -18,10 +18,10 @@ func (o *OpAlterColumn) Start(ctx context.Context, conn *sql.DB, stateSchema str
 	return op.Start(ctx, conn, stateSchema, s, cbs...)
 }
 
-func (o *OpAlterColumn) Complete(ctx context.Context, conn *sql.DB) error {
+func (o *OpAlterColumn) Complete(ctx context.Context, conn *sql.DB, s *schema.Schema) error {
 	op := o.innerOperation()
 
-	return op.Complete(ctx, conn)
+	return op.Complete(ctx, conn, s)
 }
 
 func (o *OpAlterColumn) Rollback(ctx context.Context, conn *sql.DB) error {

diff --git a/pkg/migrations/op_change_type.go b/pkg/migrations/op_change_type.go
@@ -75,7 +75,7 @@ func (o *OpChangeType) Start(ctx context.Context, conn *sql.DB, stateSchema stri
 	return nil
 }
 
-func (o *OpChangeType) Complete(ctx context.Context, conn *sql.DB) error {
+func (o *OpChangeType) Complete(ctx context.Context, conn *sql.DB, s *schema.Schema) error {
 	// Remove the up function and trigger
 	_, err := conn.ExecContext(ctx, fmt.Sprintf("DROP FUNCTION IF EXISTS %s CASCADE",
 		pq.QuoteIdentifier(TriggerFunctionName(o.Table, o.Column))))

diff --git a/pkg/migrations/op_create_index.go b/pkg/migrations/op_create_index.go
@@ -23,7 +23,7 @@ func (o *OpCreateIndex) Start(ctx context.Context, conn *sql.DB, stateSchema str
 	return err
 }
 
-func (o *OpCreateIndex) Complete(ctx context.Context, conn *sql.DB) error {
+func (o *OpCreateIndex) Complete(ctx context.Context, conn *sql.DB, s *schema.Schema) error {
 	// No-op
 	return nil
 }

diff --git a/pkg/migrations/op_create_table.go b/pkg/migrations/op_create_table.go
@@ -53,7 +53,7 @@ func (o *OpCreateTable) Start(ctx context.Context, conn *sql.DB, stateSchema str
 	return nil
 }
 
-func (o *OpCreateTable) Complete(ctx context.Context, conn *sql.DB) error {
+func (o *OpCreateTable) Complete(ctx context.Context, conn *sql.DB, s *schema.Schema) error {
 	tempName := TemporaryName(o.Name)
 	_, err := conn.ExecContext(ctx, fmt.Sprintf("ALTER TABLE IF EXISTS %s RENAME TO %s",
 		pq.QuoteIdentifier(tempName),

diff --git a/pkg/migrations/op_drop_column.go b/pkg/migrations/op_drop_column.go
@@ -34,7 +34,7 @@ func (o *OpDropColumn) Start(ctx context.Context, conn *sql.DB, stateSchema stri
 	return nil
 }
 
-func (o *OpDropColumn) Complete(ctx context.Context, conn *sql.DB) error {
+func (o *OpDropColumn) Complete(ctx context.Context, conn *sql.DB, s *schema.Schema) error {
 	_, err := conn.ExecContext(ctx, fmt.Sprintf("ALTER TABLE %s DROP COLUMN %s", o.Table, o.Column))
 	if err != nil {
 		return err

diff --git a/pkg/migrations/op_drop_constraint.go b/pkg/migrations/op_drop_constraint.go
@@ -66,7 +66,7 @@ func (o *OpDropConstraint) Start(ctx context.Context, conn *sql.DB, stateSchema
 	return nil
 }
 
-func (o *OpDropConstraint) Complete(ctx context.Context, conn *sql.DB) error {
+func (o *OpDropConstraint) Complete(ctx context.Context, conn *sql.DB, s *schema.Schema) error {
 	// Remove the up function and trigger
 	_, err := conn.ExecContext(ctx, fmt.Sprintf("DROP FUNCTION IF EXISTS %s CASCADE",
 		pq.QuoteIdentifier(TriggerFunctionName(o.Table, o.Column))))

diff --git a/pkg/migrations/op_drop_index.go b/pkg/migrations/op_drop_index.go
@@ -17,7 +17,7 @@ func (o *OpDropIndex) Start(ctx context.Context, conn *sql.DB, stateSchema strin
 	return nil
 }
 
-func (o *OpDropIndex) Complete(ctx context.Context, conn *sql.DB) error {
+func (o *OpDropIndex) Complete(ctx context.Context, conn *sql.DB, s *schema.Schema) error {
 	// drop the index concurrently
 	_, err := conn.ExecContext(ctx, fmt.Sprintf("DROP INDEX CONCURRENTLY IF EXISTS %s", o.Name))
 

diff --git a/pkg/migrations/op_drop_table.go b/pkg/migrations/op_drop_table.go
@@ -18,7 +18,7 @@ func (o *OpDropTable) Start(ctx context.Context, conn *sql.DB, stateSchema strin
 	return nil
 }
 
-func (o *OpDropTable) Complete(ctx context.Context, conn *sql.DB) error {
+func (o *OpDropTable) Complete(ctx context.Context, conn *sql.DB, s *schema.Schema) error {
 	_, err := conn.ExecContext(ctx, fmt.Sprintf("DROP TABLE IF EXISTS %s", pq.QuoteIdentifier(o.Name)))
 
 	return err

diff --git a/pkg/migrations/op_raw_sql.go b/pkg/migrations/op_raw_sql.go
@@ -19,7 +19,7 @@ func (o *OpRawSQL) Start(ctx context.Context, conn *sql.DB, stateSchema string,
 	return nil
 }
 
-func (o *OpRawSQL) Complete(ctx context.Context, conn *sql.DB) error {
+func (o *OpRawSQL) Complete(ctx context.Context, conn *sql.DB, s *schema.Schema) error {
 	return nil
 }
 

diff --git a/pkg/migrations/op_rename_column.go b/pkg/migrations/op_rename_column.go
@@ -25,7 +25,7 @@ func (o *OpRenameColumn) Start(ctx context.Context, conn *sql.DB, stateSchema st
 	return nil
 }
 
-func (o *OpRenameColumn) Complete(ctx context.Context, conn *sql.DB) error {
+func (o *OpRenameColumn) Complete(ctx context.Context, conn *sql.DB, s *schema.Schema) error {
 	// rename the column in the underlying table
 	_, err := conn.ExecContext(ctx, fmt.Sprintf("ALTER TABLE %s RENAME COLUMN %s TO %s",
 		pq.QuoteIdentifier(o.Table),

diff --git a/pkg/migrations/op_rename_table.go b/pkg/migrations/op_rename_table.go
@@ -17,7 +17,7 @@ func (o *OpRenameTable) Start(ctx context.Context, conn *sql.DB, stateSchema str
 	return s.RenameTable(o.From, o.To)
 }
 
-func (o *OpRenameTable) Complete(ctx context.Context, conn *sql.DB) error {
+func (o *OpRenameTable) Complete(ctx context.Context, conn *sql.DB, s *schema.Schema) error {
 	_, err := conn.ExecContext(ctx, fmt.Sprintf("ALTER TABLE IF EXISTS %s RENAME TO %s",
 		pq.QuoteIdentifier(o.From),
 		pq.QuoteIdentifier(o.To)))

diff --git a/pkg/migrations/op_set_check.go b/pkg/migrations/op_set_check.go
@@ -80,7 +80,7 @@ func (o *OpSetCheckConstraint) Start(ctx context.Context, conn *sql.DB, stateSch
 	return nil
 }
 
-func (o *OpSetCheckConstraint) Complete(ctx context.Context, conn *sql.DB) error {
+func (o *OpSetCheckConstraint) Complete(ctx context.Context, conn *sql.DB, s *schema.Schema) error {
 	// Validate the check constraint
 	_, err := conn.ExecContext(ctx, fmt.Sprintf("ALTER TABLE IF EXISTS %s VALIDATE CONSTRAINT %s",
 		pq.QuoteIdentifier(o.Table),

diff --git a/pkg/migrations/op_set_fk.go b/pkg/migrations/op_set_fk.go
@@ -80,7 +80,7 @@ func (o *OpSetForeignKey) Start(ctx context.Context, conn *sql.DB, stateSchema s
 	return nil
 }
 
-func (o *OpSetForeignKey) Complete(ctx context.Context, conn *sql.DB) error {
+func (o *OpSetForeignKey) Complete(ctx context.Context, conn *sql.DB, s *schema.Schema) error {
 	// Validate the foreign key constraint
 	_, err := conn.ExecContext(ctx, fmt.Sprintf("ALTER TABLE IF EXISTS %s VALIDATE CONSTRAINT %s",
 		pq.QuoteIdentifier(o.Table),

diff --git a/pkg/migrations/op_set_notnull.go b/pkg/migrations/op_set_notnull.go
@@ -79,7 +79,7 @@ func (o *OpSetNotNull) Start(ctx context.Context, conn *sql.DB, stateSchema stri
 	return nil
 }
 
-func (o *OpSetNotNull) Complete(ctx context.Context, conn *sql.DB) error {
+func (o *OpSetNotNull) Complete(ctx context.Context, conn *sql.DB, s *schema.Schema) error {
 	// Validate the NOT NULL constraint on the old column.
 	// The constraint must be valid because:
 	// * Existing NULL values in the old column were rewritten using the `up` SQL during backfill.

diff --git a/pkg/migrations/op_set_replica_identity.go b/pkg/migrations/op_set_replica_identity.go
@@ -29,7 +29,7 @@ func (o *OpSetReplicaIdentity) Start(ctx context.Context, conn *sql.DB, stateSch
 	return err
 }
 
-func (o *OpSetReplicaIdentity) Complete(ctx context.Context, conn *sql.DB) error {
+func (o *OpSetReplicaIdentity) Complete(ctx context.Context, conn *sql.DB, s *schema.Schema) error {
 	// No-op
 	return nil
 }

diff --git a/pkg/migrations/op_set_unique.go b/pkg/migrations/op_set_unique.go
@@ -80,7 +80,7 @@ func (o *OpSetUnique) Start(ctx context.Context, conn *sql.DB, stateSchema strin
 	return nil
 }
 
-func (o *OpSetUnique) Complete(ctx context.Context, conn *sql.DB) error {
+func (o *OpSetUnique) Complete(ctx context.Context, conn *sql.DB, s *schema.Schema) error {
 	// Create a unique constraint using the unique index
 	_, err := conn.ExecContext(ctx, fmt.Sprintf("ALTER TABLE IF EXISTS %s ADD CONSTRAINT %s UNIQUE USING INDEX %s",
 		pq.QuoteIdentifier(o.Table),

diff --git a/pkg/roll/execute.go b/pkg/roll/execute.go
@@ -105,9 +105,15 @@ func (m *Roll) Complete(ctx context.Context) error {
 		}
 	}
 
+	// read the current schema
+	schema, err := m.state.ReadSchema(ctx, m.schema)
+	if err != nil {
+		return fmt.Errorf("unable to read schema: %w", err)
+	}
+
 	// execute operations
 	for _, op := range migration.Operations {
-		err := op.Complete(ctx, m.pgConn)
+		err := op.Complete(ctx, m.pgConn, schema)
 		if err != nil {
 			return fmt.Errorf("unable to execute complete operation: %w", err)
 		}