Repository: xataio/pgstream
Pull request #174 (2025-03-14 10:40:35)

[PostgreSQL processor] Handle errors gracefully
This PR updates the query flushing behaviour of the postgres batch writer to ensure that if one of the queries within a batch fails with a non internal error (constraint violation, relation doesn't exist), the error is logged as dataloss, and the rest of the queries are retried and checkpointed, instead of stopping the processor at the first query error.

This aligns with the behaviour of all the other processors.



Patch:
diff --git a/Makefile b/Makefile
index 38644e0..de8e6b8 100644
--- a/Makefile
+++ b/Makefile
@@ -10,7 +10,7 @@ test:
 
 .PHONY: integration-test
 integration-test:
-	@PGSTREAM_INTEGRATION_TESTS=true go test -timeout 90s github.com/xataio/pgstream/pkg/stream/integration
+	@PGSTREAM_INTEGRATION_TESTS=true go test -timeout 180s github.com/xataio/pgstream/pkg/stream/integration
 
 .PHONY: license-check
 license-check:
diff --git a/go.sum b/go.sum
index 1862f72..86a975f 100644
--- a/go.sum
+++ b/go.sum
@@ -102,8 +102,6 @@ github.com/elastic/elastic-transport-go/v8 v8.6.0 h1:Y2S/FBjx1LlCv5m6pWAF2kDJAHo
 github.com/elastic/elastic-transport-go/v8 v8.6.0/go.mod h1:YLHer5cj0csTzNFXoNQ8qhtGY1GTvSqPnKWKaqQE3Hk=
 github.com/elastic/go-elasticsearch/v8 v8.14.0 h1:1ywU8WFReLLcxE1WJqii3hTtbPUE2hc38ZK/j4mMFow=
 github.com/elastic/go-elasticsearch/v8 v8.14.0/go.mod h1:WRvnlGkSuZyp83M2U8El/LGXpCjYLrvlkSgkAH4O5I4=
-github.com/eminano/greenmask v0.0.0-20250210141616-781d3aca2262 h1:IAhjExmPmCH/UWLv56J/50wrjlwJRTbwI8gBhb5Rl/A=
-github.com/eminano/greenmask v0.0.0-20250210141616-781d3aca2262/go.mod h1:gBBDtACvj/I3nzyEvCGqrWd2ox1XVlh5+Xv5JXKcKDI=
 github.com/eminano/greenmask v0.0.0-20250307113752-035ee2b102e6 h1:m/MuqhOhjJM2oWIjUMkyvrCuCmb/IMC8Ig3/uSM0M+U=
 github.com/eminano/greenmask v0.0.0-20250307113752-035ee2b102e6/go.mod h1:gBBDtACvj/I3nzyEvCGqrWd2ox1XVlh5+Xv5JXKcKDI=
 github.com/expr-lang/expr v1.16.9 h1:WUAzmR0JNI9JCiF0/ewwHB1gmcGw5wW7nWt8gc6PpCI=
diff --git a/internal/postgres/errors.go b/internal/postgres/errors.go
index bfd6004..67ba417 100644
--- a/internal/postgres/errors.go
+++ b/internal/postgres/errors.go
@@ -5,6 +5,7 @@ package postgres
 import (
 	"errors"
 	"fmt"
+	"strings"
 
 	"github.com/jackc/pgx/v5"
 	"github.com/jackc/pgx/v5/pgconn"
@@ -23,6 +24,14 @@ func (e *ErrRelationDoesNotExist) Error() string {
 	return fmt.Sprintf("relation does not exist: %s", e.Details)
 }
 
+type ErrConstraintViolation struct {
+	Details string
+}
+
+func (e *ErrConstraintViolation) Error() string {
+	return fmt.Sprintf("constraint violation: %s", e.Details)
+}
+
 func mapError(err error) error {
 	if pgconn.Timeout(err) {
 		return ErrConnTimeout
@@ -39,6 +48,12 @@ func mapError(err error) error {
 				Details: pgErr.Message,
 			}
 		}
+		// Class 23 — Integrity Constraint Violation
+		if strings.HasPrefix(pgErr.Code, "23") {
+			return &ErrConstraintViolation{
+				Details: pgErr.Message,
+			}
+		}
 	}
 
 	return err
diff --git a/pkg/wal/processor/postgres/postgres_batch_writer.go b/pkg/wal/processor/postgres/postgres_batch_writer.go
index 0f23f27..bc4b062 100644
--- a/pkg/wal/processor/postgres/postgres_batch_writer.go
+++ b/pkg/wal/processor/postgres/postgres_batch_writer.go
@@ -195,18 +195,64 @@ func (w *BatchWriter) flushQueries(ctx context.Context, queries []*query) error
 	if len(queries) == 0 {
 		return nil
 	}
+
+	var err error
+	for {
+		queries, err = w.execQueries(ctx, queries)
+		if err != nil {
+			return err
+		}
+
+		if len(queries) == 0 {
+			return nil
+		}
+	}
+}
+
+func (w *BatchWriter) execQueries(ctx context.Context, queries []*query) ([]*query, error) {
+	retryQueries := []*query{}
 	err := w.pgConn.ExecInTx(ctx, func(tx pglib.Tx) error {
-		for _, q := range queries {
+		for i, q := range queries {
 			if _, err := tx.Exec(ctx, q.sql, q.args...); err != nil {
+				w.logger.Error(err, "executing sql query", loglib.Fields{
+					"sql":      q.sql,
+					"args":     q.args,
+					"severity": "DATALOSS",
+				})
+				// if a query returns an error, it will abort the tx. Log it as
+				// dataloss and remove it from the list of queries to be
+				// retried.
+				retryQueries = removeIndex(queries, i)
 				return err
 			}
 		}
 		return nil
 	})
-	if err != nil {
-		return err
+	if err != nil && w.isInternalError(err) {
+		// if there was an internal error in the tx, there's no point in
+		// retrying, return error and stop processing.
+		return nil, err
 	}
 
-	queries = queries[:0]
-	return nil
+	// if there were no errors or no internal errors in the tx, return the
+	// queries to retry (none if the tx was successful)
+	return retryQueries, nil
+}
+
+func (w *BatchWriter) isInternalError(err error) bool {
+	var errRelationDoesNotExist *pglib.ErrRelationDoesNotExist
+	var errConstraintViolation *pglib.ErrConstraintViolation
+	switch {
+	case errors.As(err, &errRelationDoesNotExist),
+		errors.As(err, &errConstraintViolation):
+		return false
+	default:
+		return true
+	}
+}
+
+func removeIndex(s []*query, index int) []*query {
+	ret := make([]*query, 0)
+	ret = append(ret, s[:index]...)
+	return append(ret, s[index+1:]...)
 }


Test patch:
diff --git a/pkg/wal/processor/postgres/postgres_batch_writer_test.go b/pkg/wal/processor/postgres/postgres_batch_writer_test.go
index 8008b57..92ffe8b 100644
--- a/pkg/wal/processor/postgres/postgres_batch_writer_test.go
+++ b/pkg/wal/processor/postgres/postgres_batch_writer_test.go
@@ -5,6 +5,7 @@ package postgres
 import (
 	"context"
 	"errors"
+	"fmt"
 	"testing"
 	"time"
 
@@ -239,6 +240,112 @@ func TestBatchWriter_sendBatch(t *testing.T) {
 	}
 }
 
+func TestBatchWriter_flushQueries(t *testing.T) {
+	t.Parallel()
+
+	testQuerySQL := "INSERT INTO test(id, name) VALUES($1, $2)"
+	args1 := []any{1, "alice"}
+	args2 := []any{1, "bob"}
+
+	testQuery := func(args []any) *query {
+		return &query{
+			sql:  testQuerySQL,
+			args: args,
+		}
+	}
+	execCalls := uint(0)
+
+	tests := []struct {
+		name    string
+		pgconn  *pgmocks.Querier
+		queries []*query
+
+		wantExecCalls uint
+		wantErr       error
+	}{
+		{
+			name: "ok - no queries",
+			pgconn: &pgmocks.Querier{
+				ExecInTxFn: func(ctx context.Context, f func(tx pglib.Tx) error) error {
+					return errors.New("ExecInTxFn: should not be called")
+				},
+			},
+			queries: []*query{},
+
+			wantExecCalls: 0,
+			wantErr:       nil,
+		},
+		{
+			name: "ok",
+			pgconn: &pgmocks.Querier{
+				ExecInTxFn: func(ctx context.Context, f func(tx pglib.Tx) error) error {
+					mockTx := pgmocks.Tx{
+						ExecFn: func(ctx context.Context, query string, args ...any) (pglib.CommandTag, error) {
+							execCalls++
+							require.Equal(t, testQuerySQL, query)
+							require.Len(t, args, 2)
+							switch args[1] {
+							case args1[1]:
+								return pglib.CommandTag{}, nil
+							case args2[1]:
+								return pglib.CommandTag{}, &pglib.ErrConstraintViolation{}
+							default:
+								return pglib.CommandTag{}, fmt.Errorf("unexpected argument in sql query: %v", args[1])
+							}
+						},
+					}
+					return f(&mockTx)
+				},
+			},
+			queries: []*query{testQuery(args1), testQuery(args2)},
+
+			wantExecCalls: 3,
+			wantErr:       nil,
+		},
+		{
+			name: "error - internal error in tx exec",
+			pgconn: &pgmocks.Querier{
+				ExecInTxFn: func(ctx context.Context, f func(tx pglib.Tx) error) error {
+					mockTx := pgmocks.Tx{
+						ExecFn: func(ctx context.Context, query string, args ...any) (pglib.CommandTag, error) {
+							execCalls++
+							require.Equal(t, testQuerySQL, query)
+							require.Len(t, args, 2)
+							switch args[1] {
+							case args1[1]:
+								return pglib.CommandTag{}, nil
+							case args2[1]:
+								return pglib.CommandTag{}, errTest
+							default:
+								return pglib.CommandTag{}, fmt.Errorf("unexpected argument in sql query: %v", args[1])
+							}
+						},
+					}
+					return f(&mockTx)
+				},
+			},
+			queries: []*query{testQuery(args1), testQuery(args2)},
+
+			wantExecCalls: 2,
+			wantErr:       errTest,
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			bw := &BatchWriter{
+				logger: loglib.NewNoopLogger(),
+				pgConn: tc.pgconn,
+			}
+
+			err := bw.flushQueries(context.Background(), tc.queries)
+			require.ErrorIs(t, err, tc.wantErr)
+			require.Equal(t, tc.wantExecCalls, execCalls)
+			execCalls = 0
+		})
+	}
+}
+
 func TestBatchWriter(t *testing.T) {
 	t.Parallel()