Repository: xataio/pgstream
Pull request #54 (2024-07-04 13:18:54)

Use postgres connection pools
This PR updates the postgres library to add connection pools, as well as updating all existing postgres dependencies to use the postgres library, favouring the use of the pool when concurrent/multiple access is required (this includes the subscription store and the schema log store). This prevents conn busy errors.


Patch:
diff --git a/.gitignore b/.gitignore
index 9404647..7b36554 100644
--- a/.gitignore
+++ b/.gitignore
@@ -3,4 +3,4 @@ pgstream
 tools/webhook/webhook
 
 # misc
-.DS_Store
\ No newline at end of file
+.DS_Store
diff --git a/brand-kit/README.md b/brand-kit/README.md
index 1751ec5..42dab9b 100644
--- a/brand-kit/README.md
+++ b/brand-kit/README.md
@@ -11,4 +11,4 @@ Our logo combines a symbol and wordmark. For dark backgrounds, use the inverted
 
 The project banner is a key branding element that can be prominently featured at the top of project documentation, websites, social media profiles, and even on swag stickers. It serves as a visual representation of our project's identity and can be used to create a strong connection with our brand. Ensure that the banner is displayed at its original size to maintain clarity and visibility across various applications.
 
-![Banner](banner/pgstream-banner.svg)
\ No newline at end of file
+![Banner](banner/pgstream-banner.svg)
diff --git a/go.mod b/go.mod
index 1e942a8..8a87a73 100644
--- a/go.mod
+++ b/go.mod
@@ -41,6 +41,7 @@ require (
 	github.com/jackc/pgio v1.0.0 // indirect
 	github.com/jackc/pgpassfile v1.0.0 // indirect
 	github.com/jackc/pgservicefile v0.0.0-20221227161230-091c0ba34f0a // indirect
+	github.com/jackc/puddle/v2 v2.2.1 // indirect
 	github.com/klauspost/compress v1.17.0 // indirect
 	github.com/labstack/gommon v0.4.2 // indirect
 	github.com/lib/pq v1.10.9 // indirect
diff --git a/internal/postgres/mocks/mock_pg_conn.go b/internal/postgres/mocks/mock_pg_conn.go
deleted file mode 100644
index c8eb128..0000000
--- a/internal/postgres/mocks/mock_pg_conn.go
+++ /dev/null
@@ -1,24 +0,0 @@
-// SPDX-License-Identifier: Apache-2.0
-
-package mocks
-
-import (
-	"context"
-
-	"github.com/jackc/pgx/v5/pgconn"
-	"github.com/xataio/pgstream/internal/postgres"
-)
-
-type Conn struct {
-	QueryRowFn func(ctx context.Context, query string, args ...any) postgres.Row
-	ExecFn     func(context.Context, string, ...any) (pgconn.CommandTag, error)
-	CloseFn    func(context.Context) error
-}
-
-func (m *Conn) QueryRow(ctx context.Context, query string, args ...any) postgres.Row {
-	return m.QueryRowFn(ctx, query, args...)
-}
-
-func (m *Conn) Close(ctx context.Context) error {
-	return m.CloseFn(ctx)
-}
diff --git a/internal/postgres/mocks/mock_pg_querier.go b/internal/postgres/mocks/mock_pg_querier.go
new file mode 100644
index 0000000..ac46d93
--- /dev/null
+++ b/internal/postgres/mocks/mock_pg_querier.go
@@ -0,0 +1,32 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package mocks
+
+import (
+	"context"
+
+	"github.com/xataio/pgstream/internal/postgres"
+)
+
+type Querier 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(context.Context, string, ...any) (postgres.CommandTag, error)
+	CloseFn    func(context.Context) error
+}
+
+func (m *Querier) QueryRow(ctx context.Context, query string, args ...any) postgres.Row {
+	return m.QueryRowFn(ctx, query, args...)
+}
+
+func (m *Querier) Query(ctx context.Context, query string, args ...any) (postgres.Rows, error) {
+	return m.QueryFn(ctx, query, args...)
+}
+
+func (m *Querier) Exec(ctx context.Context, query string, args ...any) (postgres.CommandTag, error) {
+	return m.ExecFn(ctx, query, args...)
+}
+
+func (m *Querier) Close(ctx context.Context) error {
+	return m.CloseFn(ctx)
+}
diff --git a/internal/postgres/pg_conn.go b/internal/postgres/pg_conn.go
index 1cc334b..bbeefc0 100644
--- a/internal/postgres/pg_conn.go
+++ b/internal/postgres/pg_conn.go
@@ -7,6 +7,7 @@ import (
 	"fmt"
 
 	"github.com/jackc/pgx/v5"
+	"github.com/jackc/pgx/v5/pgconn"
 )
 
 type Conn struct {
@@ -17,6 +18,14 @@ type Row interface {
 	pgx.Row
 }
 
+type Rows interface {
+	pgx.Rows
+}
+
+type CommandTag struct {
+	pgconn.CommandTag
+}
+
 func NewConn(ctx context.Context, url string) (*Conn, error) {
 	pgCfg, err := pgx.ParseConfig(url)
 	if err != nil {
@@ -35,6 +44,15 @@ func (c *Conn) QueryRow(ctx context.Context, query string, args ...any) Row {
 	return c.conn.QueryRow(ctx, query, args...)
 }
 
+func (c *Conn) Query(ctx context.Context, query string, args ...any) (Rows, error) {
+	return c.conn.Query(ctx, query, args...)
+}
+
+func (c *Conn) Exec(ctx context.Context, query string, args ...any) (CommandTag, error) {
+	tag, err := c.conn.Exec(ctx, query, args...)
+	return CommandTag{tag}, err
+}
+
 func (c *Conn) Close(ctx context.Context) error {
 	return mapError(c.conn.Close(ctx))
 }
diff --git a/internal/postgres/pg_conn_pool.go b/internal/postgres/pg_conn_pool.go
new file mode 100644
index 0000000..46f7e7e
--- /dev/null
+++ b/internal/postgres/pg_conn_pool.go
@@ -0,0 +1,46 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package postgres
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/jackc/pgx/v5/pgxpool"
+)
+
+type Pool struct {
+	*pgxpool.Pool
+}
+
+func NewConnPool(ctx context.Context, url string) (*Pool, error) {
+	pgCfg, err := pgxpool.ParseConfig(url)
+	if err != nil {
+		return nil, fmt.Errorf("failed parsing postgres connection string: %w", mapError(err))
+	}
+
+	pool, err := pgxpool.NewWithConfig(ctx, pgCfg)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create a postgres connection pool: %w", mapError(err))
+	}
+
+	return &Pool{Pool: pool}, nil
+}
+
+func (c *Pool) QueryRow(ctx context.Context, query string, args ...any) Row {
+	return c.Pool.QueryRow(ctx, query, args...)
+}
+
+func (c *Pool) Query(ctx context.Context, query string, args ...any) (Rows, error) {
+	return c.Pool.Query(ctx, query, args...)
+}
+
+func (c *Pool) Exec(ctx context.Context, query string, args ...any) (CommandTag, error) {
+	tag, err := c.Pool.Exec(ctx, query, args...)
+	return CommandTag{tag}, err
+}
+
+func (c *Pool) Close(_ context.Context) error {
+	c.Pool.Close()
+	return nil
+}
diff --git a/internal/postgres/pg_querier.go b/internal/postgres/pg_querier.go
new file mode 100644
index 0000000..8a6c826
--- /dev/null
+++ b/internal/postgres/pg_querier.go
@@ -0,0 +1,12 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package postgres
+
+import "context"
+
+type Querier 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)
+	Close(ctx context.Context) error
+}
diff --git a/pkg/schemalog/postgres/pg_schemalog_store.go b/pkg/schemalog/postgres/pg_schemalog_store.go
index 309db51..94fae09 100644
--- a/pkg/schemalog/postgres/pg_schemalog_store.go
+++ b/pkg/schemalog/postgres/pg_schemalog_store.go
@@ -7,42 +7,32 @@ import (
 	"errors"
 	"fmt"
 
+	pglib "github.com/xataio/pgstream/internal/postgres"
 	"github.com/xataio/pgstream/pkg/schemalog"
 
 	"github.com/jackc/pgx/v5"
-	"github.com/jackc/pgx/v5/pgconn"
 )
 
 // Store is a postgres implementation of the schemalog.Store interface
 type Store struct {
-	querier querier
+	querier pglib.Querier
 }
 
 type Config struct {
 	URL string
 }
 
