review feedbacks:

- decouple offload and restore (so restore derive based on range_offloaded flag)
- always encrypt non-nil ranges regardless of type
- update schema to include inserted_at timestamp
- make sure metadata_qrep_partitions also does not store serialized partition in plaintext
- add dynconf PEERDB_OFFLOAD_PARTITION_RANGES to allow disabling offload/encryption behind a FF
- update tests

Author: jgao54

flow/activities/flowable.go

@@ -596,11 +596,16 @@ 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, fmt.Errorf("failed to offload partition ranges: %w", err)
+		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 {
+			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,
@@ -669,7 +674,7 @@ 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 {
+	if err := connmetadata.RestoreOffloadedPartitionRanges(ctx, a.CatalogPool, runUUID, partitions.Partitions); err != nil {
 		return fmt.Errorf("failed to rehydrate partition ranges: %w", err)
 	}
 

flow/connectors/external_metadata/store.go

@@ -6,6 +6,8 @@ import (
 	"errors"
 	"fmt"
 	"log/slog"
+	"maps"
+	"slices"
 	"strconv"
 	"time"
 
@@ -27,7 +29,7 @@ import (
 const (
 	lastSyncStateTableName       = "metadata_last_sync_state"
 	qrepTableName                = "metadata_qrep_partitions"
-	qrepPartitionRangesTableName = "metadata_qrep_partition_ranges"
+	qrepPartitionRangesTableName = "metadata_qrep_offloaded_partition_ranges"
 )
 
 type PostgresMetadata struct {
@@ -262,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.Clone(partition).(*protos.QRepPartition)
+		partition.Range = nil
+	}
 	pbytes, err := protojson.Marshal(partition)
 	if err != nil {
 		return fmt.Errorf("failed to marshal partition to json: %w", err)
@@ -288,20 +299,16 @@ func (p *PostgresMetadata) IsQRepPartitionSynced(ctx context.Context, req *proto
 	return exists, nil
 }
 
-// OffloadSensitivePartitionRanges encrypts sensitive partition ranges, persists them to
-// the catalog, and replaces each partition's range in-place with an empty range.
-// Used in conjunction with RestoreSensitivePartitionRanges.
-func OffloadSensitivePartitionRanges(
+// 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 {
-	if !partitionsContainSensitiveRanges(partitions) {
-		return nil
-	}
-
 	key, err := internal.PeerDBCurrentEncKey(ctx)
 	if err != nil {
 		return fmt.Errorf("failed to load current encryption key: %w", err)
@@ -322,6 +329,9 @@ func OffloadSensitivePartitionRanges(
 
 	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)
@@ -332,9 +342,8 @@ func OffloadSensitivePartitionRanges(
 		}
 		insertRows = append(insertRows, []any{parentMirrorName, runUUID, partition.PartitionId, key.ID, encrypted})
 
-		partition.Range = &protos.PartitionRange{
-			Range: &protos.PartitionRange_StringRange{StringRange: &protos.StringPartitionRange{}},
-		}
+		partition.Range = nil
+		partition.RangeOffloaded = true
 	}
 
 	if _, err := tx.CopyFrom(ctx,
@@ -351,29 +360,29 @@ func OffloadSensitivePartitionRanges(
 	return nil
 }
 
-// RestoreSensitivePartitionRanges hydrates sensitive partitions in-place with the decrypted
-// values fetched from the catalog. Used in conjunction with OffloadSensitivePartitionRanges
-func RestoreSensitivePartitionRanges(
+// 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 {
-	if !partitionsContainSensitiveRanges(partitions) {
-		return nil
-	}
-
-	partitionIDs := make([]string, 0, len(partitions))
 	partitionByIDs := make(map[string]*protos.QRepPartition, len(partitions))
 	for _, partition := range partitions {
-		partitionIDs = append(partitionIDs, partition.PartitionId)
+		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, partitionIDs,
+		runUUID, slices.Collect(maps.Keys(partitionByIDs)),
 	)
 	if err != nil {
 		return fmt.Errorf("failed to query partition ranges: %w", err)
@@ -417,19 +426,6 @@ func RestoreSensitivePartitionRanges(
 	return nil
 }
 
-func partitionsContainSensitiveRanges(partitions []*protos.QRepPartition) bool {
-	for _, partition := range partitions {
-		if partition == nil || partition.Range == nil {
-			continue
-		}
-		// String ranges may contain PII data like email, so deem it as sensitive.
-		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 {

flow/connectors/external_metadata/store_test.go

@@ -2,10 +2,13 @@ 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"
@@ -20,56 +23,87 @@ func TestOffloadRestoreSensitivePartitionRanges(t *testing.T) {
 	pool, err := internal.GetCatalogConnectionPoolFromEnv(ctx)
 	require.NoError(t, err)
 
-	parentMirrorName := "test_offload_restore_sensitive_partition_ranges"
-	runUUID := uuid.NewString()
-	t.Cleanup(func() {
-		_, _ = pool.Exec(ctx, `DELETE FROM `+qrepPartitionRangesTableName+` WHERE parent_mirror_name=$1`, parentMirrorName)
-	})
-
-	p1Start, p1End, p2Start, p2End := "bar", "foo", "hello", "world"
-	partitions := []*protos.QRepPartition{
+	for _, tc := range []struct {
+		name   string
+		ranges []*protos.PartitionRange
+	}{
 		{
-			PartitionId: uuid.NewString(),
-			Range: &protos.PartitionRange{Range: &protos.PartitionRange_StringRange{
-				StringRange: &protos.StringPartitionRange{Start: p1Start, End: p1End},
-			}},
+			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}}},
+			},
 		},
 		{
-			PartitionId: uuid.NewString(),
-			Range: &protos.PartitionRange{Range: &protos.PartitionRange_StringRange{
-				StringRange: &protos.StringPartitionRange{Start: p2Start, End: p2End, EndInclusive: true},
-			}},
+			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)
+			})
 
-	require.NoError(t, OffloadSensitivePartitionRanges(ctx, pool, parentMirrorName, runUUID, partitions))
-	for _, p := range partitions {
-		require.Empty(t, p.Range.GetStringRange().Start)
-		require.Empty(t, p.Range.GetStringRange().End)
-	}
+			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)
+			}
 
-	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)
-		for _, plaintext := range []string{p1Start, p1End, p2Start, p2End} {
-			require.NotContains(t, string(res.Payload), plaintext)
-		}
-	}
+			require.NoError(t, OffloadPartitionRanges(ctx, pool, parentMirrorName, runUUID, partitions))
+			for _, p := range partitions {
+				require.Nil(t, p.Range)
+				require.True(t, p.RangeOffloaded)
+			}
 
-	require.NoError(t, RestoreSensitivePartitionRanges(ctx, pool, runUUID, partitions))
-	require.Equal(t, p1Start, partitions[0].Range.GetStringRange().Start)
-	require.Equal(t, p1End, partitions[0].Range.GetStringRange().End)
-	require.False(t, partitions[0].Range.GetStringRange().EndInclusive)
-	require.Equal(t, p2Start, partitions[1].Range.GetStringRange().Start)
-	require.Equal(t, p2End, partitions[1].Range.GetStringRange().End)
-	require.True(t, partitions[1].Range.GetStringRange().EndInclusive)
+			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))
 	}