Repository: xataio/pgstream
Pull request #382 (2025-06-03 10:11:24)

[Postgres Writer] Fix updates/deletes when no PK
This PR updates the DML query adapter to leverage the identity columns available in replication to build update/delete queries, instead of only relying on the primary keys, which are only available if the metadata injector is enabled. 

Fixes https://github.com/xataio/pgstream/issues/381


Patch:
diff --git a/pkg/wal/processor/postgres/postgres_wal_adapter.go b/pkg/wal/processor/postgres/postgres_wal_adapter.go
index 5d2c5cc..abb4b6d 100644
--- a/pkg/wal/processor/postgres/postgres_wal_adapter.go
+++ b/pkg/wal/processor/postgres/postgres_wal_adapter.go
@@ -48,5 +48,10 @@ func (a *adapter) walEventToQueries(ctx context.Context, e *wal.Event) ([]*query
 		return a.ddlAdapter.walDataToQueries(ctx, e.Data)
 	}
 
-	return []*query{a.dmlAdapter.walDataToQuery(e.Data)}, nil
+	q, err := a.dmlAdapter.walDataToQuery(e.Data)
+	if err != nil {
+		return nil, err
+	}
+
+	return []*query{q}, nil
 }
diff --git a/pkg/wal/processor/postgres/postgres_wal_dml_adapter.go b/pkg/wal/processor/postgres/postgres_wal_dml_adapter.go
index a8ecffb..f595102 100644
--- a/pkg/wal/processor/postgres/postgres_wal_dml_adapter.go
+++ b/pkg/wal/processor/postgres/postgres_wal_dml_adapter.go
@@ -20,7 +20,10 @@ const (
 	onConflictDoNothing
 )
 
-var errUnsupportedOnConflictAction = errors.New("unsupported on conflict action")
+var (
+	errUnsupportedOnConflictAction = errors.New("unsupported on conflict action")
+	errUnableToBuildQuery          = errors.New("unable to build query, no primary keys of previous values available")
+)
 
 type dmlAdapter struct {
 	onConflictAction onConflictAction
@@ -36,18 +39,18 @@ func newDMLAdapter(action string) (*dmlAdapter, error) {
 	}, nil
 }
 
-func (a *dmlAdapter) walDataToQuery(d *wal.Data) *query {
+func (a *dmlAdapter) walDataToQuery(d *wal.Data) (*query, error) {
 	switch d.Action {
 	case "T":
-		return a.buildTruncateQuery(d)
+		return a.buildTruncateQuery(d), nil
 	case "D":
 		return a.buildDeleteQuery(d)
 	case "I":
-		return a.buildInsertQuery(d)
+		return a.buildInsertQuery(d), nil
 	case "U":
 		return a.buildUpdateQuery(d)
 	default:
-		return &query{}
+		return &query{}, nil
 	}
 }
 
@@ -59,15 +62,17 @@ func (a *dmlAdapter) buildTruncateQuery(d *wal.Data) *query {
 	}
 }
 
-func (a *dmlAdapter) buildDeleteQuery(d *wal.Data) *query {
-	primaryKeyCols := a.extractPrimaryKeyColumns(d.Metadata.InternalColIDs, d.Identity)
-	whereQuery, whereValues := a.buildWhereQuery(primaryKeyCols, 0)
+func (a *dmlAdapter) buildDeleteQuery(d *wal.Data) (*query, error) {
+	whereQuery, whereValues, err := a.buildWhereQuery(d, 0)
+	if err != nil {
+		return nil, fmt.Errorf("building delete query: %w", err)
+	}
 	return &query{
 		table:  d.Table,
 		schema: d.Schema,
 		sql:    fmt.Sprintf("DELETE FROM %s %s", quotedTableName(d.Schema, d.Table), whereQuery),
 		args:   whereValues,
-	}
+	}, nil
 }
 
 func (a *dmlAdapter) buildInsertQuery(d *wal.Data) *query {
@@ -92,19 +97,37 @@ func (a *dmlAdapter) buildInsertQuery(d *wal.Data) *query {
 	}
 }
 
