diff --git a/cmd/pg-schema-diff/plan_cmd.go b/cmd/pg-schema-diff/plan_cmd.go index 36eabb2..b46bfef 100644 --- a/cmd/pg-schema-diff/plan_cmd.go +++ b/cmd/pg-schema-diff/plan_cmd.go @@ -94,6 +94,8 @@ type ( lockTimeoutModifiers []string insertStatements []string outputFormat outputFormat + + templateDb string } timeoutModifier struct { @@ -117,6 +119,8 @@ type ( statementTimeoutModifiers []timeoutModifier lockTimeoutModifiers []timeoutModifier insertStatements []insertStatement + + templateDb string } ) @@ -173,6 +177,13 @@ func createPlanFlags(cmd *cobra.Command) *planFlags { cmd.Flags().Var(&flags.outputFormat, "output-format", "Change the output format for what is printed. Defaults to pretty-printed human-readable output. (options: pretty, json)") + cmd.Flags().StringVar( + &flags.templateDb, + "template-db", + "template0", + "Template database to use when creating temporary databases", + ) + return flags } @@ -250,6 +261,7 @@ func parsePlanConfig(p planFlags) (planConfig, error) { statementTimeoutModifiers: statementTimeoutModifiers, lockTimeoutModifiers: lockTimeoutModifiers, insertStatements: insertStatements, + templateDb: p.templateDb, }, nil } @@ -384,7 +396,7 @@ func generatePlan(ctx context.Context, logger log.Logger, connConfig *pgx.ConnCo copiedConfig := connConfig.Copy() copiedConfig.Database = dbName return openDbWithPgxConfig(copiedConfig) - }, tempdb.WithRootDatabase(connConfig.Database)) + }, tempdb.WithRootDatabase(connConfig.Database), tempdb.WithTemplateDatabase(planConfig.templateDb)) if err != nil { return diff.Plan{}, err } diff --git a/cmd/pg-schema-diff/plan_cmd_test.go b/cmd/pg-schema-diff/plan_cmd_test.go index b323411..eb6292f 100644 --- a/cmd/pg-schema-diff/plan_cmd_test.go +++ b/cmd/pg-schema-diff/plan_cmd_test.go @@ -5,6 +5,7 @@ import ( "testing" "time" + "github.com/spf13/cobra" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) @@ -32,11 +33,11 @@ func TestParseTimeoutModifierStr(t *testing.T) { }, { opt: "timeout=15m", - expectedErrContains: "could not find key", + expectedErrContains: "could not find key", // "pattern" missing }, { opt: `pattern="some pattern"`, - expectedErrContains: "could not find key", + expectedErrContains: "could not find key", // "timeout" missing }, { opt: `pattern="normal" timeout=5m some-unknown-key=5m`, @@ -80,19 +81,19 @@ func TestParseInsertStatementStr(t *testing.T) { }, { opt: "statement=no-index timeout=5m6s lock_timeout=1m11s", - expectedErrContains: "could not find key", + expectedErrContains: "could not find key", // "index" missing }, { opt: "index=0 timeout=5m6s lock_timeout=1m11s", - expectedErrContains: "could not find key", + expectedErrContains: "could not find key", // "statement" missing }, { opt: "index=0 statement=no-timeout lock_timeout=1m11s", - expectedErrContains: "could not find key", + expectedErrContains: "could not find key", // "timeout" missing }, { opt: "index=0 statement=no-lock-timeout-timeout timeout=5m6s", - expectedErrContains: "could not find key", + expectedErrContains: "could not find key", // "lock_timeout" missing }, { opt: "index=not-an-int statement=some-statement timeout=5m6s lock_timeout=1m11s", @@ -118,3 +119,34 @@ func TestParseInsertStatementStr(t *testing.T) { }) } } + +func TestPlanFlagsTemplateDBDefault(t *testing.T) { + cmd := &cobra.Command{} + flags := createPlanFlags(cmd) + + err := cmd.ParseFlags([]string{ + "--schema-dir=/no/such/dir", + }) + require.NoError(t, err) + + planCfg, err := parsePlanConfig(*flags) + require.NoError(t, err, "parsePlanConfig should not fail with a dummy --schema-dir") + + assert.Equal(t, "template0", planCfg.templateDb) +} + +func TestPlanFlagsTemplateDBOverride(t *testing.T) { + cmd := &cobra.Command{} + flags := createPlanFlags(cmd) + + err := cmd.ParseFlags([]string{ + "--template-db=template1", + "--schema-dir=/no/such/dir", + }) + require.NoError(t, err) + + planCfg, err := parsePlanConfig(*flags) + require.NoError(t, err, "parsePlanConfig should not fail with dummy --schema-dir") + + assert.Equal(t, "template1", planCfg.templateDb) +} diff --git a/pkg/tempdb/factory.go b/pkg/tempdb/factory.go index 6311ae6..04405b1 100644 --- a/pkg/tempdb/factory.go +++ b/pkg/tempdb/factory.go @@ -22,6 +22,7 @@ const ( DefaultOnInstanceMetadataSchema = "pgschemadiff_tmp_metadata" DefaultOnInstanceMetadataTable = "metadata" + DefaultTemplateDatabase = "template0" DefaultStatementTimeout = 3 * time.Second ) @@ -57,11 +58,12 @@ type ( ) type ( onInstanceFactoryOptions struct { - dbPrefix string - metadataSchema string - metadataTable string - logger log.Logger - rootDatabase string + dbPrefix string + metadataSchema string + metadataTable string + logger log.Logger + rootDatabase string + templateDatabase string } OnInstanceFactoryOpt func(*onInstanceFactoryOptions) @@ -102,6 +104,13 @@ func WithRootDatabase(db string) OnInstanceFactoryOpt { } } +// WithTemplateDatabase sets the template DB that CREATE DATABASE will use. +func WithTemplateDatabase(templateDB string) OnInstanceFactoryOpt { + return func(opts *onInstanceFactoryOptions) { + opts.templateDatabase = templateDB + } +} + type ( CreateConnPoolForDbFn func(ctx context.Context, dbName string) (*sql.DB, error) @@ -126,11 +135,12 @@ type ( // when the temporary database was created, e.g., to create a TTL func NewOnInstanceFactory(ctx context.Context, createConnPoolForDb CreateConnPoolForDbFn, opts ...OnInstanceFactoryOpt) (_ Factory, _retErr error) { options := onInstanceFactoryOptions{ - dbPrefix: DefaultOnInstanceDbPrefix, - metadataSchema: DefaultOnInstanceMetadataSchema, - metadataTable: DefaultOnInstanceMetadataTable, - rootDatabase: "postgres", - logger: log.SimpleLogger(), + dbPrefix: DefaultOnInstanceDbPrefix, + metadataSchema: DefaultOnInstanceMetadataSchema, + metadataTable: DefaultOnInstanceMetadataTable, + rootDatabase: "postgres", + logger: log.SimpleLogger(), + templateDatabase: DefaultTemplateDatabase, } for _, opt := range opts { opt(&options) @@ -175,8 +185,15 @@ func (o *onInstanceFactory) Create(ctx context.Context) (_ *Database, _retErr er defer rootConn.Close() tempDbName := o.options.dbPrefix + strings.ReplaceAll(dbUUID.String(), "-", "_") + + createDbSql := fmt.Sprintf( + "CREATE DATABASE %s TEMPLATE %s;", + tempDbName, + pgx.Identifier{o.options.templateDatabase}.Sanitize(), + ) + // Create the temporary database using template0, the default Postgres template with no user-defined objects. - if _, err = rootConn.ExecContext(ctx, fmt.Sprintf("CREATE DATABASE %s TEMPLATE template0;", tempDbName)); err != nil { + if _, err = rootConn.ExecContext(ctx, createDbSql); err != nil { return nil, fmt.Errorf("creating temporary database: %w", err) } defer util.DoOnErrOrPanic(&_retErr, func() { diff --git a/pkg/tempdb/factory_test.go b/pkg/tempdb/factory_test.go index d271fb6..f684adc 100644 --- a/pkg/tempdb/factory_test.go +++ b/pkg/tempdb/factory_test.go @@ -215,3 +215,114 @@ func (suite *onInstanceTempDbFactorySuite) TestDropTempDB_CannotDropNonTempDb() func TestOnInstanceFactorySuite(t *testing.T) { suite.Run(t, new(onInstanceTempDbFactorySuite)) } + +func (suite *onInstanceTempDbFactorySuite) TestCreate_UnknownTemplateDatabase() { + rootDbName := "some_other_root" + rootDb, err := suite.engine.CreateDatabaseWithName(rootDbName) + suite.Require().NoError(err, "failed to create the root DB") + suite.T().Cleanup(func() { + suite.Require().NoError(rootDb.DropDB()) + }) + + factory := suite.mustBuildFactory( + WithRootDatabase(rootDbName), + WithTemplateDatabase("non_existent_template_db"), + ) + defer func() { + suite.Require().NoError(factory.Close()) + }() + + _, err = factory.Create(context.Background()) + suite.ErrorContains( + err, + "template database \"non_existent_template_db\" does not exist", + "Expected an error about non-existent template DB", + ) +} + +func (suite *onInstanceTempDbFactorySuite) TestCreate_UsesCustomTemplateDatabase() { + // 1) Create the template DB + templateDbName := "mytemplatedb" + templateDb, err := suite.engine.CreateDatabaseWithName(templateDbName) + suite.Require().NoError(err, "failed to create the custom template database") + + // 2) Connect to it and create a table, then close all connections + templateDbPool, err := suite.getConnPoolForDb(templateDbName) + suite.Require().NoError(err, "could not get conn pool for template DB") + + conn, err := templateDbPool.Conn(context.Background()) + suite.Require().NoError(err, "could not get a connection from template DB pool") + + _, err = conn.ExecContext(context.Background(), ` + CREATE TABLE template_table ( + id SERIAL PRIMARY KEY, + name TEXT + ); + `) + suite.Require().NoError(err, "failed to create table in template DB") + + suite.NoError(conn.Close()) + suite.NoError(templateDbPool.Close()) + + // 3) Terminate any lingering sessions on "mytemplatedb" + rootConnPool, err := suite.getConnPoolForDb("postgres") + suite.Require().NoError(err, "failed to get conn pool for 'postgres'") + + rootConn, err := rootConnPool.Conn(context.Background()) + suite.Require().NoError(err, "failed to get root connection") + + _, err = rootConn.ExecContext(context.Background(), ` + SELECT pg_terminate_backend(pid) + FROM pg_stat_activity + WHERE datname = $1 + AND pid <> pg_backend_pid() + `, templateDbName) + suite.Require().NoError(err, "failed to terminate leftover connections to the template DB") + + suite.NoError(rootConn.Close()) + suite.NoError(rootConnPool.Close()) + + // 4) Create the "root" DB for CREATE DATABASE statements + rootDbName := "mytemplate_root" + rootDb, err := suite.engine.CreateDatabaseWithName(rootDbName) + suite.Require().NoError(err, "failed to create the root DB for create statements") + + suite.T().Cleanup(func() { + suite.NoError(rootDb.DropDB()) + suite.NoError(templateDb.DropDB()) + }) + + // 5) Build the factory that uses our template DB + factory := suite.mustBuildFactory( + WithRootDatabase(rootDbName), + WithTemplateDatabase(templateDbName), + ) + + defer func() { + suite.Require().NoError(factory.Close()) + }() + + // 6) Create a new DB from our template + newDb, err := factory.Create(context.Background()) + suite.Require().NoError(err, "should create a DB from mytemplatedb without error") + + // 7) Verify the table is inherited in the newly created DB + newConn, err := newDb.ConnPool.Conn(context.Background()) + suite.Require().NoError(err, "could not get conn from newly created DB pool") + + var count int + err = newConn.QueryRowContext(context.Background(), ` + SELECT count(*) + FROM pg_class c + JOIN pg_namespace n ON c.relnamespace = n.oid + WHERE c.relkind = 'r' + AND c.relname = 'template_table' + AND n.nspname = 'public'; + `).Scan(&count) + + suite.Require().NoError(err, "failed to check existence of 'template_table'") + suite.Equal(1, count, "expected 'template_table' to exist in the new DB") + + suite.NoError(newConn.Close()) + suite.NoError(newDb.Close(context.Background())) +}