treat string partition ranges as sensitive data (#4501)

Before introducing parallel snapshotting for arbitrary string watermark
column, in order to avoid sensitive data from being accidentally exposed
in db/Temporal as plaintext, this PR will offload partition ranges to db
(encrypted) before return the `GetQRepPartitions` Activity result, and
then restored in the `ReplicateQRepPartitions` Activity, to avoid
crossing Temporal boundary. Long term, offloading all partition ranges
to DB is also ideal, in order to not having to worry about Temporal
payload limit.

This PR introduces a new table `metadata_qrep_partition_ranges` for
storing partition ranges. I also evaluated building on top of existing
`peerdb_stats.qrep_partitions` but that table is intended for monitoring
so adding operational logic to it is not ideal. This PR also makes sure
that sensitive ranges are removed before writing to
`peerdb_stats.qrep_partitions`, which are not being used afaik.

Open to feedback on what we should deem as sensitive that should be
encrypted (e.g. all string vs. non-UUIDs vs. only encrypt email-shaped
string). Currently all string ranges are encrypted.

Testing: add integration test coverage for offload/restore correctness
logic; as well as e2e test for the workflow, and that table is cleanup
when mirror is dropped.

Fixes: DBI-869

Author: jgao54

flow/activities/flowable.go

@@ -596,6 +596,17 @@ 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 {
+		shouldOffload, err := internal.PeerDBOffloadPartitionRanges(ctx, config.Env)
+		if err != nil {
+			return nil, fmt.Errorf("failed to read offload partition ranges setting: %w", err)
+		}
+		if shouldOffload && config.InitialCopyOnly {
+			if err := connmetadata.OffloadPartitionRanges(
+				ctx, a.CatalogPool, config.ParentMirrorName, runUUID, partitions,
+			); err != nil {
+				return nil, fmt.Errorf("failed to offload partition ranges: %w", err)
+			}
+		}
 		if err := monitoring.InitializeQRepRun(
 			ctx,
 			logger,
@@ -663,6 +674,10 @@ func (a *FlowableActivity) ReplicateQRepPartitions(ctx context.Context,
 		return a.Alerter.LogFlowError(ctx, config.FlowJobName, err)
 	}
 
+	if err := connmetadata.RestoreOffloadedPartitionRanges(ctx, a.CatalogPool, runUUID, partitions.Partitions); err != nil {
+		return 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)),
@@ -1623,6 +1638,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

@@ -6,6 +6,8 @@ import (
 	"errors"
 	"fmt"
 	"log/slog"
+	"maps"
+	"slices"
 	"strconv"
 	"time"
 
@@ -14,6 +16,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 +27,9 @@ import (
 )
 
 const (
-	lastSyncStateTableName = "metadata_last_sync_state"
-	qrepTableName          = "metadata_qrep_partitions"
+	lastSyncStateTableName       = "metadata_last_sync_state"
+	qrepTableName                = "metadata_qrep_partitions"
+	qrepPartitionRangesTableName = "metadata_qrep_offloaded_partition_ranges"
 )
 
 type PostgresMetadata struct {
@@ -260,6 +264,15 @@ func (p *PostgresMetadata) FinishQRepPartition(
 	jobName string,
 	startTime time.Time,
 ) error {
+	if partition.RangeOffloaded {
+		// The range was offloaded but has since been rehydrated, so drop it before persisting
+		// to metadata_qrep_partitions.
+		// Clone rather than mutate in place because this partition is shared with the pull
+		// goroutine, cloning keeps us from relying on implicit ordering to avoid a write race.
+		// TODO: drop the sync_partition column entirely; nothing reads it back.
+		partition = proto.CloneOf(partition)
+		partition.Range = nil
+	}
 	pbytes, err := protojson.Marshal(partition)
 	if err != nil {
 		return fmt.Errorf("failed to marshal partition to json: %w", err)
@@ -286,6 +299,133 @@ func (p *PostgresMetadata) IsQRepPartitionSynced(ctx context.Context, req *proto
 	return exists, nil
 }
 
+// OffloadPartitionRanges encrypts partition ranges and persists them to the catalog,
+// then clears each partition's range in-place and sets RangeOffloaded so the range can
+// be rehydrated later. Used in conjunction with RestoreOffloadedPartitionRanges.
+func OffloadPartitionRanges(
+	ctx context.Context,
+	pool shared.CatalogPool,
+	parentMirrorName string,
+	runUUID string,
+	partitions []*protos.QRepPartition,
+) error {
+	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)
+	}
+
+	insertRows := make([][]any, 0, len(partitions))
+	for _, partition := range partitions {
+		if partition.Range == nil {
+			continue
+		}
+		payload, err := proto.Marshal(partition.Range)
+		if err != nil {
+			return fmt.Errorf("failed to marshal partition range: %w", err)
+		}
+		encrypted, err := key.Encrypt(payload)
+		if err != nil {
+			return fmt.Errorf("failed to encrypt partition range: %w", err)
+		}
+		insertRows = append(insertRows, []any{parentMirrorName, runUUID, partition.PartitionId, key.ID, encrypted})
+
+		partition.Range = nil
+		partition.RangeOffloaded = true
+	}
+
+	if _, err := tx.CopyFrom(ctx,
+		pgx.Identifier{qrepPartitionRangesTableName},
+		[]string{"parent_mirror_name", "run_uuid", "partition_uuid", "enc_key_id", "range_payload"},
+		pgx.CopyFromRows(insertRows),
+	); err != nil {
+		return fmt.Errorf("failed to persist encrypted partition ranges: %w", err)
+	}
+
+	if err := tx.Commit(ctx); err != nil {
+		return fmt.Errorf("failed to commit partition ranges: %w", err)
+	}
+	return nil
+}
+
+// RestoreOffloadedPartitionRanges hydrates offloaded partition ranges in-place with the decrypted
+// values fetched from the catalog. Used in conjunction with OffloadPartitionRanges.
+func RestoreOffloadedPartitionRanges(
+	ctx context.Context,
+	pool shared.CatalogPool,
+	runUUID string,
+	partitions []*protos.QRepPartition,
+) error {
+	partitionByIDs := make(map[string]*protos.QRepPartition, len(partitions))
+	for _, partition := range partitions {
+		if !partition.RangeOffloaded {
+			continue
+		}
+		partitionByIDs[partition.PartitionId] = partition
+	}
+	if len(partitionByIDs) == 0 {
+		return nil
+	}
+
+	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, slices.Collect(maps.Keys(partitionByIDs)),
+	)
+	if err != nil {
+		return fmt.Errorf("failed to query partition ranges: %w", err)
+	}
+	results, err := pgx.CollectRows(rows, pgx.RowToStructByPos[struct {
+		PartitionID string
+		EncKeyID    string
+		Payload     []byte
+	}])
+	if err != nil {
+		return fmt.Errorf("failed to read encrypted partition ranges: %w", err)
+	}
+
+	for _, res := range results {
+		partition, ok := partitionByIDs[res.PartitionID]
+		if !ok {
+			return fmt.Errorf("received unexpected partition range for partition %s", res.PartitionID)
+		}
+
+		payload, err := internal.Decrypt(ctx, res.EncKeyID, res.Payload)
+		if err != nil {
+			return fmt.Errorf("failed to decrypt partition range for %s: %w", res.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", res.PartitionID, err)
+		}
+		partition.Range = &partitionRange
+		delete(partitionByIDs, res.PartitionID)
+	}
+
+	if len(partitionByIDs) > 0 {
+		missing := make([]string, 0, len(partitionByIDs))
+		for id := range partitionByIDs {
+			missing = append(missing, id)
+		}
+		return fmt.Errorf("missing encrypted partition ranges for partitions: %v", missing)
+	}
+
+	return nil
+}
+
 func (p *PostgresMetadata) SyncFlowCleanup(ctx context.Context, jobName string) error {
 	tx, err := p.pool.Begin(ctx)
 	if err != nil {
@@ -313,5 +453,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/connectors/external_metadata/store_test.go

@@ -0,0 +1,109 @@
+package connmetadata
+
+import (
+	"testing"
+	"time"
+
+	"github.com/google/uuid"
+	"github.com/jackc/pgx/v5"
+	"github.com/stretchr/testify/require"
+	"google.golang.org/protobuf/proto"
+	"google.golang.org/protobuf/types/known/timestamppb"
+
+	"github.com/PeerDB-io/peerdb/flow/generated/protos"
+	"github.com/PeerDB-io/peerdb/flow/internal"
+)
+
+func TestOffloadRestoreSensitivePartitionRanges(t *testing.T) {
+	const encKeyID = "test_enc_key"
+	t.Setenv("PEERDB_CURRENT_ENC_KEY_ID", encKeyID)
+	t.Setenv("PEERDB_ENC_KEYS", `[{"id":"`+encKeyID+`","value":"cGVlcmRiX2NpX3Rlc3RfZW5jX2tleV8zMl9ieXRlcyE="}]`)
+
+	ctx := t.Context()
+	pool, err := internal.GetCatalogConnectionPoolFromEnv(ctx)
+	require.NoError(t, err)
+
+	for _, tc := range []struct {
+		name   string
+		ranges []*protos.PartitionRange
+	}{
+		{
+			name: "int",
+			ranges: []*protos.PartitionRange{
+				{Range: &protos.PartitionRange_IntRange{IntRange: &protos.IntPartitionRange{Start: 1, End: 100}}},
+				{Range: &protos.PartitionRange_IntRange{IntRange: &protos.IntPartitionRange{Start: 101, End: 200}}},
+			},
+		},
+		{
+			name: "timestamp",
+			ranges: []*protos.PartitionRange{
+				{Range: &protos.PartitionRange_TimestampRange{TimestampRange: &protos.TimestampPartitionRange{
+					Start: timestamppb.New(time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC)),
+					End:   timestamppb.New(time.Date(2026, 6, 30, 23, 59, 59, 0, time.UTC)),
+				}}},
+				{Range: &protos.PartitionRange_TimestampRange{TimestampRange: &protos.TimestampPartitionRange{
+					Start: timestamppb.New(time.Date(2026, 7, 1, 0, 0, 0, 0, time.UTC)),
+					End:   timestamppb.New(time.Date(2026, 12, 31, 23, 59, 59, 0, time.UTC)),
+				}}},
+			},
+		},
+		{
+			name: "string_uuid",
+			ranges: []*protos.PartitionRange{
+				{Range: &protos.PartitionRange_StringRange{
+					StringRange: &protos.StringPartitionRange{Start: uuid.NewString(), End: uuid.NewString()},
+				}},
+				{Range: &protos.PartitionRange_StringRange{
+					StringRange: &protos.StringPartitionRange{Start: uuid.NewString(), End: uuid.NewString(), EndInclusive: true},
+				}},
+			},
+		},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			parentMirrorName := "test_offload_restore_partition_ranges_" + tc.name
+			runUUID := uuid.NewString()
+			t.Cleanup(func() {
+				_, _ = pool.Exec(ctx, `DELETE FROM `+qrepPartitionRangesTableName+` WHERE parent_mirror_name=$1`, parentMirrorName)
+			})
+
+			partitions := make([]*protos.QRepPartition, len(tc.ranges))
+			expectedRanges := make([]*protos.PartitionRange, len(tc.ranges))
+			for i, r := range tc.ranges {
+				partitions[i] = &protos.QRepPartition{PartitionId: uuid.NewString(), Range: r}
+				expectedRanges[i] = proto.Clone(r).(*protos.PartitionRange)
+			}
+
+			require.NoError(t, OffloadPartitionRanges(ctx, pool, parentMirrorName, runUUID, partitions))
+			for _, p := range partitions {
+				require.Nil(t, p.Range)
+				require.True(t, p.RangeOffloaded)
+			}
+
+			rows, err := pool.Query(ctx,
+				`SELECT enc_key_id, range_payload FROM `+qrepPartitionRangesTableName+` WHERE parent_mirror_name=$1`,
+				parentMirrorName)
+			require.NoError(t, err)
+			results, err := pgx.CollectRows(rows, pgx.RowToStructByPos[struct {
+				KeyID   string
+				Payload []byte
+			}])
+			require.NoError(t, err)
+			require.Len(t, results, len(partitions))
+			for _, res := range results {
+				require.Equal(t, encKeyID, res.KeyID)
+				require.NotEmpty(t, res.Payload)
+				for _, r := range expectedRanges {
+					unencrypted, err := proto.Marshal(r)
+					require.NoError(t, err)
+					require.NotContains(t, string(res.Payload), string(unencrypted))
+				}
+			}
+
+			require.NoError(t, RestoreOffloadedPartitionRanges(ctx, pool, runUUID, partitions))
+			for i, p := range partitions {
+				require.Truef(t, proto.Equal(expectedRanges[i], p.Range),
+					"partition %d range mismatch after restore: want %v, got %v", i, expectedRanges[i], p.Range)
+			}
+		})
+	}
+}

flow/connectors/utils/monitoring/monitoring.go

@@ -332,7 +332,11 @@ func addPartitionToQRepRun(ctx context.Context, tx pgx.Tx, flowJobName string,
 	}
 
 	var rangeStart, rangeEnd *string
-	if partition.Range != nil {
+	if partition.RangeOffloaded {
+		internal.LoggerFromCtx(ctx).Info(
+			"omitting partition_start/partition_end from qrep_partitions because it's offloaded",
+			slog.String("partitionId", partition.PartitionId))
+	} else if partition.Range != nil {
 		switch x := partition.Range.Range.(type) {
 		case *protos.PartitionRange_IntRange:
 			rangeStart, rangeEnd = new(strconv.FormatInt(x.IntRange.Start, 10)), new(strconv.FormatInt(x.IntRange.End, 10))
@@ -369,7 +373,7 @@ func addPartitionToQRepRun(ctx context.Context, tx pgx.Tx, flowJobName string,
 		default:
 			return fmt.Errorf("unknown range type: %v", x)
 		}
-	} else if !partition.FullTablePartition {
+	} else if !partition.FullTablePartition && len(partition.ChildTableRanges) == 0 {
 		internal.LoggerFromCtx(ctx).Warn("[monitoring]: partition "+partition.PartitionId+" has nil range",
 			slog.String(string(shared.FlowNameKey), parentMirrorName))
 	}