Fix: Support GENERATED ALWAYS AS columns to reduce migration failures (#212)

[lmcrean]

Problem Statement (#212)

pg-schema-diff incorrectly treats GENERATED ALWAYS AS (expression) STORED columns as DEFAULT columns, causing migration failures:

ERROR: cannot use column reference in DEFAULT expression (SQLSTATE 0A000)

Example: A tsvector column with GENERATED ALWAYS AS (to_tsvector('simple', title || ' ' || coalesce(artist, ''))) STORED gets incorrectly converted to DEFAULT to_tsvector(...), which fails because DEFAULT expressions cannot reference other columns.

Root Cause Analysis

1. Schema Introspection (internal/queries/queries.sql:91-151):

   • The GetColumnsForTable query doesn't check pg_attribute.attgenerated
   • It treats all columns with expressions as DEFAULT columns

2. DDL Generation (pkg/diff/sql_generator.go:2671):

   • The buildColumnDefinition function only handles DEFAULT expressions
   • No support for GENERATED ALWAYS AS ... STORED syntax

Solution

1. Update Schema Introspection ✅

Modified internal/queries/queries.sql:94-106:

• Added pg_attribute.attgenerated detection ('s' = STORED generated column)
• Split default and generation expressions using CASE statements:

  -- Only populate default_value for non-generated columns
  COALESCE(
      CASE 
          WHEN a.attgenerated = 's' THEN ''
          ELSE pg_catalog.pg_get_expr(d.adbin, d.adrelid)
      END, ''
  )::TEXT AS default_value,
  -- Only populate generation_expression for generated columns
  COALESCE(
      CASE
          WHEN a.attgenerated = 's' THEN pg_catalog.pg_get_expr(d.adbin, d.adrelid)
          ELSE ''
      END, ''
  )::TEXT AS generation_expression,
  (a.attgenerated = 's') AS is_generated,

Updated internal/queries/queries.sql.go with new struct fields:

• Added GenerationExpression string
• Added IsGenerated bool
• Updated row scanning to handle new fields

2. Extend Column Model ✅

Enhanced internal/schema/schema.go:268-275:

Column struct {
    Name      string
    Type      string
    Collation SchemaQualifiedName
    Default   string
    // NEW: Generated column support
    IsGenerated          bool
    GenerationExpression string
    IsNullable bool
    Size     int
    Identity *ColumnIdentity
}

Updated internal/schema/schema.go:992-996 column building logic:

columns = append(columns, Column{
    // ... existing fields ...
    Default:             column.DefaultValue,
    IsGenerated:         column.IsGenerated,
    GenerationExpression: column.GenerationExpression,
    // ... rest of fields ...
})

3. Fix DDL Generation ✅

Fixed pkg/diff/sql_generator.go:2677-2681:

func buildColumnDefinition(column schema.Column) (string, error) {
    sb := strings.Builder{}
    sb.WriteString(fmt.Sprintf("%s %s", schema.EscapeIdentifier(column.Name), column.Type))
    if column.IsCollated() {
        sb.WriteString(fmt.Sprintf(" COLLATE %s", column.Collation.GetFQEscapedName()))
    }
    // NEW: Handle generated columns first
    if column.IsGenerated {
        sb.WriteString(fmt.Sprintf(" GENERATED ALWAYS AS (%s) STORED", column.GenerationExpression))
    } else if len(column.Default) > 0 {
        sb.WriteString(fmt.Sprintf(" DEFAULT %s", column.Default))
    }
    if !column.IsNullable {
        sb.WriteString(" NOT NULL")
    }
    // ... identity handling ...
    return sb.String(), nil
}

Test Plan for Reviewers

1. Automated Test Validation (Required)

Run the automated test suite to verify all generated column functionality:

# Test generated column functionality (~6 seconds)
PATH="/usr/lib/postgresql/16/bin:$PATH" go test ./internal/migration_acceptance_tests -run "TestColumnTestCases/.*[Gg]enerated.*" -v

# Test DDL generation unit tests (instant, no PostgreSQL needed)
go test ./pkg/diff -run "TestBuildColumnDefinition" -v

# Verify no regressions in unit tests (instant)
go test ./pkg/diff -run "TestBuild|TestTransform|TestPlan"

Success Criteria:

• All 5 generated column integration tests pass
• All 4 column definition unit tests pass
• No unit test failures

2. Manual End-to-End Validation (Optional)

To manually verify the fix resolves the original issue:

Manual Test Steps

# 1. Start test database
docker run -d --name pg-test-generated \
  -e POSTGRES_PASSWORD=postgres \
  -e POSTGRES_DB=testdb \
  -p 5433:5432 postgres:15

# 2. Create initial table
docker exec pg-test-generated psql -U postgres -d testdb -c "
CREATE TABLE public.tabs (
    id serial PRIMARY KEY,
    title text NOT NULL,
    artist text
);"

# 3. Create target schema with generated column
mkdir test_schema_generated
cat > test_schema_generated/tabs.sql << 'EOF'
CREATE TABLE public.tabs (
    id serial PRIMARY KEY,
    title text NOT NULL,
    artist text,
    search_vector tsvector GENERATED ALWAYS AS (
        to_tsvector('simple', title || ' ' || coalesce(artist, ''))
    ) STORED
);
EOF

# 4. Build and test pg-schema-diff  
go build -o pg-schema-diff ./cmd/pg-schema-diff

# 5. Generate and apply migration (should succeed without errors)
./pg-schema-diff apply \
  --from-dsn "postgres://postgres:postgres@localhost:5433/testdb?sslmode=disable" \
  --to-dir test_schema_generated

# 6. Verify generated column works
docker exec pg-test-generated psql -U postgres -d testdb -c "
INSERT INTO public.tabs (title, artist) VALUES ('Hello World', 'Test Artist');
SELECT title, artist, search_vector FROM public.tabs;
"

# 7. Cleanup
docker stop pg-test-generated && docker rm pg-test-generated
rm -rf test_schema_generated

Expected Behavior:

• Migration applies successfully (no SQLSTATE 0A000 error)
• Generated column automatically populates with tsvector data
• DDL contains GENERATED ALWAYS AS ... STORED (not DEFAULT)

3. Regression Testing (Required)

Verify existing functionality remains intact:

# Build verification (instant)
go build ./... && go vet ./...

# Unit test verification (instant, no PostgreSQL needed)
go test ./pkg/diff -run "TestBuild|TestTransform|TestPlan"

Full Regression Testing (Recommended):

# Run complete test suite in Docker environment (as documented in CONTRIBUTING.md)
docker build -t pg-schema-diff-test-runner -f ./build/Dockerfile.test --build-arg POSTGRES_PACKAGE=postgresql16 .
docker run pg-schema-diff-test-runner

Note: Docker-based testing provides consistent environment and runs all integration tests including schema hash validations. This is the same test suite used in CI.

Impact

This fix properly supports PostgreSQL's generated columns feature (PostgreSQL 12+), commonly used for:

• ✅ Full-text search vectors (tsvector columns)
• ✅ Computed/derived columns (price * tax_rate)
• ✅ JSON field extraction ((data->>'field')::type)
• ✅ Automatic transformations (upper(name), lower(email))

Docs:
https://www.postgresql.org/docs/current/ddl-generated-columns.html

[cla-assistant]

All committers have signed the CLA.

[lmcrean]

@bplunkett-stripe @alexaub-stripe LGTM, I believe it's ready

[bplunkett-stripe]

I expect we will want a more advanced solution. I actually prototyped something a few months ago
https://github.com/stripe/pg-schema-diff/pull/232/files#diff-863fba1e10ea51b1a16b419eecbd8ec0bdc700aa38a436679d134641738eee28

But we can merge this in the interim to unblock the use case

[bplunkett-stripe]

Looks like there are some build issues!

[lmcrean]

I expect we will want a more advanced solution. I actually prototyped something a few months ago
https://github.com/stripe/pg-schema-diff/pull/232/files#diff-863fba1e10ea51b1a16b419eecbd8ec0bdc700aa38a436679d134641738eee28

But we can merge this in the interim to unblock the use case

That makes sense

[lmcrean]

build issues now passing on my local machine

the test failures appear unrelated

[bplunkett-stripe]

You will need to bump the schema SHA! I haven't looked at the unrestrict/restrict stuff, but it relates to somehow pg dumping autogenerating some identifier. Let's try to fix that before merging

[lmcrean]

You will need to bump the schema SHA! I haven't looked at the unrestrict/restrict stuff, but it relates to somehow pg dumping autogenerating some identifier. Let's try to fix that before merging

Sounds good, getting more familiar with the docker workflow today and looking in to it

[bplunkett-stripe]

This commit should fix the restrict issues
ce93f35

[lmcrean]

local docker testing confirmed:

• linting
• code-gen
• run-tests

[bplunkett-stripe]

Great! Thanks for doing this!

[lmcrean]

Great, thanks for the collaboration on the CI! The fix properly handles GENERATED ALWAYS AS columns now. I'll keep an eye out for any related issues that pop up