Add hazard for new NOT NULL columns requiring backfill

Adding a new NOT NULL column usually needs a backfill.
This change adds a hazard and shows a suggested online migration path
(add nullable -> backfill -> set NOT NULL).

Author: catatsuy

README.md

@@ -152,6 +152,9 @@ Apply the schema. Any hazards in the generated plan must be approved
 pg-schema-diff apply --from-dsn "postgres://postgres:postgres@localhost:5432/postgres" --to-dir schema --allow-hazards INDEX_BUILD
 ```
 
+# Stateful migrations
+- Adding a brand new `NOT NULL` column without a constant `DEFAULT` requires backfilling existing rows. pg-schema-diff will emit the `NEW_NOT_NULL_COLUMN_REQUIRES_BACKFILL` hazard and annotate the plan with an online sequence: add the column as nullable, backfill in batches, then rerun pg-schema-diff to flip to `NOT NULL` via the online CHECK/VALIDATE/SET NOT NULL flow.
+
 # Using Library
 Docs to use the library can be found [here](https://pkg.go.dev/github.com/stripe/pg-schema-diff). Check out [the CLI](https://github.com/stripe/pg-schema-diff/tree/main/cmd/pg-schema-diff)
 for an example implementation with the library

cmd/pg-schema-diff/apply_cmd_test.go

@@ -3,6 +3,7 @@ package main
 import (
 	"github.com/stripe/pg-schema-diff/internal/pgdump"
 	"github.com/stripe/pg-schema-diff/internal/pgengine"
+	"github.com/stripe/pg-schema-diff/pkg/diff"
 )
 
 func (suite *cmdTestSuite) TestApplyCmd() {
@@ -87,3 +88,21 @@ func (suite *cmdTestSuite) TestApplyCmd() {
 		})
 	}
 }
+
+func (suite *cmdTestSuite) TestFailIfHazardsNotAllowed() {
+	plan := diff.Plan{
+		Statements: []diff.Statement{{
+			DDL: "ALTER TABLE public.t ADD COLUMN city_name text NOT NULL",
+			Hazards: []diff.MigrationHazard{{
+				Type: diff.MigrationHazardTypeNewNotNullColumnRequiresBackfill,
+			}},
+		}},
+	}
+
+	err := failIfHazardsNotAllowed(plan, nil)
+	suite.Error(err)
+	suite.Contains(err.Error(), string(diff.MigrationHazardTypeNewNotNullColumnRequiresBackfill))
+
+	err = failIfHazardsNotAllowed(plan, []string{string(diff.MigrationHazardTypeNewNotNullColumnRequiresBackfill)})
+	suite.NoError(err)
+}

internal/migration_acceptance_tests/acceptance_test.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	stdlog "log"
 	"os"
+	"strings"
 	"testing"
 
 	"github.com/google/uuid"
@@ -62,6 +63,9 @@ type (
 		//
 		// If no expectedDBSchemaDDL is specified, the newSchemaDDL will be used
 		expectedDBSchemaDDL []string
+		// expectedHazardMessages is a list of substrings that must be found within the set of hazard messages generated
+		// by the plan.
+		expectedHazardMessages []string
 	}
 )
 
@@ -178,6 +182,19 @@ func runTest(t *testing.T, tc acceptanceTestCase) {
 	}
 	assert.ElementsMatch(t, tc.expectedHazardTypes, getUniqueHazardTypesFromStatements(plan.Statements), prettySprintPlan(plan))
 
+	if len(tc.expectedHazardMessages) > 0 {
+		var hazardMsgs []string
+		for _, stmt := range plan.Statements {
+			for _, hazard := range stmt.Hazards {
+				hazardMsgs = append(hazardMsgs, hazard.Message)
+			}
+		}
+		joinedHazardMsgs := strings.Join(hazardMsgs, "\n")
+		for _, expectedHazardMsg := range tc.expectedHazardMessages {
+			assert.Contains(t, joinedHazardMsgs, expectedHazardMsg, prettySprintPlan(plan))
+		}
+	}
+
 	// Apply the plan
 	require.NoError(t, applyPlan(oldDb, plan), prettySprintPlan(plan))
 

internal/migration_acceptance_tests/column_cases_test.go

@@ -85,6 +85,9 @@ var columnAcceptanceTestCases = []acceptanceTestCase{
             );
 			`,
 		},
