diff --git a/flow/connectors/mysql/qrep.go b/flow/connectors/mysql/qrep.go index 40ebcd861..853b38ab3 100644 --- a/flow/connectors/mysql/qrep.go +++ b/flow/connectors/mysql/qrep.go @@ -20,6 +20,28 @@ import ( "github.com/PeerDB-io/peerdb/flow/shared/types" ) +func (c *MySqlConnector) GetDataTypeOfWatermarkColumn( + ctx context.Context, + watermarkTableName string, + watermarkColumn string, +) (types.QValueKind, error) { + if watermarkColumn == "" { + return "", errors.New("watermark column is not specified in the config") + } + + query := fmt.Sprintf("SELECT `%s` FROM %s LIMIT 0", watermarkColumn, watermarkTableName) + rs, err := c.Execute(ctx, query) + if err != nil { + return "", fmt.Errorf("failed to execute query for watermark column type: %w", err) + } + + if len(rs.Fields) == 0 { + return "", errors.New("no fields returned from select query: " + query) + } + + return qkindFromMysql(rs.Fields[0]) +} + func (c *MySqlConnector) GetQRepPartitions( ctx context.Context, config *protos.QRepConfig, @@ -46,7 +68,7 @@ func (c *MySqlConnector) GetQRepPartitions( whereClause := "" if last != nil && last.Range != nil { - whereClause = fmt.Sprintf("WHERE %s > $1", quotedWatermarkColumn) + whereClause = fmt.Sprintf("WHERE %s > ?", quotedWatermarkColumn) } parsedWatermarkTable, err := utils.ParseSchemaTable(config.WatermarkTable) if err != nil { @@ -64,7 +86,7 @@ func (c *MySqlConnector) GetQRepPartitions( case *protos.PartitionRange_UintRange: minVal = lastRange.UintRange.End case *protos.PartitionRange_TimestampRange: - minVal = lastRange.TimestampRange.End.AsTime() + minVal = lastRange.TimestampRange.End.AsTime().String() } c.logger.Info(fmt.Sprintf("count query: %s - minVal: %v", countQuery, minVal)) @@ -99,62 +121,115 @@ func (c *MySqlConnector) GetQRepPartitions( if totalRows%numRowsPerPartition != 0 { numPartitions++ } + + watermarkQKind, err := c.GetDataTypeOfWatermarkColumn(ctx, parsedWatermarkTable.MySQL(), config.WatermarkColumn) + if err != nil { + return nil, fmt.Errorf("failed to get data type of watermark column %s: %w", config.WatermarkColumn, err) + } + c.logger.Info(fmt.Sprintf("total rows: %d, num partitions: %d, num rows per partition: %d", totalRows, numPartitions, numRowsPerPartition)) var rs *mysql.Result - if minVal != nil { - // Query to get partitions using window functions - partitionsQuery := fmt.Sprintf( - `SELECT bucket, MIN(%[2]s) AS start, MAX(%[2]s) AS end - FROM ( - SELECT NTILE(%[1]d) OVER (ORDER BY %[2]s) AS bucket, %[2]s - FROM %[3]s WHERE %[2]s > $1 - ) AS subquery + + switch watermarkQKind { + case types.QValueKindInt8, types.QValueKindInt16, types.QValueKindInt32, types.QValueKindInt64, + types.QValueKindUInt8, types.QValueKindUInt16, types.QValueKindUInt32, types.QValueKindUInt64: + if minVal != nil { + partitionsQuery := fmt.Sprintf( + `WITH stats AS ( + SELECT MIN(%[2]s) AS min_watermark, + 1.0 * (MAX(%[2]s) - MIN(%[2]s)) / (%[1]d) AS range_size + FROM %[3]s WHERE %[2]s > ? + ) + SELECT FLOOR((w.%[2]s - s.min_watermark) / s.range_size) AS bucket, + MIN(w.%[2]s) AS start, MAX(w.%[2]s) AS end + FROM %[3]s AS w + CROSS JOIN stats AS s + WHERE w.%[2]s > ? GROUP BY bucket - ORDER BY start`, - numPartitions, - quotedWatermarkColumn, - parsedWatermarkTable.MySQL(), - ) - c.logger.Info("partitions query", slog.String("query", partitionsQuery), slog.Any("minVal", minVal)) - rs, err = c.Execute(ctx, partitionsQuery, minVal) - } else { - partitionsQuery := fmt.Sprintf( - `SELECT bucket, MIN(%[2]s) AS start, MAX(%[2]s) AS end - FROM ( - SELECT NTILE(%[1]d) OVER (ORDER BY %[2]s) AS bucket, %[2]s FROM %[3]s - ) AS subquery + ORDER BY start;`, + numPartitions, + quotedWatermarkColumn, + parsedWatermarkTable.MySQL(), + ) + c.logger.Info("partitions query", slog.String("query", partitionsQuery), slog.Any("minVal", minVal)) + rs, err = c.Execute(ctx, partitionsQuery, minVal, minVal) + } else { + partitionsQuery := fmt.Sprintf( + `WITH stats AS ( + SELECT MIN(%[2]s) AS min_watermark, + 1.0 * (MAX(%[2]s) - MIN(%[2]s)) / (%[1]d) AS range_size + FROM %[3]s + ) + SELECT FLOOR((w.%[2]s - s.min_watermark) / s.range_size) AS bucket, + MIN(w.%[2]s) AS start, MAX(w.%[2]s) AS end + FROM %[3]s AS w + CROSS JOIN stats AS s GROUP BY bucket - ORDER BY start`, - numPartitions, - quotedWatermarkColumn, - parsedWatermarkTable.MySQL(), - ) - c.logger.Info("partitions query", slog.String("query", partitionsQuery)) - rs, err = c.Execute(ctx, partitionsQuery) - } - if err != nil { - return nil, fmt.Errorf("failed to query for partitions: %w", err) + ORDER BY start;`, + numPartitions, + quotedWatermarkColumn, + parsedWatermarkTable.MySQL(), + ) + c.logger.Info("partitions query", slog.String("query", partitionsQuery)) + rs, err = c.Execute(ctx, partitionsQuery) + } + if err != nil { + return nil, fmt.Errorf("failed to query for partitions: %w", err) + } + case types.QValueKindTimestamp, types.QValueKindTimestampTZ: + if minVal != nil { + partitionsQuery := fmt.Sprintf( + `WITH stats AS ( + SELECT MIN(%[2]s) AS min_watermark, + 1.0 * (TIMESTAMPDIFF(MICROSECOND, MAX(%[2]s), MIN(%[2]s)) / (%[1]d)) AS range_size + FROM %[3]s WHERE %[2]s > ? + ) + SELECT FLOOR(TIMESTAMPDIFF(MICROSECOND, w.%[2]s, s.min_watermark) / s.range_size) AS bucket, + MIN(w.%[2]s) AS start, MAX(w.%[2]s) AS end + FROM %[3]s AS w + CROSS JOIN stats AS s + WHERE w.%[2]s > ? + GROUP BY bucket + ORDER BY start;`, + numPartitions, + quotedWatermarkColumn, + parsedWatermarkTable.MySQL(), + ) + c.logger.Info("partitions query", slog.String("query", partitionsQuery), slog.Any("minVal", minVal)) + rs, err = c.Execute(ctx, partitionsQuery, minVal) + } else { + partitionsQuery := fmt.Sprintf( + `WITH stats AS ( + SELECT MIN(%[2]s) AS min_watermark, + 1.0 * (TIMESTAMPDIFF(MICROSECOND, MAX(%[2]s), MIN(%[2]s)) / (%[1]d)) AS range_size + FROM %[3]s + ) + SELECT FLOOR(TIMESTAMPDIFF(MICROSECOND, w.%[2]s, s.min_watermark) / s.range_size) AS bucket, + MIN(w.%[2]s) AS start, MAX(w.%[2]s) AS end + FROM %[3]s AS w + CROSS JOIN stats AS s + GROUP BY bucket + ORDER BY start;`, + numPartitions, + quotedWatermarkColumn, + parsedWatermarkTable.MySQL(), + ) + c.logger.Info("partitions query", slog.String("query", partitionsQuery)) + rs, err = c.Execute(ctx, partitionsQuery) + } + if err != nil { + return nil, fmt.Errorf("failed to query for partitions: %w", err) + } } - qk1, err := qkindFromMysql(rs.Fields[1]) - if err != nil { - return nil, err - } - qk2, err := qkindFromMysql(rs.Fields[2]) - if err != nil { - return nil, err - } - if qk1 != qk2 { - return nil, fmt.Errorf("low/high of partition range should be same type, got low:%s, high:%s", qk1, qk2) - } partitionHelper := utils.NewPartitionHelper(c.logger) for _, row := range rs.Values { - val1, err := QValueFromMysqlFieldValue(qk1, row[1]) + val1, err := QValueFromMysqlFieldValue(watermarkQKind, row[1]) if err != nil { return nil, err } - val2, err := QValueFromMysqlFieldValue(qk2, row[2]) + val2, err := QValueFromMysqlFieldValue(watermarkQKind, row[2]) if err != nil { return nil, err } @@ -169,7 +244,7 @@ func (c *MySqlConnector) GetQRepPartitions( func (c *MySqlConnector) PullQRepRecords( ctx context.Context, config *protos.QRepConfig, - last *protos.QRepPartition, + partition *protos.QRepPartition, stream *model.QRecordStream, ) (int64, int64, error) { tableSchema, err := c.getTableSchemaForTable(ctx, config.Env, @@ -241,7 +316,7 @@ func (c *MySqlConnector) PullQRepRecords( return nil } - if last.FullTablePartition { + if partition.FullTablePartition { // this is a full table partition, so just run the query if err := c.ExecuteSelectStreaming(ctx, config.Query, &rs, onRow, onResult); err != nil { return 0, 0, err @@ -251,7 +326,7 @@ func (c *MySqlConnector) PullQRepRecords( var rangeEnd string // Depending on the type of the range, convert the range into the correct type - switch x := last.Range.Range.(type) { + switch x := partition.Range.Range.(type) { case *protos.PartitionRange_IntRange: rangeStart = strconv.FormatInt(x.IntRange.Start, 10) rangeEnd = strconv.FormatInt(x.IntRange.End, 10) diff --git a/flow/connectors/utils/partition.go b/flow/connectors/utils/partition.go index f60394603..db87dea5c 100644 --- a/flow/connectors/utils/partition.go +++ b/flow/connectors/utils/partition.go @@ -2,6 +2,7 @@ package utils import ( "cmp" + "errors" "fmt" "log/slog" "time" @@ -14,51 +15,135 @@ import ( "github.com/PeerDB-io/peerdb/flow/generated/protos" ) -// Function to compare two values -func compareValues(prevEnd any, start any) int { - switch v := start.(type) { - case int64: - return cmp.Compare(prevEnd.(int64), v) - case uint64: - return cmp.Compare(prevEnd.(uint64), v) - case time.Time: - return prevEnd.(time.Time).Compare(v) - case pgtype.TID: - pe := prevEnd.(pgtype.TID) - if c := cmp.Compare(pe.BlockNumber, v.BlockNumber); c != 0 { +type PartitionRangeType string + +const ( + PartitionEndRangeType PartitionRangeType = "end" + PartitionStartRangeType PartitionRangeType = "start" +) + +type PartitionRangeForComparison struct { + partitionRange *protos.PartitionRange + rangeTypeToCompare PartitionRangeType +} + +// Function to compare the end of a partition with the start of another +func comparePartitionRanges( + previousPartition PartitionRangeForComparison, + currentPartition PartitionRangeForComparison, +) int { + if previousPartition.partitionRange == nil || currentPartition.partitionRange == nil { + slog.Warn("one of the partition ranges is nil, cannot compare") + return 0 + } + switch pr := previousPartition.partitionRange.Range.(type) { + case *protos.PartitionRange_IntRange: + cr, ok := currentPartition.partitionRange.Range.(*protos.PartitionRange_IntRange) + if !ok { + return 0 + } + getVal := func(r *protos.IntPartitionRange, t PartitionRangeType) int64 { + if t == PartitionEndRangeType { + return r.End + } + return r.Start + } + prevVal := getVal(pr.IntRange, previousPartition.rangeTypeToCompare) + currVal := getVal(cr.IntRange, currentPartition.rangeTypeToCompare) + return cmp.Compare(prevVal, currVal) + case *protos.PartitionRange_UintRange: + cr, ok := currentPartition.partitionRange.Range.(*protos.PartitionRange_UintRange) + if !ok { + return 0 + } + getVal := func(r *protos.UIntPartitionRange, t PartitionRangeType) uint64 { + if t == PartitionEndRangeType { + return r.End + } + return r.Start + } + prevVal := getVal(pr.UintRange, previousPartition.rangeTypeToCompare) + currVal := getVal(cr.UintRange, currentPartition.rangeTypeToCompare) + return cmp.Compare(prevVal, currVal) + case *protos.PartitionRange_TimestampRange: + cr, ok := currentPartition.partitionRange.Range.(*protos.PartitionRange_TimestampRange) + if !ok { + return 0 + } + getTime := func(r *protos.TimestampPartitionRange, t PartitionRangeType) time.Time { + if t == PartitionEndRangeType { + return r.End.AsTime() + } + return r.Start.AsTime() + } + prevVal := getTime(pr.TimestampRange, previousPartition.rangeTypeToCompare) + currVal := getTime(cr.TimestampRange, currentPartition.rangeTypeToCompare) + return prevVal.Compare(currVal) + case *protos.PartitionRange_TidRange: + cr, ok := currentPartition.partitionRange.Range.(*protos.PartitionRange_TidRange) + if !ok { + return 0 + } + getTuple := func(r *protos.TIDPartitionRange, t PartitionRangeType) *protos.TID { + if t == PartitionEndRangeType { + return r.End + } + return r.Start + } + prevTuple := getTuple(pr.TidRange, previousPartition.rangeTypeToCompare) + currTuple := getTuple(cr.TidRange, currentPartition.rangeTypeToCompare) + if c := cmp.Compare(prevTuple.BlockNumber, currTuple.BlockNumber); c != 0 { return c } - return cmp.Compare(pe.OffsetNumber, v.OffsetNumber) - case uint32: // xmin - return cmp.Compare(prevEnd.(uint32), v) + return cmp.Compare(prevTuple.OffsetNumber, currTuple.OffsetNumber) default: return 0 } } // Function to adjust start value -func adjustStartValue(prevEnd any, start any) any { - switch start.(type) { - case int64: - return prevEnd.(int64) + 1 - case int32: - return int32(prevEnd.(int64) + 1) - case time.Time: - // postgres & mysql timestamps have microsecond precision - return prevEnd.(time.Time).Add(1 * time.Microsecond) - case pgtype.TID: - pe := prevEnd.(pgtype.TID) - if pe.OffsetNumber < 0xFFFF { - pe.OffsetNumber++ - } else { - pe.BlockNumber++ - pe.OffsetNumber = 0 - } - return pe - case uint32: - return prevEnd.(uint32) + 1 +func adjustStartValueOfPartition(prevRange *protos.PartitionRange, currentRange *protos.PartitionRange) { + if prevRange == nil || currentRange == nil { + return + } + + switch cr := currentRange.Range.(type) { + case *protos.PartitionRange_IntRange: + if pr, ok := prevRange.Range.(*protos.PartitionRange_IntRange); ok { + cr.IntRange.Start = pr.IntRange.End + 1 + } + return + + case *protos.PartitionRange_UintRange: + if pr, ok := prevRange.Range.(*protos.PartitionRange_UintRange); ok { + cr.UintRange.Start = pr.UintRange.End + 1 + } + return + + case *protos.PartitionRange_TimestampRange: + if pr, ok := prevRange.Range.(*protos.PartitionRange_TimestampRange); ok { + cr.TimestampRange.Start = timestamppb.New(pr.TimestampRange.End.AsTime().Add(1 * time.Microsecond)) + } + return + + case *protos.PartitionRange_TidRange: + if pr, ok := prevRange.Range.(*protos.PartitionRange_TidRange); ok { + start := &protos.TID{ + BlockNumber: pr.TidRange.End.BlockNumber, + OffsetNumber: pr.TidRange.End.OffsetNumber, + } + if start.OffsetNumber < 0xFFFF { + start.OffsetNumber++ + } else { + start.BlockNumber++ + start.OffsetNumber = 0 + } + cr.TidRange.Start = start + } + return + default: - return start + return } } @@ -144,66 +229,115 @@ func NewPartitionHelper(logger log.Logger) *PartitionHelper { func (p *PartitionHelper) AddPartition(start any, end any) error { p.logger.Info("adding partition", slog.Any("start", start), slog.Any("end", end)) + currentPartition, err := p.getPartitionForStartAndEnd(start, end) + if err != nil { + return fmt.Errorf("error getting current partition from start and end: %w", err) + } + + prevPartition, err := p.getPartitionForStartAndEnd(p.prevStart, p.prevEnd) + if err != nil { + return fmt.Errorf("error getting previous partition from prevStart and prevEnd: %w", err) + } // Skip partition if it's fully contained within the previous one // If it's not fully contained but overlaps, adjust the start - if p.prevEnd != nil { - if compareValues(p.prevEnd, start) >= 0 { + if prevPartition != nil { + prevEndCompareStart := comparePartitionRanges( + PartitionRangeForComparison{ + partitionRange: prevPartition.Range, + rangeTypeToCompare: PartitionEndRangeType, + }, + PartitionRangeForComparison{ + partitionRange: currentPartition.Range, + rangeTypeToCompare: PartitionStartRangeType, + }) + if prevEndCompareStart >= 0 { + prevEndCompareEnd := comparePartitionRanges( + PartitionRangeForComparison{ + partitionRange: prevPartition.Range, + rangeTypeToCompare: PartitionEndRangeType, + }, + PartitionRangeForComparison{ + partitionRange: currentPartition.Range, + rangeTypeToCompare: PartitionEndRangeType, + }, + ) // If end is also less than or equal to prevEnd, skip this partition - if compareValues(p.prevEnd, end) >= 0 { + if prevEndCompareEnd >= 0 { // log the skipped partition p.logger.Info("skipping partition, fully contained within previous partition", slog.Any("start", start), slog.Any("end", end), slog.Any("prevStart", p.prevStart), slog.Any("prevEnd", p.prevEnd)) return nil } // If end is greater than prevEnd, adjust the start - start = adjustStartValue(p.prevEnd, start) + adjustStartValueOfPartition(prevPartition.Range, currentPartition.Range) } } + if err := p.updatePartitionHelper(currentPartition); err != nil { + return fmt.Errorf("error adjusting start value: %w", err) + } + + return nil +} + +func (p *PartitionHelper) getPartitionForStartAndEnd(start any, end any) (*protos.QRepPartition, error) { + if start == nil || end == nil { + return nil, nil + } switch v := start.(type) { case int64: - p.partitions = append(p.partitions, createIntPartition(v, end.(int64))) - p.prevStart = v - p.prevEnd = end + return createIntPartition(v, end.(int64)), nil case uint64: - p.partitions = append(p.partitions, createUIntPartition(v, end.(uint64))) - p.prevStart = v - p.prevEnd = end.(uint64) + return createUIntPartition(v, end.(uint64)), nil case int32: - p.partitions = append(p.partitions, createIntPartition(int64(v), int64(end.(int32)))) - p.prevStart = int64(v) - p.prevEnd = int64(end.(int32)) + return createIntPartition(int64(v), int64(end.(int32))), nil case uint32: - p.partitions = append(p.partitions, createUIntPartition(uint64(v), uint64(end.(uint32)))) - p.prevStart = uint64(v) - p.prevEnd = uint64(end.(uint32)) + return createUIntPartition(uint64(v), uint64(end.(uint32))), nil case int16: - p.partitions = append(p.partitions, createIntPartition(int64(v), int64(end.(int16)))) - p.prevStart = int64(v) - p.prevEnd = int64(end.(int16)) + return createIntPartition(int64(v), int64(end.(int16))), nil case uint16: - p.partitions = append(p.partitions, createUIntPartition(uint64(v), uint64(end.(uint16)))) - p.prevStart = uint64(v) - p.prevEnd = uint64(end.(uint16)) + return createUIntPartition(uint64(v), uint64(end.(uint16))), nil case int8: - p.partitions = append(p.partitions, createIntPartition(int64(v), int64(end.(int8)))) - p.prevStart = int64(v) - p.prevEnd = int64(end.(int8)) + return createIntPartition(int64(v), int64(end.(int8))), nil case uint8: - p.partitions = append(p.partitions, createUIntPartition(uint64(v), uint64(end.(uint8)))) - p.prevStart = uint64(v) - p.prevEnd = uint64(end.(uint8)) + return createUIntPartition(uint64(v), uint64(end.(uint8))), nil case time.Time: - p.partitions = append(p.partitions, createTimePartition(v, end.(time.Time))) - p.prevStart = v - p.prevEnd = end + return createTimePartition(v, end.(time.Time)), nil case pgtype.TID: - p.partitions = append(p.partitions, createTIDPartition(v, end.(pgtype.TID))) - p.prevStart = v - p.prevEnd = end + return createTIDPartition(v, end.(pgtype.TID)), nil + default: + return nil, fmt.Errorf("unsupported type: %T", v) + } +} + +func (p *PartitionHelper) updatePartitionHelper(partition *protos.QRepPartition) error { + if partition == nil { + return errors.New("partition is nil") + } + p.partitions = append(p.partitions, partition) + + switch r := partition.Range.Range.(type) { + case *protos.PartitionRange_IntRange: + p.prevStart = r.IntRange.Start + p.prevEnd = r.IntRange.End + case *protos.PartitionRange_UintRange: + p.prevStart = r.UintRange.Start + p.prevEnd = r.UintRange.End + case *protos.PartitionRange_TimestampRange: + p.prevStart = r.TimestampRange.Start.AsTime() + p.prevEnd = r.TimestampRange.End.AsTime() + case *protos.PartitionRange_TidRange: + p.prevStart = pgtype.TID{ + BlockNumber: r.TidRange.Start.BlockNumber, + OffsetNumber: uint16(r.TidRange.Start.OffsetNumber), + } + p.prevEnd = pgtype.TID{ + BlockNumber: r.TidRange.End.BlockNumber, + OffsetNumber: uint16(r.TidRange.End.OffsetNumber), + } default: - return fmt.Errorf("unsupported type: %T", v) + return fmt.Errorf("unsupported partition range type: %T", r) } return nil diff --git a/flow/e2e/api/api_test.go b/flow/e2e/api/api_test.go index 0b19572ea..994268a20 100644 --- a/flow/e2e/api/api_test.go +++ b/flow/e2e/api/api_test.go @@ -882,10 +882,10 @@ func (s Suite) TestCustomSync() { } func (s Suite) TestQRep() { - if _, ok := s.source.(*e2e.PostgresSource); !ok { - s.t.Skip("only run with pg as mysql qrep isn't really supported") - } - + peerType, err := s.GetPeerType(s.t.Context(), &protos.PeerInfoRequest{ + PeerName: s.source.GeneratePeer(s.t).Name, + }) + require.NoError(s.t, err) tblName := "qrepapi" schemaQualified := e2e.AttachSchema(s, tblName) require.NoError(s.t, s.source.Exec(s.t.Context(), @@ -893,22 +893,24 @@ func (s Suite) TestQRep() { require.NoError(s.t, s.source.Exec(s.t.Context(), fmt.Sprintf("INSERT INTO %s(id, val) values (1,'first')", schemaQualified))) - sourcePeer := s.Source().GeneratePeer(s.t) qrepConfig := e2e.CreateQRepWorkflowConfig( s.t, - "qrepapiflow", + "qrepapiflow"+"_"+peerType.PeerType, schemaQualified, - schemaQualified+"dst", + tblName, fmt.Sprintf("SELECT * FROM %s WHERE id BETWEEN {{.start}} AND {{.end}}", schemaQualified), - sourcePeer.Name, + s.ch.Peer().Name, "", true, "", "", ) + qrepConfig.SourceName = s.source.GeneratePeer(s.t).Name qrepConfig.WatermarkColumn = "id" - - _, err := s.CreateQRepFlow(s.t.Context(), &protos.CreateQRepFlowRequest{ + qrepConfig.InitialCopyOnly = false + qrepConfig.WaitBetweenBatchesSeconds = 5 + qrepConfig.NumRowsPerPartition = 1 + _, err = s.CreateQRepFlow(s.t.Context(), &protos.CreateQRepFlowRequest{ QrepConfig: qrepConfig, CreateCatalogEntry: true, }) @@ -917,19 +919,23 @@ func (s Suite) TestQRep() { tc := e2e.NewTemporalClient(s.t) env, err := e2e.GetPeerflow(s.t.Context(), s.pg.PostgresConnector.Conn(), tc, qrepConfig.FlowJobName) require.NoError(s.t, err) - e2e.EnvWaitForFinished(s.t, env, 3*time.Minute) - require.NoError(s.t, env.Error(s.t.Context())) + e2e.EnvWaitForEqualTables(env, s.ch, "qrep initial load", tblName, "id,val") + + require.NoError(s.t, s.source.Exec(s.t.Context(), + fmt.Sprintf("INSERT INTO %s(id, val) values (2,'second')", schemaQualified))) + + e2e.EnvWaitForEqualTables(env, s.ch, "insert post qrep initial load", tblName, "id,val") statusResponse, err := s.MirrorStatus(s.t.Context(), &protos.MirrorStatusRequest{ FlowJobName: qrepConfig.FlowJobName, IncludeFlowInfo: true, ExcludeBatches: false, }) require.NoError(s.t, err) - require.Equal(s.t, protos.FlowStatus_STATUS_COMPLETED, statusResponse.CurrentFlowState) qStatus := statusResponse.GetQrepStatus() require.NotNil(s.t, qStatus) - require.Len(s.t, qStatus.Partitions, 1) - require.Equal(s.t, int64(1), qStatus.Partitions[0].RowsInPartition) - require.Equal(s.t, int64(1), qStatus.Partitions[0].RowsSynced) + require.Len(s.t, qStatus.Partitions, 2) + + env.Cancel(s.t.Context()) + e2e.RequireEnvCanceled(s.t, env) } diff --git a/flow/e2e/clickhouse/peer_flow_ch_test.go b/flow/e2e/clickhouse/peer_flow_ch_test.go index 2a40a055d..f5d1c8133 100644 --- a/flow/e2e/clickhouse/peer_flow_ch_test.go +++ b/flow/e2e/clickhouse/peer_flow_ch_test.go @@ -1502,3 +1502,117 @@ func (s ClickHouseSuite) Test_NullEngine() { env.Cancel(s.t.Context()) e2e.RequireEnvCanceled(s.t, env) } + +func (s ClickHouseSuite) Test_Partition_Key_Integer() { + srcTableName := "test_partition_key_integer" + srcFullName := s.attachSchemaSuffix(srcTableName) + dstTableName := "test_partition_key_integer" + + require.NoError(s.t, s.source.Exec(s.t.Context(), + fmt.Sprintf(`CREATE TABLE IF NOT EXISTS %s ( + id INT PRIMARY KEY, + myname TEXT NOT NULL, + updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP)`, srcFullName))) + + for i := 1; i <= 100; i++ { + if _, ok := s.source.(*e2e.PostgresSource); ok { + require.NoError(s.t, s.source.Exec(s.t.Context(), + fmt.Sprintf(`INSERT INTO %s (id,myname,updated_at) + VALUES (%d,'init_%d',CURRENT_TIMESTAMP + INTERVAL '%d seconds')`, + srcFullName, i, i, i))) + } else { + require.NoError(s.t, s.source.Exec(s.t.Context(), + fmt.Sprintf(`INSERT INTO %s (id,myname,updated_at) + VALUES (%d,'init_%d',CURRENT_TIMESTAMP + INTERVAL %d SECOND)`, + srcFullName, i, i, i))) + } + } + + connectionGen := e2e.FlowConnectionGenerationConfig{ + FlowJobName: s.attachSuffix("clickhouse_partition_key_integer"), + TableMappings: []*protos.TableMapping{{ + SourceTableIdentifier: srcFullName, + DestinationTableIdentifier: dstTableName, + PartitionKey: "id", + }}, + Destination: s.Peer().Name, + } + + flowConnConfig := connectionGen.GenerateFlowConnectionConfigs(s) + flowConnConfig.DoInitialSnapshot = true + flowConnConfig.SnapshotMaxParallelWorkers = 4 + flowConnConfig.SnapshotNumRowsPerPartition = 10 + tc := e2e.NewTemporalClient(s.t) + env := e2e.ExecutePeerflow(s.t.Context(), tc, peerflow.CDCFlowWorkflow, flowConnConfig, nil) + e2e.SetupCDCFlowStatusQuery(s.t, env, flowConnConfig) + + e2e.EnvWaitForEqualTablesWithNames(env, s, "waiting on initial", srcTableName, dstTableName, "id,myname,updated_at") + + countRow := s.Conn().QueryRow(s.t.Context(), + `SELECT COUNT(*) FROM peerdb_stats.qrep_partitions WHERE parent_mirror_name = $1`, + flowConnConfig.FlowJobName) + + var partitionCount int + require.NoError(s.t, countRow.Scan(&partitionCount), "failed to get partition count") + require.GreaterOrEqual(s.t, partitionCount, 10, "expected at least 10 partitions to be created") + + env.Cancel(s.t.Context()) + e2e.RequireEnvCanceled(s.t, env) +} + +func (s ClickHouseSuite) Test_Partition_Key_Timestamp() { + srcTableName := "test_partition_key_timestamp" + srcFullName := s.attachSchemaSuffix(srcTableName) + dstTableName := "test_partition_key_timestamp" + + require.NoError(s.t, s.source.Exec(s.t.Context(), + fmt.Sprintf(`CREATE TABLE IF NOT EXISTS %s ( + id INT PRIMARY KEY, + myname TEXT NOT NULL, + updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP)`, srcFullName))) + + for i := 1; i <= 100; i++ { + if _, ok := s.source.(*e2e.PostgresSource); ok { + require.NoError(s.t, s.source.Exec(s.t.Context(), + fmt.Sprintf(`INSERT INTO %s (id,myname,updated_at) + VALUES (%d,'init_%d',CURRENT_TIMESTAMP + INTERVAL '%d seconds')`, + srcFullName, i, i, i))) + } else { + require.NoError(s.t, s.source.Exec(s.t.Context(), + fmt.Sprintf(`INSERT INTO %s (id,myname,updated_at) + VALUES (%d,'init_%d',CURRENT_TIMESTAMP + INTERVAL %d SECOND)`, + srcFullName, i, i, i))) + } + } + + connectionGen := e2e.FlowConnectionGenerationConfig{ + FlowJobName: s.attachSuffix("clickhouse_partition_key_timestamp"), + TableMappings: []*protos.TableMapping{{ + SourceTableIdentifier: srcFullName, + DestinationTableIdentifier: dstTableName, + PartitionKey: "updated_at", + }}, + Destination: s.Peer().Name, + } + + flowConnConfig := connectionGen.GenerateFlowConnectionConfigs(s) + flowConnConfig.DoInitialSnapshot = true + flowConnConfig.SnapshotMaxParallelWorkers = 4 + flowConnConfig.SnapshotNumRowsPerPartition = 10 + tc := e2e.NewTemporalClient(s.t) + env := e2e.ExecutePeerflow(s.t.Context(), tc, peerflow.CDCFlowWorkflow, flowConnConfig, nil) + e2e.SetupCDCFlowStatusQuery(s.t, env, flowConnConfig) + + e2e.EnvWaitForEqualTablesWithNames(env, s, "waiting on initial", srcTableName, dstTableName, "id,myname,updated_at") + + countRow := s.Conn().QueryRow(s.t.Context(), + `SELECT COUNT(*) FROM peerdb_stats.qrep_partitions WHERE parent_mirror_name = $1`, + flowConnConfig.FlowJobName) + + var partitionCount int + require.NoError(s.t, countRow.Scan(&partitionCount), "failed to get partition count") + require.GreaterOrEqual(s.t, partitionCount, 10, "expected at least 10 partitions to be created") + + env.Cancel(s.t.Context()) + e2e.RequireEnvCanceled(s.t, env) +}