diff --git a/flow/activities/snapshot_activity.go b/flow/activities/snapshot_activity.go index ee9c3a1e4..8b82400fa 100644 --- a/flow/activities/snapshot_activity.go +++ b/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{} diff --git a/flow/connectors/bigquery/qrep_object_pull.go b/flow/connectors/bigquery/qrep_object_pull.go index 8765272f0..229652a28 100644 --- a/flow/connectors/bigquery/qrep_object_pull.go +++ b/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. diff --git a/flow/workflows/snapshot_flow.go b/flow/workflows/snapshot_flow.go index c37831280..66e88b7a1 100644 --- a/flow/workflows/snapshot_flow.go +++ b/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) diff --git a/protos/flow.proto b/protos/flow.proto index 006a588a4..6b6541d00 100644 --- a/protos/flow.proto +++ b/protos/flow.proto @@ -555,6 +555,7 @@ message IsQRepPartitionSyncedInput { message ExportTxSnapshotOutput { reserved 2; string snapshot_name = 1; + string snapshot_staging_path = 3; } enum DynconfValueType {