-type querier interface {
-	QueryRow(ctx context.Context, sql string, args ...any) pgx.Row
-	Exec(ctx context.Context, sql string, args ...any) (tag pgconn.CommandTag, err error)
-	Close(ctx context.Context)
-}
-
 func NewStore(ctx context.Context, cfg Config) (*Store, error) {
-	pgCfg, err := pgx.ParseConfig(cfg.URL)
+	pool, err := pglib.NewConnPool(ctx, cfg.URL)
 	if err != nil {
 		return nil, err
 	}
-	pgConn, err := pgx.ConnectConfig(ctx, pgCfg)
-	if err != nil {
-		return nil, fmt.Errorf("create postgres client: %w", err)
-	}
 	return &Store{
-		querier: &pgxConn{Conn: pgConn},
+		querier: pool,
 	}, nil
 }
 
-func NewStoreWithQuerier(querier querier) *Store {
+func NewStoreWithQuerier(querier pglib.Querier) *Store {
 	return &Store{
 		querier: querier,
 	}
@@ -80,16 +70,6 @@ func (s *Store) Close() error {
 	return nil
 }
 
-// pgxConn is a wrapper around the pgx.Conn to implement the same interface all
-// other pg packages use for pool/connections
-type pgxConn struct {
-	*pgx.Conn
-}
-
-func (c *pgxConn) Close(ctx context.Context) {
-	c.Conn.Close(ctx)
-}
-
 func mapError(err error) error {
 	if errors.Is(err, pgx.ErrNoRows) {
 		return schemalog.ErrNoRows
diff --git a/pkg/wal/processor/webhook/subscription/store/postgres/pg_subscription_store.go b/pkg/wal/processor/webhook/subscription/store/postgres/pg_subscription_store.go
index fa3e09f..363caf4 100644
--- a/pkg/wal/processor/webhook/subscription/store/postgres/pg_subscription_store.go
+++ b/pkg/wal/processor/webhook/subscription/store/postgres/pg_subscription_store.go
@@ -6,14 +6,13 @@ import (
 	"context"
 	"fmt"
 
-	"github.com/jackc/pgx/v5"
-
+	pglib "github.com/xataio/pgstream/internal/postgres"
 	loglib "github.com/xataio/pgstream/pkg/log"
 	"github.com/xataio/pgstream/pkg/wal/processor/webhook/subscription"
 )
 
 type Store struct {
-	conn   *pgx.Conn
+	conn   pglib.Querier
 	logger loglib.Logger
 }
 
@@ -22,17 +21,12 @@ type Option func(*Store)
 const subscriptionsTable = "webhook_subscriptions"
 
 func NewSubscriptionStore(ctx context.Context, url string, opts ...Option) (*Store, error) {
-	pgCfg, err := pgx.ParseConfig(url)
-	if err != nil {
-		return nil, err
-	}
-	pgConn, err := pgx.ConnectConfig(ctx, pgCfg)
+	pgpool, err := pglib.NewConnPool(ctx, url)
 	if err != nil {
-		return nil, fmt.Errorf("create postgres client: %w", err)
+		return nil, fmt.Errorf("create postgres connection pool: %w", err)
 	}
-
 	ss := &Store{
-		conn: pgConn,
+		conn: pgpool,
 	}
 
 	for _, opt := range opts {
diff --git a/pkg/wal/replication/postgres/pg_replication_handler.go b/pkg/wal/replication/postgres/pg_replication_handler.go
index 4c8c348..40aab93 100644
--- a/pkg/wal/replication/postgres/pg_replication_handler.go
+++ b/pkg/wal/replication/postgres/pg_replication_handler.go
@@ -18,16 +18,11 @@ type Handler struct {
 
 	pgReplicationConn     pgReplicationConn
 	pgReplicationSlotName string
-	pgConnBuilder         func() (pgConn, error)
+	pgConnBuilder         func() (pglib.Querier, error)
 
 	lsnParser replication.LSNParser
 }
 
-type pgConn interface {
-	QueryRow(ctx context.Context, query string, args ...any) pglib.Row
-	Close(context.Context) error
-}
-
 type pgReplicationConn interface {
 	IdentifySystem(ctx context.Context) (pglib.IdentifySystemResult, error)
 	StartReplication(ctx context.Context, cfg pglib.ReplicationConfig) error
@@ -63,7 +58,7 @@ var pluginArguments = []string{
 
 // NewHandler returns a new postgres replication handler for the database on input.
 func NewHandler(ctx context.Context, cfg Config, opts ...Option) (*Handler, error) {
-	connBuilder := func() (pgConn, error) {
+	connBuilder := func() (pglib.Querier, error) {
 		return pglib.NewConn(ctx, cfg.PostgresURL)
 	}
 
@@ -221,7 +216,7 @@ func (h *Handler) Close() error {
 // getRestartLSN returns the absolute earliest possible LSN we can support. If
 // the consumer's LSN is earlier than this, we cannot (easily) catch the
 // consumer back up.
-func (h *Handler) getRestartLSN(ctx context.Context, conn pgConn, slotName string) (replication.LSN, error) {
+func (h *Handler) getRestartLSN(ctx context.Context, conn pglib.Querier, slotName string) (replication.LSN, error) {
 	var restartLSN string
 	err := conn.QueryRow(
 		ctx,
@@ -237,7 +232,7 @@ func (h *Handler) getRestartLSN(ctx context.Context, conn pgConn, slotName strin
 
 // getLastSyncedLSN gets the `confirmed_flush_lsn` from PG. This is the last LSN
 // that the consumer confirmed it had completed.
-func (h *Handler) getLastSyncedLSN(ctx context.Context, conn pgConn) (replication.LSN, error) {
+func (h *Handler) getLastSyncedLSN(ctx context.Context, conn pglib.Querier) (replication.LSN, error) {
 	var confirmedFlushLSN string
 	err := conn.QueryRow(ctx, `select confirmed_flush_lsn from pg_replication_slots where slot_name=$1`, h.pgReplicationSlotName).Scan(&confirmedFlushLSN)
 	if err != nil {


Test patch:
diff --git a/pkg/schemalog/postgres/helper_test.go b/pkg/schemalog/postgres/helper_test.go
index c74c4ff..3eb9455 100644
--- a/pkg/schemalog/postgres/helper_test.go
+++ b/pkg/schemalog/postgres/helper_test.go
@@ -3,33 +3,12 @@
 package postgres
 
 import (
-	"context"
 	"fmt"
 
-	"github.com/jackc/pgx/v5"
-	"github.com/jackc/pgx/v5/pgconn"
 	"github.com/rs/xid"
 	"github.com/xataio/pgstream/pkg/schemalog"
 )
 
-type mockQuerier struct {
-	queryRowFn func(context.Context, string, ...any) pgx.Row
-	execFn     func(context.Context, string, ...any) (pgconn.CommandTag, error)
-	closeFn    func(context.Context)
-}
-
-func (m *mockQuerier) QueryRow(ctx context.Context, sql string, args ...any) pgx.Row {
-	return m.queryRowFn(ctx, sql, args...)
-}
-
-func (m *mockQuerier) Exec(ctx context.Context, sql string, args ...any) (pgconn.CommandTag, error) {
-	return m.execFn(ctx, sql, args...)
-}
-
-func (m *mockQuerier) Close(ctx context.Context) {
-	m.closeFn(ctx)
-}
-
 type mockRow struct {
 	logEntry *schemalog.LogEntry
 	scanFn   func(args ...any) error
diff --git a/pkg/schemalog/postgres/pg_schemalog_store_test.go b/pkg/schemalog/postgres/pg_schemalog_store_test.go
index 6eddf17..221fbd3 100644
--- a/pkg/schemalog/postgres/pg_schemalog_store_test.go
+++ b/pkg/schemalog/postgres/pg_schemalog_store_test.go
@@ -12,6 +12,8 @@ import (
 	"github.com/jackc/pgx/v5/pgconn"
 	"github.com/rs/xid"
 	"github.com/stretchr/testify/require"
+	pglib "github.com/xataio/pgstream/internal/postgres"
+	pgmocks "github.com/xataio/pgstream/internal/postgres/mocks"
 	"github.com/xataio/pgstream/pkg/schemalog"
 )
 
@@ -27,7 +29,7 @@ func TestStore_Fetch(t *testing.T) {
 
 	tests := []struct {
 		name      string
-		querier   querier
+		querier   pglib.Querier
 		ackedOnly bool
 
 		wantLogEntry *schemalog.LogEntry
@@ -35,8 +37,8 @@ func TestStore_Fetch(t *testing.T) {
 	}{
 		{
 			name: "ok - without acked",
-			querier: &mockQuerier{
-				queryRowFn: func(_ context.Context, query string, args ...any) pgx.Row {
+			querier: &pgmocks.Querier{
+				QueryRowFn: func(_ context.Context, query string, args ...any) pglib.Row {
 					require.Len(t, args, 1)
 					require.Equal(t, args[0], testSchema)
 					require.Equal(t,
@@ -51,8 +53,8 @@ func TestStore_Fetch(t *testing.T) {
 		},
 		{
 			name: "ok - with acked",
-			querier: &mockQuerier{
-				queryRowFn: func(_ context.Context, query string, args ...any) pgx.Row {
+			querier: &pgmocks.Querier{
+				QueryRowFn: func(_ context.Context, query string, args ...any) pglib.Row {
 					require.Len(t, args, 1)
 					require.Equal(t, args[0], testSchema)
 					require.Equal(t,
@@ -68,8 +70,8 @@ func TestStore_Fetch(t *testing.T) {
 		},
 		{
 			name: "error - querying rows",
-			querier: &mockQuerier{
-				queryRowFn: func(context.Context, string, ...any) pgx.Row {
+			querier: &pgmocks.Querier{
+				QueryRowFn: func(_ context.Context, query string, args ...any) pglib.Row {
 					return &mockRow{scanFn: func(...any) error { return errTest }}
 				},
 			},
@@ -106,22 +108,22 @@ func TestStore_Ack(t *testing.T) {
 
 	tests := []struct {
 		name     string
-		querier  querier
+		querier  pglib.Querier
 		logEntry *schemalog.LogEntry
 
 		wantErr error
 	}{
 		{
 			name: "ok",
-			querier: &mockQuerier{
-				execFn: func(_ context.Context, query string, args ...any) (pgconn.CommandTag, error) {
+			querier: &pgmocks.Querier{
+				ExecFn: func(_ context.Context, query string, args ...any) (pglib.CommandTag, error) {
 					require.Len(t, args, 2)
 					require.Equal(t, args[0], testID.String())
 					require.Equal(t, args[1], testSchema)
 					require.Equal(t,
 						fmt.Sprintf(`update %s.%s set acked = true where id = $1 and schema_name = $2`, schemalog.SchemaName, schemalog.TableName),
 						query)
-					return pgconn.NewCommandTag("1"), nil
+					return pglib.CommandTag{CommandTag: pgconn.NewCommandTag("1")}, nil
 				},
 			},
 			logEntry: testLogEntry,
@@ -130,9 +132,9 @@ func TestStore_Ack(t *testing.T) {
 		},
 		{
 			name: "error - executing update query",
-			querier: &mockQuerier{
-				execFn: func(_ context.Context, query string, args ...any) (pgconn.CommandTag, error) {
-					return pgconn.NewCommandTag(""), errTest
+			querier: &pgmocks.Querier{
+				ExecFn: func(_ context.Context, query string, args ...any) (pglib.CommandTag, error) {
+					return pglib.CommandTag{CommandTag: pgconn.NewCommandTag("")}, errTest
 				},
 			},
 			logEntry: testLogEntry,
diff --git a/pkg/wal/replication/postgres/pg_replication_handler_test.go b/pkg/wal/replication/postgres/pg_replication_handler_test.go
index 39fc45f..289f736 100644
--- a/pkg/wal/replication/postgres/pg_replication_handler_test.go
+++ b/pkg/wal/replication/postgres/pg_replication_handler_test.go
@@ -23,7 +23,7 @@ func TestHandler_StartReplication(t *testing.T) {
 	tests := []struct {
 		name            string
 		replicationConn pgReplicationConn
-		connBuilder     func() (pgConn, error)
+		connBuilder     func() (pglib.Querier, error)
 		slotName        string
 
 		wantErr error
@@ -48,8 +48,8 @@ func TestHandler_StartReplication(t *testing.T) {
 					return nil
 				},
 			},
-			connBuilder: func() (pgConn, error) {
-				return &pgmocks.Conn{
+			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)
@@ -89,8 +89,8 @@ func TestHandler_StartReplication(t *testing.T) {
 					return nil
 				},
 			},
-			connBuilder: func() (pgConn, error) {
-				return &pgmocks.Conn{
+			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)
@@ -130,8 +130,8 @@ func TestHandler_StartReplication(t *testing.T) {
 					return nil
 				},
 			},
-			connBuilder: func() (pgConn, error) {
-				return &pgmocks.Conn{
+			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], defaultSlot)
@@ -171,7 +171,7 @@ func TestHandler_StartReplication(t *testing.T) {
 					}, nil
 				},
 			},
-			connBuilder: func() (pgConn, error) {
+			connBuilder: func() (pglib.Querier, error) {
 				return nil, errTest
 			},
 			slotName: testSlot,
@@ -188,8 +188,8 @@ func TestHandler_StartReplication(t *testing.T) {
 					}, nil
 				},
 			},
-			connBuilder: func() (pgConn, error) {
-				return &pgmocks.Conn{
+			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)
@@ -217,8 +217,8 @@ func TestHandler_StartReplication(t *testing.T) {
 					}, nil
 				},
 			},
-			connBuilder: func() (pgConn, error) {
-				return &pgmocks.Conn{
+			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)
@@ -251,8 +251,8 @@ func TestHandler_StartReplication(t *testing.T) {
 					return errTest
 				},
 			},
-			connBuilder: func() (pgConn, error) {
-				return &pgmocks.Conn{
+			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)
@@ -289,8 +289,8 @@ func TestHandler_StartReplication(t *testing.T) {
 					return errTest
 				},
 			},
-			connBuilder: func() (pgConn, error) {
-				return &pgmocks.Conn{
+			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)
@@ -410,15 +410,15 @@ func TestHandler_GetReplicationLag(t *testing.T) {
 
 	tests := []struct {
 		name        string
-		connBuilder func() (pgConn, error)
+		connBuilder func() (pglib.Querier, error)
 
 		wantLag int64
 		wantErr error
 	}{
 		{
 			name: "ok",
-			connBuilder: func() (pgConn, error) {
-				return &pgmocks.Conn{
+			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)
@@ -438,7 +438,7 @@ func TestHandler_GetReplicationLag(t *testing.T) {
 		},
 		{
 			name: "error - building connection",
-			connBuilder: func() (pgConn, error) {
+			connBuilder: func() (pglib.Querier, error) {
 				return nil, errTest
 			},
 
@@ -447,8 +447,8 @@ func TestHandler_GetReplicationLag(t *testing.T) {
 		},
 		{
 			name: "error - getting lag",
-			connBuilder: func() (pgConn, error) {
-				return &pgmocks.Conn{
+			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 }}
 					},