Skip to content

Commit e252314

Browse files
jgao54jgao54
authored
support default primary key detection for mysql for int/temporal types (#4440)
When custom partition key is not specified, determine default partition key based on table's primary key. If the primary key(s) starts with an integer or temporal type, perform parallel snapshot instead of full table snapshot. Testing: - add unit & e2e tests - validated parallel snapshotting locally
1 parent 98993d4 commit e252314

6 files changed

Lines changed: 284 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: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1590,3 +1590,64 @@ func (s ClickHouseSuite) Test_MySQL_BinlogIncident() {
15901590
env.Cancel(s.t.Context())
15911591
RequireEnvCanceled(s.t, env)
15921592
}
1593+
1594+
func (s ClickHouseSuite) Test_MySQL_Default_Partition_Key_Parallel_Snapshot() {
1595+
if _, ok := s.source.(*MySqlSource); !ok {
1596+
s.t.Skip("only applies to mysql")
1597+
}
1598+
1599+
srcTable := "test_default_partition_key_parallel"
1600+
srcFullName := s.attachSchemaSuffix(srcTable)
1601+
dstTable := "test_default_partition_key_parallel_dst"
1602+
1603+
const numRows = 100
1604+
const numPartitions = 8
1605+
require.NoError(s.t, s.source.Exec(s.t.Context(),
1606+
fmt.Sprintf("CREATE TABLE %s (id BIGINT PRIMARY KEY, created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP)", srcFullName)))
1607+
for i := 1; i <= numRows; i++ {
1608+
require.NoError(s.t, s.source.Exec(s.t.Context(),
1609+
fmt.Sprintf("INSERT INTO %s (id) VALUES (%d)", srcFullName, i)))
1610+
}
1611+
1612+
connectionGen := FlowConnectionGenerationConfig{
1613+
FlowJobName: s.attachSuffix("default_partition_key_parallel"),
1614+
// PartitionKey is not specified to test default partition key detection
1615+
TableMappings: TableMappings(s, srcTable, dstTable),
1616+
Destination: s.Peer().Name,
1617+
}
1618+
flowConnConfig := connectionGen.GenerateFlowConnectionConfigs(s)
1619+
flowConnConfig.DoInitialSnapshot = true
1620+
flowConnConfig.SnapshotMaxParallelWorkers = 4
1621+
flowConnConfig.SnapshotNumPartitionsOverride = numPartitions
1622+
1623+
tc := NewTemporalClient(s.t)
1624+
env := ExecutePeerflow(s.t, tc, flowConnConfig)
1625+
SetupCDCFlowStatusQuery(s.t, env, flowConnConfig)
1626+
1627+
EnvWaitForEqualTablesWithNames(env, s, "waiting on initial", srcTable, dstTable, "id,created_at")
1628+
1629+
partitionRows, err := s.catalog.Query(s.t.Context(),
1630+
`SELECT partition_start, partition_end, COALESCE(rows_in_partition, 0)
1631+
FROM peerdb_stats.qrep_partitions WHERE parent_mirror_name = $1`,
1632+
flowConnConfig.FlowJobName)
1633+
require.NoError(s.t, err)
1634+
defer partitionRows.Close()
1635+
1636+
var partitionCount int32
1637+
var totalRows int64
1638+
for partitionRows.Next() {
1639+
var partitionStart, partitionEnd *string
1640+
var rowsInPartition int64
1641+
require.NoError(s.t, partitionRows.Scan(&partitionStart, &partitionEnd, &rowsInPartition))
1642+
require.NotNil(s.t, partitionStart)
1643+
require.NotNil(s.t, partitionEnd)
1644+
totalRows += rowsInPartition
1645+
partitionCount++
1646+
}
1647+
require.NoError(s.t, partitionRows.Err())
1648+
require.EqualValues(s.t, numPartitions, partitionCount)
1649+
require.EqualValues(s.t, numRows, totalRows)
1650+
1651+
env.Cancel(s.t.Context())
1652+
RequireEnvCanceled(s.t, env)
1653+
}

‎flow/internal/dynamicconf.go‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -468,6 +468,14 @@ var DynamicSettings = [...]*protos.DynamicSetting{
468468
ApplyMode: protos.DynconfApplyMode_APPLY_MODE_AFTER_RESUME,
469469
TargetForSetting: protos.DynconfTarget_ALL,
470470
},
471+
{
472+
Name: "PEERDB_MYSQL_DEFAULT_PARTITION_KEY_ENABLED",
473+
Description: "Enables automatic detection of a default partition key from primary key for MySQL initial load",
474+
DefaultValue: "true",
475+
ValueType: protos.DynconfValueType_BOOL,
476+
ApplyMode: protos.DynconfApplyMode_APPLY_MODE_NEW_MIRROR,
477+
TargetForSetting: protos.DynconfTarget_ALL,
478+
},
471479
{
472480
Name: "PEERDB_PG_AUTOMATED_SCHEMA_DUMP",
473481
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
858866
return dynamicConfBool(ctx, env, "PEERDB_POSTGRES_APPLY_CTID_BLOCK_PARTITIONING_OVERRIDE")
859867
}
860868

869+
func PeerDBMySQLDefaultPartitionKeyEnabled(ctx context.Context, env map[string]string) (bool, error) {
870+
return dynamicConfBool(ctx, env, "PEERDB_MYSQL_DEFAULT_PARTITION_KEY_ENABLED")
871+
}
872+
861873
func PeerDBPGAutomatedSchemaDump(ctx context.Context, env map[string]string) (bool, error) {
862874
return dynamicConfBool(ctx, env, "PEERDB_PG_AUTOMATED_SCHEMA_DUMP")
863875
}

‎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)