feat: add --pre-plan-file flag

Author: aleclarson

cmd/pg-schema-diff/apply_cmd.go

@@ -122,6 +122,14 @@ func runPlan(ctx context.Context, connConfig *pgx.ConnConfig, plan diff.Plan) er
 	}
 	defer conn.Close()
 
+	if plan.PrePlanDDL != "" {
+		fmt.Println(header("Executing pre-plan DDL"))
+		fmt.Printf("%s\n\n", plan.PrePlanDDL)
+		if _, err := conn.ExecContext(ctx, plan.PrePlanDDL); err != nil {
+			return fmt.Errorf("executing pre-plan DDL: %w", err)
+		}
+	}
+
 	// Due to the way *sql.Db works, when a statement_timeout is set for the session, it will NOT reset
 	// by default when it's returned to the pool.
 	//

cmd/pg-schema-diff/plan_cmd.go

@@ -75,6 +75,7 @@ type (
 
 	schemaSourceFlags struct {
 		schemaDirs        []string
+		prePlanFile       string
 		targetDatabaseDSN string
 	}
 
@@ -147,6 +148,11 @@ func schemaSourceFlagsVar(cmd *cobra.Command, p *schemaSourceFlags) {
 	if err := cmd.MarkFlagDirname("schema-dir"); err != nil {
 		panic(err)
 	}
+	cmd.Flags().StringVar(&p.prePlanFile, "pre-plan-file", "", "File path to a file containing DDL statements to prepend to the generated plan.")
+	if err := cmd.MarkFlagFilename("pre-plan-file"); err != nil {
+		panic(err)
+	}
+
 	cmd.Flags().StringVar(&p.targetDatabaseDSN, "schema-source-dsn", "", "DSN for the database to use as the schema source. Use to generate a diff between the target database and the schema in this database.")
 
 	cmd.MarkFlagsMutuallyExclusive("schema-dir", "schema-source-dsn")
@@ -222,7 +228,7 @@ func parsePlanConfig(p planFlags) (planConfig, error) {
 func parseSchemaSource(p schemaSourceFlags) (schemaSourceFactory, error) {
 	if len(p.schemaDirs) > 0 {
 		return func() (diff.SchemaSource, io.Closer, error) {
-			schemaSource, err := diff.DirSchemaSource(p.schemaDirs)
+			schemaSource, err := diff.DirSchemaSource(p.schemaDirs, p.prePlanFile)
 			if err != nil {
 				return nil, nil, err
 			}

pkg/diff/plan.go

@@ -59,6 +59,9 @@ type Plan struct {
 	// plan on running them later, you should verify that the current schema hash matches the current schema hash.
 	// To get the current schema hash, you can use schema.GetPublicSchemaHash(ctx, conn)
 	CurrentSchemaHash string
+	// PrePlanDDL is a string containing DDL statements that should be executed before the plan is applied.
+	// This can be used for setup operations or preliminary changes that need to occur before the main migration.
+	PrePlanDDL string
 }
 
 // ApplyStatementTimeoutModifier applies the given timeout to all statements that match the given regex

pkg/diff/plan_generator.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"database/sql"
 	"fmt"
+	"os"
 	"strings"
 	"time"
 
@@ -106,7 +107,7 @@ func WithGetSchemaOpts(getSchemaOpts ...externalschema.GetSchemaOpt) PlanOpt {
 // newDDL:  		DDL encoding the new schema
 // opts:  			Additional options to configure the plan generation
 func GeneratePlan(ctx context.Context, queryable sqldb.Queryable, tempdbFactory tempdb.Factory, newDDL []string, opts ...PlanOpt) (Plan, error) {
-	return Generate(ctx, queryable, DDLSchemaSource(newDDL), append(opts, WithTempDbFactory(tempdbFactory), WithIncludeSchemas("public"))...)
+	return Generate(ctx, queryable, DDLSchemaSource(newDDL, ""), append(opts, WithTempDbFactory(tempdbFactory), WithIncludeSchemas("public"))...)
 }
 
 // Generate generates a migration plan to migrate the database to the target schema
@@ -155,9 +156,20 @@ func Generate(
 		return Plan{}, fmt.Errorf("generating current schema hash: %w", err)
 	}
 
+	prePlanDDL := ""
+	// Prepend pre-plan file statements if available
+	if ddlSource, ok := targetSchema.(*ddlSchemaSource); ok && ddlSource.prePlanFile != "" {
+		content, err := os.ReadFile(ddlSource.prePlanFile)
+		if err != nil {
+			return Plan{}, fmt.Errorf("reading pre-plan file: %w", err)
+		}
+		prePlanDDL = string(content)
+	}
+
 	plan := Plan{
 		Statements:        statements,
 		CurrentSchemaHash: hash,
+		PrePlanDDL:        prePlanDDL,
 	}
 
 	if planOptions.validatePlan {
@@ -216,11 +228,11 @@ func assertValidPlan(ctx context.Context,
 	// on the database.
 	setMaxConnectionsIfNotSet(tempDb.ConnPool, tempDbMaxConnections)
 
-	if err := setSchemaForEmptyDatabase(ctx, tempDb, currentSchema, planOptions); err != nil {
+	if err := setSchemaForEmptyDatabase(ctx, tempDb, currentSchema, planOptions, plan.PrePlanDDL); err != nil {
 		return fmt.Errorf("inserting schema in temporary database: %w", err)
 	}
 
-	if err := executeStatementsIgnoreTimeouts(ctx, tempDb.ConnPool, plan.Statements); err != nil {
+	if err := executeStatementsIgnoreTimeouts(ctx, tempDb.ConnPool, plan.Statements, plan.PrePlanDDL); err != nil {
 		return fmt.Errorf("running migration plan: %w", err)
 	}
 
@@ -238,7 +250,7 @@ func setMaxConnectionsIfNotSet(db *sql.DB, defaultMax int) {
 	}
 }
 
-func setSchemaForEmptyDatabase(ctx context.Context, emptyDb *tempdb.Database, targetSchema schema.Schema, options *planOptions) error {
+func setSchemaForEmptyDatabase(ctx context.Context, emptyDb *tempdb.Database, targetSchema schema.Schema, options *planOptions, prePlanDDL string) error {
 	// We can't create invalid indexes. We'll mark them valid in the schema, which should be functionally
 	// equivalent for the sake of DDL and other statements.
 	//
@@ -261,7 +273,7 @@ func setSchemaForEmptyDatabase(ctx context.Context, emptyDb *tempdb.Database, ta
 	if err != nil {
 		return fmt.Errorf("building schema diff: %w", err)
 	}
-	if err := executeStatementsIgnoreTimeouts(ctx, emptyDb.ConnPool, statements); err != nil {
+	if err := executeStatementsIgnoreTimeouts(ctx, emptyDb.ConnPool, statements, prePlanDDL); err != nil {
 		return fmt.Errorf("executing statements: %w\n%# v", err, pretty.Formatter(statements))
 	}
 	return nil
@@ -290,7 +302,7 @@ func assertMigratedSchemaMatchesTarget(migratedSchema, targetSchema schema.Schem
 
 // executeStatementsIgnoreTimeouts executes the statements using the sql connection but ignores any provided timeouts.
 // This function is currently used to validate migration plans.
-func executeStatementsIgnoreTimeouts(ctx context.Context, connPool *sql.DB, statements []Statement) error {
+func executeStatementsIgnoreTimeouts(ctx context.Context, connPool *sql.DB, statements []Statement, prePlanDDL string) error {
 	conn, err := connPool.Conn(ctx)
 	if err != nil {
 		return fmt.Errorf("getting connection from pool: %w", err)
@@ -301,6 +313,11 @@ func executeStatementsIgnoreTimeouts(ctx context.Context, connPool *sql.DB, stat
 	if _, err := conn.ExecContext(ctx, fmt.Sprintf("SET SESSION statement_timeout = %d", (10*time.Second).Milliseconds())); err != nil {
 		return fmt.Errorf("setting statement timeout: %w", err)
 	}
+	if prePlanDDL != "" {
+		if _, err := conn.ExecContext(ctx, prePlanDDL); err != nil {
+			return fmt.Errorf("executing pre-plan DDL: %w", err)
+		}
+	}
 	// Due to the way *sql.Db works, when a statement_timeout is set for the session, it will NOT reset
 	// by default when it's returned to the pool.
 	//

pkg/diff/schema_source.go

@@ -33,12 +33,13 @@ type (
 
 	ddlSchemaSource struct {
 		ddl []ddlStatement
+		prePlanFile string
 	}
 )
 
 // DirSchemaSource returns a SchemaSource that returns a schema based on the provided directories. You must provide a tempDBFactory
 // via the WithTempDbFactory option.
-func DirSchemaSource(dirs []string) (SchemaSource, error) {
+func DirSchemaSource(dirs []string, prePlanFile string) (SchemaSource, error) {
 	var ddl []ddlStatement
 	for _, dir := range dirs {
 		stmts, err := getDDLFromPath(dir)
@@ -50,6 +51,7 @@ func DirSchemaSource(dirs []string) (SchemaSource, error) {
 	}
 	return &ddlSchemaSource{
 		ddl: ddl,
+		prePlanFile: prePlanFile,
 	}, nil
 }
 
@@ -86,7 +88,7 @@ func getDDLFromPath(path string) ([]ddlStatement, error) {
 
 // DDLSchemaSource returns a SchemaSource that returns a schema based on the provided DDL. You must provide a tempDBFactory
 // via the WithTempDbFactory option.
-func DDLSchemaSource(stmts []string) SchemaSource {
+func DDLSchemaSource(stmts []string, prePlanFile string) SchemaSource {
 	var ddl []ddlStatement
 	for _, stmt := range stmts {
 		ddl = append(ddl, ddlStatement{
@@ -96,7 +98,7 @@ func DDLSchemaSource(stmts []string) SchemaSource {
 		)
 	}
 
-	return &ddlSchemaSource{ddl: ddl}
+	return &ddlSchemaSource{ddl: ddl, prePlanFile: prePlanFile}
 }
 
 func (s *ddlSchemaSource) GetSchema(ctx context.Context, deps schemaSourcePlanDeps) (schema.Schema, error) {
@@ -114,6 +116,17 @@ func (s *ddlSchemaSource) GetSchema(ctx context.Context, deps schemaSourcePlanDe
 		}
 	}(tempDb.ContextualCloser)
 
+	if s.prePlanFile != "" {
+		prePlanDDL, err := os.ReadFile(s.prePlanFile)
+		if err != nil {
+			return schema.Schema{}, fmt.Errorf("opening pre-plan file: %w", err)
+		}
+
+		if _, err := tempDb.ConnPool.ExecContext(ctx, string(prePlanDDL)); err != nil {
+			return schema.Schema{}, fmt.Errorf("running pre-plan DDL: %w", err)
+		}
+	}
+
 	for _, ddlStmt := range s.ddl {
 		if _, err := tempDb.ConnPool.ExecContext(ctx, ddlStmt.stmt); err != nil {
 			debugInfo := ""