Skip to content

Commit f1df47e

Browse files
jgao54jgao54
committed
support default primary key detection for mysql for integer and temporal types
1 parent b5c24ef commit f1df47e

6 files changed

Lines changed: 287 additions & 4 deletions

File tree

‎flow/activities/snapshot_activity.go‎

Lines changed: 33 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -194,8 +194,40 @@ func (a *SnapshotActivity) GetDefaultPartitionKeyForTables(
194194
}
195195
defer connClose(ctx)
196196

197+
srcType, err := connectors.LoadPeerType(ctx, a.CatalogPool, input.SourceName)
198+
if err != nil {
199+
return nil, a.Alerter.LogFlowError(ctx, input.FlowJobName, fmt.Errorf("failed to load source peer type: %w", err))
200+
}
201+
202+
var tableSchemaMapping map[string]*protos.TableSchema
203+
if srcType == protos.DBType_MYSQL {
204+
mysqlDefaultPartitionKeyEnabled, err := internal.PeerDBMySQLDefaultPartitionKeyEnabled(ctx, input.Env)
205+
if err != nil {
206+
return nil, a.Alerter.LogFlowError(ctx, input.FlowJobName,
207+
fmt.Errorf("failed to check if mysql default partition key detection is enabled: %w", err))
208+
}
209+
210+
if mysqlDefaultPartitionKeyEnabled {
211+
dstTableNames := make([]string, 0, len(input.TableMappings))
212+
for _, tm := range input.TableMappings {
213+
dstTableNames = append(dstTableNames, tm.DestinationTableIdentifier)
214+
}
215+
schemasByDstTable, err := internal.LoadTableSchemasFromCatalog(ctx, a.CatalogPool.Pool, input.FlowJobName, dstTableNames)
216+
if err != nil {
217+
return nil, a.Alerter.LogFlowError(ctx, input.FlowJobName, fmt.Errorf("failed to load table schemas from catalog: %w", err))
218+
}
219+
tableSchemaMapping = make(map[string]*protos.TableSchema, len(input.TableMappings))
220+
for _, tm := range input.TableMappings {
221+
if schema, ok := schemasByDstTable[tm.DestinationTableIdentifier]; ok {
222+
tableSchemaMapping[tm.SourceTableIdentifier] = schema
223+
}
224+
}
225+
}
226+
}
227+
197228
output, err := conn.GetDefaultPartitionKeyForTables(ctx, &protos.GetDefaultPartitionKeyForTablesInput{
198-
TableMappings: input.TableMappings,
229+
TableMappings: input.TableMappings,
230+
TableSchemaMapping: tableSchemaMapping,
199231
})
200232
if err != nil {
201233
return nil, a.Alerter.LogFlowError(ctx, input.FlowJobName, fmt.Errorf("failed to check if tables can parallel load: %w", err))

‎flow/connectors/mysql/qrep.go‎

Lines changed: 54 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -168,13 +168,64 @@ func (c *MySqlConnector) GetQRepPartitions(
168168
return partitionHelper.GetPartitions(), nil
169169
}
170170

171+
func supportsRangePartition(qkind types.QValueKind) bool {
172+
switch qkind {
173+
// integer types
174+
case types.QValueKindInt8, types.QValueKindInt16, types.QValueKindInt32, types.QValueKindInt64,
175+
types.QValueKindUInt8, types.QValueKindUInt16, types.QValueKindUInt32, types.QValueKindUInt64:
176+
return true
177+
// temporal types
178+
case types.QValueKindDate, types.QValueKindTimestamp:
179+
return true
180+
default:
181+
return false
182+
}
183+
}
184+
171185
func (c *MySqlConnector) GetDefaultPartitionKeyForTables(
172186
ctx context.Context,
173187
input *protos.GetDefaultPartitionKeyForTablesInput,
174188
) (*protos.GetDefaultPartitionKeyForTablesOutput, error) {
175-
return &protos.GetDefaultPartitionKeyForTablesOutput{
176-
TableDefaultPartitionKeyMapping: nil,
177-
}, nil
189+
c.logger.Info("Evaluating if tables can perform parallel load")
190+
191+
output := &protos.GetDefaultPartitionKeyForTablesOutput{
192+
TableDefaultPartitionKeyMapping: make(map[string]string, len(input.TableMappings)),
193+
}
194+
for _, tm := range input.TableMappings {
195+
source := tm.SourceTableIdentifier
196+
schema, ok := input.TableSchemaMapping[source]
197+
if !ok {
198+
c.logger.Warn("[mysql] table schema not found, defaulting to full table snapshot",
199+
slog.String("table", source))
200+
continue
201+
}
202+
if len(schema.PrimaryKeyColumns) == 0 {
203+
c.logger.Info("[mysql] table has no primary key, defaulting to full table snapshot",
204+
slog.String("table", source))
205+
continue
206+
}
207+
pkColumn := schema.PrimaryKeyColumns[0]
208+
var pkQKind types.QValueKind
209+
for _, col := range schema.Columns {
210+
if col.Name == pkColumn {
211+
pkQKind = types.QValueKind(col.Type)
212+
break
213+
}
214+
}
215+
if !supportsRangePartition(pkQKind) {
216+
c.logger.Info("[mysql] primary key type does not support range partitioning, defaulting to full table snapshot",
217+
slog.String("table", source),
218+
slog.String("column", pkColumn),
219+
slog.String("qkind", string(pkQKind)))
220+
continue
221+
}
222+
c.logger.Info("[mysql] using primary key as default partition key",
223+
slog.String("table", source),
224+
slog.String("column", pkColumn),
225+
slog.String("qkind", string(pkQKind)))
226+
output.TableDefaultPartitionKeyMapping[source] = pkColumn
227+
}
228+
return output, nil
178229
}
179230

180231
func buildSelectedColumns(cols []*protos.FieldDescription, exclude []string) string {

‎flow/connectors/mysql/qrep_test.go‎

Lines changed: 123 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,13 @@
11
package connmysql
22

33
import (
4+
"context"
5+
"log/slog"
46
"testing"
57

8+
"github.com/stretchr/testify/require"
9+
"go.temporal.io/sdk/log"
10+
611
"github.com/PeerDB-io/peerdb/flow/generated/protos"
712
"github.com/PeerDB-io/peerdb/flow/shared/types"
813
)
@@ -82,3 +87,121 @@ func TestBuildSelectedColumns(t *testing.T) {
8287
})
8388
}
8489
}
90+
91+
func TestGetDefaultPartitionKeyForTables(t *testing.T) {
92+
c := &MySqlConnector{logger: log.NewStructuredLogger(slog.Default())}
93+
94+
tableMapping := func(name string) *protos.TableMapping {
95+
return &protos.TableMapping{SourceTableIdentifier: name, DestinationTableIdentifier: name}
96+
}
97+
98+
fieldDesc := func(name string, qkind types.QValueKind) *protos.FieldDescription {
99+
return &protos.FieldDescription{Name: name, Type: string(qkind)}
100+
}
101+
102+
testCases := []struct {
103+
schemas map[string]*protos.TableSchema
104+
expected map[string]string
105+
name string
106+
tableMappings []*protos.TableMapping
107+
}{
108+
{
109+
name: "integer primary key is supported",
110+
tableMappings: []*protos.TableMapping{tableMapping("db.ints")},
111+
schemas: map[string]*protos.TableSchema{
112+
"db.ints": {
113+
PrimaryKeyColumns: []string{"id"},
114+
Columns: []*protos.FieldDescription{fieldDesc("id", types.QValueKindInt64)},
115+
},
116+
},
117+
expected: map[string]string{"db.ints": "id"},
118+
},
119+
{
120+
name: "timestamp primary key is supported",
121+
tableMappings: []*protos.TableMapping{tableMapping("db.ts")},
122+
schemas: map[string]*protos.TableSchema{
123+
"db.ts": {
124+
PrimaryKeyColumns: []string{"created_at"},
125+
Columns: []*protos.FieldDescription{fieldDesc("created_at", types.QValueKindTimestamp)},
126+
},
127+
},
128+
expected: map[string]string{"db.ts": "created_at"},
129+
},
130+
{
131+
name: "composite primary key with valid first column",
132+
tableMappings: []*protos.TableMapping{tableMapping("db.composite")},
133+
schemas: map[string]*protos.TableSchema{
134+
"db.composite": {
135+
PrimaryKeyColumns: []string{"id", "created_at"},
136+
Columns: []*protos.FieldDescription{
137+
fieldDesc("id", types.QValueKindInt32),
138+
fieldDesc("created_at", types.QValueKindTimestamp),
139+
},
140+
},
141+
},
142+
expected: map[string]string{"db.composite": "id"},
143+
},
144+
{
145+
name: "composite primary key with invalid first column",
146+
tableMappings: []*protos.TableMapping{tableMapping("db.composite2")},
147+
schemas: map[string]*protos.TableSchema{
148+
"db.composite2": {
149+
PrimaryKeyColumns: []string{"name", "id"},
150+
Columns: []*protos.FieldDescription{
151+
fieldDesc("name", types.QValueKindString),
152+
fieldDesc("id", types.QValueKindInt32),
153+
},
154+
},
155+
},
156+
expected: map[string]string{},
157+
},
158+
{
159+
name: "no primary key",
160+
tableMappings: []*protos.TableMapping{tableMapping("db.nopk")},
161+
schemas: map[string]*protos.TableSchema{
162+
"db.nopk": {
163+
PrimaryKeyColumns: nil,
164+
Columns: []*protos.FieldDescription{fieldDesc("id", types.QValueKindInt64)},
165+
},
166+
},
167+
expected: map[string]string{},
168+
},
169+
{
170+
name: "multiple table schemas",
171+
tableMappings: []*protos.TableMapping{
172+
tableMapping("db.a"),
173+
tableMapping("db.b"),
174+
tableMapping("db.c"),
175+
},
176+
schemas: map[string]*protos.TableSchema{
177+
"db.a": {
178+
PrimaryKeyColumns: []string{"id"},
179+
Columns: []*protos.FieldDescription{fieldDesc("id", types.QValueKindInt16)},
180+
},
181+
"db.b": {
182+
PrimaryKeyColumns: []string{"uuid"},
183+
Columns: []*protos.FieldDescription{
184+
fieldDesc("uuid", types.QValueKindString),
185+
},
186+
},
187+
"db.c": {
188+
PrimaryKeyColumns: []string{"date"},
189+
Columns: []*protos.FieldDescription{fieldDesc("date", types.QValueKindDate)},
190+
},
191+
},
192+
expected: map[string]string{"db.a": "id", "db.c": "date"},
193+
},
194+
}
195+
196+
for _, tc := range testCases {
197+
t.Run(tc.name, func(t *testing.T) {
198+
output, err := c.GetDefaultPartitionKeyForTables(context.Background(), &protos.GetDefaultPartitionKeyForTablesInput{
199+
TableMappings: tc.tableMappings,
200+
TableSchemaMapping: tc.schemas,
201+
})
202+
require.NoError(t, err)
203+
require.NotNil(t, output)
204+
require.Equal(t, tc.expected, output.TableDefaultPartitionKeyMapping)
205+
})
206+
}
207+
}

