Repository: xataio/pgstream
Pull request #293 (2025-04-23 11:28:23)

Use connection pool in transformer validator
This PR updates the transformer validator to use a connection pool instead of reusing the same connection for the different schema table queries. This should prevent `conn busy` errors during the validation.


Patch:
diff --git a/pkg/stream/stream_run.go b/pkg/stream/stream_run.go
index f2daa84..b86a645 100644
--- a/pkg/stream/stream_run.go
+++ b/pkg/stream/stream_run.go
@@ -233,7 +233,11 @@ func Run(ctx context.Context, logger loglib.Logger, config *Config, instrumentat
 			pgURL = config.Listener.Snapshot.Generator.URL
 		}
 		if pgURL != "" {
-			pgValidator := transformer.NewPostgresTransformerValidator(pgURL)
+			pgValidator, err := transformer.NewPostgresTransformerValidator(ctx, pgURL)
+			if err != nil {
+				return fmt.Errorf("error creating postgres transformer validator: %w", err)
+			}
+			defer pgValidator.Close()
 			opts = append(opts, transformer.WithValidator(pgValidator.Validate))
 		}
 		transformer, err := transformer.New(ctx, config.Processor.Transformer, processor, opts...)
diff --git a/pkg/wal/processor/transformer/wal_transformer_validator.go b/pkg/wal/processor/transformer/wal_transformer_validator.go
index 7f87894..cfeaeaa 100644
--- a/pkg/wal/processor/transformer/wal_transformer_validator.go
+++ b/pkg/wal/processor/transformer/wal_transformer_validator.go
@@ -13,29 +13,23 @@ import (
 )
 
 type PostgresTransformerValidator struct {
-	connBuilder connBuilder
+	conn pglib.Querier
 }
 
-type connBuilder func(context.Context) (pglib.Querier, error)
-
-func NewPostgresTransformerValidator(pgURL string) *PostgresTransformerValidator {
-	return &PostgresTransformerValidator{
-		connBuilder: func(ctx context.Context) (pglib.Querier, error) {
-			return pglib.NewConn(ctx, pgURL)
-		},
+func NewPostgresTransformerValidator(ctx context.Context, pgURL string) (*PostgresTransformerValidator, error) {
+	pool, err := pglib.NewConnPool(ctx, pgURL)
+	if err != nil {
+		return nil, err
 	}
+	return &PostgresTransformerValidator{
+		conn: pool,
+	}, nil
 }
 
 func (v *PostgresTransformerValidator) Validate(ctx context.Context, transformerMap map[string]ColumnTransformers) error {
-	conn, err := v.connBuilder(ctx)
-	if err != nil {
-		return fmt.Errorf("creating postgres connection: %w", err)
-	}
-	defer conn.Close(context.Background())
-
 	for schemaTable, columnTransformers := range transformerMap {
 		query := fmt.Sprintf("SELECT * FROM %s LIMIT 0", schemaTable)
-		rows, err := conn.Query(ctx, query)
+		rows, err := v.conn.Query(ctx, query)
 		if err != nil {
 			return fmt.Errorf("querying table rows: %w", err)
 		}
@@ -71,6 +65,10 @@ func (v *PostgresTransformerValidator) Validate(ctx context.Context, transformer
 	return nil
 }
 
+func (v *PostgresTransformerValidator) Close() error {
+	return v.conn.Close(context.Background())
+}
+
 func pgTypeCompatibleWithTransformerType(compatibleTypes []transformers.SupportedDataType, pgType uint32) bool {
 	switch pgType {
 	case pgtype.TextOID, pgtype.VarcharOID, pgtype.BPCharOID:


Test patch:
diff --git a/pkg/wal/processor/transformer/wal_transformer_validator_test.go b/pkg/wal/processor/transformer/wal_transformer_validator_test.go
index 2372fe4..57698f4 100644
--- a/pkg/wal/processor/transformer/wal_transformer_validator_test.go
+++ b/pkg/wal/processor/transformer/wal_transformer_validator_test.go
@@ -74,9 +74,7 @@ func TestPostgresTransformerValidator(t *testing.T) {
 	}
 
 	testPGValidator := PostgresTransformerValidator{
-		connBuilder: func(context.Context) (pglib.Querier, error) {
-			return testQuerier, nil
-		},
+		conn: testQuerier,
 	}
 
 	tests := []struct {