Repository: xataio/pgstream (go)

Issue:
[Replication] Improve error log when replication slot doesn't exist
Currently we get an error `no rows` when starting pgstream without a replication slot created. Instead this should have a proper description so that it's easy to identify.

Pull request:
Improve errors
This PR updates the errors for unmarshaling yaml configuration and no replication slot so that they're easier to understand. 

Closes #227 

Patch:
diff --git a/cmd/config/config.go b/cmd/config/config.go
index 4ee01e9..d131d2e 100644
--- a/cmd/config/config.go
+++ b/cmd/config/config.go
@@ -52,7 +52,7 @@ func ParseStreamConfig() (*stream.Config, error) {
 		}
 		yamlCfg := YAMLConfig{}
 		if err := viper.Unmarshal(&yamlCfg); err != nil {
-			return nil, fmt.Errorf("unmarshaling config file %q: %w", cfgFile, err)
+			return nil, fmt.Errorf("invalid format in config file %q: %w", cfgFile, err)
 		}
 		// yaml.Unmarshal is used to override the viper.Umarshal to be able to
 		// parse the transformers configuration with support for case sensitive
@@ -60,7 +60,7 @@ func ParseStreamConfig() (*stream.Config, error) {
 		// https://github.com/spf13/viper/issues/260
 		err = yaml.Unmarshal(buf, &yamlCfg.Modifiers)
 		if err != nil {
-			return nil, fmt.Errorf("unmarshaling modifiers config from %q: %w", cfgFile, err)
+			return nil, fmt.Errorf("invalid format for modifiers config in file %q: %w", cfgFile, err)
 		}
 		return yamlCfg.toStreamConfig()
 	default:
diff --git a/cmd/config/config_env.go b/cmd/config/config_env.go
index afd675b..601def3 100644
--- a/cmd/config/config_env.go
+++ b/cmd/config/config_env.go
@@ -421,7 +421,7 @@ func parseTransformerConfig() (*transformer.Config, error) {
 	}{}
 	err = yaml.Unmarshal(buf, &yamlConfig)
 	if err != nil {
-		return nil, fmt.Errorf("in file %q: %w", filename, err)
+		return nil, fmt.Errorf("invalid format for transformations config in file %q: %w", filename, err)
 	}
 
 	return yamlConfig.Transformations.parseTransformationConfig()
diff --git a/docs/README.md b/docs/README.md
index 470594b..da3dbd5 100644
--- a/docs/README.md
+++ b/docs/README.md
@@ -958,14 +958,15 @@ The rules for the transformers are defined in a dedicated yaml file, with the fo
 
 ```yaml
 transformations:
-  validation_mode: <validation_mode>
-  table_transformers:
-    - schema: <schema_name>
-      table: <table_name>
-      column_transformers:
-        <column_name>:
-          name: <transformer_name>
-          parameters:
+  validation_mode: <validation_mode> # Validation mode for the transformation rules. Can be one of strict, relaxed or table_level. Defaults to relaxed if not provided
+  table_transformers: # List of table transformations
+    - schema: <schema_name> # Name of the table schema
+      table: <table_name> # Name of the table
+      validation_mode: <validation_mode> # To be used when the global validation_mode is set to `table_level`. Can be one of strict or relaxed
+      column_transformers: # List of column transformations
+        <column_name>: # Name of the column to which the transformation will be applied
+          name: <transformer_name> # Name of the transformer to be applied to the column. If no transformer needs to be applied on strict validation mode, it can be left empty or use `noop`
+          parameters: # Transformer parameters as defined in the supported transformers documentation
             <transformer_parameter>: <transformer_parameter_value>
 ```
 
diff --git a/pkg/wal/replication/postgres/pg_replication_handler.go b/pkg/wal/replication/postgres/pg_replication_handler.go
index 1ef73ab..03e5cd9 100644
--- a/pkg/wal/replication/postgres/pg_replication_handler.go
+++ b/pkg/wal/replication/postgres/pg_replication_handler.go
@@ -96,6 +96,12 @@ func NewHandler(ctx context.Context, cfg Config, opts ...Option) (*Handler, erro
 		opt(h)
 	}
 
+	// make sure the replication slot we are going to use exists
+	if err := h.verifyReplicationSlotExists(ctx); err != nil {
+		h.logger.Error(err, "verifying replication slot")
+		return nil, err
+	}
+
 	return h, nil
 }
 
@@ -261,6 +267,24 @@ func (h *Handler) getLastSyncedLSN(ctx context.Context, conn pglib.Querier) (rep
 	return h.lsnParser.FromString(confirmedFlushLSN)
 }
 