‎flow/e2e/clickhouse_mysql_test.go‎

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1270,3 +1270,67 @@ func (s ClickHouseSuite) Test_MySQL_Column_Position_Shifting_DDL_Error() {
12701270
})
12711271
}
12721272
}
1273+
1274+
func (s ClickHouseSuite) Test_MySQL_Default_Partition_Key_Parallel_Snapshot() {
1275+
if _, ok := s.source.(*MySqlSource); !ok {
1276+
s.t.Skip("only applies to mysql")
1277+
}
1278+
1279+
srcTable := "test_default_partition_key_parallel"
1280+
srcFullName := s.attachSchemaSuffix(srcTable)
1281+
dstTable := "test_default_partition_key_parallel_dst"
1282+
1283+
const numRows = 100
1284+
const numPartitions = 8
1285+
require.NoError(s.t, s.source.Exec(s.t.Context(),
1286+
fmt.Sprintf("CREATE TABLE %s (id BIGINT PRIMARY KEY, created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP)", srcFullName)))
1287+
for i := 1; i <= numRows; i++ {
1288+
require.NoError(s.t, s.source.Exec(s.t.Context(),
1289+
fmt.Sprintf("INSERT INTO %s (id) VALUES (%d)", srcFullName, i)))
1290+
}
1291+
1292+
connectionGen := FlowConnectionGenerationConfig{
1293+
FlowJobName: s.attachSuffix("default_partition_key_parallel"),
1294+
TableMappings: []*protos.TableMapping{{
1295+
SourceTableIdentifier: srcFullName,
1296+
DestinationTableIdentifier: dstTable,
1297+
// PartitionKey is not specified to test default partition key detection
1298+
}},
1299+
Destination: s.Peer().Name,
1300+
}
1301+
flowConnConfig := connectionGen.GenerateFlowConnectionConfigs(s)
1302+
flowConnConfig.DoInitialSnapshot = true
1303+
flowConnConfig.SnapshotMaxParallelWorkers = 4
1304+
flowConnConfig.SnapshotNumPartitionsOverride = numPartitions
1305+
1306+
tc := NewTemporalClient(s.t)
1307+
env := ExecutePeerflow(s.t, tc, flowConnConfig)
1308+
SetupCDCFlowStatusQuery(s.t, env, flowConnConfig)
1309+
1310+
EnvWaitForEqualTablesWithNames(env, s, "waiting on initial", srcTable, dstTable, "id,created_at")
1311+
1312+
partitionRows, err := s.catalog.Query(s.t.Context(),
1313+
`SELECT partition_start, partition_end, COALESCE(rows_in_partition, 0)
1314+
FROM peerdb_stats.qrep_partitions WHERE parent_mirror_name = $1`,
1315+
flowConnConfig.FlowJobName)
1316+
require.NoError(s.t, err)
1317+
defer partitionRows.Close()
1318+
1319+
var partitionCount int32
1320+
var totalRows int64
1321+
for partitionRows.Next() {
1322+
var partitionStart, partitionEnd *string
1323+
var rowsInPartition int64
1324+
require.NoError(s.t, partitionRows.Scan(&partitionStart, &partitionEnd, &rowsInPartition))
1325+
require.NotNil(s.t, partitionStart)
1326+
require.NotNil(s.t, partitionEnd)
1327+
totalRows += rowsInPartition
1328+
partitionCount++
1329+
}
1330+
require.NoError(s.t, partitionRows.Err())
1331+
require.EqualValues(s.t, numPartitions, partitionCount)
1332+
require.EqualValues(s.t, numRows, totalRows)
1333+
1334+
env.Cancel(s.t.Context())
1335+
RequireEnvCanceled(s.t, env)
1336+
}

