Merge pull request #174 from xataio/pg-processor-handle-errors-gracefully

[PostgreSQL processor] Handle errors gracefully

Author: eminano

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:

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=

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

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:]...)
 }

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()