Skip to content

Commit 122e4c1

Browse files
Amogh-BharadwajAmogh-Bharadwaj
committed
fix(test): use dedicated source db
1 parent 4529a84 commit 122e4c1

1 file changed

Lines changed: 86 additions & 26 deletions

File tree

‎flow/e2e/pg_schema_dump_test.go‎

Lines changed: 86 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ package e2e
33
import (
44
"context"
55
"fmt"
6+
"testing"
67
"time"
78

89
"github.com/jackc/pgx/v5"
@@ -12,8 +13,66 @@ import (
1213
"github.com/PeerDB-io/peerdb/flow/internal"
1314
)
1415

16+
// setupDedicatedPgDumpSource creates a fresh database on the primary PG instance
17+
// to serve as the source for schema-dump tests. pg_dump dumps a whole database,
18+
// so sharing the source DB with other parallel tests caused the dump to include
19+
// every concurrent test's e2e_test_<suffix> schema, blowing past the workflow
20+
// SETUP timeout. A dedicated source DB keeps the dump scoped to this test.
21+
//
22+
// Returns a connection to the new DB, the registered source peer name, and the
23+
// schema name to use within it. All resources are cleaned up via t.Cleanup.
24+
func setupDedicatedPgDumpSource(t *testing.T, suffix string) (*pgx.Conn, string, string) {
25+
t.Helper()
26+
27+
srcCfg := internal.GetAncillaryPostgresConfigFromEnv()
28+
srcDBName := "e2e_pgdump_src_" + suffix
29+
srcSchema := "e2e_test_" + suffix
30+
31+
bootstrapStr := fmt.Sprintf("host=%s port=%d user=%s password=%s dbname=%s",
32+
srcCfg.Host, srcCfg.Port, srcCfg.User, srcCfg.Password, srcCfg.Database)
33+
bootstrap, err := pgx.Connect(t.Context(), bootstrapStr)
34+
require.NoError(t, err)
35+
_, err = bootstrap.Exec(t.Context(), "CREATE DATABASE "+srcDBName)
36+
require.NoError(t, err)
37+
bootstrap.Close(t.Context())
38+
39+
srcConnStr := fmt.Sprintf("host=%s port=%d user=%s password=%s dbname=%s",
40+
srcCfg.Host, srcCfg.Port, srcCfg.User, srcCfg.Password, srcDBName)
41+
srcConn, err := pgx.Connect(t.Context(), srcConnStr)
42+
require.NoError(t, err)
43+
_, err = srcConn.Exec(t.Context(), "CREATE SCHEMA "+srcSchema)
44+
require.NoError(t, err)
45+
46+
srcPeerCfg := internal.GetAncillaryPostgresConfigFromEnv()
47+
srcPeerCfg.Database = srcDBName
48+
srcPeerName := "pgdump_src_" + suffix
49+
CreatePeer(t, &protos.Peer{
50+
Name: srcPeerName,
51+
Type: protos.DBType_POSTGRES,
52+
Config: &protos.Peer_PostgresConfig{PostgresConfig: srcPeerCfg},
53+
})
54+
55+
t.Cleanup(func() {
56+
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
57+
defer cancel()
58+
srcConn.Close(ctx)
59+
dropConn, err := pgx.Connect(ctx, bootstrapStr)
60+
if err != nil {
61+
t.Logf("failed to connect for source DB cleanup: %v", err)
62+
return
63+
}
64+
defer dropConn.Close(ctx)
65+
if _, err := dropConn.Exec(ctx, "DROP DATABASE IF EXISTS "+srcDBName+" WITH (FORCE)"); err != nil {
66+
t.Logf("failed to drop source database %s: %v", srcDBName, err)
67+
}
68+
})
69+
70+
return srcConn, srcPeerName, srcSchema
71+
}
72+
1573
func (s PeerFlowE2ETestSuitePG) Test_PG_Schema_Dump_And_CDC() {
16-
srcSchema := "e2e_test_" + s.suffix
74+
srcConn, srcPeerName, srcSchema := setupDedicatedPgDumpSource(s.t, s.suffix)
75+
1776
dstDBName := "e2e_pgdump_" + s.suffix
1877

1978
// create destination database on the same PG instance
@@ -44,9 +103,6 @@ func (s PeerFlowE2ETestSuitePG) Test_PG_Schema_Dump_And_CDC() {
44103
require.NoError(s.t, err)
45104
s.t.Cleanup(func() { dstConn.Close(s.t.Context()) })
46105

47-
// create the source schema in the destination database (pg_dump will create it, but we need it for the peer)
48-
// Actually, pg_dump --schema-only will create the schema, so we don't need to pre-create it.
49-
50106
// create a destination peer pointing to the new database
51107
dstPeerCfg := internal.GetAncillaryPostgresConfigFromEnv()
52108
dstPeerCfg.Database = dstDBName
@@ -63,13 +119,13 @@ func (s PeerFlowE2ETestSuitePG) Test_PG_Schema_Dump_And_CDC() {
63119
// --- set up rich schema on source ---
64120

65121
// create custom enum type in the source schema
66-
_, err = s.Conn().Exec(s.t.Context(), fmt.Sprintf(
122+
_, err = srcConn.Exec(s.t.Context(), fmt.Sprintf(
67123
"CREATE TYPE %s.color AS ENUM ('red', 'green', 'blue')", srcSchema))
68124
require.NoError(s.t, err)
69125

70126
// create parent table with various column types
71127
parentTable := srcSchema + ".parent_tbl"
72-
_, err = s.Conn().Exec(s.t.Context(), fmt.Sprintf(`
128+
_, err = srcConn.Exec(s.t.Context(), fmt.Sprintf(`
73129
CREATE TABLE %s (
74130
id SERIAL PRIMARY KEY,
75131
name TEXT NOT NULL,
@@ -81,13 +137,13 @@ func (s PeerFlowE2ETestSuitePG) Test_PG_Schema_Dump_And_CDC() {
81137
require.NoError(s.t, err)
82138

83139
// create unique index on parent
84-
_, err = s.Conn().Exec(s.t.Context(), fmt.Sprintf(
140+
_, err = srcConn.Exec(s.t.Context(), fmt.Sprintf(
85141
"CREATE UNIQUE INDEX idx_parent_name ON %s (name)", parentTable))
86142
require.NoError(s.t, err)
87143

88144
// create child table with foreign key referencing parent
89145
childTable := srcSchema + ".child_tbl"
90-
_, err = s.Conn().Exec(s.t.Context(), fmt.Sprintf(`
146+
_, err = srcConn.Exec(s.t.Context(), fmt.Sprintf(`
91147
CREATE TABLE %s (
92148
id SERIAL PRIMARY KEY,
93149
parent_id INT NOT NULL REFERENCES %s(id),
@@ -97,13 +153,13 @@ func (s PeerFlowE2ETestSuitePG) Test_PG_Schema_Dump_And_CDC() {
97153
require.NoError(s.t, err)
98154

99155
// create btree index on child
100-
_, err = s.Conn().Exec(s.t.Context(), fmt.Sprintf(
156+
_, err = srcConn.Exec(s.t.Context(), fmt.Sprintf(
101157
"CREATE INDEX idx_child_parent ON %s (parent_id)", childTable))
102158
require.NoError(s.t, err)
103159

104160
// insert initial data for snapshot
105161
for i := 1; i <= 5; i++ {
106-
_, err = s.Conn().Exec(s.t.Context(), fmt.Sprintf(
162+
_, err = srcConn.Exec(s.t.Context(), fmt.Sprintf(
107163
"INSERT INTO %s (name, color, score, metadata) VALUES ($1, $2, $3, $4)",
108164
parentTable),
109165
fmt.Sprintf("item_%d", i),
@@ -115,7 +171,7 @@ func (s PeerFlowE2ETestSuitePG) Test_PG_Schema_Dump_And_CDC() {
115171
}
116172

117173
for i := 1; i <= 10; i++ {
118-
_, err = s.Conn().Exec(s.t.Context(), fmt.Sprintf(
174+
_, err = srcConn.Exec(s.t.Context(), fmt.Sprintf(
119175
"INSERT INTO %s (parent_id, value, tags) VALUES ($1, $2, $3)",
120176
childTable),
121177
(i%5)+1,
@@ -145,7 +201,7 @@ func (s PeerFlowE2ETestSuitePG) Test_PG_Schema_Dump_And_CDC() {
145201
DestinationTableIdentifier: dstChild,
146202
},
147203
},
148-
SourceName: GeneratePostgresPeer(s.t).Name,
204+
SourceName: srcPeerName,
149205
MaxBatchSize: 100,
150206
DoInitialSnapshot: true,
151207
System: protos.TypeSystem_PG,
@@ -223,7 +279,7 @@ func (s PeerFlowE2ETestSuitePG) Test_PG_Schema_Dump_And_CDC() {
223279

224280
// insert more parent rows
225281
for i := 6; i <= 8; i++ {
226-
_, err = s.Conn().Exec(s.t.Context(), fmt.Sprintf(
282+
_, err = srcConn.Exec(s.t.Context(), fmt.Sprintf(
227283
"INSERT INTO %s (name, color, score, metadata) VALUES ($1, $2, $3, $4)",
228284
parentTable),
229285
fmt.Sprintf("item_%d", i),
@@ -236,7 +292,7 @@ func (s PeerFlowE2ETestSuitePG) Test_PG_Schema_Dump_And_CDC() {
236292

237293
// insert more child rows
238294
for i := 11; i <= 15; i++ {
239-
_, err = s.Conn().Exec(s.t.Context(), fmt.Sprintf(
295+
_, err = srcConn.Exec(s.t.Context(), fmt.Sprintf(
240296
"INSERT INTO %s (parent_id, value, tags) VALUES ($1, $2, $3)",
241297
childTable),
242298
(i%5)+1,
@@ -263,14 +319,14 @@ func (s PeerFlowE2ETestSuitePG) Test_PG_Schema_Dump_And_CDC() {
263319
// verify data integrity: compare actual row content
264320
// query source and destination and compare
265321
var srcParentCount, dstParentCount int64
266-
err = s.Conn().QueryRow(s.t.Context(), fmt.Sprintf("SELECT COUNT(*) FROM %s", srcParent)).Scan(&srcParentCount)
322+
err = srcConn.QueryRow(s.t.Context(), fmt.Sprintf("SELECT COUNT(*) FROM %s", srcParent)).Scan(&srcParentCount)
267323
require.NoError(s.t, err)
268324
err = dstConn.QueryRow(s.t.Context(), fmt.Sprintf("SELECT COUNT(*) FROM %s", dstParent)).Scan(&dstParentCount)
269325
require.NoError(s.t, err)
270326
require.Equal(s.t, srcParentCount, dstParentCount, "parent table row counts should match")
271327

272328
var srcChildCount, dstChildCount int64
273-
err = s.Conn().QueryRow(s.t.Context(), fmt.Sprintf("SELECT COUNT(*) FROM %s", srcChild)).Scan(&srcChildCount)
329+
err = srcConn.QueryRow(s.t.Context(), fmt.Sprintf("SELECT COUNT(*) FROM %s", srcChild)).Scan(&srcChildCount)
274330
require.NoError(s.t, err)
275331
err = dstConn.QueryRow(s.t.Context(), fmt.Sprintf("SELECT COUNT(*) FROM %s", dstChild)).Scan(&dstChildCount)
276332
require.NoError(s.t, err)
@@ -290,9 +346,10 @@ func (s PeerFlowE2ETestSuitePG) Test_PG_Schema_Dump_And_CDC() {
290346
// Also verifies initial load + CDC into the dumped table on the destination
291347
// cluster, so we know the table is usable end-to-end.
292348
func (s PeerFlowE2ETestSuitePG) Test_PG_Schema_Dump_No_Owner_No_Privileges() {
349+
srcConn, srcPeerName, srcSchema := setupDedicatedPgDumpSource(s.t, s.suffix)
350+
293351
dstCfg := internal.GetSecondaryPostgresConfigFromEnv()
294352

295-
srcSchema := "e2e_test_" + s.suffix
296353
roleName := "peerdb_owner_role_" + s.suffix
297354
tableName := "owned_tbl"
298355
qualified := fmt.Sprintf("%s.%s", srcSchema, tableName)
@@ -310,28 +367,31 @@ func (s PeerFlowE2ETestSuitePG) Test_PG_Schema_Dump_No_Owner_No_Privileges() {
310367
"SELECT EXISTS(SELECT 1 FROM pg_roles WHERE rolname=$1)", roleName).Scan(&roleExistsOnDst))
311368
require.False(s.t, roleExistsOnDst, "role %s unexpectedly exists on destination", roleName)
312369

313-
// create role + owned/granted table on source with seed rows
314-
_, err = s.Conn().Exec(s.t.Context(), fmt.Sprintf("CREATE ROLE %s LOGIN PASSWORD 'pw'", roleName))
370+
// create role + owned/granted table on source with seed rows.
371+
// the role is cluster-wide; the table lives in the dedicated source DB
372+
// and will go away when that DB is dropped during cleanup.
373+
_, err = srcConn.Exec(s.t.Context(), fmt.Sprintf("CREATE ROLE %s LOGIN PASSWORD 'pw'", roleName))
315374
require.NoError(s.t, err)
316-
_, err = s.Conn().Exec(s.t.Context(), fmt.Sprintf(
375+
_, err = srcConn.Exec(s.t.Context(), fmt.Sprintf(
317376
"CREATE TABLE %s (id SERIAL PRIMARY KEY, val TEXT)", qualified))
318377
require.NoError(s.t, err)
319-
_, err = s.Conn().Exec(s.t.Context(), fmt.Sprintf("ALTER TABLE %s OWNER TO %s", qualified, roleName))
378+
_, err = srcConn.Exec(s.t.Context(), fmt.Sprintf("ALTER TABLE %s OWNER TO %s", qualified, roleName))
320379
require.NoError(s.t, err)
321-
_, err = s.Conn().Exec(s.t.Context(), fmt.Sprintf("GRANT SELECT, INSERT ON %s TO %s", qualified, roleName))
380+
_, err = srcConn.Exec(s.t.Context(), fmt.Sprintf("GRANT SELECT, INSERT ON %s TO %s", qualified, roleName))
322381
require.NoError(s.t, err)
323382

324383
for i := 1; i <= 5; i++ {
325-
_, err = s.Conn().Exec(s.t.Context(),
384+
_, err = srcConn.Exec(s.t.Context(),
326385
fmt.Sprintf("INSERT INTO %s (val) VALUES ($1)", qualified),
327386
fmt.Sprintf("snap_%d", i))
328387
require.NoError(s.t, err)
329388
}
330389

390+
// role is cluster-wide so it outlives the dedicated source DB; drop it
391+
// against the shared source connection after the source DB is gone.
331392
s.t.Cleanup(func() {
332393
cleanupCtx, cleanupCancel := context.WithTimeout(context.Background(), 30*time.Second)
333394
defer cleanupCancel()
334-
_, _ = s.Conn().Exec(cleanupCtx, fmt.Sprintf("DROP TABLE IF EXISTS %s", qualified))
335395
_, _ = s.Conn().Exec(cleanupCtx, "DROP ROLE IF EXISTS "+roleName)
336396

337397
dropConn, err := pgx.Connect(cleanupCtx, dstConnStr)
@@ -358,7 +418,7 @@ func (s PeerFlowE2ETestSuitePG) Test_PG_Schema_Dump_No_Owner_No_Privileges() {
358418
SourceTableIdentifier: qualified,
359419
DestinationTableIdentifier: qualified,
360420
}},
361-
SourceName: GeneratePostgresPeer(s.t).Name,
421+
SourceName: srcPeerName,
362422
MaxBatchSize: 100,
363423
DoInitialSnapshot: true,
364424
System: protos.TypeSystem_PG,
@@ -382,7 +442,7 @@ func (s PeerFlowE2ETestSuitePG) Test_PG_Schema_Dump_No_Owner_No_Privileges() {
382442

383443
// CDC: insert more rows on source and wait for them on dst
384444
for i := 6; i <= 10; i++ {
385-
_, err = s.Conn().Exec(s.t.Context(),
445+
_, err = srcConn.Exec(s.t.Context(),
386446
fmt.Sprintf("INSERT INTO %s (val) VALUES ($1)", qualified),
387447
fmt.Sprintf("cdc_%d", i))
388448
EnvNoError(s.t, env, err)

0 commit comments

Comments
 (0)