-func (a *dmlAdapter) buildUpdateQuery(d *wal.Data) *query {
+func (a *dmlAdapter) buildUpdateQuery(d *wal.Data) (*query, error) {
 	setQuery, setValues := a.buildSetQuery(d.Columns)
-	primaryKeyCols := a.extractPrimaryKeyColumns(d.Metadata.InternalColIDs, d.Columns)
-	whereQuery, whereValues := a.buildWhereQuery(primaryKeyCols, len(d.Columns))
+	whereQuery, whereValues, err := a.buildWhereQuery(d, len(d.Columns))
+	if err != nil {
+		return nil, fmt.Errorf("building update query: %w", err)
+	}
+
 	return &query{
 		table:  d.Table,
 		schema: d.Schema,
 		sql:    fmt.Sprintf("UPDATE %s %s %s", quotedTableName(d.Schema, d.Table), setQuery, whereQuery),
 		args:   append(setValues, whereValues...),
-	}
+	}, nil
 }
 
-func (a *dmlAdapter) buildWhereQuery(cols []wal.Column, placeholderOffset int) (string, []any) {
+func (a *dmlAdapter) buildWhereQuery(d *wal.Data, placeholderOffset int) (string, []any, error) {
+	var cols []wal.Column
+	switch {
+	case len(d.Identity) > 0:
+		// if we have the previous values (replica identity), add them to the where query
+		cols = d.Identity
+	case len(d.Metadata.InternalColIDs) > 0:
+		// if we don't have previous values we have to rely on the primary keys
+		primaryKeyCols := a.extractPrimaryKeyColumns(d.Metadata.InternalColIDs, d.Columns)
+		cols = primaryKeyCols
+	default:
+		// without a where clause in the query we'd be updating/deleting all table
+		// rows, so we need to error to prevent that from happening
+		return "", nil, errUnableToBuildQuery
+	}
+
 	whereQuery := "WHERE"
 	whereValues := make([]any, 0, len(cols))
 	for i, c := range cols {
@@ -114,7 +137,7 @@ func (a *dmlAdapter) buildWhereQuery(cols []wal.Column, placeholderOffset int) (
 		whereQuery = fmt.Sprintf("%s %s = $%d", whereQuery, pglib.QuoteIdentifier(c.Name), i+placeholderOffset+1)
 		whereValues = append(whereValues, c.Value)
 	}
-	return whereQuery, whereValues
+	return whereQuery, whereValues, nil
 }
 
 func (a *dmlAdapter) buildSetQuery(cols []wal.Column) (string, []any) {


Test patch:
diff --git a/pkg/wal/processor/postgres/postgres_wal_dml_adapter_test.go b/pkg/wal/processor/postgres/postgres_wal_dml_adapter_test.go
index f1b4c88..7d74ea2 100644
--- a/pkg/wal/processor/postgres/postgres_wal_dml_adapter_test.go
+++ b/pkg/wal/processor/postgres/postgres_wal_dml_adapter_test.go
@@ -30,6 +30,7 @@ func TestDMLAdapter_walDataToQuery(t *testing.T) {
 		action  onConflictAction
 
 		wantQuery *query
+		wantErr   error
 	}{
 		{
 			name: "truncate",
@@ -56,7 +57,6 @@ func TestDMLAdapter_walDataToQuery(t *testing.T) {
 				Table:  testTable,
 				Identity: []wal.Column{
 					{ID: columnID(1), Name: "id", Value: 1},
-					{ID: columnID(2), Name: "name", Value: "alice"},
 				},
 				Metadata: wal.Metadata{
 					InternalColIDs: []string{columnID(1)},
@@ -92,6 +92,39 @@ func TestDMLAdapter_walDataToQuery(t *testing.T) {
 				args:   []any{1, "alice"},
 			},
 		},
+		{
+			name: "delete with full identity",
+			walData: &wal.Data{
+				Action: "D",
+				Schema: testSchema,
+				Table:  testTable,
+				Identity: []wal.Column{
+					{ID: columnID(1), Name: "id", Value: 1},
+					{ID: columnID(2), Name: "name", Value: "alice"},
+				},
+				Metadata: wal.Metadata{},
+			},
+
+			wantQuery: &query{
+				schema: testSchema,
+				table:  testTable,
+				sql:    fmt.Sprintf("DELETE FROM %s WHERE \"id\" = $1 AND \"name\" = $2", quotedTestTable),
+				args:   []any{1, "alice"},
+			},
+		},
+		{
+			name: "error - delete",
+			walData: &wal.Data{
+				Action:   "D",
+				Schema:   testSchema,
+				Table:    testTable,
+				Identity: []wal.Column{},
+				Metadata: wal.Metadata{},
+			},
+
+			wantQuery: nil,
+			wantErr:   errUnableToBuildQuery,
+		},
 		{
 			name: "insert",
 			walData: &wal.Data{
@@ -185,7 +218,7 @@ func TestDMLAdapter_walDataToQuery(t *testing.T) {
 			},
 		},
 		{
-			name: "update",
+			name: "update - primary key",
 			walData: &wal.Data{
 				Action: "U",
 				Schema: testSchema,
@@ -206,6 +239,70 @@ func TestDMLAdapter_walDataToQuery(t *testing.T) {
 				args:   []any{1, "alice", 1},
 			},
 		},
+		{
+			name: "update - default identity",
+			walData: &wal.Data{
+				Action: "U",
+				Schema: testSchema,
+				Table:  testTable,
+				Columns: []wal.Column{
+					{ID: columnID(1), Name: "id", Value: 1},
+					{ID: columnID(2), Name: "name", Value: "alice"},
+				},
+				Identity: []wal.Column{
+					{ID: columnID(1), Name: "id", Value: 1},
+				},
+				Metadata: wal.Metadata{},
+			},
+
+			wantQuery: &query{
+				schema: testSchema,
+				table:  testTable,
+				sql:    fmt.Sprintf("UPDATE %s SET \"id\" = $1, \"name\" = $2 WHERE \"id\" = $3", quotedTestTable),
+				args:   []any{1, "alice", 1},
+			},
+		},
+		{
+			name: "update - full identity",
+			walData: &wal.Data{
+				Action: "U",
+				Schema: testSchema,
+				Table:  testTable,
+				Columns: []wal.Column{
+					{ID: columnID(1), Name: "id", Value: 1},
+					{ID: columnID(2), Name: "name", Value: "alice"},
+				},
+				Identity: []wal.Column{
+					{ID: columnID(1), Name: "id", Value: 1},
+					{ID: columnID(2), Name: "name", Value: "a"},
+				},
+				Metadata: wal.Metadata{},
+			},
+
+			wantQuery: &query{
+				schema: testSchema,
+				table:  testTable,
+				sql:    fmt.Sprintf("UPDATE %s SET \"id\" = $1, \"name\" = $2 WHERE \"id\" = $3 AND \"name\" = $4", quotedTestTable),
+				args:   []any{1, "alice", 1, "a"},
+			},
+		},
+		{
+			name: "error - update",
+			walData: &wal.Data{
+				Action: "U",
+				Schema: testSchema,
+				Table:  testTable,
+				Columns: []wal.Column{
+					{ID: columnID(1), Name: "id", Value: 1},
+					{ID: columnID(2), Name: "name", Value: "alice"},
+				},
+				Identity: []wal.Column{},
+				Metadata: wal.Metadata{},
+			},
+
+			wantQuery: nil,
+			wantErr:   errUnableToBuildQuery,
+		},
 		{
 			name: "unknown",
 			walData: &wal.Data{
@@ -232,7 +329,8 @@ func TestDMLAdapter_walDataToQuery(t *testing.T) {
 			a := &dmlAdapter{
 				onConflictAction: tc.action,
 			}
-			query := a.walDataToQuery(tc.walData)
+			query, err := a.walDataToQuery(tc.walData)
+			require.ErrorIs(t, err, tc.wantErr)
 			require.Equal(t, tc.wantQuery, query)
 		})
 	}