support parallel snapshotting for uuid partitioning (#4482)

Add parallel snapshotting support for MySQL when default primary key is
a UUID string column.

We best-effort detect uuid by evaluating the min and max values in the
table; if both min and max value match uuid format (and have a
_consistent hex casing_), we can guarantee correctness (see description
for `DetectUuidWithHexCasing` for a bit more details). Original plan was
to sample first and last N records, but it doesn't add a ton of value
here, so just checking min/max values for simplicity.

If string column is deemed non-UUID, we fallback to full table partition
for now; will follow-up with partitioning for arbitrary string.

Testing: add unit and integration tests to ensure uuids PKs are parallel
snapshotted while arbitrary string PKs still correctly replicate with
full table partition.

Fixes: DBI-868

Author: jgao54

flow/connectors/mysql/qrep.go

@@ -38,13 +38,7 @@ func (c *MySqlConnector) GetQRepPartitions(
 ) ([]*protos.QRepPartition, error) {
 	if config.WatermarkColumn == "" || config.NumPartitionsOverride == 1 {
 		// if no watermark column is specified, return a single partition
-		return []*protos.QRepPartition{
-			{
-				PartitionId:        utils.FullTablePartitionID,
-				Range:              nil,
-				FullTablePartition: true,
-			},
-		}, nil
+		return utils.FullTablePartition(), nil
 	}
 
 	if config.NumPartitionsOverride == 0 && config.NumRowsPerPartition == 0 {
@@ -101,6 +95,10 @@ func (c *MySqlConnector) GetQRepPartitions(
 		case *protos.PartitionRange_TimestampRange:
 			time := lastRange.TimestampRange.End.AsTime()
 			minVal = "'" + time.Format("2006-01-02 15:04:05.999999") + "'"
+		case *protos.PartitionRange_StringRange:
+			// resuming from last partition range is only possible for standalone QRepFlowWorkflow;
+			// this is a legacy feature and string partitioning is not supported
+			return nil, errors.New("resuming QRep by a string partition range is not supported")
 		case *protos.PartitionRange_NullRange:
 			// this case should never happen because we only add null partition for InitialCopyOnly replication
 			// (so there shouldn't be a resume scenario with null range)
@@ -155,7 +153,21 @@ func (c *MySqlConnector) GetQRepPartitions(
 	if err != nil {
 		return nil, fmt.Errorf("failed to convert partition maximum to qvalue: %w", err)
 	}
-	if err := partitionHelper.AddPartitionsWithRange(val1.Value(), val2.Value(), numPartitions); err != nil {
+
+	start, startIsString := val1.Value().(string)
+	end, endIsString := val2.Value().(string)
+	if startIsString && endIsString {
+		if isUUID, casing := detectUuidWithHexCasing(start, end); isUUID {
+			uuidPartitions, err := buildUuidStringPartitions(start, end, casing, numPartitions)
+			if err != nil {
+				return nil, fmt.Errorf("failed to add uuid string partitions: %w", err)
+			}
+			partitionHelper.AddPartitions(uuidPartitions)
+		} else {
+			c.logger.Info("string watermark column is not uuid, falling back to full table partition")
+			return utils.FullTablePartition(), nil
+		}
+	} else if err := partitionHelper.AddPartitionsWithRange(val1.Value(), val2.Value(), numPartitions); err != nil {
 		return nil, fmt.Errorf("failed to add partitions: %w", err)
 	}
 
@@ -177,6 +189,9 @@ func supportsRangePartition(qkind types.QValueKind) bool {
 	// temporal types
 	case types.QValueKindDate, types.QValueKindTimestamp:
 		return true
+	// string types
+	case types.QValueKindString:
+		return true
 	default:
 		return false
 	}
@@ -354,6 +369,19 @@ func (c *MySqlConnector) PullQRepRecords(
 		case *protos.PartitionRange_TimestampRange:
 			rangeStart = "'" + x.TimestampRange.Start.AsTime().Format("2006-01-02 15:04:05.999999") + "'"
 			rangeEnd = "'" + x.TimestampRange.End.AsTime().Format("2006-01-02 15:04:05.999999") + "'"
+		case *protos.PartitionRange_StringRange:
+			rangeStart = "'" + mysql.Escape(x.StringRange.Start) + "'"
+			rangeEnd = "'" + mysql.Escape(x.StringRange.End) + "'"
+			if config.Query != "" {
+				// custom query is only possible for standalone QRepFlowWorkflow;
+				// this is a legacy feature and string partitioning is not supported
+				return 0, 0, errors.New("can't construct a string range partition for custom queries")
+			}
+			if !x.StringRange.EndInclusive {
+				queryTemplate = fmt.Sprintf(
+					"SELECT %[1]s FROM %[2]s WHERE %[3]s >= {{.start}} AND %[3]s < {{.end}}",
+					selectedColumns, parsedSrcTable.MySQL(), common.QuoteMySQLIdentifier(config.WatermarkColumn))
+			}
 		case *protos.PartitionRange_NullRange:
 			if config.Query != "" {
 				return 0, 0, errors.New("can't construct a null range partition for custom queries")

flow/connectors/mysql/qrep_partition.go

@@ -0,0 +1,120 @@
+package connmysql
+
+import (
+	"fmt"
+	"math/big"
+	"regexp"
+	"strings"
+
+	"github.com/google/uuid"
+
+	"github.com/PeerDB-io/peerdb/flow/generated/protos"
+	"github.com/PeerDB-io/peerdb/flow/shared"
+)
+
+// String watermark partitioning is currently only supported by the MySQL connector,
+// so these helpers live in the MySQL connector subtree to prevent accidental usage
+// by other connectors.
+
+var (
+	uuidLowerRe = regexp.MustCompile(`^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$`)
+	uuidUpperRe = regexp.MustCompile(`^[0-9A-F]{8}-[0-9A-F]{4}-[0-9A-F]{4}-[0-9A-F]{4}-[0-9A-F]{12}$`)
+)
+
+type hexCasing int
+
+const (
+	hexCasingUnknown hexCasing = iota
+	hexCasingLower
+	hexCasingUpper
+)
+
+// detectUuidWithHexCasing best-effort classifies a string watermark column by
+// inspecting only its min and max bounds: it determines whether both are
+// canonical UUIDs, and if so, the shared casing of their hex letters.
+//
+// Because only the bounds are examined, the classification can be wrong in
+// the following rare cases:
+//   - non-UUID rows can exist between UUID-shaped min/max bounds;
+//   - non-boundary UUIDs contain a mix of upper and lower case
+//   - min/max bounds both don't contain hex, so casing blindly defaults to lower
+//
+// These cases may lead to partition skew, but will not affect correctness.
+func detectUuidWithHexCasing(minVal string, maxVal string) (bool, hexCasing) {
+	switch {
+	case uuidLowerRe.MatchString(minVal) && uuidLowerRe.MatchString(maxVal):
+		return true, hexCasingLower
+	case uuidUpperRe.MatchString(minVal) && uuidUpperRe.MatchString(maxVal):
+		return true, hexCasingUpper
+	default:
+		return false, hexCasingUnknown
+	}
+}
+
+// buildUuidStringPartitions splits a UUID string watermark column into partitions.
+// The original minVal/maxVal are used for first/last partitions to ensure correctness.
+func buildUuidStringPartitions(
+	minVal string, maxVal string, casing hexCasing, numPartitions int64,
+) ([]*protos.QRepPartition, error) {
+	minInt, err := uuidToBigInt(minVal)
+	if err != nil {
+		return nil, fmt.Errorf("failed to convert min uuid to bigint: %w", err)
+	}
+	maxInt, err := uuidToBigInt(maxVal)
+	if err != nil {
+		return nil, fmt.Errorf("failed to convert max uuid to bigint: %w", err)
+	}
+	if minInt.Cmp(maxInt) > 0 {
+		return nil, fmt.Errorf("min uuid (%s) greater than max uuid (%s)", minVal, maxVal)
+	}
+
+	var partitions []*protos.QRepPartition
+	start := minVal
+	step := shared.BigIntDivCeil(new(big.Int).Sub(maxInt, minInt), big.NewInt(numPartitions))
+	for value := new(big.Int).Add(minInt, step); value.Cmp(maxInt) < 0; value.Add(value, step) {
+		end, err := bigIntToUUID(value, casing)
+		if err != nil {
+			return nil, fmt.Errorf("failed to convert bigint to uuid: %w", err)
+		}
+		partitions = append(partitions, createStringPartition(start, end, false))
+		start = end
+	}
+	partitions = append(partitions, createStringPartition(start, maxVal, true))
+	return partitions, nil
+}
+
+func uuidToBigInt(s string) (*big.Int, error) {
+	u, err := uuid.Parse(s)
+	if err != nil {
+		return nil, err
+	}
+	return new(big.Int).SetBytes(u[:]), nil
+}
+
+func bigIntToUUID(n *big.Int, casing hexCasing) (string, error) {
+	if n.BitLen() > 128 {
+		return "", fmt.Errorf("value does not fit in a UUID (%d bits)", n.BitLen())
+	}
+	var b [16]byte
+	n.FillBytes(b[:])
+	s := uuid.UUID(b).String()
+	if casing == hexCasingUpper {
+		s = strings.ToUpper(s)
+	}
+	return s, nil
+}
+
+func createStringPartition(start string, end string, endInclusive bool) *protos.QRepPartition {
+	return &protos.QRepPartition{
+		PartitionId: uuid.NewString(),
+		Range: &protos.PartitionRange{
+			Range: &protos.PartitionRange_StringRange{
+				StringRange: &protos.StringPartitionRange{
+					Start:        start,
+					End:          end,
+					EndInclusive: endInclusive,
+				},
+			},
+		},
+	}
+}

flow/connectors/mysql/qrep_partition_test.go

@@ -0,0 +1,167 @@
+package connmysql
+
+import (
+	"regexp"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/PeerDB-io/peerdb/flow/generated/protos"
+)
+
+func TestUUIDToBigIntRoundTrip(t *testing.T) {
+	cases := []struct {
+		name   string
+		uuid   string
+		casing hexCasing
+	}{
+		{"lower case v4", "f47ac10b-58cc-4372-a567-0e02b2c3d479", hexCasingLower},
+		{"upper case v7", "017F22E2-79B0-7CC3-98C4-DC0C0C07398F", hexCasingUpper},
+		{"zero", "00000000-0000-0000-0000-000000000000", hexCasingLower},
+		{"max", "ffffffff-ffff-ffff-ffff-ffffffffffff", hexCasingLower},
+		{"all digits", "01234567-8901-2345-6789-012345678901", hexCasingLower},
+	}
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			num, err := uuidToBigInt(tc.uuid)
+			require.NoError(t, err)
+			uuid, err := bigIntToUUID(num, tc.casing)
+			require.NoError(t, err)
+			assert.Equal(t, tc.uuid, uuid)
+		})
+	}
+}
+
+func TestUUIDToBigIntInvalid(t *testing.T) {
+	for _, s := range []string{
+		"",
+		"not-a-uuid",
+		"zzzzzzzz-zzzz-zzzz-zzzz-zzzzzzzzzzzz",
+		"01234567-8901-2345-6789-012345678901a", // 37 chars
+	} {
+		_, err := uuidToBigInt(s)
+		assert.Error(t, err)
+	}
+}
+
+func TestDetectUUIDCase(t *testing.T) {
+	cases := []struct {
+		name         string
+		minVal       string
+		maxVal       string
+		expectIsUUID bool
+		expectCasing hexCasing
+	}{
+		{"both lower", "00000000-0000-0000-0000-000000000000", "f47ac10b-58cc-4372-a567-0e02b2c3d479", true, hexCasingLower},
+		{"both upper", "00000000-0000-0000-0000-000000000000", "F47AC10B-58CC-4372-A567-0E02B2C3D479", true, hexCasingUpper},
+		{"all digits", "01234567-8901-2345-6789-012345678901", "09234567-8901-2345-6789-012345678901", true, hexCasingLower},
+		{"digits min, upper max", "01234567-8901-2345-6789-012345678901", "F47AC10B-58CC-4372-A567-0E02B2C3D479", true, hexCasingUpper},
+		{"lower min, digits max", "f47ac10b-58cc-4372-a567-0e02b2c3d479", "01234567-8901-2345-6789-012345678901", true, hexCasingLower},
+		{"mixed case across bounds", "f47ac10b-58cc-4372-a567-0e02b2c3d479", "F47AC10B-58CC-4372-A567-0E02B2C3D479", false, hexCasingUnknown},
+		{"mixed case within a bound", "00000000-0000-0000-0000-000000000000", "F47ac10b-58cc-4372-a567-0e02b2c3d479", false, hexCasingUnknown},
+		{"non-uuid", "apple", "banana", false, hexCasingUnknown},
+	}
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			isUUID, casing := detectUuidWithHexCasing(tc.minVal, tc.maxVal)
+			assert.Equal(t, tc.expectIsUUID, isUUID)
+			assert.Equal(t, tc.expectCasing, casing)
+		})
+	}
+}
+
+func TestBuildUuidStringPartitions(t *testing.T) {
+	cases := []struct {
+		name         string
+		minV         string
+		maxV         string
+		expectedCase *regexp.Regexp
+	}{
+		{"lower", "018f6e7a-1b2c-7def-8a3b-1c2d3e4f5a6b", "f47ac10b-58cc-4372-a567-0e02b2c3d479", uuidLowerRe},
+		{"upper", "018F6E7A-1B2C-7DEF-8A3B-1C2D3E4F5A6B", "F47AC10B-58CC-4372-A567-0E02B2C3D479", uuidUpperRe},
+	}
+	const numPartitions = 32
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			isUUID, casing := detectUuidWithHexCasing(tc.minV, tc.maxV)
+			require.True(t, isUUID)
+			parts, err := buildUuidStringPartitions(tc.minV, tc.maxV, casing, numPartitions)
+			require.NoError(t, err)
+			require.Len(t, parts, numPartitions)
+
+			prevEnd := ""
+			for i, part := range parts {
+				r, ok := part.Range.Range.(*protos.PartitionRange_StringRange)
+				require.True(t, ok)
+				sr := r.StringRange
+				if i == 0 {
+					assert.Equal(t, tc.minV, sr.Start)
+				} else {
+					assert.Equal(t, prevEnd, sr.Start)
+					assert.Regexp(t, tc.expectedCase, sr.Start)
+				}
+				startInt, err := uuidToBigInt(sr.Start)
+				require.NoError(t, err)
+				endInt, err := uuidToBigInt(sr.End)
+				require.NoError(t, err)
+				assert.Negative(t, startInt.Cmp(endInt))
+
+				isLast := i == len(parts)-1
+				assert.Equal(t, isLast, sr.EndInclusive)
+				if isLast {
+					assert.Equal(t, tc.maxV, sr.End)
+				}
+				prevEnd = sr.End
+			}
+		})
+	}
+}
+
+func TestBuildUuidStringPartitionsSinglePartition(t *testing.T) {
+	uuid1 := "018f6e7a-1b2c-7def-8a3b-1c2d3e4f5a6b"
+	uuid2 := "f47ac10b-58cc-4372-a567-0e02b2c3d479"
+	cases := []struct {
+		name          string
+		minV          string
+		maxV          string
+		numPartitions int64
+	}{
+		{"one partition requested", uuid1, uuid2, 1},
+		{"min equals max", uuid1, uuid1, 8},
+	}
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			isUUID, casing := detectUuidWithHexCasing(tc.minV, tc.maxV)
+			require.True(t, isUUID)
+			parts, err := buildUuidStringPartitions(tc.minV, tc.maxV, casing, tc.numPartitions)
+			require.NoError(t, err)
+			require.Len(t, parts, 1)
+			sr, ok := parts[0].Range.Range.(*protos.PartitionRange_StringRange)
+			require.True(t, ok)
+			assert.Equal(t, tc.minV, sr.StringRange.Start)
+			assert.Equal(t, tc.maxV, sr.StringRange.End)
+			assert.True(t, sr.StringRange.EndInclusive)
+		})
+	}
+}
+
+func TestBuildUuidStringPartitionsInvalid(t *testing.T) {
+	cases := []struct {
+		name string
+		minV string
+		maxV string
+	}{
+		{"arbitrary string", "apple", "banana"},
+		{"inverted", "f47ac10b-58cc-4372-a567-0e02b2c3d479", "018f6e7a-1b2c-7def-8a3b-1c2d3e4f5a6b"},
+		{"too short", "018f6e7a-1b2c", "f47ac10b-58cc"},
+		{"too long", "018f6e7a-1b2c-7def-8a3b-1c2d3e4f5a6bb", "f47ac10b-58cc-4372-a567-0e02b2c3d4799"},
+	}
+	const numPartitions = 32
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			_, err := buildUuidStringPartitions(tc.minV, tc.maxV, hexCasingLower, numPartitions)
+			require.Error(t, err)
+		})
+	}
+}