[FEATURE] Add support for --from-dir and providing a dedicated temporary database server (#199)

Author: bplunkett-stripe

README.md

@@ -110,7 +110,7 @@ echo "CREATE TABLE bar (id varchar(255), message TEXT NOT NULL);" > schema/bar.s
 Apply the schema to a fresh database. [The connection string spec can be found here](https://www.postgresql.org/docs/current/libpq-connect.html#LIBPQ-CONNSTRING).
 Setting the `PGPASSWORD` env var will override any password set in the connection string and is recommended.
 ```bash
-pg-schema-diff apply --dsn "postgres://postgres:postgres@localhost:5432/postgres" --schema-dir schema 
+pg-schema-diff apply --from-dsn "postgres://postgres:postgres@localhost:5432/postgres" --to-dir schema 
 ```
 
 ## 2. Updating schema
@@ -121,7 +121,7 @@ echo "CREATE INDEX message_idx ON bar(message)" >> schema/bar.sql
 
 Apply the schema. Any hazards in the generated plan must be approved
 ```bash
-pg-schema-diff apply --dsn "postgres://postgres:postgres@localhost:5432/postgres" --schema-dir schema --allow-hazards INDEX_BUILD
+pg-schema-diff apply --from-dsn "postgres://postgres:postgres@localhost:5432/postgres" --to-dir schema --allow-hazards INDEX_BUILD
 ```
 
 # Using Library

cmd/pg-schema-diff/apply_cmd.go

@@ -20,37 +20,51 @@ func buildApplyCmd() *cobra.Command {
 		Short: "Migrate your database to the match the inputted schema (apply the schema to the database)",
 	}
 
-	connFlags := createConnFlags(cmd)
-	planFlags := createPlanFlags(cmd)
+	connFlags := createConnectionFlags(cmd, "from-", " The database to migrate")
+	toSchemaFlags := createSchemaSourceFlags(cmd, "to-")
+	planOptsFlags := createPlanOptionsFlags(cmd)
 	allowedHazardsTypesStrs := cmd.Flags().StringSlice("allow-hazards", nil,
 		"Specify the hazards that are allowed. Order does not matter, and duplicates are ignored. If the"+
 			" migration plan contains unwanted hazards (hazards not in this list), then the migration will fail to run"+
 			" (example: --allowed-hazards DELETES_DATA,INDEX_BUILD)")
 	skipConfirmPrompt := cmd.Flags().Bool("skip-confirm-prompt", false, "Skips prompt asking for user to confirm before applying")
 	cmd.RunE = func(cmd *cobra.Command, args []string) error {
 		logger := log.SimpleLogger()
-		connConfig, err := parseConnConfig(*connFlags, logger)
+
+		connConfig, err := parseConnectionFlags(connFlags)
+		if err != nil {
+			return err
+		}
+		fromSchema := dsnSchemaSource(connConfig)
+
+		toSchema, err := parseSchemaSource(*toSchemaFlags)
 		if err != nil {
 			return err
 		}
 
-		planConfig, err := parsePlanConfig(*planFlags)
+		planOptions, err := parsePlanOptions(*planOptsFlags)
 		if err != nil {
 			return err
 		}
 
 		cmd.SilenceUsage = true
 
-		plan, err := generatePlan(context.Background(), logger, connConfig, planConfig)
+		plan, err := generatePlan(cmd.Context(), generatePlanParameters{
+			fromSchema:       fromSchema,
+			toSchema:         toSchema,
+			tempDbConnConfig: connConfig,
+			planOptions:      planOptions,
+			logger:           logger,
+		})
 		if err != nil {
 			return err
 		} else if len(plan.Statements) == 0 {
-			fmt.Println("Schema matches expected. No plan generated")
+			cmd.Println("Schema matches expected. No plan generated")
 			return nil
 		}
 
-		fmt.Println(header("Review plan"))
-		fmt.Print(planToPrettyS(plan), "\n\n")
+		cmd.Println(header("Review plan"))
+		cmd.Print(planToPrettyS(plan), "\n\n")
 
 		if err := failIfHazardsNotAllowed(plan, *allowedHazardsTypesStrs); err != nil {
 			return err
@@ -67,10 +81,10 @@ func buildApplyCmd() *cobra.Command {
 			}
 		}
 
-		if err := runPlan(context.Background(), connConfig, plan); err != nil {
+		if err := runPlan(cmd.Context(), cmd, connConfig, plan); err != nil {
 			return err
 		}
-		fmt.Println("Schema applied successfully")
+		cmd.Println("Schema applied successfully")
 		return nil
 	}
 
@@ -109,7 +123,7 @@ func failIfHazardsNotAllowed(plan diff.Plan, allowedHazardsTypesStrs []string) e
 	return nil
 }
 
-func runPlan(ctx context.Context, connConfig *pgx.ConnConfig, plan diff.Plan) error {
+func runPlan(ctx context.Context, cmd *cobra.Command, connConfig *pgx.ConnConfig, plan diff.Plan) error {
 	connPool, err := openDbWithPgxConfig(connConfig)
 	if err != nil {
 		return err
@@ -129,8 +143,8 @@ func runPlan(ctx context.Context, connConfig *pgx.ConnConfig, plan diff.Plan) er
 	// must be executed within its own transaction block. Postgres will error if you try to set a TRANSACTION-level
 	// timeout for it. SESSION-level statement_timeouts are respected by `ADD INDEX CONCURRENTLY`
 	for i, stmt := range plan.Statements {
-		fmt.Println(header(fmt.Sprintf("Executing statement %d", getDisplayableStmtIdx(i))))
-		fmt.Printf("%s\n\n", statementToPrettyS(stmt))
+		cmd.Println(header(fmt.Sprintf("Executing statement %d", getDisplayableStmtIdx(i))))
+		cmd.Printf("%s\n\n", statementToPrettyS(stmt))
 		start := time.Now()
 		if _, err := conn.ExecContext(ctx, fmt.Sprintf("SET SESSION statement_timeout = %d", stmt.Timeout.Milliseconds())); err != nil {
 			return fmt.Errorf("setting statement timeout: %w", err)
@@ -141,9 +155,9 @@ func runPlan(ctx context.Context, connConfig *pgx.ConnConfig, plan diff.Plan) er
 		if _, err := conn.ExecContext(ctx, stmt.ToSQL()); err != nil {
 			return fmt.Errorf("executing migration statement. the database maybe be in a dirty state: %s: %w", stmt, err)
 		}
-		fmt.Printf("Finished executing statement. Duration: %s\n", time.Since(start))
+		cmd.Printf("Finished executing statement. Duration: %s\n", time.Since(start))
 	}
-	fmt.Println(header("Complete"))
+	cmd.Println(header("Complete"))
 
 	return nil
 }

cmd/pg-schema-diff/apply_cmd_test.go

@@ -0,0 +1,89 @@
+package main
+
+import (
+	"github.com/stripe/pg-schema-diff/internal/pgdump"
+	"github.com/stripe/pg-schema-diff/internal/pgengine"
+)
+
+func (suite *cmdTestSuite) TestApplyCmd() {
+	// Non-comprehensive set of tests for the plan command. Not totally comprehensive to avoid needing to avoid
+	// hindering developer velocity when updating the command.
+	type testCase struct {
+		name string
+		// fromDbArg is an optional argument to override the default "--from-dsn" arg.
+		fromDbArg func(db *pgengine.DB) []string
+		args      []string
+		// dynamicArgs is function that can be used to build args that are dynamic, i.e.,
+		// saving schemas to a randomly generated temporary directory.
+		dynamicArgs []dArgGenerator
+
+		outputContains []string
+		// expectedSchema is the schema that is expected to be in the database after the migration.
+		// If nil, the expected schema will be the fromDDL.
+		expectedSchemaDDL []string
+		// expectErrContains is a list of substrings that are expected to be contained in the error returned by
+		// cmd.RunE. This is DISTINCT from stdErr.
+		expectErrContains []string
+	}
+	for _, tc := range []testCase{
+		{
+			name:        "to dir",
+			dynamicArgs: []dArgGenerator{tempSchemaDirDArg("to-dir", []string{"CREATE TABLE foobar();"})},
+
+			expectedSchemaDDL: []string{"CREATE TABLE foobar();"},
+		},
+		{
+			name:        "to dsn",
+			dynamicArgs: []dArgGenerator{tempDsnDArg(suite.pgEngine, "to-dsn", []string{"CREATE TABLE foobar();"})},
+
+			expectedSchemaDDL: []string{"CREATE TABLE foobar();"},
+		},
+		{
+			name: "from empty dsn",
+			fromDbArg: func(db *pgengine.DB) []string {
+				tempSetPqEnvVarsForDb(suite.T(), db)
+				return []string{"--from-empty-dsn"}
+			},
+			dynamicArgs: []dArgGenerator{tempSchemaDirDArg("to-dir", []string{"CREATE TABLE foobar();"})},
+
+			expectedSchemaDDL: []string{"CREATE TABLE foobar();"},
+		},
+		{
+			name:              "no to schema provided",
+			expectErrContains: []string{"must be set"},
+		},
+		{
+			name:              "two to schemas provided",
+			args:              []string{"--to-dir", "some-other-dir", "--to-dsn", "some-dsn"},
+			expectErrContains: []string{"only one of"},
+		},
+	} {
+		suite.Run(tc.name, func() {
+			fromDb := tempDbWithSchema(suite.T(), suite.pgEngine, nil)
+			if tc.fromDbArg == nil {
+				tc.fromDbArg = func(db *pgengine.DB) []string {
+					return []string{"--from-dsn", db.GetDSN()}
+				}
+			}
+			args := append([]string{
+				"apply",
+				"--skip-confirm-prompt",
+			}, tc.fromDbArg(fromDb)...)
+			args = append(args, tc.args...)
+			suite.runCmdWithAssertions(runCmdWithAssertionsParams{
+				args:              args,
+				dynamicArgs:       tc.dynamicArgs,
+				outputContains:    tc.outputContains,
+				expectErrContains: tc.expectErrContains,
+			})
+			// The migration should have been successful. Assert it was.
+			expectedDb := tempDbWithSchema(suite.T(), suite.pgEngine, tc.expectedSchemaDDL)
+			expectedDbDump, err := pgdump.GetDump(expectedDb, pgdump.WithSchemaOnly())
+			suite.Require().NoError(err)
+			fromDbDump, err := pgdump.GetDump(fromDb, pgdump.WithSchemaOnly())
+			suite.Require().NoError(err)
+
+			suite.Equal(expectedDbDump, fromDbDump)
+		})
+	}
+}

cmd/pg-schema-diff/flags.go

@@ -7,29 +7,51 @@ import (
 	"github.com/go-logfmt/logfmt"
 	"github.com/jackc/pgx/v4"
 	"github.com/spf13/cobra"
-	"github.com/stripe/pg-schema-diff/pkg/log"
 )
 
-type connFlags struct {
-	dsn string
+type connectionFlags struct {
+	// dsn is the connection string for the database.
+	dsn         string
+	dsnFlagName string
+
+	// isEmptyDsnUsingPq indicates to connect via DSN using the pq environment variables and defaults.
+	isEmptyDsnUsingPq         bool
+	isEmptyDsnUsingPqFlagName string
 }
 
-func createConnFlags(cmd *cobra.Command) *connFlags {
-	flags := &connFlags{}
+func createConnectionFlags(cmd *cobra.Command, prefix string, additionalHelp string) *connectionFlags {
+	var c connectionFlags
+
+	c.dsnFlagName = prefix + "dsn"
+	dsnFlagHelp := "Connection string for the database (DB password can be specified through PGPASSWORD environment variable)."
+	if additionalHelp != "" {
+		dsnFlagHelp += " " + additionalHelp
+	}
+	cmd.Flags().StringVar(&c.dsn, c.dsnFlagName, "", dsnFlagHelp)
+
+	c.isEmptyDsnUsingPqFlagName = prefix + "empty-dsn"
+	isEmptyDsnUsingPqFlagHelp := "Connect with an empty DSN using the pq environment variables and defaults."
+	if additionalHelp != "" {
+		isEmptyDsnUsingPqFlagHelp += " " + additionalHelp
+	}
+	cmd.Flags().BoolVar(&c.isEmptyDsnUsingPq, c.isEmptyDsnUsingPqFlagName, false, isEmptyDsnUsingPqFlagHelp)
 
-	cmd.Flags().StringVar(&flags.dsn, "dsn", "", "Connection string for the database (DB password can be specified through PGPASSWORD environment variable)")
-	// Don't mark dsn as a required flag.
-	// Allow users to use the "PGHOST" etc environment variables like `psql`.
+	return &c
+}
 
-	return flags
+func (c *connectionFlags) IsSet() bool {
+	return c.dsn != "" || c.isEmptyDsnUsingPq
 }
 
-func parseConnConfig(c connFlags, logger log.Logger) (*pgx.ConnConfig, error) {
-	if c.dsn == "" {
-		logger.Warnf("DSN flag not set. Using libpq environment variables and default values.")
+func parseConnectionFlags(flags *connectionFlags) (*pgx.ConnConfig, error) {
+	if !flags.isEmptyDsnUsingPq && flags.dsn == "" {
+		return nil, fmt.Errorf("must specify either --%s or --%s", flags.dsnFlagName, flags.isEmptyDsnUsingPqFlagName)
 	}
-
-	return pgx.ParseConfig(c.dsn)
+	connConfig, err := pgx.ParseConfig(flags.dsn)
+	if err != nil {
+		return nil, fmt.Errorf("could not parse connection string %q: %w", flags.dsn, err)
+	}
+	return connConfig, nil
 }
 
 // logFmtToMap parses all LogFmt key/value pairs from the provided string into a

cmd/pg-schema-diff/main.go

@@ -6,21 +6,19 @@ import (
 	"github.com/spf13/cobra"
 )
 
-// rootCmd represents the base command when called without any subcommands
-var rootCmd = &cobra.Command{
-	Use:   "pg-schema-diff",
-	Short: "Diff two Postgres schemas and generate the SQL to get from one to the other",
-}
-
-func init() {
+func buildRootCmd() *cobra.Command {
+	rootCmd := &cobra.Command{
+		Use:   "pg-schema-diff",
+		Short: "Diff two Postgres schemas and generate the SQL to get from one to the other",
+	}
 	rootCmd.AddCommand(buildPlanCmd())
 	rootCmd.AddCommand(buildApplyCmd())
 	rootCmd.AddCommand(buildVersionCmd())
+	return rootCmd
 }
 
 func main() {
-	err := rootCmd.Execute()
-	if err != nil {
+	if err := buildRootCmd().Execute(); err != nil {
 		os.Exit(1)
 	}
 }