Backport: fix UPDATE failures on GENERATED ALWAYS AS IDENTITY columns (#820) (#822)

Author: tsg

pkg/schemalog/schema.go

@@ -276,6 +276,14 @@ func (c *Column) IsGenerated() bool {
 	return c.Generated || c.Identity != ""
 }
 
+// IsAlwaysIdentity reports whether the column is defined as
+// GENERATED ALWAYS AS IDENTITY. Such columns reject explicit values in UPDATE
+// SET clauses — only DEFAULT is accepted. The DDL injector encodes this as
+// "ALWAYS" in the schema log (see migrations/postgres/core/2_create_emit_ddl_function_and_triggers.up.sql).
+func (c *Column) IsAlwaysIdentity() bool {
+	return c.Identity == "ALWAYS"
+}
+
 func (c *Column) IsSerial() bool {
 	return c.HasSequence() &&
 		(strings.ToUpper(c.DataType) == "INTEGER" ||

pkg/stream/integration/pg_pg_integration_test.go

@@ -416,6 +416,93 @@ func Test_PostgresToPostgres_SchemaObjects(t *testing.T) {
 	}, 20*time.Second, 200*time.Millisecond)
 }
 
+// Test_PostgresToPostgres_AlwaysIdentityUpdate verifies that UPDATE events on a
+// table with a GENERATED ALWAYS AS IDENTITY column are replicated successfully.
+// Postgres rejects UPDATE ... SET <always-identity-col> = <value> with
+// "column can only be updated to DEFAULT", so the column must be filtered out
+// of the SET clause on the target.
+func Test_PostgresToPostgres_AlwaysIdentityUpdate(t *testing.T) {
+	if os.Getenv("PGSTREAM_INTEGRATION_TESTS") == "" {
+		t.Skip("skipping integration test...")
+	}
+
+	cfg := &stream.Config{
+		Listener:  testPostgresListenerCfg(),
+		Processor: testPostgresProcessorCfg(pgurl, withoutBulkIngestion),
+	}
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	targetConn, err := pglib.NewConn(ctx, targetPGURL)
+	require.NoError(t, err)
+	defer targetConn.Close(ctx)
+
+	runStream(t, ctx, cfg)
+
+	testTable := "pg2pg_always_identity_update_test"
+	defer execQuery(t, ctx, fmt.Sprintf("DROP TABLE IF EXISTS %s", testTable))
+
+	// Table has a separate primary key plus a GENERATED ALWAYS AS IDENTITY
+	// column. UPDATEs on `name` must not include `request_id` in the SET
+	// clause, or Postgres rejects them.
+	execQuery(t, ctx, fmt.Sprintf(`
+		CREATE TABLE %s (
+			id bigint PRIMARY KEY,
+			request_id bigint GENERATED ALWAYS AS IDENTITY,
+			name text
+		)
+	`, testTable))
+
+	// Wait for the table schema to land on the target before querying it.
+	require.Eventually(t, func() bool {
+		columns := getInformationSchemaColumns(t, ctx, targetConn, testTable)
+		return len(columns) == 3
+	}, 20*time.Second, 200*time.Millisecond, "table schema not replicated")
+
+	execQuery(t, ctx, fmt.Sprintf("INSERT INTO %s(id, name) VALUES(1, 'Alice')", testTable))
+	execQuery(t, ctx, fmt.Sprintf("INSERT INTO %s(id, name) VALUES(2, 'Bob')", testTable))
+
+	require.Eventually(t, func() bool {
+		rows := getIDNameRows(t, ctx, targetConn,
+			fmt.Sprintf("SELECT id, name FROM %s ORDER BY id", testTable))
+		return len(rows) == 2 && rows[0].name == "Alice" && rows[1].name == "Bob"
+	}, 20*time.Second, 200*time.Millisecond, "initial inserts not replicated")
+
+	// Capture the original request_id values on the target so we can verify
+	// they are preserved across the UPDATE (since the SET clause must not
+	// touch the always-identity column).
+	var aliceReqIDBefore, bobReqIDBefore int64
+	err = targetConn.QueryRow(ctx, []any{&aliceReqIDBefore},
+		fmt.Sprintf("SELECT request_id FROM %s WHERE id = 1", testTable))
+	require.NoError(t, err)
+	err = targetConn.QueryRow(ctx, []any{&bobReqIDBefore},
+		fmt.Sprintf("SELECT request_id FROM %s WHERE id = 2", testTable))
+	require.NoError(t, err)
+
+	// Perform UPDATEs that, pre-fix, produced
+	// "column \"request_id\" can only be updated to DEFAULT" on the target.
+	execQuery(t, ctx, fmt.Sprintf("UPDATE %s SET name = 'Alice2' WHERE id = 1", testTable))
+	execQuery(t, ctx, fmt.Sprintf("UPDATE %s SET name = 'Bob2' WHERE id = 2", testTable))
+
+	require.Eventually(t, func() bool {
+		rows := getIDNameRows(t, ctx, targetConn,
+			fmt.Sprintf("SELECT id, name FROM %s ORDER BY id", testTable))
+		return len(rows) == 2 && rows[0].name == "Alice2" && rows[1].name == "Bob2"
+	}, 20*time.Second, 200*time.Millisecond, "UPDATE not replicated — likely rejected by always-identity rule")
+
+	var aliceReqIDAfter, bobReqIDAfter int64
+	err = targetConn.QueryRow(ctx, []any{&aliceReqIDAfter},
+		fmt.Sprintf("SELECT request_id FROM %s WHERE id = 1", testTable))
+	require.NoError(t, err)
+	err = targetConn.QueryRow(ctx, []any{&bobReqIDAfter},
+		fmt.Sprintf("SELECT request_id FROM %s WHERE id = 2", testTable))
+	require.NoError(t, err)
+
+	require.Equal(t, aliceReqIDBefore, aliceReqIDAfter, "request_id should be unchanged by UPDATE")
+	require.Equal(t, bobReqIDBefore, bobReqIDAfter, "request_id should be unchanged by UPDATE")
+}
+
 func Test_PostgresToPostgres_Sequences(t *testing.T) {
 	if os.Getenv("PGSTREAM_INTEGRATION_TESTS") == "" {
 		t.Skip("skipping integration test...")

pkg/wal/processor/postgres/helper_test.go

@@ -22,17 +22,25 @@ func (m *mockAdapter) close() error {
 }
 
 type mockSchemaObserver struct {
-	getGeneratedColumnNamesFn func(ctx context.Context, schema, table string) (map[string]struct{}, error)
-	getSequenceColumnsFn      func(ctx context.Context, schema, table string) (map[string]string, error)
-	isMaterializedViewFn      func(schema, table string) bool
-	updateFn                  func(logEntry *schemalog.LogEntry)
-	closeFn                   func() error
+	getGeneratedColumnNamesFn      func(ctx context.Context, schema, table string) (map[string]struct{}, error)
+	getAlwaysIdentityColumnNamesFn func(ctx context.Context, schema, table string) (map[string]struct{}, error)
+	getSequenceColumnsFn           func(ctx context.Context, schema, table string) (map[string]string, error)
+	isMaterializedViewFn           func(schema, table string) bool
+	updateFn                       func(logEntry *schemalog.LogEntry)
+	closeFn                        func() error
 }
 
 func (m *mockSchemaObserver) getGeneratedColumnNames(ctx context.Context, schema, table string) (map[string]struct{}, error) {
 	return m.getGeneratedColumnNamesFn(ctx, schema, table)
 }
 
+func (m *mockSchemaObserver) getAlwaysIdentityColumnNames(ctx context.Context, schema, table string) (map[string]struct{}, error) {
+	if m.getAlwaysIdentityColumnNamesFn == nil {
+		return nil, nil
+	}
+	return m.getAlwaysIdentityColumnNamesFn(ctx, schema, table)
+}
+
 func (m *mockSchemaObserver) getSequenceColumns(ctx context.Context, schema, table string) (map[string]string, error) {
 	return m.getSequenceColumnsFn(ctx, schema, table)
 }

pkg/wal/processor/postgres/postgres_schema_observer.go

@@ -21,6 +21,10 @@ type pgSchemaObserver struct {
 	pgConn pglib.Querier
 	// generatedTableColumns is a map of schema.table to a list of generated column names.
 	generatedTableColumns *synclib.Map[string, map[string]struct{}]
+	// alwaysIdentityTableColumns is a map of schema.table to a set of column names
+	// defined as GENERATED ALWAYS AS IDENTITY. These must be filtered from UPDATE
+	// SET clauses since Postgres rejects explicit values for them.
+	alwaysIdentityTableColumns *synclib.Map[string, map[string]struct{}]
 	// materializedViews is a map of schema name to a set of materialized view names.
 	materializedViews *synclib.Map[string, map[string]struct{}]
 	// columnTableSequences is a map of schema.table to a map of sequence column names.
@@ -37,11 +41,12 @@ func newPGSchemaObserver(ctx context.Context, pgURL string, logger loglib.Logger
 		return nil, err
 	}
 	return &pgSchemaObserver{
-		pgConn:                pgConn,
-		generatedTableColumns: synclib.NewMap[string, map[string]struct{}](),
-		materializedViews:     synclib.NewMap[string, map[string]struct{}](),
-		columnTableSequences:  synclib.NewMap[string, map[string]string](),
-		logger:                logger,
+		pgConn:                     pgConn,
+		generatedTableColumns:      synclib.NewMap[string, map[string]struct{}](),
+		alwaysIdentityTableColumns: synclib.NewMap[string, map[string]struct{}](),
+		materializedViews:          synclib.NewMap[string, map[string]struct{}](),
+		columnTableSequences:       synclib.NewMap[string, map[string]string](),
+		logger:                     logger,
 	}, nil
 }
 
@@ -66,6 +71,25 @@ func (o *pgSchemaObserver) getGeneratedColumnNames(ctx context.Context, schema,
 	return colNames, nil
 }
 
+// getAlwaysIdentityColumnNames returns the set of GENERATED ALWAYS AS IDENTITY
+// column names for the given schema.table. If not cached, it queries postgres.
+func (o *pgSchemaObserver) getAlwaysIdentityColumnNames(ctx context.Context, schema, table string) (map[string]struct{}, error) {
+	key := pglib.QuoteQualifiedIdentifier(schema, table)
+
+	columns, found := o.alwaysIdentityTableColumns.Get(key)
+	if found {
+		return columns, nil
+	}
+
+	colNames, err := o.queryAlwaysIdentityColumnNames(ctx, schema, table)
+	if err != nil {
+		return nil, err
+	}
+
+	o.alwaysIdentityTableColumns.Set(key, colNames)
+	return colNames, nil
+}
+
 // isMaterializedView will return true if the input schema.table is a
 // materialized view. It uses an internal cache to reduce the number of calls to
 // postgres. If the value is not in the cache, it will query postgres.
@@ -114,18 +138,28 @@ func (o *pgSchemaObserver) update(logEntry *schemalog.LogEntry) {
 }
 
 // updateGeneratedColumnNames will update the internal cache with the table
-// columns for the schema log on input.
+// columns for the schema log on input. Identity columns are added to
+// generatedColumns via IsGenerated() (preserved historical behavior so live
+// INSERTs let the target auto-generate ids and the sequence increments
+// naturally). GENERATED ALWAYS AS IDENTITY columns are additionally tracked in
+// alwaysIdentityTableColumns so UPDATE SET clauses can drop them even on
+// cache paths where generatedColumns is empty (e.g. populated via SQL query).
 func (o *pgSchemaObserver) updateGeneratedColumnNames(logEntry *schemalog.LogEntry) {
 	for _, table := range logEntry.Schema.Tables {
 		key := pglib.QuoteQualifiedIdentifier(logEntry.SchemaName, table.Name)
 		generatedColumns := make(map[string]struct{}, len(table.Columns))
+		alwaysIdentityColumns := make(map[string]struct{}, len(table.Columns))
 		for _, c := range table.Columns {
+			if c.IsAlwaysIdentity() {
+				alwaysIdentityColumns[pglib.QuoteIdentifier(c.Name)] = struct{}{}
+			}
 			if c.IsGenerated() {
 				generatedColumns[pglib.QuoteIdentifier(c.Name)] = struct{}{}
 			}
 		}
 
 		o.generatedTableColumns.Set(key, generatedColumns)
+		o.alwaysIdentityTableColumns.Set(key, alwaysIdentityColumns)
 	}
 }
 
@@ -183,6 +217,34 @@ func (o *pgSchemaObserver) queryGeneratedColumnNames(ctx context.Context, schema
 	return columnNames, nil
 }
 
+const alwaysIdentityTableColumnsQuery = `SELECT attname FROM pg_attribute
+		WHERE attnum > 0
+		AND attrelid = (SELECT c.oid FROM pg_class c JOIN pg_namespace n ON c.relnamespace=n.oid WHERE c.relname=$1 and n.nspname=$2)
+		AND attidentity = 'a'`
+
+func (o *pgSchemaObserver) queryAlwaysIdentityColumnNames(ctx context.Context, schemaName, tableName string) (map[string]struct{}, error) {
+	columnNames := map[string]struct{}{}
+	rows, err := o.pgConn.Query(ctx, alwaysIdentityTableColumnsQuery, tableName, schemaName)
+	if err != nil {
+		return nil, fmt.Errorf("getting table always-identity column names for table %s.%s: %w", schemaName, tableName, err)
+	}
+	defer rows.Close()
+
+	for rows.Next() {
+		var columnName string
+		if err := rows.Scan(&columnName); err != nil {
+			return nil, fmt.Errorf("scanning table always-identity column name: %w", err)
+		}
+		columnNames[pglib.QuoteIdentifier(columnName)] = struct{}{}
+	}
+
+	if err := rows.Err(); err != nil {
+		return nil, err
+	}
+
+	return columnNames, nil
+}
+
 const materializedViewsQuery = `SELECT matviewname FROM pg_matviews WHERE schemaname = $1`
 
 func (o *pgSchemaObserver) queryMaterializedViews(ctx context.Context, schemaName string) (map[string]struct{}, error) {

pkg/wal/processor/postgres/postgres_schema_observer_test.go

@@ -153,6 +153,112 @@ func TestPGSchemaObserver_getGeneratedColumnNames(t *testing.T) {
 	}
 }
 
+func TestPGSchemaObserver_updateGeneratedColumnNames(t *testing.T) {
+	t.Parallel()
+
+	// The DDL injector encodes attidentity as the human-readable strings
+	// "ALWAYS" / "BY DEFAULT" in the schema log JSON, not the raw 'a'/'d'.
+	const always = "ALWAYS"
+	const byDefault = "BY DEFAULT"
+	quotedTable := `"test_schema"."test_table"`
+
+	tests := []struct {
+		name string
+		cols []schemalog.Column
+
+		wantGenerated      map[string]struct{}
+		wantAlwaysIdentity map[string]struct{}
+	}{
+		{
+			name: "plain columns",
+			cols: []schemalog.Column{
+				{Name: "id", DataType: "integer"},
+				{Name: "name", DataType: "text"},
+			},
+
+			wantGenerated:      map[string]struct{}{},
+			wantAlwaysIdentity: map[string]struct{}{},
+		},
+		{
+			name: "truly generated column",
+			cols: []schemalog.Column{
+				{Name: "id", DataType: "integer"},
+				{Name: "full_name", DataType: "text", Generated: true},
+			},
+
+			wantGenerated:      map[string]struct{}{`"full_name"`: {}},
+			wantAlwaysIdentity: map[string]struct{}{},
+		},
+		{
+			// GENERATED ALWAYS AS IDENTITY lands in BOTH caches: filtered
+			// from all DML via generatedColumns (historical behavior — lets
+			// target sequence auto-increment on INSERT) and additionally
+			// tracked in alwaysIdentityColumns so UPDATE SET drops it even
+			// when generatedColumns is empty (e.g. SQL-query cache path).
+			name: "always identity column",
+			cols: []schemalog.Column{
+				{Name: "id", DataType: "bigint", Identity: always},
+				{Name: "name", DataType: "text"},
+			},
+
+			wantGenerated:      map[string]struct{}{`"id"`: {}},
+			wantAlwaysIdentity: map[string]struct{}{`"id"`: {}},
+		},
+		{
+			// GENERATED BY DEFAULT AS IDENTITY also lands in generatedColumns
+			// via IsGenerated() (Identity != ""). It is NOT in the
+			// alwaysIdentity map since PG accepts explicit values for it in
+			// UPDATE SET.
+			name: "by-default identity column",
+			cols: []schemalog.Column{
+				{Name: "id", DataType: "bigint", Identity: byDefault},
+				{Name: "name", DataType: "text"},
+			},
+
+			wantGenerated:      map[string]struct{}{`"id"`: {}},
+			wantAlwaysIdentity: map[string]struct{}{},
+		},
+		{
+			name: "mix of generated and always identity",
+			cols: []schemalog.Column{
+				{Name: "id", DataType: "bigint", Identity: always},
+				{Name: "request_id", DataType: "bigint", Identity: always},
+				{Name: "full_name", DataType: "text", Generated: true},
+				{Name: "name", DataType: "text"},
+			},
+
+			wantGenerated:      map[string]struct{}{`"id"`: {}, `"request_id"`: {}, `"full_name"`: {}},
+			wantAlwaysIdentity: map[string]struct{}{`"id"`: {}, `"request_id"`: {}},
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+
+			obs := &pgSchemaObserver{
+				generatedTableColumns:      synclib.NewMap[string, map[string]struct{}](),
+				alwaysIdentityTableColumns: synclib.NewMap[string, map[string]struct{}](),
+				logger:                     loglib.NewNoopLogger(),
+			}
+
+			obs.updateGeneratedColumnNames(&schemalog.LogEntry{
+				SchemaName: testSchema,
+				Schema: schemalog.Schema{
+					Tables: []schemalog.Table{
+						{Name: "test_table", Columns: tc.cols},
+					},
+				},
+			})
+
+			require.Equal(t, map[string]map[string]struct{}{quotedTable: tc.wantGenerated},
+				obs.generatedTableColumns.GetMap())
+			require.Equal(t, map[string]map[string]struct{}{quotedTable: tc.wantAlwaysIdentity},
+				obs.alwaysIdentityTableColumns.GetMap())
+		})
+	}
+}
+
 func TestPGSchemaObserver_isMaterializedView(t *testing.T) {
 	t.Parallel()