Repository: xataio/pgstream
Pull request #334 (2025-05-13 11:32:44)

Add instrumentation to postgres lib
This PR adds instrumentation to the postgres lib utilities, including the querier, the querier builder and the pdump/pgrestore functions. 

Relates to #323 


Patch:
diff --git a/internal/postgres/instrumentation/instrumented_pg_dump_restore.go b/internal/postgres/instrumentation/instrumented_pg_dump_restore.go
new file mode 100644
index 0000000..2ccd164
--- /dev/null
+++ b/internal/postgres/instrumentation/instrumented_pg_dump_restore.go
@@ -0,0 +1,51 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package instrumentation
+
+import (
+	"context"
+
+	pglib "github.com/xataio/pgstream/internal/postgres"
+	"github.com/xataio/pgstream/pkg/otel"
+	"go.opentelemetry.io/otel/attribute"
+	"go.opentelemetry.io/otel/trace"
+)
+
+type PGDumpRestore struct {
+	pgdumpFn    pglib.PGDumpFn
+	pgrestoreFn pglib.PGRestoreFn
+	tracer      trace.Tracer
+}
+
+func NewPGDumpFn(pgdumpFn pglib.PGDumpFn, instrumentation *otel.Instrumentation) pglib.PGDumpFn {
+	pgdr := &PGDumpRestore{
+		pgdumpFn: pgdumpFn,
+		tracer:   instrumentation.Tracer,
+	}
+	return pgdr.PGDump
+}
+
+func NewPGRestoreFn(pgrestoreFn pglib.PGRestoreFn, instrumentation *otel.Instrumentation) pglib.PGRestoreFn {
+	pgdr := &PGDumpRestore{
+		pgrestoreFn: pgrestoreFn,
+		tracer:      instrumentation.Tracer,
+	}
+	return pgdr.PGRestore
+}
+
+func (i *PGDumpRestore) PGDump(ctx context.Context, opts pglib.PGDumpOptions) (dump []byte, err error) {
+	ctx, span := otel.StartSpan(ctx, i.tracer, "pgdump", trace.WithAttributes([]attribute.KeyValue{
+		{Key: "schemas", Value: attribute.StringSliceValue(opts.Schemas)},
+		{Key: "tables", Value: attribute.StringSliceValue(opts.Tables)},
+		{Key: "exclude_tables", Value: attribute.StringSliceValue(opts.ExcludeTables)},
+		{Key: "clean", Value: attribute.BoolValue(opts.Clean)},
+	}...))
+	defer otel.CloseSpan(span, err)
+	return i.pgdumpFn(ctx, opts)
+}
+
+func (i *PGDumpRestore) PGRestore(ctx context.Context, opts pglib.PGRestoreOptions, dump []byte) (out string, err error) {
+	ctx, span := otel.StartSpan(ctx, i.tracer, "pgrestore")
+	defer otel.CloseSpan(span, err)
+	return i.pgrestoreFn(ctx, opts, dump)
+}
diff --git a/internal/postgres/instrumentation/instrumented_querier.go b/internal/postgres/instrumentation/instrumented_querier.go
new file mode 100644
index 0000000..275f03d
--- /dev/null
+++ b/internal/postgres/instrumentation/instrumented_querier.go
@@ -0,0 +1,181 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package instrumentation
+
+import (
+	"context"
+	"fmt"
+	"strings"
+	"time"
+
+	pglib "github.com/xataio/pgstream/internal/postgres"
+	"github.com/xataio/pgstream/pkg/otel"
+
+	"go.opentelemetry.io/otel/attribute"
+	"go.opentelemetry.io/otel/metric"
+	"go.opentelemetry.io/otel/trace"
+)
+
+type Querier struct {
+	inner   pglib.Querier
+	tracer  trace.Tracer
+	meter   metric.Meter
+	metrics *metrics
+}
+
+type metrics struct {
+	queryLatency metric.Int64Histogram
+}
+
+const (
+	queryTypeAttributeKey = "query_type"
+	queryAttributeKey     = "query"
+	unknownQueryType      = "unknown"
+	txQueryType           = "tx"
+)
+
+func NewQuerier(q pglib.Querier, instrumentation *otel.Instrumentation) (pglib.Querier, error) {
+	if instrumentation == nil {
+		return q, nil
+	}
+
+	querier := &Querier{
+		inner:   q,
+		tracer:  instrumentation.Tracer,
+		meter:   instrumentation.Meter,
+		metrics: &metrics{},
+	}
+
+	if err := querier.initMetrics(); err != nil {
+		return nil, fmt.Errorf("initialising postgres querier metrics: %w", err)
+	}
+
+	return querier, nil
+}
+
+func (i *Querier) Query(ctx context.Context, query string, args ...any) (rows pglib.Rows, err error) {
+	queryAttrs := i.queryAttributes(query)
+	ctx, span := otel.StartSpan(ctx, i.tracer, "querier.Query", trace.WithAttributes(queryAttrs...))
+	defer otel.CloseSpan(span, err)
+
+	if i.meter != nil {
+		startTime := time.Now()
+		defer func() {
+			i.metrics.queryLatency.Record(ctx, int64(time.Since(startTime).Milliseconds()), metric.WithAttributes(queryAttrs...))
+		}()
+	}
+	return i.inner.Query(ctx, query, args...)
+}
+
+func (i *Querier) QueryRow(ctx context.Context, query string, args ...any) pglib.Row {
+	queryAttrs := i.queryAttributes(query)
+	ctx, span := otel.StartSpan(ctx, i.tracer, "querier.QueryRow", trace.WithAttributes(queryAttrs...))
+	defer otel.CloseSpan(span, nil)
+
+	if i.meter != nil {
+		startTime := time.Now()
+		defer func() {
+			i.metrics.queryLatency.Record(ctx, int64(time.Since(startTime).Milliseconds()), metric.WithAttributes(queryAttrs...))
+		}()
+	}
+	return i.inner.QueryRow(ctx, query, args...)
+}
+
+func (i *Querier) Exec(ctx context.Context, query string, args ...any) (tag pglib.CommandTag, err error) {
+	queryAttrs := i.queryAttributes(query)
+	ctx, span := otel.StartSpan(ctx, i.tracer, "querier.Exec", trace.WithAttributes(queryAttrs...))
+	defer otel.CloseSpan(span, err)
+
+	if i.meter != nil {
+		startTime := time.Now()
+		defer func() {
+			i.metrics.queryLatency.Record(ctx, int64(time.Since(startTime).Milliseconds()), metric.WithAttributes(queryAttrs...))
+		}()
+	}
+	return i.inner.Exec(ctx, query, args...)
+}
+
+func (i *Querier) ExecInTx(ctx context.Context, fn func(tx pglib.Tx) error) (err error) {
+	queryAttrs := i.queryAttributes(txQueryType)
+	ctx, span := otel.StartSpan(ctx, i.tracer, "querier.ExecInTx", trace.WithAttributes(queryAttrs...))
+	defer otel.CloseSpan(span, err)
+
+	if i.meter != nil {
+		startTime := time.Now()
+		defer func() {
+			i.metrics.queryLatency.Record(ctx, int64(time.Since(startTime).Milliseconds()), metric.WithAttributes(queryAttrs...))
+		}()
+	}
+	return i.inner.ExecInTx(ctx, fn)
+}
+
+func (i *Querier) ExecInTxWithOptions(ctx context.Context, fn func(tx pglib.Tx) error, txOpts pglib.TxOptions) (err error) {
+	queryAttrs := i.queryAttributes(txQueryType)
+	ctx, span := otel.StartSpan(ctx, i.tracer, "querier.ExecInTxWithOptions", trace.WithAttributes(queryAttrs...))
+	defer otel.CloseSpan(span, err)
+
+	if i.meter != nil {
+		startTime := time.Now()
+		defer func() {
+			i.metrics.queryLatency.Record(ctx, int64(time.Since(startTime).Milliseconds()), metric.WithAttributes(queryAttrs...))
+		}()
+	}
+
+	return i.inner.ExecInTxWithOptions(ctx, fn, txOpts)
+}
+
+func (i *Querier) Ping(ctx context.Context) (err error) {
+	ctx, span := otel.StartSpan(ctx, i.tracer, "querier.Ping")
+	defer otel.CloseSpan(span, err)
+	return i.inner.Ping(ctx)
+}
+
+func (i *Querier) Close(ctx context.Context) (err error) {
+	ctx, span := otel.StartSpan(ctx, i.tracer, "querier.Close")
+	defer otel.CloseSpan(span, err)
+	return i.inner.Close(ctx)
+}
+
+func (i *Querier) initMetrics() error {
+	if i.meter == nil {
+		return nil
+	}
+
+	var err error
+	i.metrics.queryLatency, err = i.meter.Int64Histogram("pgstream.postgres.querier.latency",
+		metric.WithUnit("ms"),
+		metric.WithDescription("Distribution of the time taken to perform a query"))
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func (i *Querier) queryAttributes(query string) []attribute.KeyValue {
+	var qt string
+	switch query {
+	case txQueryType:
+		qt = txQueryType
+	case "":
+		qt = unknownQueryType
+	default:
+		qt = strings.ToUpper(strings.Split(query, " ")[0])
+	}
+
+	attrs := []attribute.KeyValue{
+		{
+			Key:   queryTypeAttributeKey,
+			Value: attribute.StringValue(qt),
+		},
+	}
+
+	if qt == txQueryType || qt == unknownQueryType {
+		return attrs
+	}
+
+	return append(attrs, attribute.KeyValue{
+		Key:   queryAttributeKey,
+		Value: attribute.StringValue(query),
+	})
+}
diff --git a/internal/postgres/instrumentation/instrumented_querier_builder.go b/internal/postgres/instrumentation/instrumented_querier_builder.go
new file mode 100644
index 0000000..fb1e594
--- /dev/null
+++ b/internal/postgres/instrumentation/instrumented_querier_builder.go
@@ -0,0 +1,20 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package instrumentation
+
+import (
+	"context"
+
+	pglib "github.com/xataio/pgstream/internal/postgres"
+	"github.com/xataio/pgstream/pkg/otel"
+)
+
+func NewQuerierBuilder(b pglib.QuerierBuilder, i *otel.Instrumentation) (pglib.QuerierBuilder, error) {
+	return func(ctx context.Context, url string) (pglib.Querier, error) {
+		querier, err := b(ctx, url)
+		if err != nil {
+			return nil, err
+		}
+		return NewQuerier(querier, i)
+	}, nil
+}
diff --git a/internal/postgres/pg_dump.go b/internal/postgres/pg_dump.go
index e9c55a1..72c74e1 100644
--- a/internal/postgres/pg_dump.go
+++ b/internal/postgres/pg_dump.go
@@ -4,6 +4,7 @@ package postgres
 
 import (
 	"bytes"
+	"context"
 	"fmt"
 	"os/exec"
 )
@@ -70,7 +71,7 @@ func (opts *PGDumpOptions) ToArgs() []string {
 }
 
 // Func RunPGDump runs pg_dump command with the given options and returns the result.
-func RunPGDump(opts PGDumpOptions) ([]byte, error) {
+func RunPGDump(_ context.Context, opts PGDumpOptions) ([]byte, error) {
 	cmd := exec.Command(pgDumpCmd, opts.ToArgs()...) //nolint:gosec
 
 	// TODO: add streaming support when it needs to be used for large data dumps
diff --git a/internal/postgres/pg_conn_builder.go b/internal/postgres/pg_querier_builder.go
similarity index 100%
rename from internal/postgres/pg_conn_builder.go
rename to internal/postgres/pg_querier_builder.go
diff --git a/internal/postgres/pg_restore.go b/internal/postgres/pg_restore.go
index 5a0b179..8564414 100644
--- a/internal/postgres/pg_restore.go
+++ b/internal/postgres/pg_restore.go
@@ -5,6 +5,7 @@ package postgres
 import (
 	"bufio"
 	"bytes"
+	"context"
 	"errors"
 	"fmt"
 	"io"
@@ -54,7 +55,7 @@ func (opts PGRestoreOptions) toPSQLArgs() []string {
 
 // Func RunPGRestore runs pg_restore command with the given options and returns
 // the result.
-func RunPGRestore(opts PGRestoreOptions, dump []byte) (string, error) {
+func RunPGRestore(_ context.Context, opts PGRestoreOptions, dump []byte) (string, error) {
 	var cmd *exec.Cmd
 	switch opts.Format {
 	case "c":
diff --git a/internal/postgres/pg_utils.go b/internal/postgres/pg_utils.go
index e2e0708..9960f1b 100644
--- a/internal/postgres/pg_utils.go
+++ b/internal/postgres/pg_utils.go
@@ -2,7 +2,11 @@
 
 package postgres
 
-import "github.com/lib/pq"
+import (
+	"context"
+
+	"github.com/lib/pq"
+)
 
 func QuoteIdentifier(s string) string {
 	return pq.QuoteIdentifier(s)
@@ -11,3 +15,8 @@ func QuoteIdentifier(s string) string {
 func QuoteQualifiedIdentifier(schema, table string) string {
 	return pq.QuoteIdentifier(schema) + "." + pq.QuoteIdentifier(table)
 }
+
+type (
+	PGDumpFn    func(context.Context, PGDumpOptions) ([]byte, error)
+	PGRestoreFn func(context.Context, PGRestoreOptions, []byte) (string, error)
+)
diff --git a/pkg/snapshot/generator/postgres/schema/pgdumprestore/snapshot_pg_dump_restore_generator.go b/pkg/snapshot/generator/postgres/schema/pgdumprestore/snapshot_pg_dump_restore_generator.go
index a8acd75..dc8cc3d 100644
--- a/pkg/snapshot/generator/postgres/schema/pgdumprestore/snapshot_pg_dump_restore_generator.go
+++ b/pkg/snapshot/generator/postgres/schema/pgdumprestore/snapshot_pg_dump_restore_generator.go
@@ -21,8 +21,8 @@ import (
 type SnapshotGenerator struct {
 	sourceURL      string
 	targetURL      string
-	pgDumpFn       pgdumpFn
-	pgRestoreFn    pgrestoreFn
+	pgDumpFn       pglib.PGDumpFn
+	pgRestoreFn    pglib.PGRestoreFn
 	schemalogStore schemalog.Store
 	connBuilder    pglib.QuerierBuilder
 	cleanTargetDB  bool
@@ -35,11 +35,6 @@ type Config struct {
 	CleanTargetDB bool
 }
 
-type (
-	pgdumpFn    func(pglib.PGDumpOptions) ([]byte, error)
-	pgrestoreFn func(pglib.PGRestoreOptions, []byte) (string, error)
-)
-
 type Option func(s *SnapshotGenerator)
 
 const (
@@ -93,7 +88,7 @@ func (s *SnapshotGenerator) CreateSnapshot(ctx context.Context, ss *snapshot.Sna
 		return fmt.Errorf("preparing pg_dump options: %w", err)
 	}
 
-	dump, err := s.pgDumpFn(*pgdumpOpts)
+	dump, err := s.pgDumpFn(ctx, *pgdumpOpts)
 	if err != nil {
 		return err
 	}
@@ -107,7 +102,7 @@ func (s *SnapshotGenerator) CreateSnapshot(ctx context.Context, ss *snapshot.Sna
 		}
 	}
 
-	_, err = s.pgRestoreFn(s.pgrestoreOptions(), dump)
+	_, err = s.pgRestoreFn(ctx, s.pgrestoreOptions(), dump)
 	pgrestoreErr := &pglib.PGRestoreErrors{}
 	if err != nil {
 		switch {


Test patch:
diff --git a/internal/postgres/pg_dump_pg_restore_integration_test.go b/internal/postgres/pg_dump_pg_restore_integration_test.go
index 2cbd1e8..6e0d939 100644
--- a/internal/postgres/pg_dump_pg_restore_integration_test.go
+++ b/internal/postgres/pg_dump_pg_restore_integration_test.go
@@ -58,7 +58,7 @@ func Test_pgdump_pgrestore(t *testing.T) {
 			ExcludeTables:    []string{testSchema + "." + testTable2},
 		}
 
-		dump, err := RunPGDump(pgdumpOpts)
+		dump, err := RunPGDump(context.TODO(), pgdumpOpts)
 		require.NoError(t, err)
 
 		pgrestoreOpts := PGRestoreOptions{
@@ -68,7 +68,7 @@ func Test_pgdump_pgrestore(t *testing.T) {
 			Format:           format,
 		}
 
-		_, err = RunPGRestore(pgrestoreOpts, dump)
+		_, err = RunPGRestore(context.TODO(), pgrestoreOpts, dump)
 		require.NoError(t, err)
 
 		// schema only pgdump, no data should be available but the schema and
diff --git a/pkg/snapshot/generator/postgres/schema/pgdumprestore/snapshot_pg_dump_restore_generator_test.go b/pkg/snapshot/generator/postgres/schema/pgdumprestore/snapshot_pg_dump_restore_generator_test.go
index 6bc8ee4..7513047 100644
--- a/pkg/snapshot/generator/postgres/schema/pgdumprestore/snapshot_pg_dump_restore_generator_test.go
+++ b/pkg/snapshot/generator/postgres/schema/pgdumprestore/snapshot_pg_dump_restore_generator_test.go
@@ -31,8 +31,8 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 		snapshot       *snapshot.Snapshot
 		conn           pglib.Querier
 		connBuilder    pglib.QuerierBuilder
-		pgdumpFn       pgdumpFn
-		pgrestoreFn    pgrestoreFn
+		pgdumpFn       pglib.PGDumpFn
+		pgrestoreFn    pglib.PGRestoreFn
 		schemalogStore schemalog.Store
 
 		wantErr error
@@ -65,7 +65,7 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 					}, nil
 				},
 			},
-			pgdumpFn: func(po pglib.PGDumpOptions) ([]byte, error) {
+			pgdumpFn: func(_ context.Context, po pglib.PGDumpOptions) ([]byte, error) {
 				require.Equal(t, pglib.PGDumpOptions{
 					ConnectionString: "source-url",
 					Format:           "p",
@@ -76,7 +76,7 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 				}, po)
 				return testDump, nil
 			},
-			pgrestoreFn: func(po pglib.PGRestoreOptions, dump []byte) (string, error) {
+			pgrestoreFn: func(_ context.Context, po pglib.PGRestoreOptions, dump []byte) (string, error) {
 				require.Equal(t, pglib.PGRestoreOptions{
 					ConnectionString: "target-url",
 					SchemaOnly:       true,
@@ -100,7 +100,7 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 					return pglib.CommandTag{}, nil
 				},
 			},
-			pgdumpFn: func(po pglib.PGDumpOptions) ([]byte, error) {
+			pgdumpFn: func(_ context.Context, po pglib.PGDumpOptions) ([]byte, error) {
 				require.Equal(t, pglib.PGDumpOptions{
 					ConnectionString: "source-url",
 					Format:           "p",
@@ -109,7 +109,7 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 				}, po)
 				return testDump, nil
 			},
-			pgrestoreFn: func(po pglib.PGRestoreOptions, dump []byte) (string, error) {
+			pgrestoreFn: func(_ context.Context, po pglib.PGRestoreOptions, dump []byte) (string, error) {
 				require.Equal(t, pglib.PGRestoreOptions{
 					ConnectionString: "target-url",
 					SchemaOnly:       true,
@@ -132,10 +132,10 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 					return pglib.CommandTag{}, errors.New("ExecFn: should not be called")
 				},
 			},
-			pgdumpFn: func(po pglib.PGDumpOptions) ([]byte, error) {
+			pgdumpFn: func(_ context.Context, po pglib.PGDumpOptions) ([]byte, error) {
 				return nil, errors.New("pgdumpFn: should not be called")
 			},
-			pgrestoreFn: func(po pglib.PGRestoreOptions, dump []byte) (string, error) {
+			pgrestoreFn: func(_ context.Context, po pglib.PGRestoreOptions, dump []byte) (string, error) {
 				return "", errors.New("pgrestoreFn: should not be called")
 			},
 
@@ -152,10 +152,10 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 					return nil, errTest
 				},
 			},
-			pgdumpFn: func(po pglib.PGDumpOptions) ([]byte, error) {
+			pgdumpFn: func(_ context.Context, po pglib.PGDumpOptions) ([]byte, error) {
 				return nil, errors.New("pgdumpFn: should not be called")
 			},
-			pgrestoreFn: func(po pglib.PGRestoreOptions, dump []byte) (string, error) {
+			pgrestoreFn: func(_ context.Context, po pglib.PGRestoreOptions, dump []byte) (string, error) {
 				return "", errors.New("pgrestoreFn: should not be called")
 			},
 
@@ -180,10 +180,10 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 					}, nil
 				},
 			},
-			pgdumpFn: func(po pglib.PGDumpOptions) ([]byte, error) {
+			pgdumpFn: func(_ context.Context, po pglib.PGDumpOptions) ([]byte, error) {
 				return nil, errors.New("pgdumpFn: should not be called")
 			},
-			pgrestoreFn: func(po pglib.PGRestoreOptions, dump []byte) (string, error) {
+			pgrestoreFn: func(_ context.Context, po pglib.PGRestoreOptions, dump []byte) (string, error) {
 				return "", errors.New("pgrestoreFn: should not be called")
 			},
 
@@ -216,10 +216,10 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 					}, nil
 				},
 			},
-			pgdumpFn: func(po pglib.PGDumpOptions) ([]byte, error) {
+			pgdumpFn: func(_ context.Context, po pglib.PGDumpOptions) ([]byte, error) {
 				return nil, errors.New("pgdumpFn: should not be called")
 			},
-			pgrestoreFn: func(po pglib.PGRestoreOptions, dump []byte) (string, error) {
+			pgrestoreFn: func(_ context.Context, po pglib.PGRestoreOptions, dump []byte) (string, error) {
 				return "", errors.New("pgrestoreFn: should not be called")
 			},
 
@@ -251,10 +251,10 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 					}, nil
 				},
 			},
-			pgdumpFn: func(po pglib.PGDumpOptions) ([]byte, error) {
+			pgdumpFn: func(_ context.Context, po pglib.PGDumpOptions) ([]byte, error) {
 				return nil, errTest
 			},
-			pgrestoreFn: func(po pglib.PGRestoreOptions, dump []byte) (string, error) {
+			pgrestoreFn: func(_ context.Context, po pglib.PGRestoreOptions, dump []byte) (string, error) {
 				return "", errors.New("pgrestoreFn: should not be called")
 			},
 
@@ -286,10 +286,10 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 					}, nil
 				},
 			},
-			pgdumpFn: func(po pglib.PGDumpOptions) ([]byte, error) {
+			pgdumpFn: func(_ context.Context, po pglib.PGDumpOptions) ([]byte, error) {
 				return testDump, nil
 			},
-			pgrestoreFn: func(po pglib.PGRestoreOptions, dump []byte) (string, error) {
+			pgrestoreFn: func(_ context.Context, po pglib.PGRestoreOptions, dump []byte) (string, error) {
 				return "", errTest
 			},
 
@@ -321,10 +321,10 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 					}, nil
 				},
 			},
-			pgdumpFn: func(po pglib.PGDumpOptions) ([]byte, error) {
+			pgdumpFn: func(_ context.Context, po pglib.PGDumpOptions) ([]byte, error) {
 				return testDump, nil
 			},
-			pgrestoreFn: func(po pglib.PGRestoreOptions, dump []byte) (string, error) {
+			pgrestoreFn: func(_ context.Context, po pglib.PGRestoreOptions, dump []byte) (string, error) {
 				return "", pglib.NewPGRestoreErrors(errTest)
 			},
 
@@ -356,10 +356,10 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 					}, nil
 				},
 			},
-			pgdumpFn: func(po pglib.PGDumpOptions) ([]byte, error) {
+			pgdumpFn: func(_ context.Context, po pglib.PGDumpOptions) ([]byte, error) {
 				return testDump, nil
 			},
-			pgrestoreFn: func(po pglib.PGRestoreOptions, dump []byte) (string, error) {
+			pgrestoreFn: func(_ context.Context, po pglib.PGRestoreOptions, dump []byte) (string, error) {
 				return "", pglib.NewPGRestoreErrors(&pglib.ErrRelationAlreadyExists{})
 			},
 
@@ -393,10 +393,10 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 					},
 				}, nil
 			},
-			pgdumpFn: func(po pglib.PGDumpOptions) ([]byte, error) {
+			pgdumpFn: func(_ context.Context, po pglib.PGDumpOptions) ([]byte, error) {
 				return testDump, nil
 			},
-			pgrestoreFn: func(po pglib.PGRestoreOptions, dump []byte) (string, error) {
+			pgrestoreFn: func(_ context.Context, po pglib.PGRestoreOptions, dump []byte) (string, error) {
 				return "", errors.New("pgrestoreFn: should not be called")
 			},
 
@@ -428,10 +428,10 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 					}, nil
 				},
 			},
-			pgdumpFn: func(po pglib.PGDumpOptions) ([]byte, error) {
+			pgdumpFn: func(_ context.Context, po pglib.PGDumpOptions) ([]byte, error) {
 				return testDump, nil
 			},
-			pgrestoreFn: func(po pglib.PGRestoreOptions, dump []byte) (string, error) {
+			pgrestoreFn: func(_ context.Context, po pglib.PGRestoreOptions, dump []byte) (string, error) {
 				return "", errors.New("pgrestoreFn: should not be called")
 			},
 
@@ -464,10 +464,10 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 					}, nil
 				},
 			},
-			pgdumpFn: func(po pglib.PGDumpOptions) ([]byte, error) {
+			pgdumpFn: func(_ context.Context, po pglib.PGDumpOptions) ([]byte, error) {
 				return testDump, nil
 			},
-			pgrestoreFn: func(po pglib.PGRestoreOptions, dump []byte) (string, error) {
+			pgrestoreFn: func(_ context.Context, po pglib.PGRestoreOptions, dump []byte) (string, error) {
 				require.Equal(t, testDump, dump)
 				return "", nil
 			},