Use unique staging paths for BigQuery source (#4411)

A very good write-up of the situation is in
[DBI-790](https://linear.app/clickhouse/issue/DBI-790/).

To resolve the conflicts between the previous snapshot cleanup and the
new snapshot sync, add a subfolder with the workflow's run id. Use the
same temporal session mechanism that's used to pass PG snapshot name to
the QRep flows, it's great that we have it. It won't help us on the
future snapshot+CDC path but we'll get there when we get there, the
whole "return string as an `any` transaction" situation is already a
mess.

Validated locally and as a hotfix on the customer service that hit this.

Closes DBI-790

Author: ilidemi

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

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