‎flow/internal/dynamicconf.go‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -459,6 +459,14 @@ var DynamicSettings = [...]*protos.DynamicSetting{
459459
ApplyMode: protos.DynconfApplyMode_APPLY_MODE_AFTER_RESUME,
460460
TargetForSetting: protos.DynconfTarget_ALL,
461461
},
462+
{
463+
Name: "PEERDB_MYSQL_DEFAULT_PARTITION_KEY_ENABLED",
464+
Description: "Enables automatic detection of a default partition key from primary key for MySQL initial load",
465+
DefaultValue: "true",
466+
ValueType: protos.DynconfValueType_BOOL,
467+
ApplyMode: protos.DynconfApplyMode_APPLY_MODE_NEW_MIRROR,
468+
TargetForSetting: protos.DynconfTarget_ALL,
469+
},
462470
{
463471
Name: "PEERDB_PG_AUTOMATED_SCHEMA_DUMP",
464472
Description: "For PG-to-PG mirrors, run pg_dump --schema-only from source into psql on destination " +
@@ -845,6 +853,10 @@ func PeerDBPostgresApplyCtidBlockPartitioning(ctx context.Context, env map[strin
845853
return dynamicConfBool(ctx, env, "PEERDB_POSTGRES_APPLY_CTID_BLOCK_PARTITIONING_OVERRIDE")
846854
}
847855

856+
func PeerDBMySQLDefaultPartitionKeyEnabled(ctx context.Context, env map[string]string) (bool, error) {
857+
return dynamicConfBool(ctx, env, "PEERDB_MYSQL_DEFAULT_PARTITION_KEY_ENABLED")
858+
}
859+
848860
func PeerDBPGAutomatedSchemaDump(ctx context.Context, env map[string]string) (bool, error) {
849861
return dynamicConfBool(ctx, env, "PEERDB_PG_AUTOMATED_SCHEMA_DUMP")
850862
}

‎protos/flow.proto‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -666,6 +666,7 @@ message AdditionalContextMetadata{
666666

667667
message GetDefaultPartitionKeyForTablesInput {
668668
repeated TableMapping table_mappings = 1;
669+
map<string, TableSchema> table_schema_mapping = 2;
669670
}
670671

671672
message GetDefaultPartitionKeyForTablesOutput {

0 commit comments

Comments
 (0)