diff --git a/flow/activities/snapshot_activity.go b/flow/activities/snapshot_activity.go index 8b82400fa..96478d761 100644 --- a/flow/activities/snapshot_activity.go +++ b/flow/activities/snapshot_activity.go @@ -194,8 +194,40 @@ func (a *SnapshotActivity) GetDefaultPartitionKeyForTables( } defer connClose(ctx) + srcType, err := connectors.LoadPeerType(ctx, a.CatalogPool, input.SourceName) + if err != nil { + return nil, a.Alerter.LogFlowError(ctx, input.FlowJobName, fmt.Errorf("failed to load source peer type: %w", err)) + } + + var tableSchemaMapping map[string]*protos.TableSchema + if srcType == protos.DBType_MYSQL { + mysqlDefaultPartitionKeyEnabled, err := internal.PeerDBMySQLDefaultPartitionKeyEnabled(ctx, input.Env) + if err != nil { + return nil, a.Alerter.LogFlowError(ctx, input.FlowJobName, + fmt.Errorf("failed to check if mysql default partition key detection is enabled: %w", err)) + } + + if mysqlDefaultPartitionKeyEnabled { + dstTableNames := make([]string, 0, len(input.TableMappings)) + for _, tm := range input.TableMappings { + dstTableNames = append(dstTableNames, tm.DestinationTableIdentifier) + } + schemasByDstTable, err := internal.LoadTableSchemasFromCatalog(ctx, a.CatalogPool.Pool, input.FlowJobName, dstTableNames) + if err != nil { + return nil, a.Alerter.LogFlowError(ctx, input.FlowJobName, fmt.Errorf("failed to load table schemas from catalog: %w", err)) + } + tableSchemaMapping = make(map[string]*protos.TableSchema, len(input.TableMappings)) + for _, tm := range input.TableMappings { + if schema, ok := schemasByDstTable[tm.DestinationTableIdentifier]; ok { + tableSchemaMapping[tm.SourceTableIdentifier] = schema + } + } + } + } + output, err := conn.GetDefaultPartitionKeyForTables(ctx, &protos.GetDefaultPartitionKeyForTablesInput{ - TableMappings: input.TableMappings, + TableMappings: input.TableMappings, + TableSchemaMapping: tableSchemaMapping, }) if err != nil { return nil, a.Alerter.LogFlowError(ctx, input.FlowJobName, fmt.Errorf("failed to check if tables can parallel load: %w", err)) diff --git a/flow/connectors/mysql/qrep.go b/flow/connectors/mysql/qrep.go index bb841b80f..86793762a 100644 --- a/flow/connectors/mysql/qrep.go +++ b/flow/connectors/mysql/qrep.go @@ -168,13 +168,64 @@ func (c *MySqlConnector) GetQRepPartitions( return partitionHelper.GetPartitions(), nil } +func supportsRangePartition(qkind types.QValueKind) bool { + switch qkind { + // integer types + case types.QValueKindInt8, types.QValueKindInt16, types.QValueKindInt32, types.QValueKindInt64, + types.QValueKindUInt8, types.QValueKindUInt16, types.QValueKindUInt32, types.QValueKindUInt64: + return true + // temporal types + case types.QValueKindDate, types.QValueKindTimestamp: + return true + default: + return false + } +} + func (c *MySqlConnector) GetDefaultPartitionKeyForTables( ctx context.Context, input *protos.GetDefaultPartitionKeyForTablesInput, ) (*protos.GetDefaultPartitionKeyForTablesOutput, error) { - return &protos.GetDefaultPartitionKeyForTablesOutput{ - TableDefaultPartitionKeyMapping: nil, - }, nil + c.logger.Info("Evaluating if tables can perform parallel load") + + output := &protos.GetDefaultPartitionKeyForTablesOutput{ + TableDefaultPartitionKeyMapping: make(map[string]string, len(input.TableMappings)), + } + for _, tm := range input.TableMappings { + source := tm.SourceTableIdentifier + schema, ok := input.TableSchemaMapping[source] + if !ok { + c.logger.Warn("[mysql] table schema not found, defaulting to full table snapshot", + slog.String("table", source)) + continue + } + if len(schema.PrimaryKeyColumns) == 0 { + c.logger.Info("[mysql] table has no primary key, defaulting to full table snapshot", + slog.String("table", source)) + continue + } + pkColumn := schema.PrimaryKeyColumns[0] + var pkQKind types.QValueKind + for _, col := range schema.Columns { + if col.Name == pkColumn { + pkQKind = types.QValueKind(col.Type) + break + } + } + if !supportsRangePartition(pkQKind) { + c.logger.Info("[mysql] primary key type does not support range partitioning, defaulting to full table snapshot", + slog.String("table", source), + slog.String("column", pkColumn), + slog.String("qkind", string(pkQKind))) + continue + } + c.logger.Info("[mysql] using primary key as default partition key", + slog.String("table", source), + slog.String("column", pkColumn), + slog.String("qkind", string(pkQKind))) + output.TableDefaultPartitionKeyMapping[source] = pkColumn + } + return output, nil } func buildSelectedColumns(cols []*protos.FieldDescription, exclude []string) string { diff --git a/flow/connectors/mysql/qrep_test.go b/flow/connectors/mysql/qrep_test.go index 206dea6b7..eb0b8e42b 100644 --- a/flow/connectors/mysql/qrep_test.go +++ b/flow/connectors/mysql/qrep_test.go @@ -1,8 +1,13 @@ package connmysql import ( + "context" + "log/slog" "testing" + "github.com/stretchr/testify/require" + "go.temporal.io/sdk/log" + "github.com/PeerDB-io/peerdb/flow/generated/protos" "github.com/PeerDB-io/peerdb/flow/shared/types" ) @@ -82,3 +87,121 @@ func TestBuildSelectedColumns(t *testing.T) { }) } } + +func TestGetDefaultPartitionKeyForTables(t *testing.T) { + c := &MySqlConnector{logger: log.NewStructuredLogger(slog.Default())} + + tableMapping := func(name string) *protos.TableMapping { + return &protos.TableMapping{SourceTableIdentifier: name, DestinationTableIdentifier: name} + } + + fieldDesc := func(name string, qkind types.QValueKind) *protos.FieldDescription { + return &protos.FieldDescription{Name: name, Type: string(qkind)} + } + + testCases := []struct { + schemas map[string]*protos.TableSchema + expected map[string]string + name string + tableMappings []*protos.TableMapping + }{ + { + name: "integer primary key is supported", + tableMappings: []*protos.TableMapping{tableMapping("db.ints")}, + schemas: map[string]*protos.TableSchema{ + "db.ints": { + PrimaryKeyColumns: []string{"id"}, + Columns: []*protos.FieldDescription{fieldDesc("id", types.QValueKindInt64)}, + }, + }, + expected: map[string]string{"db.ints": "id"}, + }, + { + name: "timestamp primary key is supported", + tableMappings: []*protos.TableMapping{tableMapping("db.ts")}, + schemas: map[string]*protos.TableSchema{ + "db.ts": { + PrimaryKeyColumns: []string{"created_at"}, + Columns: []*protos.FieldDescription{fieldDesc("created_at", types.QValueKindTimestamp)}, + }, + }, + expected: map[string]string{"db.ts": "created_at"}, + }, + { + name: "composite primary key with valid first column", + tableMappings: []*protos.TableMapping{tableMapping("db.composite")}, + schemas: map[string]*protos.TableSchema{ + "db.composite": { + PrimaryKeyColumns: []string{"id", "created_at"}, + Columns: []*protos.FieldDescription{ + fieldDesc("id", types.QValueKindInt32), + fieldDesc("created_at", types.QValueKindTimestamp), + }, + }, + }, + expected: map[string]string{"db.composite": "id"}, + }, + { + name: "composite primary key with invalid first column", + tableMappings: []*protos.TableMapping{tableMapping("db.composite2")}, + schemas: map[string]*protos.TableSchema{ + "db.composite2": { + PrimaryKeyColumns: []string{"name", "id"}, + Columns: []*protos.FieldDescription{ + fieldDesc("name", types.QValueKindString), + fieldDesc("id", types.QValueKindInt32), + }, + }, + }, + expected: map[string]string{}, + }, + { + name: "no primary key", + tableMappings: []*protos.TableMapping{tableMapping("db.nopk")}, + schemas: map[string]*protos.TableSchema{ + "db.nopk": { + PrimaryKeyColumns: nil, + Columns: []*protos.FieldDescription{fieldDesc("id", types.QValueKindInt64)}, + }, + }, + expected: map[string]string{}, + }, + { + name: "multiple table schemas", + tableMappings: []*protos.TableMapping{ + tableMapping("db.a"), + tableMapping("db.b"), + tableMapping("db.c"), + }, + schemas: map[string]*protos.TableSchema{ + "db.a": { + PrimaryKeyColumns: []string{"id"}, + Columns: []*protos.FieldDescription{fieldDesc("id", types.QValueKindInt16)}, + }, + "db.b": { + PrimaryKeyColumns: []string{"uuid"}, + Columns: []*protos.FieldDescription{ + fieldDesc("uuid", types.QValueKindString), + }, + }, + "db.c": { + PrimaryKeyColumns: []string{"date"}, + Columns: []*protos.FieldDescription{fieldDesc("date", types.QValueKindDate)}, + }, + }, + expected: map[string]string{"db.a": "id", "db.c": "date"}, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + output, err := c.GetDefaultPartitionKeyForTables(context.Background(), &protos.GetDefaultPartitionKeyForTablesInput{ + TableMappings: tc.tableMappings, + TableSchemaMapping: tc.schemas, + }) + require.NoError(t, err) + require.NotNil(t, output) + require.Equal(t, tc.expected, output.TableDefaultPartitionKeyMapping) + }) + } +} diff --git a/flow/e2e/clickhouse_mysql_test.go b/flow/e2e/clickhouse_mysql_test.go index 83e71cf9f..1d4a530c3 100644 --- a/flow/e2e/clickhouse_mysql_test.go +++ b/flow/e2e/clickhouse_mysql_test.go @@ -1590,3 +1590,64 @@ func (s ClickHouseSuite) Test_MySQL_BinlogIncident() { env.Cancel(s.t.Context()) RequireEnvCanceled(s.t, env) } + +func (s ClickHouseSuite) Test_MySQL_Default_Partition_Key_Parallel_Snapshot() { + if _, ok := s.source.(*MySqlSource); !ok { + s.t.Skip("only applies to mysql") + } + + srcTable := "test_default_partition_key_parallel" + srcFullName := s.attachSchemaSuffix(srcTable) + dstTable := "test_default_partition_key_parallel_dst" + + const numRows = 100 + const numPartitions = 8 + require.NoError(s.t, s.source.Exec(s.t.Context(), + fmt.Sprintf("CREATE TABLE %s (id BIGINT PRIMARY KEY, created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP)", srcFullName))) + for i := 1; i <= numRows; i++ { + require.NoError(s.t, s.source.Exec(s.t.Context(), + fmt.Sprintf("INSERT INTO %s (id) VALUES (%d)", srcFullName, i))) + } + + connectionGen := FlowConnectionGenerationConfig{ + FlowJobName: s.attachSuffix("default_partition_key_parallel"), + // PartitionKey is not specified to test default partition key detection + TableMappings: TableMappings(s, srcTable, dstTable), + 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,created_at") + + 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) +} diff --git a/flow/internal/dynamicconf.go b/flow/internal/dynamicconf.go index 9e0f03a36..482c9d793 100644 --- a/flow/internal/dynamicconf.go +++ b/flow/internal/dynamicconf.go @@ -468,6 +468,14 @@ var DynamicSettings = [...]*protos.DynamicSetting{ ApplyMode: protos.DynconfApplyMode_APPLY_MODE_AFTER_RESUME, TargetForSetting: protos.DynconfTarget_ALL, }, + { + Name: "PEERDB_MYSQL_DEFAULT_PARTITION_KEY_ENABLED", + Description: "Enables automatic detection of a default partition key from primary key for MySQL initial load", + DefaultValue: "true", + ValueType: protos.DynconfValueType_BOOL, + ApplyMode: protos.DynconfApplyMode_APPLY_MODE_NEW_MIRROR, + TargetForSetting: protos.DynconfTarget_ALL, + }, { Name: "PEERDB_PG_AUTOMATED_SCHEMA_DUMP", Description: "For PG-to-PG mirrors, run pg_dump --schema-only from source into psql on destination " + @@ -858,6 +866,10 @@ func PeerDBPostgresApplyCtidBlockPartitioning(ctx context.Context, env map[strin return dynamicConfBool(ctx, env, "PEERDB_POSTGRES_APPLY_CTID_BLOCK_PARTITIONING_OVERRIDE") } +func PeerDBMySQLDefaultPartitionKeyEnabled(ctx context.Context, env map[string]string) (bool, error) { + return dynamicConfBool(ctx, env, "PEERDB_MYSQL_DEFAULT_PARTITION_KEY_ENABLED") +} + func PeerDBPGAutomatedSchemaDump(ctx context.Context, env map[string]string) (bool, error) { return dynamicConfBool(ctx, env, "PEERDB_PG_AUTOMATED_SCHEMA_DUMP") } diff --git a/protos/flow.proto b/protos/flow.proto index b9964bacc..827c0809b 100644 --- a/protos/flow.proto +++ b/protos/flow.proto @@ -666,6 +666,7 @@ message AdditionalContextMetadata{ message GetDefaultPartitionKeyForTablesInput { repeated TableMapping table_mappings = 1; + map table_schema_mapping = 2; } message GetDefaultPartitionKeyForTablesOutput {