+		expectedHazardTypes: []diff.MigrationHazardType{
+			diff.MigrationHazardTypeNewNotNullColumnRequiresBackfill,
+		},
 	},
 	{
 		name: "Add one column with serial",

internal/migration_acceptance_tests/database_schema_source_cases_test.go

@@ -156,6 +156,7 @@ var databaseSchemaSourceTestCases = []acceptanceTestCase{
 
 		expectedHazardTypes: []diff.MigrationHazardType{
 			diff.MigrationHazardTypeIndexBuild,
+			diff.MigrationHazardTypeNewNotNullColumnRequiresBackfill,
 		},
 		expectedDBSchemaDDL: []string{
 			`

internal/migration_acceptance_tests/foreign_key_constraint_cases_test.go

@@ -866,6 +866,7 @@ var foreignKeyConstraintCases = []acceptanceTestCase{
 
 			diff.MigrationHazardTypeAcquiresAccessExclusiveLock,
 			diff.MigrationHazardTypeAcquiresShareRowExclusiveLock,
+			diff.MigrationHazardTypeNewNotNullColumnRequiresBackfill,
 			diff.MigrationHazardTypeIndexBuild,
 			diff.MigrationHazardTypeIndexDropped,
 		},

internal/migration_acceptance_tests/procedure_cases_test.go

@@ -95,7 +95,10 @@ var procedureAcceptanceTestCases = []acceptanceTestCase{
             $$;
 			`,
 		},
-		expectedHazardTypes: []diff.MigrationHazardType{diff.MigrationHazardTypeHasUntrackableDependencies},
+		expectedHazardTypes: []diff.MigrationHazardType{
+			diff.MigrationHazardTypeHasUntrackableDependencies,
+			diff.MigrationHazardTypeNewNotNullColumnRequiresBackfill,
+		},
 	},
 	{
 		name: "Drop procedure and its dependencies",

internal/migration_acceptance_tests/table_cases_test.go

@@ -541,6 +541,7 @@ var tableAcceptanceTestCases = []acceptanceTestCase{
 			diff.MigrationHazardTypeAuthzUpdate,
 			diff.MigrationHazardTypeDeletesData,
 			diff.MigrationHazardTypeIndexDropped,
+			diff.MigrationHazardTypeNewNotNullColumnRequiresBackfill,
 			diff.MigrationHazardTypeIndexBuild,
 		},
 	},
@@ -569,6 +570,50 @@ var tableAcceptanceTestCases = []acceptanceTestCase{
 			diff.MigrationHazardTypeImpactsDatabasePerformance,
 		},
 	},
