Use unique staging paths for BigQuery source

flow/activities/snapshot_activity.go

@@ -25,7 +25,8 @@ type SlotSnapshotState struct {
 }
 
 type TxSnapshotState struct {
-	SnapshotName string
+	SnapshotName        string
+	SnapshotStagingPath string
 }
 
 type SnapshotActivity struct {
@@ -123,7 +124,8 @@ func (a *SnapshotActivity) MaintainTx(ctx context.Context, sessionID string, flo
 	a.SnapshotStatesMutex.Lock()
 	if exportSnapshotOutput != nil {
 		a.TxSnapshotStates[sessionID] = TxSnapshotState{
-			SnapshotName: exportSnapshotOutput.SnapshotName,
+			SnapshotName:        exportSnapshotOutput.SnapshotName,
+			SnapshotStagingPath: exportSnapshotOutput.SnapshotStagingPath,
 		}
 	} else {
 		a.TxSnapshotStates[sessionID] = TxSnapshotState{}

flow/connectors/bigquery/qrep_object_pull.go

@@ -11,6 +11,7 @@ import (
 	"cloud.google.com/go/bigquery"
 	"cloud.google.com/go/storage"
 	"github.com/google/uuid"
+	"go.temporal.io/sdk/activity"
 	"google.golang.org/api/googleapi"
 	"google.golang.org/api/iterator"
 
@@ -331,9 +332,18 @@ func (c *BigQueryConnector) ExportTxSnapshot(
 
 	_ = c.LogFlowInfo(ctx, flowName, "Starting snapshot BigQuery export to GCS staging bucket")
 
+	// Make the staging path unique per-run so it can be cleaned up asynchronously on resync.
+	activityInfo := activity.GetInfo(ctx)
+	basePath, err := parseGCSPath(cfg.SnapshotStagingPath)
+	if err != nil {
+		return nil, nil, fmt.Errorf("failed to parse snapshot staging path %q: %w", cfg.SnapshotStagingPath, err)
+	}
+	runSnapshotStagingPath := basePath.JoinPath(activityInfo.WorkflowExecution.RunID).String()
+	c.logger.Info("Run snapshot staging path", slog.String("path", runSnapshotStagingPath))
+
 	jobs := make(map[string]*bigquery.Job)
 	for _, tm := range cfg.TableMappings {
-		exportSQL, err := c.bigQueryExportQueryStatement(ctx, tm.SourceTableIdentifier, cfg.SnapshotStagingPath)
+		exportSQL, err := c.bigQueryExportQueryStatement(ctx, tm.SourceTableIdentifier, runSnapshotStagingPath)
 		if err != nil {
 			return nil, nil, fmt.Errorf("failed to build export SQL for table %s: %w", tm.SourceTableIdentifier, err)
 		}
@@ -365,7 +375,7 @@ func (c *BigQueryConnector) ExportTxSnapshot(
 		_ = c.LogFlowInfo(ctx, flowName, "Exported snapshot data to GCS for table "+sourceTableIdentifier)
 	}
 
-	return nil, cfg.SnapshotStagingPath, nil
+	return &protos.ExportTxSnapshotOutput{SnapshotStagingPath: runSnapshotStagingPath}, runSnapshotStagingPath, nil
 }
 
 // bigQueryExportQueryStatement builds the EXPORT DATA SQL statement for exporting data from BigQuery to GCS in Parquet format.

flow/workflows/snapshot_flow.go

@@ -1,6 +1,7 @@
 package peerflow
 
 import (
+	"cmp"
 	"fmt"
 	"log/slog"
 	"time"
@@ -102,6 +103,7 @@ func (s *SnapshotFlowExecution) cloneTable(
 	ctx workflow.Context,
 	boundSelector *shared.BoundSelector,
 	snapshotName string,
+	stagingPathOverride string,
 	mapping *protos.TableMapping,
 	sourcePeerType protos.DBType,
 	destinationPeerType protos.DBType,
@@ -209,7 +211,7 @@ func (s *SnapshotFlowExecution) cloneTable(
 		NumRowsPerPartition:        numRowsPerPartition,
 		NumPartitionsOverride:      numPartitionsOverride,
 		MaxParallelWorkers:         numWorkers,
-		StagingPath:                s.config.SnapshotStagingPath,
+		StagingPath:                cmp.Or(stagingPathOverride, s.config.SnapshotStagingPath),
 		SyncedAtColName:            s.config.SyncedAtColName,
 		SoftDeleteColName:          s.config.SoftDeleteColName,
 		WriteMode:                  snapshotWriteMode,
@@ -231,6 +233,7 @@ func (s *SnapshotFlowExecution) cloneTables(
 	snapshotType snapshotType,
 	slotName string,
 	snapshotName string,
+	stagingPathOverride string,
 	maxParallelClones int,
 ) error {
 	if snapshotType == SNAPSHOT_TYPE_SLOT {
@@ -273,7 +276,7 @@ func (s *SnapshotFlowExecution) cloneTables(
 		if v.PartitionKey == "" {
 			v.PartitionKey = res.TableDefaultPartitionKeyMapping[source]
 		}
-		if err := s.cloneTable(ctx, boundSelector, snapshotName, v, sourcePeerType, destinationPeerType); err != nil {
+		if err := s.cloneTable(ctx, boundSelector, snapshotName, stagingPathOverride, v, sourcePeerType, destinationPeerType); err != nil {
 			s.logger.Error("failed to start clone child workflow", slog.Any("error", err))
 			return err
 		}
@@ -316,6 +319,7 @@ func (s *SnapshotFlowExecution) cloneTablesWithSlot(
 		SNAPSHOT_TYPE_SLOT,
 		slotName,
 		snapshotName,
+		"",
 		numTablesInParallel,
 	); err != nil {
 		s.logger.Error("failed to clone tables", slog.Any("error", err))
@@ -410,6 +414,7 @@ func SnapshotFlowWorkflow(
 			SNAPSHOT_TYPE_TX,
 			"",
 			txnSnapshotState.SnapshotName,
+			txnSnapshotState.SnapshotStagingPath,
 			numTablesInParallel,
 		); err != nil {
 			return fmt.Errorf("failed to clone tables: %w", err)

protos/flow.proto

@@ -555,6 +555,7 @@ message IsQRepPartitionSyncedInput {
 message ExportTxSnapshotOutput {
   reserved 2;
   string snapshot_name = 1;
+  string snapshot_staging_path = 3;
 }
 
 enum DynconfValueType {