Skip to content

Commit cb24492

Browse files
Amogh-BharadwajAmogh-Bharadwaj
committed
chore: add tests
1 parent 481df03 commit cb24492

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),
@@ -3728,3 +3743,146 @@ func (s APITestSuite) TestSnapshotNullPartitionKey() {
37283743
EnvWaitForFinished(s.t, env, 3*time.Minute)
37293744
EnvWaitForCount(env, s.ch, "all 4 rows including nulls should be snapshotted", tableName, "*", 4)
37303745
}
3746+
3747+
func (s APITestSuite) TestGetMirrorRowCounts() {
3748+
_, ok := s.source.(*PostgresSource)
3749+
if !ok {
3750+
s.t.Skip("row count validation is only supported for PostgreSQL source and destination")
3751+
}
3752+
3753+
srcTable1 := AttachSchema(s, "rowcount_src1")
3754+
dstTable1 := AttachSchema(s, "rowcount_dst1")
3755+
srcTable2 := AttachSchema(s, "rowcount_src2")
3756+
dstTable2 := AttachSchema(s, "rowcount_dst2")
3757+
srcTable3 := AttachSchema(s, "rowcount_src3")
3758+
dstTable3 := AttachSchema(s, "rowcount_dst3")
3759+
3760+
for _, tbl := range []string{srcTable1, srcTable2, srcTable3} {
3761+
require.NoError(s.t, s.source.Exec(s.t.Context(),
3762+
fmt.Sprintf("CREATE TABLE %s(id int primary key, val text)", tbl)))
3763+
}
3764+
for _, tbl := range []string{dstTable1, dstTable2, dstTable3} {
3765+
_, err := s.pg.Conn().Exec(s.t.Context(),
3766+
fmt.Sprintf("CREATE TABLE %s(id int primary key, val text)", tbl))
3767+
require.NoError(s.t, err)
3768+
}
3769+
3770+
require.NoError(s.t, s.source.Exec(s.t.Context(),
3771+
fmt.Sprintf("INSERT INTO %s(id, val) SELECT g, 'v'||g FROM generate_series(1,3) g", srcTable1)))
3772+
require.NoError(s.t, s.source.Exec(s.t.Context(),
3773+
fmt.Sprintf("INSERT INTO %s(id, val) SELECT g, 'v'||g FROM generate_series(1,5) g", srcTable2)))
3774+
3775+
flowJobName := "get_row_counts_" + s.suffix
3776+
flowConnConfig := &protos.FlowConnectionConfigs{
3777+
FlowJobName: flowJobName,
3778+
SourceName: s.source.GeneratePeer(s.t).Name,
3779+
DestinationName: s.pg.GeneratePeer(s.t).Name,
3780+
TableMappings: []*protos.TableMapping{
3781+
{SourceTableIdentifier: srcTable1, DestinationTableIdentifier: dstTable1},
3782+
{SourceTableIdentifier: srcTable2, DestinationTableIdentifier: dstTable2},
3783+
{SourceTableIdentifier: srcTable3, DestinationTableIdentifier: dstTable3},
3784+
},
3785+
DoInitialSnapshot: true,
3786+
System: protos.TypeSystem_PG,
3787+
IdleTimeoutSeconds: 15,
3788+
Version: shared.InternalVersion_Latest,
3789+
}
3790+
3791+
response, err := s.CreateCDCFlow(s.t.Context(), &protos.CreateCDCFlowRequest{ConnectionConfigs: flowConnConfig})
3792+
require.NoError(s.t, err)
3793+
require.NotNil(s.t, response)
3794+
3795+
tc := NewTemporalClient(s.t)
3796+
env, err := GetPeerflow(s.t.Context(), s.catalog, tc, flowJobName)
3797+
require.NoError(s.t, err)
3798+
SetupCDCFlowStatusQuery(s.t, env, flowConnConfig)
3799+
EnvWaitFor(s.t, env, 3*time.Minute, "wait for initial load to finish", func() bool {
3800+
return env.GetFlowStatus(s.t) == protos.FlowStatus_STATUS_RUNNING
3801+
})
3802+
s.waitForDestinationRowCount(env, dstTable1, 3)
3803+
s.waitForDestinationRowCount(env, dstTable2, 5)
3804+
s.waitForDestinationRowCount(env, dstTable3, 0)
3805+
3806+
rowCounts, err := s.GetMirrorRowCounts(s.t.Context(), &protos.RowCountRequest{FlowJobName: flowJobName})
3807+
require.NoError(s.t, err)
3808+
require.NotNil(s.t, rowCounts)
3809+
require.Len(s.t, rowCounts.TableCounts, 3)
3810+
3811+
bySrc := make(map[string]*protos.TableRowCount, len(rowCounts.TableCounts))
3812+
for _, tc := range rowCounts.TableCounts {
3813+
bySrc[tc.SourceTable] = tc
3814+
}
3815+
require.Contains(s.t, bySrc, srcTable1)
3816+
require.Contains(s.t, bySrc, srcTable2)
3817+
require.Contains(s.t, bySrc, srcTable3)
3818+
3819+
require.Equal(s.t, dstTable1, bySrc[srcTable1].DestinationTable)
3820+
require.Equal(s.t, int64(3), bySrc[srcTable1].SourceCount)
3821+
require.Equal(s.t, int64(3), bySrc[srcTable1].DestinationCount)
3822+
require.False(s.t, bySrc[srcTable1].SourceIsApproximate)
3823+
require.False(s.t, bySrc[srcTable1].DestinationIsApproximate)
3824+
3825+
require.Equal(s.t, dstTable2, bySrc[srcTable2].DestinationTable)
3826+
require.Equal(s.t, int64(5), bySrc[srcTable2].SourceCount)
3827+
require.Equal(s.t, int64(5), bySrc[srcTable2].DestinationCount)
3828+
3829+
require.Equal(s.t, dstTable3, bySrc[srcTable3].DestinationTable)
3830+
require.Equal(s.t, int64(0), bySrc[srcTable3].SourceCount)
3831+
require.Equal(s.t, int64(0), bySrc[srcTable3].DestinationCount)
3832+
require.False(s.t, bySrc[srcTable3].SourceIsApproximate)
3833+
require.False(s.t, bySrc[srcTable3].DestinationIsApproximate)
3834+
3835+
filtered, err := s.GetMirrorRowCounts(s.t.Context(), &protos.RowCountRequest{
3836+
FlowJobName: flowJobName,
3837+
SourceTables: []string{srcTable2},
3838+
})
3839+
require.NoError(s.t, err)
3840+
require.Len(s.t, filtered.TableCounts, 1)
3841+
require.Equal(s.t, srcTable2, filtered.TableCounts[0].SourceTable)
3842+
require.Equal(s.t, int64(5), filtered.TableCounts[0].SourceCount)
3843+
require.Equal(s.t, int64(5), filtered.TableCounts[0].DestinationCount)
3844+
3845+
require.NoError(s.t, s.source.Exec(s.t.Context(),
3846+
fmt.Sprintf("INSERT INTO %s(id, val) VALUES (4,'v4'),(5,'v5')", srcTable1)))
3847+
s.waitForDestinationRowCount(env, dstTable1, 5)
3848+
3849+
afterCdc, err := s.GetMirrorRowCounts(s.t.Context(), &protos.RowCountRequest{
3850+
FlowJobName: flowJobName,
3851+
SourceTables: []string{srcTable1},
3852+
})
3853+
require.NoError(s.t, err)
3854+
require.Len(s.t, afterCdc.TableCounts, 1)
3855+
require.Equal(s.t, int64(5), afterCdc.TableCounts[0].SourceCount)
3856+
require.Equal(s.t, int64(5), afterCdc.TableCounts[0].DestinationCount)
3857+
3858+
env.Cancel(s.t.Context())
3859+
RequireEnvCanceled(s.t, env)
3860+
3861+
require.NoError(s.t, s.source.Exec(s.t.Context(), "DROP TABLE "+srcTable3))
3862+
3863+
missingCounts, err := s.GetMirrorRowCounts(s.t.Context(), &protos.RowCountRequest{
3864+
FlowJobName: flowJobName,
3865+
SourceTables: []string{srcTable3},
3866+
})
3867+
require.NoError(s.t, err)
3868+
require.Len(s.t, missingCounts.TableCounts, 1)
3869+
require.Equal(s.t, srcTable3, missingCounts.TableCounts[0].SourceTable)
3870+
require.Equal(s.t, int64(-1), missingCounts.TableCounts[0].SourceCount)
3871+
require.Equal(s.t, int64(0), missingCounts.TableCounts[0].DestinationCount)
3872+
}
3873+
3874+
func (s APITestSuite) TestGetMirrorRowCountsNonCDC() {
3875+
_, ok := s.source.(*PostgresSource)
3876+
if !ok {
3877+
s.t.Skip("row count validation is only supported for PostgreSQL source and destination")
3878+
}
3879+
3880+
_, err := s.GetMirrorRowCounts(s.t.Context(), &protos.RowCountRequest{
3881+
FlowJobName: "nonexistent_flow_" + s.suffix,
3882+
})
3883+
require.Error(s.t, err)
3884+
grpcStatus, ok := status.FromError(err)
3885+
require.True(s.t, ok)
3886+
require.Equal(s.t, codes.InvalidArgument, grpcStatus.Code())
3887+
require.Contains(s.t, grpcStatus.Message(), "not supported for query replication mirrors")
3888+
}

0 commit comments

Comments
 (0)