+func (h *Handler) verifyReplicationSlotExists(ctx context.Context) error {
+	slotExists := false
+	conn, err := h.pgConnBuilder()
+	if err != nil {
+		return fmt.Errorf("creating pg connection: %w", err)
+	}
+	defer conn.Close(context.Background())
+
+	err = conn.QueryRow(ctx, `SELECT EXISTS(SELECT 1 FROM pg_replication_slots WHERE slot_name=$1)`, h.pgReplicationSlotName).Scan(&slotExists)
+	if err != nil {
+		return fmt.Errorf("retrieving replication slot: %w", err)
+	}
+	if !slotExists {
+		return fmt.Errorf("replication slot %s does not exist", h.pgReplicationSlotName)
+	}
+	return nil
+}
+
 func mapPostgresError(err error) error {
 	if errors.Is(err, pglib.ErrConnTimeout) {
 		return replication.ErrConnTimeout


Test patch:
diff --git a/pkg/wal/replication/postgres/helper_test.go b/pkg/wal/replication/postgres/helper_test.go
index 8667a37..16661da 100644
--- a/pkg/wal/replication/postgres/helper_test.go
+++ b/pkg/wal/replication/postgres/helper_test.go
@@ -11,6 +11,7 @@ import (
 type mockRow struct {
 	lsn    string
 	lag    int64
+	exists bool
 	scanFn func(args ...any) error
 }
 
@@ -19,11 +20,17 @@ func (m *mockRow) Scan(args ...any) error {
 		return m.scanFn(args...)
 	}
 
+	if len(args) != 1 {
+		return fmt.Errorf("expected 1 argument, got %d", len(args))
+	}
+
 	switch arg := args[0].(type) {
 	case *string:
 		*arg = m.lsn
 	case *int64:
 		*arg = m.lag
+	case *bool:
+		*arg = m.exists
 	default:
 		return fmt.Errorf("unexpected argument type in scan: %T", args[0])
 	}
diff --git a/pkg/wal/replication/postgres/pg_replication_handler_test.go b/pkg/wal/replication/postgres/pg_replication_handler_test.go
index 2b6ab9b..b4f3aa7 100644
--- a/pkg/wal/replication/postgres/pg_replication_handler_test.go
+++ b/pkg/wal/replication/postgres/pg_replication_handler_test.go
@@ -497,3 +497,93 @@ func TestHandler_GetCurrentLSN(t *testing.T) {
 		})
 	}
 }
+
+func TestHandler_verifyReplicationSlotExists(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name        string
+		connBuilder func() (pglib.Querier, error)
+		slotExists  bool
+		wantErr     error
+	}{
+		{
+			name: "ok - slot exists",
+			connBuilder: func() (pglib.Querier, error) {
+				return &pgmocks.Querier{
+					QueryRowFn: func(ctx context.Context, query string, args ...any) pglib.Row {
+						require.Len(t, args, 1)
+						require.Equal(t, args[0], testSlot)
+						switch query {
+						case "SELECT EXISTS(SELECT 1 FROM pg_replication_slots WHERE slot_name=$1)":
+							return &mockRow{exists: true}
+						default:
+							return &mockRow{scanFn: func(args ...any) error { return fmt.Errorf("unexpected query: %s", query) }}
+						}
+					},
+					CloseFn: func(ctx context.Context) error { return nil },
+				}, nil
+			},
+			slotExists: true,
+			wantErr:    nil,
+		},
+		{
+			name: "error - slot does not exist",
+			connBuilder: func() (pglib.Querier, error) {
+				return &pgmocks.Querier{
+					QueryRowFn: func(ctx context.Context, query string, args ...any) pglib.Row {
+						require.Len(t, args, 1)
+						require.Equal(t, args[0], testSlot)
+						switch query {
+						case "SELECT EXISTS(SELECT 1 FROM pg_replication_slots WHERE slot_name=$1)":
+							return &mockRow{exists: false}
+						default:
+							return &mockRow{scanFn: func(args ...any) error { return fmt.Errorf("unexpected query: %s", query) }}
+						}
+					},
+					CloseFn: func(ctx context.Context) error { return nil },
+				}, nil
+			},
+			slotExists: false,
+			wantErr:    fmt.Errorf("replication slot %s does not exist", testSlot),
+		},
+		{
+			name: "error - building connection",
+			connBuilder: func() (pglib.Querier, error) {
+				return nil, errTest
+			},
+			slotExists: false,
+			wantErr:    errTest,
+		},
+		{
+			name: "error - query execution",
+			connBuilder: func() (pglib.Querier, error) {
+				return &pgmocks.Querier{
+					QueryRowFn: func(ctx context.Context, query string, args ...any) pglib.Row {
+						return &mockRow{scanFn: func(args ...any) error { return errTest }}
+					},
+					CloseFn: func(ctx context.Context) error { return nil },
+				}, nil
+			},
+			slotExists: false,
+			wantErr:    fmt.Errorf("retrieving replication slot: %w", errTest),
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+
+			h := Handler{
+				logger:                log.NewNoopLogger(),
+				pgConnBuilder:         tc.connBuilder,
+				pgReplicationSlotName: testSlot,
+			}
+
+			err := h.verifyReplicationSlotExists(context.Background())
+			if !errors.Is(err, tc.wantErr) {
+				require.Equal(t, tc.wantErr, err)
+			}
+		})
+	}
+}