support default primary key detection for mysql for int/temporal types

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

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 {

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)
+		})
+	}
+}

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)
+}

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")
 }

protos/flow.proto

@@ -666,6 +666,7 @@ message AdditionalContextMetadata{
 
 message GetDefaultPartitionKeyForTablesInput {
   repeated TableMapping table_mappings = 1;
+  map<string, TableSchema> table_schema_mapping = 2;
 }
 
 message GetDefaultPartitionKeyForTablesOutput {