diff --git a/flow/connectors/mysql/qrep.go b/flow/connectors/mysql/qrep.go index 86793762a..8c0e5532c 100644 --- a/flow/connectors/mysql/qrep.go +++ b/flow/connectors/mysql/qrep.go @@ -38,13 +38,7 @@ func (c *MySqlConnector) GetQRepPartitions( ) ([]*protos.QRepPartition, error) { if config.WatermarkColumn == "" || config.NumPartitionsOverride == 1 { // if no watermark column is specified, return a single partition - return []*protos.QRepPartition{ - { - PartitionId: utils.FullTablePartitionID, - Range: nil, - FullTablePartition: true, - }, - }, nil + return utils.FullTablePartition(), nil } if config.NumPartitionsOverride == 0 && config.NumRowsPerPartition == 0 { @@ -101,6 +95,10 @@ func (c *MySqlConnector) GetQRepPartitions( case *protos.PartitionRange_TimestampRange: time := lastRange.TimestampRange.End.AsTime() minVal = "'" + time.Format("2006-01-02 15:04:05.999999") + "'" + case *protos.PartitionRange_StringRange: + // resuming from last partition range is only possible for standalone QRepFlowWorkflow; + // this is a legacy feature and string partitioning is not supported + return nil, errors.New("resuming QRep by a string partition range is not supported") case *protos.PartitionRange_NullRange: // this case should never happen because we only add null partition for InitialCopyOnly replication // (so there shouldn't be a resume scenario with null range) @@ -155,7 +153,21 @@ func (c *MySqlConnector) GetQRepPartitions( if err != nil { return nil, fmt.Errorf("failed to convert partition maximum to qvalue: %w", err) } - if err := partitionHelper.AddPartitionsWithRange(val1.Value(), val2.Value(), numPartitions); err != nil { + + start, startIsString := val1.Value().(string) + end, endIsString := val2.Value().(string) + if startIsString && endIsString { + if isUUID, casing := detectUuidWithHexCasing(start, end); isUUID { + uuidPartitions, err := buildUuidStringPartitions(start, end, casing, numPartitions) + if err != nil { + return nil, fmt.Errorf("failed to add uuid string partitions: %w", err) + } + partitionHelper.AddPartitions(uuidPartitions) + } else { + c.logger.Info("string watermark column is not uuid, falling back to full table partition") + return utils.FullTablePartition(), nil + } + } else if err := partitionHelper.AddPartitionsWithRange(val1.Value(), val2.Value(), numPartitions); err != nil { return nil, fmt.Errorf("failed to add partitions: %w", err) } @@ -177,6 +189,9 @@ func supportsRangePartition(qkind types.QValueKind) bool { // temporal types case types.QValueKindDate, types.QValueKindTimestamp: return true + // string types + case types.QValueKindString: + return true default: return false } @@ -354,6 +369,19 @@ func (c *MySqlConnector) PullQRepRecords( case *protos.PartitionRange_TimestampRange: rangeStart = "'" + x.TimestampRange.Start.AsTime().Format("2006-01-02 15:04:05.999999") + "'" rangeEnd = "'" + x.TimestampRange.End.AsTime().Format("2006-01-02 15:04:05.999999") + "'" + case *protos.PartitionRange_StringRange: + rangeStart = "'" + mysql.Escape(x.StringRange.Start) + "'" + rangeEnd = "'" + mysql.Escape(x.StringRange.End) + "'" + if config.Query != "" { + // custom query is only possible for standalone QRepFlowWorkflow; + // this is a legacy feature and string partitioning is not supported + return 0, 0, errors.New("can't construct a string range partition for custom queries") + } + if !x.StringRange.EndInclusive { + queryTemplate = fmt.Sprintf( + "SELECT %[1]s FROM %[2]s WHERE %[3]s >= {{.start}} AND %[3]s < {{.end}}", + selectedColumns, parsedSrcTable.MySQL(), common.QuoteMySQLIdentifier(config.WatermarkColumn)) + } case *protos.PartitionRange_NullRange: if config.Query != "" { return 0, 0, errors.New("can't construct a null range partition for custom queries") diff --git a/flow/connectors/mysql/qrep_partition.go b/flow/connectors/mysql/qrep_partition.go new file mode 100644 index 000000000..6c165c2e6 --- /dev/null +++ b/flow/connectors/mysql/qrep_partition.go @@ -0,0 +1,120 @@ +package connmysql + +import ( + "fmt" + "math/big" + "regexp" + "strings" + + "github.com/google/uuid" + + "github.com/PeerDB-io/peerdb/flow/generated/protos" + "github.com/PeerDB-io/peerdb/flow/shared" +) + +// String watermark partitioning is currently only supported by the MySQL connector, +// so these helpers live in the MySQL connector subtree to prevent accidental usage +// by other connectors. + +var ( + uuidLowerRe = regexp.MustCompile(`^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$`) + uuidUpperRe = regexp.MustCompile(`^[0-9A-F]{8}-[0-9A-F]{4}-[0-9A-F]{4}-[0-9A-F]{4}-[0-9A-F]{12}$`) +) + +type hexCasing int + +const ( + hexCasingUnknown hexCasing = iota + hexCasingLower + hexCasingUpper +) + +// detectUuidWithHexCasing best-effort classifies a string watermark column by +// inspecting only its min and max bounds: it determines whether both are +// canonical UUIDs, and if so, the shared casing of their hex letters. +// +// Because only the bounds are examined, the classification can be wrong in +// the following rare cases: +// - non-UUID rows can exist between UUID-shaped min/max bounds; +// - non-boundary UUIDs contain a mix of upper and lower case +// - min/max bounds both don't contain hex, so casing blindly defaults to lower +// +// These cases may lead to partition skew, but will not affect correctness. +func detectUuidWithHexCasing(minVal string, maxVal string) (bool, hexCasing) { + switch { + case uuidLowerRe.MatchString(minVal) && uuidLowerRe.MatchString(maxVal): + return true, hexCasingLower + case uuidUpperRe.MatchString(minVal) && uuidUpperRe.MatchString(maxVal): + return true, hexCasingUpper + default: + return false, hexCasingUnknown + } +} + +// buildUuidStringPartitions splits a UUID string watermark column into partitions. +// The original minVal/maxVal are used for first/last partitions to ensure correctness. +func buildUuidStringPartitions( + minVal string, maxVal string, casing hexCasing, numPartitions int64, +) ([]*protos.QRepPartition, error) { + minInt, err := uuidToBigInt(minVal) + if err != nil { + return nil, fmt.Errorf("failed to convert min uuid to bigint: %w", err) + } + maxInt, err := uuidToBigInt(maxVal) + if err != nil { + return nil, fmt.Errorf("failed to convert max uuid to bigint: %w", err) + } + if minInt.Cmp(maxInt) > 0 { + return nil, fmt.Errorf("min uuid (%s) greater than max uuid (%s)", minVal, maxVal) + } + + var partitions []*protos.QRepPartition + start := minVal + step := shared.BigIntDivCeil(new(big.Int).Sub(maxInt, minInt), big.NewInt(numPartitions)) + for value := new(big.Int).Add(minInt, step); value.Cmp(maxInt) < 0; value.Add(value, step) { + end, err := bigIntToUUID(value, casing) + if err != nil { + return nil, fmt.Errorf("failed to convert bigint to uuid: %w", err) + } + partitions = append(partitions, createStringPartition(start, end, false)) + start = end + } + partitions = append(partitions, createStringPartition(start, maxVal, true)) + return partitions, nil +} + +func uuidToBigInt(s string) (*big.Int, error) { + u, err := uuid.Parse(s) + if err != nil { + return nil, err + } + return new(big.Int).SetBytes(u[:]), nil +} + +func bigIntToUUID(n *big.Int, casing hexCasing) (string, error) { + if n.BitLen() > 128 { + return "", fmt.Errorf("value does not fit in a UUID (%d bits)", n.BitLen()) + } + var b [16]byte + n.FillBytes(b[:]) + s := uuid.UUID(b).String() + if casing == hexCasingUpper { + s = strings.ToUpper(s) + } + return s, nil +} + +func createStringPartition(start string, end string, endInclusive bool) *protos.QRepPartition { + return &protos.QRepPartition{ + PartitionId: uuid.NewString(), + Range: &protos.PartitionRange{ + Range: &protos.PartitionRange_StringRange{ + StringRange: &protos.StringPartitionRange{ + Start: start, + End: end, + EndInclusive: endInclusive, + }, + }, + }, + } +} diff --git a/flow/connectors/mysql/qrep_partition_test.go b/flow/connectors/mysql/qrep_partition_test.go new file mode 100644 index 000000000..501f36c75 --- /dev/null +++ b/flow/connectors/mysql/qrep_partition_test.go @@ -0,0 +1,167 @@ +package connmysql + +import ( + "regexp" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/PeerDB-io/peerdb/flow/generated/protos" +) + +func TestUUIDToBigIntRoundTrip(t *testing.T) { + cases := []struct { + name string + uuid string + casing hexCasing + }{ + {"lower case v4", "f47ac10b-58cc-4372-a567-0e02b2c3d479", hexCasingLower}, + {"upper case v7", "017F22E2-79B0-7CC3-98C4-DC0C0C07398F", hexCasingUpper}, + {"zero", "00000000-0000-0000-0000-000000000000", hexCasingLower}, + {"max", "ffffffff-ffff-ffff-ffff-ffffffffffff", hexCasingLower}, + {"all digits", "01234567-8901-2345-6789-012345678901", hexCasingLower}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + num, err := uuidToBigInt(tc.uuid) + require.NoError(t, err) + uuid, err := bigIntToUUID(num, tc.casing) + require.NoError(t, err) + assert.Equal(t, tc.uuid, uuid) + }) + } +} + +func TestUUIDToBigIntInvalid(t *testing.T) { + for _, s := range []string{ + "", + "not-a-uuid", + "zzzzzzzz-zzzz-zzzz-zzzz-zzzzzzzzzzzz", + "01234567-8901-2345-6789-012345678901a", // 37 chars + } { + _, err := uuidToBigInt(s) + assert.Error(t, err) + } +} + +func TestDetectUUIDCase(t *testing.T) { + cases := []struct { + name string + minVal string + maxVal string + expectIsUUID bool + expectCasing hexCasing + }{ + {"both lower", "00000000-0000-0000-0000-000000000000", "f47ac10b-58cc-4372-a567-0e02b2c3d479", true, hexCasingLower}, + {"both upper", "00000000-0000-0000-0000-000000000000", "F47AC10B-58CC-4372-A567-0E02B2C3D479", true, hexCasingUpper}, + {"all digits", "01234567-8901-2345-6789-012345678901", "09234567-8901-2345-6789-012345678901", true, hexCasingLower}, + {"digits min, upper max", "01234567-8901-2345-6789-012345678901", "F47AC10B-58CC-4372-A567-0E02B2C3D479", true, hexCasingUpper}, + {"lower min, digits max", "f47ac10b-58cc-4372-a567-0e02b2c3d479", "01234567-8901-2345-6789-012345678901", true, hexCasingLower}, + {"mixed case across bounds", "f47ac10b-58cc-4372-a567-0e02b2c3d479", "F47AC10B-58CC-4372-A567-0E02B2C3D479", false, hexCasingUnknown}, + {"mixed case within a bound", "00000000-0000-0000-0000-000000000000", "F47ac10b-58cc-4372-a567-0e02b2c3d479", false, hexCasingUnknown}, + {"non-uuid", "apple", "banana", false, hexCasingUnknown}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + isUUID, casing := detectUuidWithHexCasing(tc.minVal, tc.maxVal) + assert.Equal(t, tc.expectIsUUID, isUUID) + assert.Equal(t, tc.expectCasing, casing) + }) + } +} + +func TestBuildUuidStringPartitions(t *testing.T) { + cases := []struct { + name string + minV string + maxV string + expectedCase *regexp.Regexp + }{ + {"lower", "018f6e7a-1b2c-7def-8a3b-1c2d3e4f5a6b", "f47ac10b-58cc-4372-a567-0e02b2c3d479", uuidLowerRe}, + {"upper", "018F6E7A-1B2C-7DEF-8A3B-1C2D3E4F5A6B", "F47AC10B-58CC-4372-A567-0E02B2C3D479", uuidUpperRe}, + } + const numPartitions = 32 + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + isUUID, casing := detectUuidWithHexCasing(tc.minV, tc.maxV) + require.True(t, isUUID) + parts, err := buildUuidStringPartitions(tc.minV, tc.maxV, casing, numPartitions) + require.NoError(t, err) + require.Len(t, parts, numPartitions) + + prevEnd := "" + for i, part := range parts { + r, ok := part.Range.Range.(*protos.PartitionRange_StringRange) + require.True(t, ok) + sr := r.StringRange + if i == 0 { + assert.Equal(t, tc.minV, sr.Start) + } else { + assert.Equal(t, prevEnd, sr.Start) + assert.Regexp(t, tc.expectedCase, sr.Start) + } + startInt, err := uuidToBigInt(sr.Start) + require.NoError(t, err) + endInt, err := uuidToBigInt(sr.End) + require.NoError(t, err) + assert.Negative(t, startInt.Cmp(endInt)) + + isLast := i == len(parts)-1 + assert.Equal(t, isLast, sr.EndInclusive) + if isLast { + assert.Equal(t, tc.maxV, sr.End) + } + prevEnd = sr.End + } + }) + } +} + +func TestBuildUuidStringPartitionsSinglePartition(t *testing.T) { + uuid1 := "018f6e7a-1b2c-7def-8a3b-1c2d3e4f5a6b" + uuid2 := "f47ac10b-58cc-4372-a567-0e02b2c3d479" + cases := []struct { + name string + minV string + maxV string + numPartitions int64 + }{ + {"one partition requested", uuid1, uuid2, 1}, + {"min equals max", uuid1, uuid1, 8}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + isUUID, casing := detectUuidWithHexCasing(tc.minV, tc.maxV) + require.True(t, isUUID) + parts, err := buildUuidStringPartitions(tc.minV, tc.maxV, casing, tc.numPartitions) + require.NoError(t, err) + require.Len(t, parts, 1) + sr, ok := parts[0].Range.Range.(*protos.PartitionRange_StringRange) + require.True(t, ok) + assert.Equal(t, tc.minV, sr.StringRange.Start) + assert.Equal(t, tc.maxV, sr.StringRange.End) + assert.True(t, sr.StringRange.EndInclusive) + }) + } +} + +func TestBuildUuidStringPartitionsInvalid(t *testing.T) { + cases := []struct { + name string + minV string + maxV string + }{ + {"arbitrary string", "apple", "banana"}, + {"inverted", "f47ac10b-58cc-4372-a567-0e02b2c3d479", "018f6e7a-1b2c-7def-8a3b-1c2d3e4f5a6b"}, + {"too short", "018f6e7a-1b2c", "f47ac10b-58cc"}, + {"too long", "018f6e7a-1b2c-7def-8a3b-1c2d3e4f5a6bb", "f47ac10b-58cc-4372-a567-0e02b2c3d4799"}, + } + const numPartitions = 32 + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + _, err := buildUuidStringPartitions(tc.minV, tc.maxV, hexCasingLower, numPartitions) + require.Error(t, err) + }) + } +} diff --git a/flow/connectors/mysql/qrep_test.go b/flow/connectors/mysql/qrep_test.go index eb0b8e42b..57a26c85a 100644 --- a/flow/connectors/mysql/qrep_test.go +++ b/flow/connectors/mysql/qrep_test.go @@ -127,6 +127,17 @@ func TestGetDefaultPartitionKeyForTables(t *testing.T) { }, expected: map[string]string{"db.ts": "created_at"}, }, + { + name: "string primary key is supported", + tableMappings: []*protos.TableMapping{tableMapping("db.uuidpk")}, + schemas: map[string]*protos.TableSchema{ + "db.uuidpk": { + PrimaryKeyColumns: []string{"id"}, + Columns: []*protos.FieldDescription{fieldDesc("id", types.QValueKindString)}, + }, + }, + expected: map[string]string{"db.uuidpk": "id"}, + }, { name: "composite primary key with valid first column", tableMappings: []*protos.TableMapping{tableMapping("db.composite")}, @@ -146,9 +157,9 @@ func TestGetDefaultPartitionKeyForTables(t *testing.T) { tableMappings: []*protos.TableMapping{tableMapping("db.composite2")}, schemas: map[string]*protos.TableSchema{ "db.composite2": { - PrimaryKeyColumns: []string{"name", "id"}, + PrimaryKeyColumns: []string{"data", "id"}, Columns: []*protos.FieldDescription{ - fieldDesc("name", types.QValueKindString), + fieldDesc("data", types.QValueKindBytes), fieldDesc("id", types.QValueKindInt32), }, }, @@ -172,6 +183,7 @@ func TestGetDefaultPartitionKeyForTables(t *testing.T) { tableMapping("db.a"), tableMapping("db.b"), tableMapping("db.c"), + tableMapping("db.d"), }, schemas: map[string]*protos.TableSchema{ "db.a": { @@ -188,8 +200,14 @@ func TestGetDefaultPartitionKeyForTables(t *testing.T) { PrimaryKeyColumns: []string{"date"}, Columns: []*protos.FieldDescription{fieldDesc("date", types.QValueKindDate)}, }, + "db.d": { + PrimaryKeyColumns: []string{"float"}, + Columns: []*protos.FieldDescription{ + fieldDesc("bad", types.QValueKindArrayFloat32), + }, + }, }, - expected: map[string]string{"db.a": "id", "db.c": "date"}, + expected: map[string]string{"db.a": "id", "db.b": "uuid", "db.c": "date"}, }, } diff --git a/flow/connectors/utils/monitoring/monitoring.go b/flow/connectors/utils/monitoring/monitoring.go index 147291bbf..98fa75e6a 100644 --- a/flow/connectors/utils/monitoring/monitoring.go +++ b/flow/connectors/utils/monitoring/monitoring.go @@ -362,6 +362,8 @@ func addPartitionToQRepRun(ctx context.Context, tx pgx.Tx, flowJobName string, rangeEnd = new(rangeEndValue.(string)) case *protos.PartitionRange_ObjectIdRange: rangeStart, rangeEnd = &x.ObjectIdRange.Start, &x.ObjectIdRange.End + case *protos.PartitionRange_StringRange: + rangeStart, rangeEnd = &x.StringRange.Start, &x.StringRange.End case *protos.PartitionRange_NullRange: // leave rangeStart and rangeEnd as nil default: diff --git a/flow/connectors/utils/partition.go b/flow/connectors/utils/partition.go index e1d1721a7..13548cdc7 100644 --- a/flow/connectors/utils/partition.go +++ b/flow/connectors/utils/partition.go @@ -17,6 +17,14 @@ import ( const FullTablePartitionID = "full-table-partition-id" +func FullTablePartition() []*protos.QRepPartition { + return []*protos.QRepPartition{{ + PartitionId: FullTablePartitionID, + Range: nil, + FullTablePartition: true, + }} +} + type PartitionRangeType string const ( @@ -394,6 +402,10 @@ func (p *PartitionHelper) updatePartitionHelper(partition *protos.QRepPartition) return nil } +func (p *PartitionHelper) AddPartitions(partitions []*protos.QRepPartition) { + p.partitions = append(p.partitions, partitions...) +} + func (p *PartitionHelper) GetPartitions() []*protos.QRepPartition { return p.partitions } diff --git a/flow/e2e/clickhouse_mysql_test.go b/flow/e2e/clickhouse_mysql_test.go index 1d4a530c3..8841764ae 100644 --- a/flow/e2e/clickhouse_mysql_test.go +++ b/flow/e2e/clickhouse_mysql_test.go @@ -10,6 +10,7 @@ import ( "testing" "time" + "github.com/google/uuid" "github.com/stretchr/testify/require" "github.com/testcontainers/testcontainers-go" "github.com/testcontainers/testcontainers-go/wait" @@ -1651,3 +1652,137 @@ func (s ClickHouseSuite) Test_MySQL_Default_Partition_Key_Parallel_Snapshot() { env.Cancel(s.t.Context()) RequireEnvCanceled(s.t, env) } + +func (s ClickHouseSuite) Test_MySQL_String_Partition_Key_UUID_Parallel_Snapshot() { + if _, ok := s.source.(*MySqlSource); !ok { + s.t.Skip("only applies to mysql") + } + + srcTable := "test_string_pk_uuid" + srcFullName := s.attachSchemaSuffix(srcTable) + dstTable := "test_string_pk_uuid_dst" + + const numRows = 100 + const numPartitions = 8 + require.NoError(s.t, s.source.Exec(s.t.Context(), + fmt.Sprintf("CREATE TABLE %s (id CHAR(36) PRIMARY KEY, val INT NOT NULL)", srcFullName))) + for i := 1; i <= numRows; i++ { + require.NoError(s.t, s.source.Exec(s.t.Context(), + fmt.Sprintf("INSERT INTO %s (id, val) VALUES ('%s', %d)", srcFullName, uuid.NewString(), i))) + } + + tableMappings := TableMappings(s, srcTable, dstTable) + // a String column can't be used directly as a Distributed sharding key (ClickHouse + // requires an integer-typed sharding expression), so hash it for the cluster suite + for _, tm := range tableMappings { + tm.ShardingKey = "cityHash64(id)" + } + connectionGen := FlowConnectionGenerationConfig{ + FlowJobName: s.attachSuffix("string_pk_uuid"), + // PartitionKey is not specified to test default partition key detection + TableMappings: tableMappings, + Destination: s.Peer().Name, + } + flowConnConfig := connectionGen.GenerateFlowConnectionConfigs(s) + flowConnConfig.DoInitialSnapshot = true + flowConnConfig.SnapshotMaxParallelWorkers = 4 + flowConnConfig.SnapshotNumPartitionsOverride = numPartitions + + tc := NewTemporalClient(s.t) + env := ExecutePeerflow(s.t, tc, flowConnConfig) + SetupCDCFlowStatusQuery(s.t, env, flowConnConfig) + + EnvWaitForEqualTablesWithNames(env, s, "waiting on initial", srcTable, dstTable, "id,val") + + partitionRows, err := s.catalog.Query(s.t.Context(), + `SELECT partition_start, partition_end, COALESCE(rows_in_partition, 0) + FROM peerdb_stats.qrep_partitions WHERE parent_mirror_name = $1`, + flowConnConfig.FlowJobName) + require.NoError(s.t, err) + defer partitionRows.Close() + + var partitionCount int32 + var totalRows int64 + for partitionRows.Next() { + var partitionStart, partitionEnd *string + var rowsInPartition int64 + require.NoError(s.t, partitionRows.Scan(&partitionStart, &partitionEnd, &rowsInPartition)) + require.NotNil(s.t, partitionStart) + require.NotNil(s.t, partitionEnd) + totalRows += rowsInPartition + partitionCount++ + } + require.NoError(s.t, partitionRows.Err()) + require.EqualValues(s.t, numPartitions, partitionCount) + require.EqualValues(s.t, numRows, totalRows) + + env.Cancel(s.t.Context()) + RequireEnvCanceled(s.t, env) +} + +func (s ClickHouseSuite) Test_MySQL_String_Partition_Key_Arbitrary_FullTable() { + if _, ok := s.source.(*MySqlSource); !ok { + s.t.Skip("only applies to mysql") + } + + srcTable := "test_string_pk_non_uuid" + srcFullName := s.attachSchemaSuffix(srcTable) + dstTable := "test_string_pk_non_uuid_dst" + + const numRows = 100 + const numPartitions = 8 + require.NoError(s.t, s.source.Exec(s.t.Context(), + fmt.Sprintf("CREATE TABLE %s (id VARCHAR(50) PRIMARY KEY, val INT NOT NULL)", srcFullName))) + for i := 1; i <= numRows; i++ { + require.NoError(s.t, s.source.Exec(s.t.Context(), + fmt.Sprintf("INSERT INTO %s (id, val) VALUES ('key_%04d', %d)", srcFullName, i, i))) + } + + tableMappings := TableMappings(s, srcTable, dstTable) + // a String column can't be used directly as a Distributed sharding key (ClickHouse + // requires an integer-typed sharding expression), so hash it for the cluster suite + for _, tm := range tableMappings { + tm.ShardingKey = "cityHash64(id)" + } + connectionGen := FlowConnectionGenerationConfig{ + FlowJobName: s.attachSuffix("string_pk_non_uuid"), + // PartitionKey is not specified to test default partition key detection + TableMappings: tableMappings, + Destination: s.Peer().Name, + } + flowConnConfig := connectionGen.GenerateFlowConnectionConfigs(s) + flowConnConfig.DoInitialSnapshot = true + flowConnConfig.SnapshotMaxParallelWorkers = 4 + flowConnConfig.SnapshotNumPartitionsOverride = numPartitions + + tc := NewTemporalClient(s.t) + env := ExecutePeerflow(s.t, tc, flowConnConfig) + SetupCDCFlowStatusQuery(s.t, env, flowConnConfig) + + EnvWaitForEqualTablesWithNames(env, s, "waiting on initial", srcTable, dstTable, "id,val") + + partitionRows, err := s.catalog.Query(s.t.Context(), + `SELECT partition_start, partition_end, COALESCE(rows_in_partition, 0) + FROM peerdb_stats.qrep_partitions WHERE parent_mirror_name = $1`, + flowConnConfig.FlowJobName) + require.NoError(s.t, err) + defer partitionRows.Close() + + var partitionCount int32 + var totalRows int64 + for partitionRows.Next() { + var partitionStart, partitionEnd *string + var rowsInPartition int64 + require.NoError(s.t, partitionRows.Scan(&partitionStart, &partitionEnd, &rowsInPartition)) + require.Nil(s.t, partitionStart) + require.Nil(s.t, partitionEnd) + totalRows += rowsInPartition + partitionCount++ + } + require.NoError(s.t, partitionRows.Err()) + require.EqualValues(s.t, 1, partitionCount) + require.EqualValues(s.t, numRows, totalRows) + + env.Cancel(s.t.Context()) + RequireEnvCanceled(s.t, env) +} diff --git a/protos/flow.proto b/protos/flow.proto index 827c0809b..af3d5f155 100644 --- a/protos/flow.proto +++ b/protos/flow.proto @@ -324,6 +324,15 @@ message ObjectIdPartitionRange { string end = 2; } +message StringPartitionRange { + string start = 1; + string end = 2; + // unlike numeric/temporal ranges, string boundaries cannot be safely + // incremented; so the partition boundaries are [start, end) except + // for the last partition being [start, end]; hence this bool is needed. + bool end_inclusive = 3; +} + message PartitionRange { // can be a timestamp range or an integer range oneof range { @@ -333,6 +342,7 @@ message PartitionRange { UIntPartitionRange uint_range = 4; ObjectIdPartitionRange object_id_range = 5; NullPartitionRange null_range = 6; + StringPartitionRange string_range = 7; } }