treat string partition ranges as sensitive data

Author: jgao54

flow/activities/flowable.go

@@ -596,6 +596,10 @@ func (a *FlowableActivity) GetQRepPartitions(ctx context.Context,
 		return nil, a.Alerter.LogFlowError(ctx, config.FlowJobName, shared.WrapError("failed to get partitions from source", err))
 	}
 	if len(partitions) > 0 {
+		// run before monitoring.InitializeQRepRun to avoid logging sensitive fields to peerdb_stats.qrep_partitions as well
+		if err := connmetadata.OffloadSensitivePartitionRanges(ctx, a.CatalogPool, config.ParentMirrorName, runUUID, partitions); err != nil {
+			return nil, a.Alerter.LogFlowError(ctx, config.FlowJobName, shared.WrapError("failed to offload partition ranges", err))
+		}
 		if err := monitoring.InitializeQRepRun(
 			ctx,
 			logger,
@@ -663,6 +667,10 @@ func (a *FlowableActivity) ReplicateQRepPartitions(ctx context.Context,
 		return a.Alerter.LogFlowError(ctx, config.FlowJobName, err)
 	}
 
+	if err := connmetadata.RestoreSensitivePartitionRanges(ctx, a.CatalogPool, runUUID, partitions.Partitions); err != nil {
+		return a.Alerter.LogFlowError(ctx, config.FlowJobName, fmt.Errorf("failed to rehydrate partition ranges: %w", err))
+	}
+
 	for i, partition := range partitions.Partitions {
 		partLogger := log.With(logger,
 			slog.Int64("batchID", int64(partitions.BatchId)),
@@ -1619,6 +1627,10 @@ func (a *FlowableActivity) QRepHasNewRows(ctx context.Context,
 		if maxValue.(time.Time).After(x.TimestampRange.End.AsTime()) {
 			return true, nil
 		}
+	case *protos.PartitionRange_StringRange:
+		// checking for new rows is only possible for standalone QRepFlowWorkflow;
+		// this is a legacy feature and string partitioning is not supported
+		return false, errors.New("checking for new rows by a string partition range is not supported")
 	default:
 		return false, fmt.Errorf("unknown range type: %v", x)
 	}

flow/connectors/external_metadata/store.go

@@ -14,6 +14,7 @@ import (
 	"github.com/jackc/pgx/v5/pgtype"
 	"go.temporal.io/sdk/log"
 	"google.golang.org/protobuf/encoding/protojson"
+	"google.golang.org/protobuf/proto"
 
 	"github.com/PeerDB-io/peerdb/flow/alerting"
 	"github.com/PeerDB-io/peerdb/flow/connectors/utils/monitoring"
@@ -24,8 +25,9 @@ import (
 )
 
 const (
-	lastSyncStateTableName = "metadata_last_sync_state"
-	qrepTableName          = "metadata_qrep_partitions"
+	lastSyncStateTableName       = "metadata_last_sync_state"
+	qrepTableName                = "metadata_qrep_partitions"
+	qrepPartitionRangesTableName = "metadata_qrep_partition_ranges"
 )
 
 type PostgresMetadata struct {
@@ -286,6 +288,151 @@ func (p *PostgresMetadata) IsQRepPartitionSynced(ctx context.Context, req *proto
 	return exists, nil
 }
 
+// OffloadSensitivePartitionRanges encrypts and persists sensitive partition ranges
+// into metadata_qrep_partition_ranges and then strip the existing ranges and
+// replacing them with an empty range so the value never crosses a Temporal boundary.
+// RestoreSensitivePartitionRanges repopulates them from the catalog before replicating.
+func OffloadSensitivePartitionRanges(
+	ctx context.Context,
+	pool shared.CatalogPool,
+	parentMirrorName string,
+	runUUID string,
+	partitions []*protos.QRepPartition,
+) error {
+	if !partitionsContainSensitiveRanges(partitions) {
+		return nil
+	}
+
+	key, err := internal.PeerDBCurrentEncKey(ctx)
+	if err != nil {
+		return fmt.Errorf("failed to load current encryption key: %w", err)
+	}
+
+	tx, err := pool.Begin(ctx)
+	if err != nil {
+		return fmt.Errorf("failed to begin transaction: %w", err)
+	}
+	defer shared.RollbackTx(tx, internal.LoggerFromCtx(ctx))
+
+	// delete any orphaned rows from previously failed attempt
+	if _, err := tx.Exec(ctx,
+		`DELETE FROM `+qrepPartitionRangesTableName+` WHERE run_uuid=$1`, runUUID,
+	); err != nil {
+		return fmt.Errorf("failed to clear existing partition ranges: %w", err)
+	}
+
+	for _, partition := range partitions {
+		payload, err := proto.Marshal(partition.Range)
+		if err != nil {
+			return fmt.Errorf("failed to marshal partition range: %w", err)
+		}
+		ciphertext, err := key.Encrypt(payload)
+		if err != nil {
+			return fmt.Errorf("failed to encrypt partition range: %w", err)
+		}
+
+		if _, err := tx.Exec(ctx,
+			`INSERT INTO `+qrepPartitionRangesTableName+
+				`(parent_mirror_name,run_uuid,partition_uuid,enc_key_id,range_payload) VALUES($1,$2,$3,$4,$5)`,
+			parentMirrorName, runUUID, partition.PartitionId, key.ID, ciphertext,
+		); err != nil {
+			return fmt.Errorf("failed to persist encrypted partition range: %w", err)
+		}
+
+		partition.Range = &protos.PartitionRange{
+			Range: &protos.PartitionRange_StringRange{StringRange: &protos.StringPartitionRange{}},
+		}
+	}
+
+	if err := tx.Commit(ctx); err != nil {
+		return fmt.Errorf("failed to commit partition ranges: %w", err)
+	}
+	return nil
+}
+
+// RestoreSensitivePartitionRanges hydrates sensitive partitions in-place with the decrypted
+// values fetched from the catalog. Used in conjunction with OffloadSensitivePartitionRanges
+func RestoreSensitivePartitionRanges(
+	ctx context.Context,
+	pool shared.CatalogPool,
+	runUUID string,
+	partitions []*protos.QRepPartition,
+) error {
+	if !partitionsContainSensitiveRanges(partitions) {
+		return nil
+	}
+
+	byID := make(map[string]*protos.QRepPartition, len(partitions))
+	partitionIDs := make([]string, 0, len(partitions))
+	for _, partition := range partitions {
+		byID[partition.PartitionId] = partition
+		partitionIDs = append(partitionIDs, partition.PartitionId)
+	}
+
+	rows, err := pool.Query(ctx,
+		`SELECT partition_uuid,enc_key_id,range_payload FROM `+qrepPartitionRangesTableName+` `+
+			`WHERE run_uuid=$1 AND partition_uuid=ANY($2)`,
+		runUUID, partitionIDs,
+	)
+	if err != nil {
+		return fmt.Errorf("failed to read encrypted partition ranges: %w", err)
+	}
+	defer rows.Close()
+
+	for rows.Next() {
+		var partitionID, encKeyID string
+		var ciphertext []byte
+		if err := rows.Scan(&partitionID, &encKeyID, &ciphertext); err != nil {
+			return fmt.Errorf("failed to scan encrypted partition range: %w", err)
+		}
+
+		partition, ok := byID[partitionID]
+		if !ok {
+			continue
+		}
+
+		payload, err := internal.Decrypt(ctx, encKeyID, ciphertext)
+		if err != nil {
+			return fmt.Errorf("failed to decrypt partition range for %s: %w", partitionID, err)
+		}
+
+		var partitionRange protos.PartitionRange
+		if err := proto.Unmarshal(payload, &partitionRange); err != nil {
+			return fmt.Errorf("failed to unmarshal partition range for %s: %w", partitionID, err)
+		}
+		partition.Range = &partitionRange
+		delete(byID, partitionID)
+	}
+	if err := rows.Err(); err != nil {
+		return fmt.Errorf("failed to read encrypted partition ranges: %w", err)
+	}
+
+	if len(byID) > 0 {
+		missing := make([]string, 0, len(byID))
+		for id := range byID {
+			missing = append(missing, id)
+		}
+		return fmt.Errorf("missing encrypted partition ranges for partitions: %v", missing)
+	}
+
+	return nil
+}
+
+// partitionsContainSensitiveRanges reports whether the partitions hold data that
+// must not cross a Temporal boundary in plaintext. String ranges may contain PII
+// data like email, so deem it as sensitive by default.
+func partitionsContainSensitiveRanges(partitions []*protos.QRepPartition) bool {
+	for _, partition := range partitions {
+		if partition == nil || partition.Range == nil {
+			return false
+		}
+		if _, ok := partition.Range.Range.(*protos.PartitionRange_StringRange); ok {
+			return true
+		}
+	}
+	return false
+}
+
 func (p *PostgresMetadata) SyncFlowCleanup(ctx context.Context, jobName string) error {
 	tx, err := p.pool.Begin(ctx)
 	if err != nil {
@@ -313,5 +460,9 @@ func SyncFlowCleanupInTx(ctx context.Context, tx pgx.Tx, jobName string) error {
 		return err
 	}
 
+	if _, err := tx.Exec(ctx, `DELETE FROM `+qrepPartitionRangesTableName+` WHERE parent_mirror_name = $1`, jobName); err != nil {
+		return err
+	}
+
 	return nil
 }

flow/e2e/clickhouse_mysql_test.go

@@ -1819,20 +1819,16 @@ func (s ClickHouseSuite) Test_MySQL_String_Partition_Key_UUID_Parallel_Snapshot(
 	EnvWaitForEqualTablesWithNames(env, s, "waiting on initial", srcTable, dstTable, "id,val")
 
 	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`,
+		`SELECT 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)
+		require.NoError(s.t, partitionRows.Scan(&rowsInPartition))
 		totalRows += rowsInPartition
 		partitionCount++
 	}
@@ -1910,3 +1906,103 @@ func (s ClickHouseSuite) Test_MySQL_String_Partition_Key_Arbitrary_FullTable() {
 	env.Cancel(s.t.Context())
 	RequireEnvCanceled(s.t, env)
 }
+
+// Test_MySQL_String_Partition_Key_Sensitive_Data_Handling verifies that string partition ranges,
+// which may carry sensitive watermark values, are never stored in plaintext in the catalog: they are
+// stripped from peerdb_stats.qrep_partitions and persisted encrypted in metadata_qrep_partition_ranges.
+func (s ClickHouseSuite) Test_MySQL_String_Partition_Key_Sensitive_Range_Handling() {
+	if _, ok := s.source.(*MySqlSource); !ok {
+		s.t.Skip("only applies to mysql")
+	}
+
+	srcTable := "test_string_pk_sensitive"
+	srcFullName := s.attachSchemaSuffix(srcTable)
+	dstTable := "test_string_pk_sensitive_dst"
+
+	const numRows = 8
+	const numPartitions = 4
+	require.NoError(s.t, s.source.Exec(s.t.Context(),
+		fmt.Sprintf("CREATE TABLE %s (id CHAR(36) PRIMARY KEY, val INT NOT NULL)", srcFullName)))
+	insertedUUIDs := make([]string, 0, numRows)
+	for i := 1; i <= numRows; i++ {
+		id := uuid.NewString()
+		insertedUUIDs = append(insertedUUIDs, id)
+		require.NoError(s.t, s.source.Exec(s.t.Context(),
+			fmt.Sprintf("INSERT INTO %s (id, val) VALUES ('%s', %d)", srcFullName, id, i)))
+	}
+
+	tableMappings := TableMappings(s, srcTable, dstTable)
+	// a String column can't be used directly as a Distributed sharding key (ClickHouse
+	// requires an integer-typed sharding expression), so hash it for the cluster suite
+	for _, tm := range tableMappings {
+		tm.ShardingKey = "cityHash64(id)"
+	}
+	connectionGen := FlowConnectionGenerationConfig{
+		FlowJobName:   s.attachSuffix("string_pk_sensitive"),
+		TableMappings: tableMappings,
+		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,val")
+
+	partitionRows, err := s.catalog.Query(s.t.Context(),
+		`SELECT partition_start, partition_end FROM peerdb_stats.qrep_partitions WHERE parent_mirror_name = $1`,
+		flowConnConfig.FlowJobName)
+	require.NoError(s.t, err)
+	defer partitionRows.Close()
+
+	var partitionCount int32
+	for partitionRows.Next() {
+		var partitionStart, partitionEnd *string
+		var rowsInPartition int64
+		require.NoError(s.t, partitionRows.Scan(&partitionStart, &partitionEnd, &rowsInPartition))
+		require.Empty(s.t, partitionStart)
+		require.Empty(s.t, partitionEnd)
+		partitionCount++
+	}
+	require.NoError(s.t, partitionRows.Err())
+	require.EqualValues(s.t, numPartitions, partitionCount)
+
+	rangeRows, err := s.catalog.Query(s.t.Context(),
+		`SELECT enc_key_id, range_payload FROM metadata_qrep_partition_ranges WHERE parent_mirror_name = $1`,
+		flowConnConfig.FlowJobName)
+	require.NoError(s.t, err)
+	defer rangeRows.Close()
+
+	var encryptedCount int32
+	for rangeRows.Next() {
+		var encKeyID string
+		var rangePayload []byte
+		require.NoError(s.t, rangeRows.Scan(&encKeyID, &rangePayload))
+		require.NotEmpty(s.t, encKeyID)
+		require.NotEmpty(s.t, rangePayload)
+		// payload must be ciphertext: none of the plaintext watermark UUIDs should appear in it
+		for _, id := range insertedUUIDs {
+			require.NotContains(s.t, string(rangePayload), id)
+		}
+		encryptedCount++
+	}
+	require.NoError(s.t, rangeRows.Err())
+	require.EqualValues(s.t, numPartitions, encryptedCount)
+
+	env.Cancel(s.t.Context())
+	RequireEnvCanceled(s.t, env)
+
+	// dropping the mirror should delete metadata_qrep_partition_ranges entries from the catalog for this mirror.
+	dropEnv := ExecuteDropFlow(s.t.Context(), tc, flowConnConfig, 0)
+	EnvWaitForFinished(s.t, dropEnv, 3*time.Minute)
+
+	var remainingRanges int64
+	require.NoError(s.t, s.catalog.QueryRow(s.t.Context(),
+		`SELECT COUNT(*) FROM metadata_qrep_partition_ranges WHERE parent_mirror_name = $1`,
+		flowConnConfig.FlowJobName).Scan(&remainingRanges))
+	require.Zero(s.t, remainingRanges)
+}

nexus/catalog/migrations/V54__qrep_partition_ranges.sql

@@ -0,0 +1,17 @@
+-- Stores QRep partition ranges offloaded to the catalog instead of being passed
+-- through Temporal. Today this is used only for string partition ranges, whose
+-- bounds carry user watermark values that could be sensitive, so they are encrypted
+-- and kept out of Temporal payloads.
+-- The table name and columns are deliberately generic so a future change can offload
+-- all partition ranges here to stay under Temporal's payload size limit, not just strings.
+CREATE TABLE IF NOT EXISTS metadata_qrep_partition_ranges (
+    parent_mirror_name TEXT NOT NULL,
+    run_uuid TEXT NOT NULL,
+    partition_uuid TEXT NOT NULL,
+    enc_key_id TEXT NOT NULL,
+    range_payload BYTEA NOT NULL,
+    PRIMARY KEY (run_uuid, partition_uuid)
+);
+
+CREATE INDEX IF NOT EXISTS idx_metadata_qrep_partition_ranges_parent_mirror_name
+    ON metadata_qrep_partition_ranges (parent_mirror_name);