[CMCSMACD-6648] Allow for a setup function in SchemaTest

cmd/schema-test/main.go

@@ -1,11 +1,13 @@
 package main
 
 import (
+	"context"
 	"fmt"
 	"log"
 	"os"
 
 	"github.com/corbaltcode/go-libraries/migrations"
+	"github.com/corbaltcode/go-libraries/pgutils"
 	_ "github.com/lib/pq"
 )
 
@@ -22,11 +24,37 @@ func main() {
 		os.Exit(1)
 	}
 
-	err := migrations.SchemaTest(&cfg, allMigrations)
+	commonSchema := "common"
+
+	postgresConnectionString := fmt.Sprintf(
+		"postgres://%s:%s@%s:%s/%s?sslmode=disable",
+		cfg.User, cfg.Password, cfg.Host, cfg.Port, cfg.Database)
+	connectionStringProvider, err := pgutils.NewConnectionStringProviderFromURLString(context.Background(), postgresConnectionString)
+	if err != nil {
+		log.Fatalf("NewConnectionStringProviderFromURLString: %v", err)
+	}
+	dbWithSearchPath, err := pgutils.ConnectDB(pgutils.ToConnector(pgutils.WithSchemaSearchPath(connectionStringProvider, commonSchema)))
+	if err != nil {
+		log.Fatalf("ConnectDB: %v", err)
+	}
+
+	commonSetup := func() error {
+		err = migrations.EnsureSchema(dbWithSearchPath, commonSchema)
+		if err != nil {
+			return fmt.Errorf("EnsureSchema: %v", err)
+		}
+		return migrations.Migrate(dbWithSearchPath, commonMigrations)
+	}
+
+	err = migrations.SchemaTestWithSetup(&cfg, allMigrations, commonSetup)
 	if err != nil {
 		log.Fatalf("First schema test failed: %s", err)
 	}
-	err = migrations.SchemaTest(&cfg, allMigrations)
+
+	// Database must be empty before calling SchemaTest a second time.
+	dbWithSearchPath.MustExec(fmt.Sprintf("DROP SCHEMA %s CASCADE", commonSchema))
+
+	err = migrations.SchemaTestWithSetup(&cfg, allMigrations, commonSetup)
 	if err != nil {
 		log.Fatalf("Second schema test failed: %s", err)
 	}

cmd/schema-test/migrations.go

@@ -2,6 +2,20 @@ package main
 
 import "github.com/corbaltcode/go-libraries/migrations"
 
+var commonMigrations = []migrations.NamedMigration{
+	{
+		Name: "Create a common table",
+		Migration: migrations.StaticMigration([]string{
+			`CREATE TABLE common (
+				id serial PRIMARY KEY
+			)`,
+		}),
+		Reverse: migrations.StaticMigration([]string{
+			`DROP TABLE common`,
+		}),
+	},
+}
+
 var allMigrations = []migrations.NamedMigration{
 	{
 		Name: "Create a table",
@@ -77,4 +91,16 @@ var allMigrations = []migrations.NamedMigration{
 			`DROP VIEW v`,
 		}),
 	},
+	{
+		Name: "Create a table referencing a common table",
+		Migration: migrations.StaticMigration([]string{
+			`CREATE TABLE dependent (
+				id serial PRIMARY KEY,
+				common_id integer REFERENCES common.common(id)
+			)`,
+		}),
+		Reverse: migrations.StaticMigration([]string{
+			`DROP TABLE dependent`,
+		}),
+	},
 }

migrations/verify.go

@@ -45,7 +45,7 @@ func dump(c *PostgresConfig) ([]byte, error) {
 	return out, err
 }
 
-func verifyNoTables(db *sqlx.DB) error {
+func verifyNoRelations(db *sqlx.DB) error {
 	// Based on the query run for the "\d" command in psql
 	// (as revealed when started with -E flag).
 	q := `SELECT 1 FROM pg_catalog.pg_class c
@@ -60,9 +60,9 @@ func verifyNoTables(db *sqlx.DB) error {
 		// Expected
 		return nil
 	} else if err != nil {
-		return fmt.Errorf("Error checking for existing tables: %w", err)
+		return fmt.Errorf("Error checking for existing relations: %w", err)
 	}
-	return errors.New("Existing tables found. You must run SchemaTest on an empty database.")
+	return errors.New("Existing relations found. You must run this on an empty database.")
 }
 
 func migrateAndRollback(emptyDBConfig *PostgresConfig, db *sqlx.DB, allMigrations []NamedMigration, migrateToIndex, rollbackThroughIndex int, repeatForward bool) error {
@@ -107,22 +107,9 @@ func migrateAndRollback(emptyDBConfig *PostgresConfig, db *sqlx.DB, allMigration
 	return err
 }
 
-// Schema test expects a new *empty* postgres database.
-// It will:
-//  1. Apply all migrations
-//  2. Reverse all migrations
-//  3. For each migration:
-//     a. Apply the migration
-//     b. Reverse the migration
-//     c. Apply the migration again
-//
-// Before and after each step it will use pg_dump to dump the database schema.
-// It will verify that:
-// A. The schema is the same after reversing as before applying.
-// B. (If re-applying) The schema is the same after applying as after re-applying.
-//
-// You must have `pg_dump` in your `PATH` to run this.
-func SchemaTest(emptyDBConfig *PostgresConfig, allMigrations []NamedMigration) error {
+// Does the same thing as SchemaTest but calls the provided setup function after verifying that the
+// database is empty.
+func SchemaTestWithSetup(emptyDBConfig *PostgresConfig, allMigrations []NamedMigration, setup func() error) error {
 	for _, v := range []string{
 		emptyDBConfig.Host,
 		emptyDBConfig.Port,
@@ -141,10 +128,14 @@ func SchemaTest(emptyDBConfig *PostgresConfig, allMigrations []NamedMigration) e
 	if err != nil {
 		return fmt.Errorf("Error connecting: %s", err)
 	}
-	err = verifyNoTables(db)
+	err = verifyNoRelations(db)
 	if err != nil {
 		return err
 	}
+	err = setup()
+	if err != nil {
+		return fmt.Errorf("setup: %s", err)
+	}
 	err = Migrate(db, []NamedMigration{})
 	if err != nil {
 		return fmt.Errorf("Setting up migrations table failed: %s", err)
@@ -170,3 +161,22 @@ func SchemaTest(emptyDBConfig *PostgresConfig, allMigrations []NamedMigration) e
 	}
 	return nil
 }
+
+// Schema test expects a new *empty* postgres database.
+// It will:
+//  1. Apply all migrations
+//  2. Reverse all migrations
+//  3. For each migration:
+//     a. Apply the migration
+//     b. Reverse the migration
+//     c. Apply the migration again
+//
+// Before and after each step it will use pg_dump to dump the database schema.
+// It will verify that:
+// A. The schema is the same after reversing as before applying.
+// B. (If re-applying) The schema is the same after applying as after re-applying.
+//
+// You must have `pg_dump` in your `PATH` to run this.
+func SchemaTest(emptyDBConfig *PostgresConfig, allMigrations []NamedMigration) error {
+	return SchemaTestWithSetup(emptyDBConfig, allMigrations, func() error { return nil })
+}