+	{
+		name: "Add NOT NULL column without default emits backfill hazard",
+		oldSchemaDDL: []string{
+			`
+            CREATE TABLE t(
+                id INT
+            );
+			`,
+		},
+		newSchemaDDL: []string{
+			`
+            CREATE TABLE t(
+                id INT,
+                city_name VARCHAR(255) NOT NULL
+            );
+			`,
+		},
+		expectedHazardTypes: []diff.MigrationHazardType{
+			diff.MigrationHazardTypeNewNotNullColumnRequiresBackfill,
+		},
+		expectedHazardMessages: []string{
+			"backfill",
+			"Add the column as NULLABLE",
+			"SET NOT NULL",
+		},
+	},
+	{
+		name: "Add NOT NULL column with constant default avoids backfill hazard",
+		oldSchemaDDL: []string{
+			`
+            CREATE TABLE t(
+                id INT
+            );
+			`,
+		},
+		newSchemaDDL: []string{
+			`
+            CREATE TABLE t(
+                id INT,
+                city_name VARCHAR(255) NOT NULL DEFAULT ''
+            );
+			`,
+		},
+	},
 }
 
 func TestTableTestCases(t *testing.T) {

pkg/diff/plan.go

@@ -10,18 +10,19 @@ import (
 type MigrationHazardType = string
 
 const (
-	MigrationHazardTypeAcquiresAccessExclusiveLock   MigrationHazardType = "ACQUIRES_ACCESS_EXCLUSIVE_LOCK"
-	MigrationHazardTypeAcquiresShareLock             MigrationHazardType = "ACQUIRES_SHARE_LOCK"
-	MigrationHazardTypeAcquiresShareRowExclusiveLock MigrationHazardType = "ACQUIRES_SHARE_ROW_EXCLUSIVE_LOCK"
-	MigrationHazardTypeCorrectness                   MigrationHazardType = "CORRECTNESS"
-	MigrationHazardTypeDeletesData                   MigrationHazardType = "DELETES_DATA"
-	MigrationHazardTypeHasUntrackableDependencies    MigrationHazardType = "HAS_UNTRACKABLE_DEPENDENCIES"
-	MigrationHazardTypeIndexBuild                    MigrationHazardType = "INDEX_BUILD"
-	MigrationHazardTypeIndexDropped                  MigrationHazardType = "INDEX_DROPPED"
-	MigrationHazardTypeImpactsDatabasePerformance    MigrationHazardType = "IMPACTS_DATABASE_PERFORMANCE"
-	MigrationHazardTypeIsUserGenerated               MigrationHazardType = "IS_USER_GENERATED"
-	MigrationHazardTypeExtensionVersionUpgrade       MigrationHazardType = "UPGRADING_EXTENSION_VERSION"
-	MigrationHazardTypeAuthzUpdate                   MigrationHazardType = "AUTHZ_UPDATE"
+	MigrationHazardTypeAcquiresAccessExclusiveLock      MigrationHazardType = "ACQUIRES_ACCESS_EXCLUSIVE_LOCK"
+	MigrationHazardTypeAcquiresShareLock                MigrationHazardType = "ACQUIRES_SHARE_LOCK"
+	MigrationHazardTypeAcquiresShareRowExclusiveLock    MigrationHazardType = "ACQUIRES_SHARE_ROW_EXCLUSIVE_LOCK"
+	MigrationHazardTypeCorrectness                      MigrationHazardType = "CORRECTNESS"
+	MigrationHazardTypeDeletesData                      MigrationHazardType = "DELETES_DATA"
+	MigrationHazardTypeHasUntrackableDependencies       MigrationHazardType = "HAS_UNTRACKABLE_DEPENDENCIES"
+	MigrationHazardTypeIndexBuild                       MigrationHazardType = "INDEX_BUILD"
+	MigrationHazardTypeIndexDropped                     MigrationHazardType = "INDEX_DROPPED"
+	MigrationHazardTypeImpactsDatabasePerformance       MigrationHazardType = "IMPACTS_DATABASE_PERFORMANCE"
+	MigrationHazardTypeIsUserGenerated                  MigrationHazardType = "IS_USER_GENERATED"
+	MigrationHazardTypeExtensionVersionUpgrade          MigrationHazardType = "UPGRADING_EXTENSION_VERSION"
+	MigrationHazardTypeAuthzUpdate                      MigrationHazardType = "AUTHZ_UPDATE"
+	MigrationHazardTypeNewNotNullColumnRequiresBackfill MigrationHazardType = "NEW_NOT_NULL_COLUMN_REQUIRES_BACKFILL"
 )
 
 // MigrationHazard represents a hazard that a statement poses to a database

pkg/diff/sql_generator.go

@@ -66,6 +66,14 @@ var (
 		Type:    MigrationHazardTypeExtensionVersionUpgrade,
 		Message: "This extension's version is being upgraded. Be sure the newer version is backwards compatible with your use case.",
 	}
+	migrationHazardNewNotNullColumnRequiresBackfill = MigrationHazard{
+		Type: MigrationHazardTypeNewNotNullColumnRequiresBackfill,
+		Message: "Adding a new NOT NULL column without a constant DEFAULT requires a backfill to populate existing rows. " +
+			"Recommended online sequence:\n" +
+			"  1) Add the column as NULLABLE (no NOT NULL)\n" +
+			"  2) Backfill existing rows in batches via application jobs or scripts\n" +
+			"  3) Re-run pg-schema-diff to enforce NOT NULL using the online CHECK/VALIDATE/SET NOT NULL flow",
+	}
 )
 
 type oldAndNew[S any] struct {
@@ -1204,11 +1212,29 @@ func (csg *columnSQLVertexGenerator) Add(column schema.Column) ([]Statement, err
 	if err != nil {
 		return nil, fmt.Errorf("building column definition: %w", err)
 	}
-	return []Statement{{
+	stmt := Statement{
 		DDL:         fmt.Sprintf("%s ADD COLUMN %s", alterTablePrefix(csg.tableName), columnDef),
 		Timeout:     statementTimeoutDefault,
 		LockTimeout: lockTimeoutDefault,
-	}}, nil
+	}
+	if columnRequiresBackfill(column) {
+		stmt.Hazards = append(stmt.Hazards, migrationHazardNewNotNullColumnRequiresBackfill)
+	}
+	return []Statement{stmt}, nil
+}
+
+func columnRequiresBackfill(column schema.Column) bool {
+	if column.IsNullable {
+		return false
+	}
+	return !columnHasSafeDefault(column)
+}
+
+func columnHasSafeDefault(column schema.Column) bool {
+	if column.Identity != nil || column.IsGenerated {
+		return true
+	}
+	return len(strings.TrimSpace(column.Default)) > 0
 }
 
 func (csg *columnSQLVertexGenerator) Delete(column schema.Column) ([]Statement, error) {