Skip to content

Commit 5c73eb9

Browse files
Amogh-BharadwajAmogh-Bharadwaj
committed
chore: add tests
1 parent 65e0afb commit 5c73eb9

2 files changed

Lines changed: 164 additions & 10 deletions

File tree

‎flow/cmd/row_counts.go‎

Lines changed: 6 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -25,29 +25,31 @@ const (
2525
countTimeoutSeconds = 30
2626
)
2727

28+
// GetMirrorRowCounts returns per-table source and destination row counts for a
29+
// Postgres-to-Postgres CDC mirror, so the caller can compare the two sides and spot drift.
30+
// Both peers must be Postgres; query replication (QRep) mirrors are not supported. Tables can
31+
// be filtered via req.SourceTables. Small tables are counted exactly while large tables use an
32+
// approximate estimate; a count of -1 signals the value could not be obtained (missing table or timeout).
2833
func (h *FlowRequestHandler) GetMirrorRowCounts(
2934
ctx context.Context,
3035
req *protos.RowCountRequest,
3136
) (*protos.RowCountResponse, APIError) {
3237
ctx, cancel := context.WithTimeout(ctx, 5*time.Minute)
3338
defer cancel()
3439

35-
// Validate mirror exists and is CDC
3640
isCdc, err := h.isCDCFlow(ctx, req.FlowJobName)
3741
if err != nil {
3842
return nil, NewInternalApiError(fmt.Errorf("unable to check flow type: %w", err))
3943
}
4044
if !isCdc {
41-
return nil, NewInvalidArgumentApiError(fmt.Errorf("row count validation is only supported for CDC mirrors"))
45+
return nil, NewInvalidArgumentApiError(fmt.Errorf("row count validation is not supported for query replication mirrors"))
4246
}
4347

44-
// Get flow config to extract table mappings and peer names
4548
config, err := h.getFlowConfigFromCatalog(ctx, req.FlowJobName)
4649
if err != nil {
4750
return nil, NewInternalApiError(fmt.Errorf("unable to get flow config: %w", err))
4851
}
4952

50-
// Verify both peers are Postgres
5153
srcType, err := connectors.LoadPeerType(ctx, h.pool, config.SourceName)
5254
if err != nil {
5355
return nil, NewInternalApiError(fmt.Errorf("unable to load source peer type: %w", err))
@@ -63,7 +65,6 @@ func (h *FlowRequestHandler) GetMirrorRowCounts(
6365
return nil, NewInvalidArgumentApiError(fmt.Errorf("destination peer %s is not PostgreSQL", config.DestinationName))
6466
}
6567

66-
// Filter table mappings if source_tables filter is specified
6768
tableMappings := config.TableMappings
6869
if len(req.SourceTables) > 0 {
6970
filterSet := make(map[string]struct{}, len(req.SourceTables))
@@ -83,7 +84,6 @@ func (h *FlowRequestHandler) GetMirrorRowCounts(
8384
return &protos.RowCountResponse{TableCounts: nil}, nil
8485
}
8586

86-
// Open connections to source and destination
8787
srcConn, srcClose, err := connectors.GetByNameAs[*connpostgres.PostgresConnector](ctx, nil, h.pool, config.SourceName)
8888
if err != nil {
8989
return nil, NewInternalApiError(fmt.Errorf("failed to connect to source peer: %w", err))
@@ -96,7 +96,6 @@ func (h *FlowRequestHandler) GetMirrorRowCounts(
9696
}
9797
defer dstClose(ctx)
9898

99-
// Build result slice
10099
results := make([]*protos.TableRowCount, len(tableMappings))
101100
for i, tm := range tableMappings {
102101
results[i] = &protos.TableRowCount{
@@ -105,8 +104,6 @@ func (h *FlowRequestHandler) GetMirrorRowCounts(
105104
}
106105
}
107106

108-
// Query source and destination in parallel,
109-
// but serialize queries within each peer since pgx.Conn is not concurrent-safe.
110107
g, gCtx := errgroup.WithContext(ctx)
111108

112109
g.Go(func() error {
@@ -162,7 +159,6 @@ func getTableRowCount(ctx context.Context, conn *pgx.Conn, tableIdentifier strin
162159
}
163160
qualifiedName := parsed.String()
164161

165-
// Get table size to decide strategy
166162
var tableSizeBytes pgtype.Int8
167163
if err := conn.QueryRow(ctx,
168164
"SELECT pg_total_relation_size(to_regclass($1))", qualifiedName,

‎flow/e2e/api_test.go‎

Lines changed: 158 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -219,6 +219,21 @@ func (s APITestSuite) waitForFlowDropped(env WorkflowRun, flowJobName string) {
219219
})
220220
}
221221

222+
func (s APITestSuite) waitForDestinationRowCount(env WorkflowRun, dstTable string, expected int64) {
223+
s.t.Helper()
224+
225+
EnvWaitFor(s.t, env, 3*time.Minute, fmt.Sprintf("wait for %d rows in %s", expected, dstTable), func() bool {
226+
var count int64
227+
if err := s.pg.Conn().QueryRow(
228+
s.t.Context(), "SELECT COUNT(*) FROM "+dstTable,
229+
).Scan(&count); err != nil {
230+
s.t.Log(err)
231+
return false
232+
}
233+
return count == expected
234+
})
235+
}
236+
222237
func testApi[TSource SuiteSource](
223238
t *testing.T,
224239
setup func(*testing.T, string) (TSource, error),
@@ -3805,3 +3820,146 @@ func (s APITestSuite) TestSnapshotNullPartitionKey() {
38053820
EnvWaitForFinished(s.t, env, 3*time.Minute)
38063821
EnvWaitForCount(env, s.ch, "all 4 rows including nulls should be snapshotted", tableName, "*", 4)
38073822
}
3823+
3824+
func (s APITestSuite) TestGetMirrorRowCounts() {
3825+
_, ok := s.source.(*PostgresSource)
3826+
if !ok {
3827+
s.t.Skip("row count validation is only supported for PostgreSQL source and destination")
3828+
}
3829+
3830+
srcTable1 := AttachSchema(s, "rowcount_src1")
3831+
dstTable1 := AttachSchema(s, "rowcount_dst1")
3832+
srcTable2 := AttachSchema(s, "rowcount_src2")
3833+
dstTable2 := AttachSchema(s, "rowcount_dst2")
3834+
srcTable3 := AttachSchema(s, "rowcount_src3")
3835+
dstTable3 := AttachSchema(s, "rowcount_dst3")
3836+
3837+
for _, tbl := range []string{srcTable1, srcTable2, srcTable3} {
3838+
require.NoError(s.t, s.source.Exec(s.t.Context(),
3839+
fmt.Sprintf("CREATE TABLE %s(id int primary key, val text)", tbl)))
3840+
}
3841+
for _, tbl := range []string{dstTable1, dstTable2, dstTable3} {
3842+
_, err := s.pg.Conn().Exec(s.t.Context(),
3843+
fmt.Sprintf("CREATE TABLE %s(id int primary key, val text)", tbl))
3844+
require.NoError(s.t, err)
3845+
}
3846+
3847+
require.NoError(s.t, s.source.Exec(s.t.Context(),
3848+
fmt.Sprintf("INSERT INTO %s(id, val) SELECT g, 'v'||g FROM generate_series(1,3) g", srcTable1)))
3849+
require.NoError(s.t, s.source.Exec(s.t.Context(),
3850+
fmt.Sprintf("INSERT INTO %s(id, val) SELECT g, 'v'||g FROM generate_series(1,5) g", srcTable2)))
3851+
3852+
flowJobName := "get_row_counts_" + s.suffix
3853+
flowConnConfig := &protos.FlowConnectionConfigs{
3854+
FlowJobName: flowJobName,
3855+
SourceName: s.source.GeneratePeer(s.t).Name,
3856+
DestinationName: s.pg.GeneratePeer(s.t).Name,
3857+
TableMappings: []*protos.TableMapping{
3858+
{SourceTableIdentifier: srcTable1, DestinationTableIdentifier: dstTable1},
3859+
{SourceTableIdentifier: srcTable2, DestinationTableIdentifier: dstTable2},
3860+
{SourceTableIdentifier: srcTable3, DestinationTableIdentifier: dstTable3},
3861+
},
3862+
DoInitialSnapshot: true,
3863+
System: protos.TypeSystem_PG,
3864+
IdleTimeoutSeconds: 15,
3865+
Version: shared.InternalVersion_Latest,
3866+
}
3867+
3868+
response, err := s.CreateCDCFlow(s.t.Context(), &protos.CreateCDCFlowRequest{ConnectionConfigs: flowConnConfig})
3869+
require.NoError(s.t, err)
3870+
require.NotNil(s.t, response)
3871+
3872+
tc := NewTemporalClient(s.t)
3873+
env, err := GetPeerflow(s.t.Context(), s.catalog, tc, flowJobName)
3874+
require.NoError(s.t, err)
3875+
SetupCDCFlowStatusQuery(s.t, env, flowConnConfig)
3876+
EnvWaitFor(s.t, env, 3*time.Minute, "wait for initial load to finish", func() bool {
3877+
return env.GetFlowStatus(s.t) == protos.FlowStatus_STATUS_RUNNING
3878+
})
3879+
s.waitForDestinationRowCount(env, dstTable1, 3)
3880+
s.waitForDestinationRowCount(env, dstTable2, 5)
3881+
s.waitForDestinationRowCount(env, dstTable3, 0)
3882+
3883+
rowCounts, err := s.GetMirrorRowCounts(s.t.Context(), &protos.RowCountRequest{FlowJobName: flowJobName})
3884+
require.NoError(s.t, err)
3885+
require.NotNil(s.t, rowCounts)
3886+
require.Len(s.t, rowCounts.TableCounts, 3)
3887+
3888+
bySrc := make(map[string]*protos.TableRowCount, len(rowCounts.TableCounts))
3889+
for _, tc := range rowCounts.TableCounts {
3890+
bySrc[tc.SourceTable] = tc
3891+
}
3892+
require.Contains(s.t, bySrc, srcTable1)
3893+
require.Contains(s.t, bySrc, srcTable2)
3894+
require.Contains(s.t, bySrc, srcTable3)
3895+
3896+
require.Equal(s.t, dstTable1, bySrc[srcTable1].DestinationTable)
3897+
require.Equal(s.t, int64(3), bySrc[srcTable1].SourceCount)
3898+
require.Equal(s.t, int64(3), bySrc[srcTable1].DestinationCount)
3899+
require.False(s.t, bySrc[srcTable1].SourceIsApproximate)
3900+
require.False(s.t, bySrc[srcTable1].DestinationIsApproximate)
3901+
3902+
require.Equal(s.t, dstTable2, bySrc[srcTable2].DestinationTable)
3903+
require.Equal(s.t, int64(5), bySrc[srcTable2].SourceCount)
3904+
require.Equal(s.t, int64(5), bySrc[srcTable2].DestinationCount)
3905+
3906+
require.Equal(s.t, dstTable3, bySrc[srcTable3].DestinationTable)
3907+
require.Equal(s.t, int64(0), bySrc[srcTable3].SourceCount)
3908+
require.Equal(s.t, int64(0), bySrc[srcTable3].DestinationCount)
3909+
require.False(s.t, bySrc[srcTable3].SourceIsApproximate)
3910+
require.False(s.t, bySrc[srcTable3].DestinationIsApproximate)
3911+
3912+
filtered, err := s.GetMirrorRowCounts(s.t.Context(), &protos.RowCountRequest{
3913+
FlowJobName: flowJobName,
3914+
SourceTables: []string{srcTable2},
3915+
})
3916+
require.NoError(s.t, err)
3917+
require.Len(s.t, filtered.TableCounts, 1)
3918+
require.Equal(s.t, srcTable2, filtered.TableCounts[0].SourceTable)
3919+
require.Equal(s.t, int64(5), filtered.TableCounts[0].SourceCount)
3920+
require.Equal(s.t, int64(5), filtered.TableCounts[0].DestinationCount)
3921+
3922+
require.NoError(s.t, s.source.Exec(s.t.Context(),
3923+
fmt.Sprintf("INSERT INTO %s(id, val) VALUES (4,'v4'),(5,'v5')", srcTable1)))
3924+
s.waitForDestinationRowCount(env, dstTable1, 5)
3925+
3926+
afterCdc, err := s.GetMirrorRowCounts(s.t.Context(), &protos.RowCountRequest{
3927+
FlowJobName: flowJobName,
3928+
SourceTables: []string{srcTable1},
3929+
})
3930+
require.NoError(s.t, err)
3931+
require.Len(s.t, afterCdc.TableCounts, 1)
3932+
require.Equal(s.t, int64(5), afterCdc.TableCounts[0].SourceCount)
3933+
require.Equal(s.t, int64(5), afterCdc.TableCounts[0].DestinationCount)
3934+
3935+
env.Cancel(s.t.Context())
3936+
RequireEnvCanceled(s.t, env)
3937+
3938+
require.NoError(s.t, s.source.Exec(s.t.Context(), "DROP TABLE "+srcTable3))
3939+
3940+
missingCounts, err := s.GetMirrorRowCounts(s.t.Context(), &protos.RowCountRequest{
3941+
FlowJobName: flowJobName,
3942+
SourceTables: []string{srcTable3},
3943+
})
3944+
require.NoError(s.t, err)
3945+
require.Len(s.t, missingCounts.TableCounts, 1)
3946+
require.Equal(s.t, srcTable3, missingCounts.TableCounts[0].SourceTable)
3947+
require.Equal(s.t, int64(-1), missingCounts.TableCounts[0].SourceCount)
3948+
require.Equal(s.t, int64(0), missingCounts.TableCounts[0].DestinationCount)
3949+
}
3950+
3951+
func (s APITestSuite) TestGetMirrorRowCountsNonCDC() {
3952+
_, ok := s.source.(*PostgresSource)
3953+
if !ok {
3954+
s.t.Skip("row count validation is only supported for PostgreSQL source and destination")
3955+
}
3956+
3957+
_, err := s.GetMirrorRowCounts(s.t.Context(), &protos.RowCountRequest{
3958+
FlowJobName: "nonexistent_flow_" + s.suffix,
3959+
})
3960+
require.Error(s.t, err)
3961+
grpcStatus, ok := status.FromError(err)
3962+
require.True(s.t, ok)
3963+
require.Equal(s.t, codes.InvalidArgument, grpcStatus.Code())
3964+
require.Contains(s.t, grpcStatus.Message(), "not supported for query replication mirrors")
3965+
}

0 commit comments

Comments
 (0)