Repository: xataio/pgstream
Pull request #408 (2025-06-17 10:08:11)

Fix placeholderOffset for update queries
When building UPDATE queries, arg count in the WHERE clause must start from the arg count used in the SET clause. We currently start it from the total column count, resulting error when there's a generated column that is skipped in the SET clause.
Also adding a test for generating update query with a generated column.

github-actions[bot]: ### Merging this branch will **not change** overall coverage

| Impacted Packages | Coverage Δ | :robot: |
|-------------------|------------|---------|
| github.com/xataio/pgstream/pkg/wal/processor/postgres | 73.42% (ø) |  |

---

<details>

<summary>Coverage by file</summary>

### Changed files (no unit tests)

| Changed File | Coverage Δ | Total | Covered | Missed | :robot: |
|--------------|------------|-------|---------|--------|---------|
| github.com/xataio/pgstream/pkg/wal/processor/postgres/postgres_wal_dml_adapter.go | 96.81% (ø) | 94 | 91 | 3 |  |

_Please note that the "Total", "Covered", and "Missed" counts above refer to ***code statements*** instead of lines of code. The value in brackets refers to the test coverage of that file in the old version of the code._

### Changed unit test files

- github.com/xataio/pgstream/pkg/wal/processor/postgres/postgres_wal_dml_adapter_test.go

</details>


Patch:
diff --git a/pkg/wal/processor/postgres/postgres_wal_dml_adapter.go b/pkg/wal/processor/postgres/postgres_wal_dml_adapter.go
index 0d0c18c..345eb01 100644
--- a/pkg/wal/processor/postgres/postgres_wal_dml_adapter.go
+++ b/pkg/wal/processor/postgres/postgres_wal_dml_adapter.go
@@ -111,7 +111,7 @@ func (a *dmlAdapter) buildUpdateQuery(d *wal.Data, generatedColumns []string) (*
 	if setQuery == "" {
 		return &query{}, nil
 	}
-	whereQuery, whereValues, err := a.buildWhereQuery(d, len(d.Columns))
+	whereQuery, whereValues, err := a.buildWhereQuery(d, len(rowColumns))
 	if err != nil {
 		return nil, fmt.Errorf("building update query: %w", err)
 	}


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 710c966..9f26a3c 100644
--- a/pkg/wal/processor/postgres/postgres_wal_dml_adapter_test.go
+++ b/pkg/wal/processor/postgres/postgres_wal_dml_adapter_test.go
@@ -19,15 +19,15 @@ func TestDMLAdapter_walDataToQuery(t *testing.T) {
 	testSchema := "test"
 	quotedTestTable := quotedTableName(testSchema, testTable)
 	quotedColumnNames := []string{`"id"`, `"name"`}
-	generatedColumns := []string{}
 	columnID := func(i int) string {
 		return fmt.Sprintf("%s-%d", testTableID, i)
 	}
 
 	tests := []struct {
-		name    string
-		walData *wal.Data
-		action  onConflictAction
+		name             string
+		walData          *wal.Data
+		action           onConflictAction
+		generatedColumns []string
 
 		wantQuery *query
 		wantErr   error
@@ -286,6 +286,31 @@ func TestDMLAdapter_walDataToQuery(t *testing.T) {
 				args:   []any{1, "alice", 1, "a"},
 			},
 		},
+		{
+			name: "update - with generated column",
+			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"},
+					{ID: columnID(3), Name: "generated_col", Value: "gen_value"},
+				},
+				Identity: []wal.Column{
+					{ID: columnID(1), Name: "id", Value: 1},
+				},
+				Metadata: wal.Metadata{},
+			},
+			generatedColumns: []string{"generated_col"},
+
+			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: "error - update",
 			walData: &wal.Data{
@@ -329,7 +354,7 @@ func TestDMLAdapter_walDataToQuery(t *testing.T) {
 			a := &dmlAdapter{
 				onConflictAction: tc.action,
 			}
-			query, err := a.walDataToQuery(tc.walData, generatedColumns)
+			query, err := a.walDataToQuery(tc.walData, tc.generatedColumns)
 			require.ErrorIs(t, err, tc.wantErr)
 			require.Equal(t, tc.wantQuery, query)
 		})