treat string partition ranges as sensitive data

Author: jgao54

flow/activities/flowable.go

@@ -596,6 +596,12 @@ 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)
+		}
 		if err := monitoring.InitializeQRepRun(
 			ctx,
 			logger,
@@ -663,6 +669,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 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 +1633,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,146 @@ 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(
+	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)
+		}
+		encrypted, 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, encrypted,
+		); 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
+	}
+
+	partitionIDs := make([]string, 0, len(partitions))
+	partitionByIDs := make(map[string]*protos.QRepPartition, len(partitions))
+	for _, partition := range partitions {
+		partitionIDs = append(partitionIDs, partition.PartitionId)
+		partitionByIDs[partition.PartitionId] = partition
+	}
+
+	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 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 partitionsContainSensitiveRanges(partitions []*protos.QRepPartition) bool {
+	for _, partition := range partitions {
+		if partition == nil || partition.Range == nil {
+			return false
+		}
+		// 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 {
@@ -313,5 +455,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,75 @@
+package connmetadata
+
+import (
+	"testing"
+
+	"github.com/google/uuid"
+	"github.com/jackc/pgx/v5"
+	"github.com/stretchr/testify/require"
+
+	"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)
+
+	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{
+		{
+			PartitionId: uuid.NewString(),
+			Range: &protos.PartitionRange{Range: &protos.PartitionRange_StringRange{
+				StringRange: &protos.StringPartitionRange{Start: p1Start, End: p1End},
+			}},
+		},
+		{
+			PartitionId: uuid.NewString(),
+			Range: &protos.PartitionRange{Range: &protos.PartitionRange_StringRange{
+				StringRange: &protos.StringPartitionRange{Start: p2Start, End: p2End, EndInclusive: true},
+			}},
+		},
+	}
+
+	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)
+	}
+
+	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, 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)
+}