Repository: xataio/pgstream
Pull request #368 (2025-05-29 15:10:11)

Fix disable triggers
This PR updates the logic to disable triggers (`set session replication role`) to set it per session, since that's the scope of the setting. This guarantees that the setting is enforced for each connection from the pool for both bulk ingest and batch postgres writers.


Patch:
diff --git a/internal/postgres/instrumentation/instrumented_tx.go b/internal/postgres/instrumentation/instrumented_tx.go
index 96f12d5..aaf189a 100644
--- a/internal/postgres/instrumentation/instrumented_tx.go
+++ b/internal/postgres/instrumentation/instrumented_tx.go
@@ -46,3 +46,9 @@ func (i *Tx) Exec(ctx context.Context, query string, args ...any) (tag pglib.Com
 	defer otel.CloseSpan(span, err)
 	return i.inner.Exec(ctx, query, args...)
 }
+
+func (i *Tx) CopyFrom(ctx context.Context, tableName string, columnNames []string, srcRows [][]any) (rowCount int64, err error) {
+	ctx, span := otel.StartSpan(ctx, i.tracer, "tx.CopyFrom")
+	defer otel.CloseSpan(span, err)
+	return i.inner.CopyFrom(ctx, tableName, columnNames, srcRows)
+}
diff --git a/internal/postgres/mocks/mock_tx.go b/internal/postgres/mocks/mock_tx.go
index ec91427..e3a6de8 100644
--- a/internal/postgres/mocks/mock_tx.go
+++ b/internal/postgres/mocks/mock_tx.go
@@ -9,9 +9,11 @@ import (
 )
 
 type Tx struct {
-	QueryRowFn func(ctx context.Context, query string, args ...any) postgres.Row
-	QueryFn    func(ctx context.Context, query string, args ...any) (postgres.Rows, error)
-	ExecFn     func(ctx context.Context, query string, args ...any) (postgres.CommandTag, error)
+	QueryRowFn    func(ctx context.Context, query string, args ...any) postgres.Row
+	QueryFn       func(ctx context.Context, query string, args ...any) (postgres.Rows, error)
+	ExecFn        func(ctx context.Context, i uint, query string, args ...any) (postgres.CommandTag, error)
+	CopyFromFn    func(ctx context.Context, tableName string, columnNames []string, srcRows [][]any) (int64, error)
+	execCallCount uint
 }
 
 func (m *Tx) QueryRow(ctx context.Context, query string, args ...any) postgres.Row {
@@ -23,5 +25,10 @@ func (m *Tx) Query(ctx context.Context, query string, args ...any) (postgres.Row
 }
 
 func (m *Tx) Exec(ctx context.Context, query string, args ...any) (postgres.CommandTag, error) {
-	return m.ExecFn(ctx, query, args...)
+	m.execCallCount++
+	return m.ExecFn(ctx, m.execCallCount, query, args...)
+}
+
+func (m *Tx) CopyFrom(ctx context.Context, tableName string, columnNames []string, srcRows [][]any) (rowCount int64, err error) {
+	return m.CopyFromFn(ctx, tableName, columnNames, srcRows)
 }
diff --git a/internal/postgres/pg_querier.go b/internal/postgres/pg_querier.go
index 1c73b33..ef0d8e9 100644
--- a/internal/postgres/pg_querier.go
+++ b/internal/postgres/pg_querier.go
@@ -28,12 +28,6 @@ type Rows interface {
 	pgx.Rows
 }
 
-type Tx interface {
-	Query(ctx context.Context, query string, args ...any) (Rows, error)
-	QueryRow(ctx context.Context, query string, args ...any) Row
-	Exec(ctx context.Context, query string, args ...any) (CommandTag, error)
-}
-
 type CommandTag struct {
 	pgconn.CommandTag
 }
diff --git a/internal/postgres/pg_tx.go b/internal/postgres/pg_tx.go
index 766af89..93e3ba1 100644
--- a/internal/postgres/pg_tx.go
+++ b/internal/postgres/pg_tx.go
@@ -8,6 +8,13 @@ import (
 	"github.com/jackc/pgx/v5"
 )
 
+type Tx interface {
+	Query(ctx context.Context, query string, args ...any) (Rows, error)
+	QueryRow(ctx context.Context, query string, args ...any) Row
+	Exec(ctx context.Context, query string, args ...any) (CommandTag, error)
+	CopyFrom(ctx context.Context, tableName string, columnNames []string, srcRows [][]any) (int64, error)
+}
+
 type TxIsolationLevel string
 
 const (
@@ -48,6 +55,21 @@ func (t *Txn) Exec(ctx context.Context, query string, args ...any) (CommandTag,
 	return CommandTag{tag}, mapError(err)
 }
 
+func (t *Txn) CopyFrom(ctx context.Context, tableName string, columnNames []string, srcRows [][]any) (int64, error) {
+	identifier, err := newIdentifier(tableName)
+	if err != nil {
+		return -1, err
+	}
+
+	// sanitize the input, removing any added quotes. The CopyFrom will sanitize
+	// them and double quotes will cause errors.
+	for i, c := range columnNames {
+		columnNames[i] = removeQuotes(c)
+	}
+
+	return t.Tx.CopyFrom(ctx, identifier, columnNames, pgx.CopyFromRows(srcRows))
+}
+
 func toTxOptions(opts TxOptions) pgx.TxOptions {
 	return pgx.TxOptions{
 		IsoLevel:   pgx.TxIsoLevel(opts.IsolationLevel),
diff --git a/pkg/wal/processor/postgres/postgres_batch_writer.go b/pkg/wal/processor/postgres/postgres_batch_writer.go
index 3226e7e..b87a1f9 100644
--- a/pkg/wal/processor/postgres/postgres_batch_writer.go
+++ b/pkg/wal/processor/postgres/postgres_batch_writer.go
@@ -167,6 +167,10 @@ func (w *BatchWriter) flushQueries(ctx context.Context, queries []*query) error
 func (w *BatchWriter) execQueries(ctx context.Context, queries []*query) ([]*query, error) {
 	retryQueries := []*query{}
 	err := w.pgConn.ExecInTx(ctx, func(tx pglib.Tx) error {
+		if err := w.setReplicationRoleToReplica(ctx, tx); err != nil {
+			return err
+		}
+
 		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{
@@ -181,7 +185,8 @@ func (w *BatchWriter) execQueries(ctx context.Context, queries []*query) ([]*que
 				return err
 			}
 		}
-		return nil
+
+		return w.resetReplicationRole(ctx, tx)
 	})
 	if err != nil && w.isInternalError(err) {
 		// if there was an internal error in the tx, there's no point in
diff --git a/pkg/wal/processor/postgres/postgres_bulk_ingest_writer.go b/pkg/wal/processor/postgres/postgres_bulk_ingest_writer.go
index b47cec2..c4c6b28 100644
--- a/pkg/wal/processor/postgres/postgres_bulk_ingest_writer.go
+++ b/pkg/wal/processor/postgres/postgres_bulk_ingest_writer.go
@@ -4,6 +4,7 @@ package postgres
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"runtime/debug"
 	"sync"
@@ -28,6 +29,8 @@ type BulkIngestWriter struct {
 
 const bulkIngestWriter = "postgres_bulk_ingest_writer"
 
+var errUnexpectedCopiedRows = errors.New("number of rows copied doesn't match the source rows")
+
 // NewBulkIngestWriter returns a postgres processor that batches and writes data
 // to the configured postgres instance in bulk using the COPY command. It uses a
 // batch sender per schema table to parallelise the bulk ingest for different
@@ -175,15 +178,25 @@ func (w *BulkIngestWriter) copyFromInsertQueries(ctx context.Context, inserts []
 		srcRows = append(srcRows, q.args)
 	}
 
-	rowsToCopy := len(inserts)
-	rowsCopied, err := w.pgConn.CopyFrom(ctx, pglib.QuoteQualifiedIdentifier(query.schema, query.table), query.columnNames, srcRows)
-	if err != nil {
-		return err
-	}
+	err := w.pgConn.ExecInTx(ctx, func(tx pglib.Tx) error {
+		if err := w.setReplicationRoleToReplica(ctx, tx); err != nil {
+			return err
+		}
 
-	if rowsCopied != int64(rowsToCopy) {
-		return fmt.Errorf("number of rows copied (%d) doesn't match the source rows (%d)", rowsCopied, rowsToCopy)
-	}
+		rowsCopied, err := tx.CopyFrom(ctx, pglib.QuoteQualifiedIdentifier(query.schema, query.table), query.columnNames, srcRows)
+		if err != nil {
+			return err
+		}
 
+		rowsToCopy := len(inserts)
+		if rowsCopied != int64(rowsToCopy) {
+			return fmt.Errorf("%w: copied (%d), expected(%d)", errUnexpectedCopiedRows, rowsCopied, rowsToCopy)
+		}
+
+		return w.resetReplicationRole(ctx, tx)
+	})
+	if err != nil {
+		return fmt.Errorf("copy from: %w", err)
+	}
 	return nil
 }
diff --git a/pkg/wal/processor/postgres/postgres_writer.go b/pkg/wal/processor/postgres/postgres_writer.go
index c9efbab..d80822c 100644
--- a/pkg/wal/processor/postgres/postgres_writer.go
+++ b/pkg/wal/processor/postgres/postgres_writer.go
@@ -47,25 +47,11 @@ func newWriter(ctx context.Context, config *Config, adapter walAdapter, writerTy
 		opt(w)
 	}
 
-	if w.disableTriggers {
-		w.logger.Info("disabling triggers on postgres instance")
-		if err := w.setReplicationRoleToReplica(ctx); err != nil {
-			return nil, err
-		}
-	}
-
 	return w, nil
 }
 
 func (w *Writer) close() error {
-	ctx := context.Background()
-	if w.disableTriggers {
-		if err := w.resetReplicationRole(ctx); err != nil {
-			w.logger.Error(err, "reseting triggers")
-		}
-	}
-
-	return w.pgConn.Close(ctx)
+	return w.pgConn.Close(context.Background())
 }
 
 func WithLogger(l loglib.Logger) WriterOption {
@@ -88,16 +74,24 @@ func WithInstrumentation(i *otel.Instrumentation) WriterOption {
 	}
 }
 
-func (w *Writer) setReplicationRoleToReplica(ctx context.Context) error {
-	_, err := w.pgConn.Exec(ctx, "SET session_replication_role = replica")
+func (w *Writer) setReplicationRoleToReplica(ctx context.Context, tx pglib.Tx) error {
+	if !w.disableTriggers {
+		return nil
+	}
+
+	_, err := tx.Exec(ctx, "SET session_replication_role = replica")
 	if err != nil {
 		return fmt.Errorf("disabling triggers on postgres instance: %w", err)
 	}
 	return nil
 }
 
-func (w *Writer) resetReplicationRole(ctx context.Context) error {
-	if _, err := w.pgConn.Exec(ctx, "SET session_replication_role = DEFAULT"); err != nil {
+func (w *Writer) resetReplicationRole(ctx context.Context, tx pglib.Tx) error {
+	if !w.disableTriggers {
+		return nil
+	}
+
+	if _, err := tx.Exec(ctx, "SET session_replication_role = DEFAULT"); err != nil {
 		return fmt.Errorf("resetting session replication role to default: %w", err)
 	}
 	return nil


Test patch:
diff --git a/pkg/snapshot/generator/postgres/data/pg_snapshot_generator_test.go b/pkg/snapshot/generator/postgres/data/pg_snapshot_generator_test.go
index e079ad1..093b72f 100644
--- a/pkg/snapshot/generator/postgres/data/pg_snapshot_generator_test.go
+++ b/pkg/snapshot/generator/postgres/data/pg_snapshot_generator_test.go
@@ -88,7 +88,7 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 						return f(&mockTx)
 					case 2:
 						mockTx := pgmocks.Tx{
-							ExecFn: func(ctx context.Context, query string, args ...any) (pglib.CommandTag, error) {
+							ExecFn: func(ctx context.Context, _ uint, query string, args ...any) (pglib.CommandTag, error) {
 								require.Equal(t, fmt.Sprintf("SET TRANSACTION SNAPSHOT '%s'", testSnapshotID), query)
 								require.Len(t, args, 0)
 								return pglib.CommandTag{}, nil
@@ -110,7 +110,7 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 						return f(&mockTx)
 					case 3:
 						mockTx := pgmocks.Tx{
-							ExecFn: func(ctx context.Context, query string, args ...any) (pglib.CommandTag, error) {
+							ExecFn: func(ctx context.Context, _ uint, query string, args ...any) (pglib.CommandTag, error) {
 								require.Equal(t, fmt.Sprintf("SET TRANSACTION SNAPSHOT '%s'", testSnapshotID), query)
 								require.Len(t, args, 0)
 								return pglib.CommandTag{}, nil
@@ -193,7 +193,7 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 								ErrFn: func() error { return nil },
 							}, nil
 						},
-						ExecFn: func(ctx context.Context, query string, args ...any) (pglib.CommandTag, error) {
+						ExecFn: func(ctx context.Context, _ uint, query string, args ...any) (pglib.CommandTag, error) {
 							require.Equal(t, fmt.Sprintf("SET TRANSACTION SNAPSHOT '%s'", testSnapshotID), query)
 							require.Len(t, args, 0)
 							return pglib.CommandTag{}, nil
@@ -240,7 +240,7 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 						return f(&mockTx)
 					case 2:
 						mockTx := pgmocks.Tx{
-							ExecFn: func(ctx context.Context, query string, args ...any) (pglib.CommandTag, error) {
+							ExecFn: func(ctx context.Context, _ uint, query string, args ...any) (pglib.CommandTag, error) {
 								require.Equal(t, fmt.Sprintf("SET TRANSACTION SNAPSHOT '%s'", testSnapshotID), query)
 								require.Len(t, args, 0)
 								return pglib.CommandTag{}, nil
@@ -262,7 +262,7 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 						return f(&mockTx)
 					case 3:
 						mockTx := pgmocks.Tx{
-							ExecFn: func(ctx context.Context, query string, args ...any) (pglib.CommandTag, error) {
+							ExecFn: func(ctx context.Context, _ uint, query string, args ...any) (pglib.CommandTag, error) {
 								require.Equal(t, fmt.Sprintf("SET TRANSACTION SNAPSHOT '%s'", testSnapshotID), query)
 								require.Len(t, args, 0)
 								return pglib.CommandTag{}, nil
@@ -321,7 +321,7 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 						return f(&mockTx)
 					case 2:
 						mockTx := pgmocks.Tx{
-							ExecFn: func(ctx context.Context, query string, args ...any) (pglib.CommandTag, error) {
+							ExecFn: func(ctx context.Context, _ uint, query string, args ...any) (pglib.CommandTag, error) {
 								require.Equal(t, fmt.Sprintf("SET TRANSACTION SNAPSHOT '%s'", testSnapshotID), query)
 								require.Len(t, args, 0)
 								return pglib.CommandTag{}, nil
@@ -343,7 +343,7 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 						return f(&mockTx)
 					case 3:
 						mockTx := pgmocks.Tx{
-							ExecFn: func(ctx context.Context, query string, args ...any) (pglib.CommandTag, error) {
+							ExecFn: func(ctx context.Context, _ uint, query string, args ...any) (pglib.CommandTag, error) {
 								require.Equal(t, fmt.Sprintf("SET TRANSACTION SNAPSHOT '%s'", testSnapshotID), query)
 								require.Len(t, args, 0)
 								return pglib.CommandTag{}, nil
@@ -421,7 +421,7 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 						return f(&mockTx)
 					case 2:
 						mockTx := pgmocks.Tx{
-							ExecFn: func(ctx context.Context, query string, args ...any) (pglib.CommandTag, error) {
+							ExecFn: func(ctx context.Context, _ uint, query string, args ...any) (pglib.CommandTag, error) {
 								require.Equal(t, fmt.Sprintf("SET TRANSACTION SNAPSHOT '%s'", testSnapshotID), query)
 								require.Len(t, args, 0)
 								return pglib.CommandTag{}, errTest
@@ -468,7 +468,7 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 						return f(&mockTx)
 					case 2:
 						mockTx := pgmocks.Tx{
-							ExecFn: func(ctx context.Context, query string, args ...any) (pglib.CommandTag, error) {
+							ExecFn: func(ctx context.Context, _ uint, query string, args ...any) (pglib.CommandTag, error) {
 								require.Equal(t, fmt.Sprintf("SET TRANSACTION SNAPSHOT '%s'", testSnapshotID), query)
 								require.Len(t, args, 0)
 								return pglib.CommandTag{}, nil
@@ -524,7 +524,7 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 						return f(&mockTx)
 					case 2:
 						mockTx := pgmocks.Tx{
-							ExecFn: func(ctx context.Context, query string, args ...any) (pglib.CommandTag, error) {
+							ExecFn: func(ctx context.Context, _ uint, query string, args ...any) (pglib.CommandTag, error) {
 								require.Equal(t, fmt.Sprintf("SET TRANSACTION SNAPSHOT '%s'", testSnapshotID), query)
 								require.Len(t, args, 0)
 								return pglib.CommandTag{}, nil
@@ -546,7 +546,7 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 						return f(&mockTx)
 					case 3:
 						mockTx := pgmocks.Tx{
-							ExecFn: func(ctx context.Context, query string, args ...any) (pglib.CommandTag, error) {
+							ExecFn: func(ctx context.Context, _ uint, query string, args ...any) (pglib.CommandTag, error) {
 								require.Equal(t, fmt.Sprintf("SET TRANSACTION SNAPSHOT '%s'", testSnapshotID), query)
 								require.Len(t, args, 0)
 								return pglib.CommandTag{}, errTest
@@ -593,7 +593,7 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 						return f(&mockTx)
 					case 2:
 						mockTx := pgmocks.Tx{
-							ExecFn: func(ctx context.Context, query string, args ...any) (pglib.CommandTag, error) {
+							ExecFn: func(ctx context.Context, _ uint, query string, args ...any) (pglib.CommandTag, error) {
 								require.Equal(t, fmt.Sprintf("SET TRANSACTION SNAPSHOT '%s'", testSnapshotID), query)
 								require.Len(t, args, 0)
 								return pglib.CommandTag{}, nil
@@ -615,7 +615,7 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 						return f(&mockTx)
 					case 3:
 						mockTx := pgmocks.Tx{
-							ExecFn: func(ctx context.Context, query string, args ...any) (pglib.CommandTag, error) {
+							ExecFn: func(ctx context.Context, _ uint, query string, args ...any) (pglib.CommandTag, error) {
 								require.Equal(t, fmt.Sprintf("SET TRANSACTION SNAPSHOT '%s'", testSnapshotID), query)
 								require.Len(t, args, 0)
 								return pglib.CommandTag{}, nil
@@ -667,7 +667,7 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 						return f(&mockTx)
 					case 2:
 						mockTx := pgmocks.Tx{
-							ExecFn: func(ctx context.Context, query string, args ...any) (pglib.CommandTag, error) {
+							ExecFn: func(ctx context.Context, _ uint, query string, args ...any) (pglib.CommandTag, error) {
 								require.Equal(t, fmt.Sprintf("SET TRANSACTION SNAPSHOT '%s'", testSnapshotID), query)
 								require.Len(t, args, 0)
 								return pglib.CommandTag{}, nil
@@ -689,7 +689,7 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 						return f(&mockTx)
 					case 3:
 						mockTx := pgmocks.Tx{
-							ExecFn: func(ctx context.Context, query string, args ...any) (pglib.CommandTag, error) {
+							ExecFn: func(ctx context.Context, _ uint, query string, args ...any) (pglib.CommandTag, error) {
 								require.Equal(t, fmt.Sprintf("SET TRANSACTION SNAPSHOT '%s'", testSnapshotID), query)
 								require.Len(t, args, 0)
 								return pglib.CommandTag{}, nil
@@ -745,7 +745,7 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 						return f(&mockTx)
 					case 2:
 						mockTx := pgmocks.Tx{
-							ExecFn: func(ctx context.Context, query string, args ...any) (pglib.CommandTag, error) {
+							ExecFn: func(ctx context.Context, _ uint, query string, args ...any) (pglib.CommandTag, error) {
 								require.Equal(t, fmt.Sprintf("SET TRANSACTION SNAPSHOT '%s'", testSnapshotID), query)
 								require.Len(t, args, 0)
 								return pglib.CommandTag{}, nil
@@ -767,7 +767,7 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 						return f(&mockTx)
 					case 3:
 						mockTx := pgmocks.Tx{
-							ExecFn: func(ctx context.Context, query string, args ...any) (pglib.CommandTag, error) {
+							ExecFn: func(ctx context.Context, _ uint, query string, args ...any) (pglib.CommandTag, error) {
 								require.Equal(t, fmt.Sprintf("SET TRANSACTION SNAPSHOT '%s'", testSnapshotID), query)
 								require.Len(t, args, 0)
 								return pglib.CommandTag{}, nil
@@ -828,7 +828,7 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 								ScanFn: func(args ...any) error { return fmt.Errorf("unexpected call to QueryRowFn: %s", query) },
 							}
 						},
-						ExecFn: func(ctx context.Context, query string, args ...any) (pglib.CommandTag, error) {
+						ExecFn: func(ctx context.Context, _ uint, query string, args ...any) (pglib.CommandTag, error) {
 							require.Equal(t, fmt.Sprintf("SET TRANSACTION SNAPSHOT '%s'", testSnapshotID), query)
 							require.Len(t, args, 0)
 							return pglib.CommandTag{}, nil
diff --git a/pkg/wal/processor/postgres/postgres_batch_writer_test.go b/pkg/wal/processor/postgres/postgres_batch_writer_test.go
index cc5f4ac..faf66a0 100644
--- a/pkg/wal/processor/postgres/postgres_batch_writer_test.go
+++ b/pkg/wal/processor/postgres/postgres_batch_writer_test.go
@@ -170,7 +170,7 @@ func TestBatchWriter_sendBatch(t *testing.T) {
 			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) {
+						ExecFn: func(ctx context.Context, _ uint, query string, args ...any) (pglib.CommandTag, error) {
 							require.Equal(t, testQuery.sql, query)
 							require.Equal(t, testQuery.args, args)
 							return pglib.CommandTag{}, nil
@@ -190,7 +190,7 @@ func TestBatchWriter_sendBatch(t *testing.T) {
 			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) {
+						ExecFn: func(ctx context.Context, _ uint, query string, args ...any) (pglib.CommandTag, error) {
 							return pglib.CommandTag{}, errTest
 						},
 					}
@@ -208,7 +208,7 @@ func TestBatchWriter_sendBatch(t *testing.T) {
 			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) {
+						ExecFn: func(ctx context.Context, _ uint, query string, args ...any) (pglib.CommandTag, error) {
 							return pglib.CommandTag{}, nil
 						},
 					}
@@ -249,7 +249,7 @@ func TestBatchWriter_flushQueries(t *testing.T) {
 
 	testQuerySQL := "INSERT INTO test(id, name) VALUES($1, $2)"
 	args1 := []any{1, "alice"}
-	args2 := []any{1, "bob"}
+	args2 := []any{2, "bob"}
 
 	testQuery := func(args []any) *query {
 		return &query{
@@ -260,9 +260,10 @@ func TestBatchWriter_flushQueries(t *testing.T) {
 	execCalls := uint(0)
 
 	tests := []struct {
-		name    string
-		pgconn  *pgmocks.Querier
-		queries []*query
+		name            string
+		pgconn          *pgmocks.Querier
+		queries         []*query
+		disableTriggers bool
 
 		wantExecCalls uint
 		wantErr       error
@@ -284,17 +285,19 @@ func TestBatchWriter_flushQueries(t *testing.T) {
 			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) {
+						ExecFn: func(ctx context.Context, i uint, query string, args ...any) (pglib.CommandTag, error) {
 							execCalls++
-							require.Equal(t, testQuerySQL, query)
-							require.Len(t, args, 2)
-							switch args[1] {
-							case args1[1]:
+							switch i {
+							case 1:
+								require.Equal(t, testQuerySQL, query)
+								require.Equal(t, args1, args)
 								return pglib.CommandTag{}, nil
-							case args2[1]:
+							case 2:
+								require.Equal(t, testQuerySQL, query)
+								require.Equal(t, args2, args)
 								return pglib.CommandTag{}, &pglib.ErrConstraintViolation{}
 							default:
-								return pglib.CommandTag{}, fmt.Errorf("unexpected argument in sql query: %v", args[1])
+								return pglib.CommandTag{}, fmt.Errorf("unexpected call to tx ExecFn: %v", args[1])
 							}
 						},
 					}
@@ -306,22 +309,68 @@ func TestBatchWriter_flushQueries(t *testing.T) {
 			wantExecCalls: 3,
 			wantErr:       nil,
 		},
+		{
+			name: "ok - disable triggers",
+			pgconn: &pgmocks.Querier{
+				ExecInTxFn: func(ctx context.Context, f func(tx pglib.Tx) error) error {
+					mockTx := pgmocks.Tx{
+						ExecFn: func(ctx context.Context, i uint, query string, args ...any) (pglib.CommandTag, error) {
+							execCalls++
+							switch i {
+							case 1:
+								require.Equal(t, "SET session_replication_role = replica", query)
+							case 2:
+								require.Equal(t, testQuerySQL, query)
+								require.Len(t, args, 2)
+								if args[0] != args1[0] && args[0] != args2[0] {
+									return pglib.CommandTag{}, fmt.Errorf("unexpected arguments in query: %v", args)
+								}
+								if args[0] == args1[0] {
+									require.Equal(t, args1, args)
+								}
+								// the second time it's called we don't return a retriable error
+								if args[0] == args2[0] {
+									require.Equal(t, args2, args)
+								}
+							case 3:
+								if query == testQuerySQL {
+									require.Equal(t, args2, args)
+									return pglib.CommandTag{}, &pglib.ErrConstraintViolation{}
+								}
+								require.Equal(t, "SET session_replication_role = DEFAULT", query)
+							default:
+								return pglib.CommandTag{}, fmt.Errorf("unexpected call to tx ExecFn: %v", args[1])
+							}
+							return pglib.CommandTag{}, nil
+						},
+					}
+					return f(&mockTx)
+				},
+			},
+			queries:         []*query{testQuery(args1), testQuery(args2)},
+			disableTriggers: true,
+
+			wantExecCalls: 6,
+			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) {
+						ExecFn: func(ctx context.Context, i uint, query string, args ...any) (pglib.CommandTag, error) {
 							execCalls++
-							require.Equal(t, testQuerySQL, query)
-							require.Len(t, args, 2)
-							switch args[1] {
-							case args1[1]:
+							switch i {
+							case 1:
+								require.Equal(t, testQuerySQL, query)
+								require.Equal(t, args1, args)
 								return pglib.CommandTag{}, nil
-							case args2[1]:
+							case 2:
+								require.Equal(t, testQuerySQL, query)
+								require.Equal(t, args2, args)
 								return pglib.CommandTag{}, errTest
 							default:
-								return pglib.CommandTag{}, fmt.Errorf("unexpected argument in sql query: %v", args[1])
+								return pglib.CommandTag{}, fmt.Errorf("unexpected call to tx ExecFn: %v", args[1])
 							}
 						},
 					}
@@ -333,14 +382,71 @@ func TestBatchWriter_flushQueries(t *testing.T) {
 			wantExecCalls: 2,
 			wantErr:       errTest,
 		},
+		{
+			name: "error - setting replication role to replica",
+			pgconn: &pgmocks.Querier{
+				ExecInTxFn: func(ctx context.Context, f func(tx pglib.Tx) error) error {
+					mockTx := pgmocks.Tx{
+						ExecFn: func(ctx context.Context, i uint, query string, args ...any) (pglib.CommandTag, error) {
+							execCalls++
+							switch i {
+							case 1:
+								require.Equal(t, "SET session_replication_role = replica", query)
+								return pglib.CommandTag{}, errTest
+							default:
+								return pglib.CommandTag{}, fmt.Errorf("unexpected call to tx ExecFn: %v", args[1])
+							}
+						},
+					}
+					return f(&mockTx)
+				},
+			},
+			queries:         []*query{testQuery(args1), testQuery(args2)},
+			disableTriggers: true,
+
+			wantExecCalls: 1,
+			wantErr:       errTest,
+		},
+		{
+			name: "error - resetting replication role to replica",
+			pgconn: &pgmocks.Querier{
+				ExecInTxFn: func(ctx context.Context, f func(tx pglib.Tx) error) error {
+					mockTx := pgmocks.Tx{
+						ExecFn: func(ctx context.Context, i uint, query string, args ...any) (pglib.CommandTag, error) {
+							execCalls++
+							switch i {
+							case 1:
+								require.Equal(t, "SET session_replication_role = replica", query)
+							case 2:
+								require.Equal(t, testQuerySQL, query)
+								require.Equal(t, args1, args)
+							case 3:
+								require.Equal(t, "SET session_replication_role = DEFAULT", query)
+								return pglib.CommandTag{}, errTest
+							default:
+								return pglib.CommandTag{}, fmt.Errorf("unexpected call to tx ExecFn: %v", args[1])
+							}
+							return pglib.CommandTag{}, nil
+						},
+					}
+					return f(&mockTx)
+				},
+			},
+			queries:         []*query{testQuery(args1)},
+			disableTriggers: true,
+
+			wantExecCalls: 3,
+			wantErr:       errTest,
+		},
 	}
 
 	for _, tc := range tests {
 		t.Run(tc.name, func(t *testing.T) {
 			bw := &BatchWriter{
 				Writer: &Writer{
-					logger: loglib.NewNoopLogger(),
-					pgConn: tc.pgconn,
+					logger:          loglib.NewNoopLogger(),
+					pgConn:          tc.pgconn,
+					disableTriggers: tc.disableTriggers,
 				},
 			}
 
diff --git a/pkg/wal/processor/postgres/postgres_bulk_ingest_writer_test.go b/pkg/wal/processor/postgres/postgres_bulk_ingest_writer_test.go
index c4dd0d4..12f9ada 100644
--- a/pkg/wal/processor/postgres/postgres_bulk_ingest_writer_test.go
+++ b/pkg/wal/processor/postgres/postgres_bulk_ingest_writer_test.go
@@ -240,9 +240,10 @@ func TestBulkIngestWriter_sendBatch(t *testing.T) {
 	}
 
 	tests := []struct {
-		name   string
-		batch  *batch.Batch[*query]
-		pgConn *pgmocks.Querier
+		name            string
+		batch           *batch.Batch[*query]
+		pgConn          *pgmocks.Querier
+		disableTriggers bool
 
 		wantErr error
 	}{
@@ -250,8 +251,13 @@ func TestBulkIngestWriter_sendBatch(t *testing.T) {
 			name:  "ok - empty batch",
 			batch: batch.NewBatch([]*query{}, nil),
 			pgConn: &pgmocks.Querier{
-				CopyFromFn: func(ctx context.Context, tableName string, columnNames []string, rows [][]any) (int64, error) {
-					return -1, errors.New("unexpected call to CopyFrom with empty batch")
+				ExecInTxFn: func(ctx context.Context, f func(tx pglib.Tx) error) error {
+					tx := &pgmocks.Tx{
+						CopyFromFn: func(ctx context.Context, tableName string, columnNames []string, rows [][]any) (int64, error) {
+							return -1, errors.New("unexpected call to CopyFrom with empty batch")
+						},
+					}
+					return f(tx)
 				},
 			},
 
@@ -261,13 +267,49 @@ func TestBulkIngestWriter_sendBatch(t *testing.T) {
 			name:  "ok",
 			batch: batch.NewBatch([]*query{testQuery}, nil),
 			pgConn: &pgmocks.Querier{
-				CopyFromFn: func(ctx context.Context, tableName string, columnNames []string, rows [][]any) (int64, error) {
-					require.Equal(t, pglib.QuoteQualifiedIdentifier(testSchema, testTable), tableName)
-					require.Equal(t, testQuery.columnNames, columnNames)
-					require.Equal(t, testQuery.args, rows[0])
-					return 1, nil
+				ExecInTxFn: func(ctx context.Context, f func(tx pglib.Tx) error) error {
+					tx := &pgmocks.Tx{
+						CopyFromFn: func(ctx context.Context, tableName string, columnNames []string, rows [][]any) (int64, error) {
+							require.Equal(t, pglib.QuoteQualifiedIdentifier(testSchema, testTable), tableName)
+							require.Equal(t, testQuery.columnNames, columnNames)
+							require.Equal(t, testQuery.args, rows[0])
+							return 1, nil
+						},
+					}
+					return f(tx)
+				},
+			},
+
+			wantErr: nil,
+		},
+		{
+			name:  "ok - disable triggers",
+			batch: batch.NewBatch([]*query{testQuery}, nil),
+			pgConn: &pgmocks.Querier{
+				ExecInTxFn: func(ctx context.Context, f func(tx pglib.Tx) error) error {
+					tx := &pgmocks.Tx{
+						ExecFn: func(ctx context.Context, i uint, s string, a ...any) (pglib.CommandTag, error) {
+							switch i {
+							case 1:
+								require.Equal(t, "SET session_replication_role = replica", s)
+							case 2:
+								require.Equal(t, "SET session_replication_role = DEFAULT", s)
+							default:
+								return pglib.CommandTag{}, fmt.Errorf("unexpected number of calls to tx ExecFn: %d", i)
+							}
+							return pglib.CommandTag{}, nil
+						},
+						CopyFromFn: func(ctx context.Context, tableName string, columnNames []string, rows [][]any) (int64, error) {
+							require.Equal(t, pglib.QuoteQualifiedIdentifier(testSchema, testTable), tableName)
+							require.Equal(t, testQuery.columnNames, columnNames)
+							require.Equal(t, testQuery.args, rows[0])
+							return 1, nil
+						},
+					}
+					return f(tx)
 				},
 			},
+			disableTriggers: true,
 
 			wantErr: nil,
 		},
@@ -275,8 +317,13 @@ func TestBulkIngestWriter_sendBatch(t *testing.T) {
 			name:  "error - copying from",
 			batch: batch.NewBatch([]*query{testQuery}, nil),
 			pgConn: &pgmocks.Querier{
-				CopyFromFn: func(ctx context.Context, tableName string, columnNames []string, rows [][]any) (int64, error) {
-					return 0, errTest
+				ExecInTxFn: func(ctx context.Context, f func(tx pglib.Tx) error) error {
+					tx := &pgmocks.Tx{
+						CopyFromFn: func(ctx context.Context, tableName string, columnNames []string, rows [][]any) (int64, error) {
+							return 0, errTest
+						},
+					}
+					return f(tx)
 				},
 			},
 
@@ -286,12 +333,72 @@ func TestBulkIngestWriter_sendBatch(t *testing.T) {
 			name:  "error - rows copied mismatch",
 			batch: batch.NewBatch([]*query{testQuery}, nil),
 			pgConn: &pgmocks.Querier{
-				CopyFromFn: func(ctx context.Context, tableName string, columnNames []string, rows [][]any) (int64, error) {
-					return 0, nil
+				ExecInTxFn: func(ctx context.Context, f func(tx pglib.Tx) error) error {
+					tx := &pgmocks.Tx{
+						CopyFromFn: func(ctx context.Context, tableName string, columnNames []string, rows [][]any) (int64, error) {
+							return 0, nil
+						},
+					}
+					return f(tx)
+				},
+			},
+
+			wantErr: errUnexpectedCopiedRows,
+		},
+		{
+			name:  "error - setting replication role to replica",
+			batch: batch.NewBatch([]*query{testQuery}, nil),
+			pgConn: &pgmocks.Querier{
+				ExecInTxFn: func(ctx context.Context, f func(tx pglib.Tx) error) error {
+					tx := &pgmocks.Tx{
+						ExecFn: func(ctx context.Context, i uint, s string, a ...any) (pglib.CommandTag, error) {
+							switch i {
+							case 1:
+								require.Equal(t, "SET session_replication_role = replica", s)
+								return pglib.CommandTag{}, errTest
+							default:
+								return pglib.CommandTag{}, fmt.Errorf("unexpected number of calls to tx ExecFn: %d", i)
+							}
+						},
+						CopyFromFn: func(ctx context.Context, tableName string, columnNames []string, rows [][]any) (int64, error) {
+							return 0, errors.New("CopyFrom should not be called")
+						},
+					}
+					return f(tx)
+				},
+			},
+			disableTriggers: true,
+
+			wantErr: errTest,
+		},
+		{
+			name:  "error - resetting replication role",
+			batch: batch.NewBatch([]*query{testQuery}, nil),
+			pgConn: &pgmocks.Querier{
+				ExecInTxFn: func(ctx context.Context, f func(tx pglib.Tx) error) error {
+					tx := &pgmocks.Tx{
+						ExecFn: func(ctx context.Context, i uint, s string, a ...any) (pglib.CommandTag, error) {
+							switch i {
+							case 1:
+								require.Equal(t, "SET session_replication_role = replica", s)
+								return pglib.CommandTag{}, nil
+							case 2:
+								require.Equal(t, "SET session_replication_role = DEFAULT", s)
+								return pglib.CommandTag{}, errTest
+							default:
+								return pglib.CommandTag{}, fmt.Errorf("unexpected number of calls to tx ExecFn: %d", i)
+							}
+						},
+						CopyFromFn: func(ctx context.Context, tableName string, columnNames []string, rows [][]any) (int64, error) {
+							return 1, nil
+						},
+					}
+					return f(tx)
 				},
 			},
+			disableTriggers: true,
 
-			wantErr: errors.New("number of rows copied (0) doesn't match the source rows (1)"),
+			wantErr: errTest,
 		},
 	}
 
@@ -301,15 +408,14 @@ func TestBulkIngestWriter_sendBatch(t *testing.T) {
 
 			writer := &BulkIngestWriter{
 				Writer: &Writer{
-					logger: loglib.NewNoopLogger(),
-					pgConn: tc.pgConn,
+					logger:          loglib.NewNoopLogger(),
+					pgConn:          tc.pgConn,
+					disableTriggers: tc.disableTriggers,
 				},
 			}
 
 			err := writer.sendBatch(context.Background(), tc.batch)
-			if errors.Is(err, tc.wantErr) {
-				require.Equal(t, tc.wantErr, err)
-			}
+			require.ErrorIs(t, err, tc.wantErr)
 		})
 	}
 }