From a2ce24c76599ac13102a8465d735827316a85e63 Mon Sep 17 00:00:00 2001 From: Vishal Gupta Date: Thu, 2 Jul 2026 13:23:09 +0000 Subject: [PATCH] Add script to rollback failed import --- .../workflow/ingestion-helper/app_test.py | 101 +++++++++- .../ingestion-helper/clients/spanner.py | 180 ++++++++++++++++++ .../ingestion-helper/clients/spanner_test.py | 137 +++++++++++++ .../ingestion-helper/routes/imports.py | 69 +++++++ .../ingestion-helper/utils/imports.py | 115 +++++++++++ 5 files changed, 601 insertions(+), 1 deletion(-) diff --git a/pipeline/workflow/ingestion-helper/app_test.py b/pipeline/workflow/ingestion-helper/app_test.py index 8150e5a6b..4261a792f 100644 --- a/pipeline/workflow/ingestion-helper/app_test.py +++ b/pipeline/workflow/ingestion-helper/app_test.py @@ -135,7 +135,9 @@ def test_seed_database_success(self): self.assertEqual(response.json()["status"], "OK") mock_spanner_client.seed_database.assert_called_once() - def test_update_import_status_success(self): + @patch('routes.imports.import_utils.get_next_refresh') + def test_update_import_status_success(self, mock_get_next_refresh): + mock_get_next_refresh.return_value = "2026-07-01T00:00:00Z" mock_spanner_client = MagicMock() mock_storage_client = MagicMock() app.dependency_overrides[get_spanner_client] = lambda: mock_spanner_client @@ -263,6 +265,103 @@ def test_update_import_version_override_success(self, mock_get_caller_identity): # Verify get_caller_identity was called exactly once outside of the loop mock_get_caller_identity.assert_called_once() + @patch('routes.imports.import_utils.retrigger_workflow') + @patch('routes.imports.import_utils.get_failed_imports') + def test_rollback_ingestion_success(self, mock_get_failed, mock_retrigger): + mock_get_failed.return_value = ["imp1", "imp2"] + mock_retrigger.return_value = "exec-123" + mock_spanner_client = MagicMock() + mock_spanner_client.revert_import.return_value = True + mock_spanner_client.reset_retry_imports.return_value = 2 + app.dependency_overrides[get_spanner_client] = lambda: mock_spanner_client + + payload = { + "jobId": "job-123", + "projectId": "test-project", + "location": "us-central1", + "workflowName": "test-workflow", + "dryRun": False + } + response = client.post("/imports/rollback", json=payload) + self.assertEqual(response.status_code, 200) + data = response.json() + self.assertEqual(data["status"], "OK") + self.assertEqual(data["revertedImports"], ["imp1", "imp2"]) + self.assertEqual(data["retriggeredExecution"], "exec-123") + self.assertFalse(data["dryRun"]) + + mock_get_failed.assert_called_once_with("test-project", "us-central1", "job-123") + self.assertEqual(mock_spanner_client.revert_import.call_count, 2) + mock_spanner_client.reset_retry_imports.assert_called_once_with(dry_run=False) + mock_retrigger.assert_called_once_with("test-project", "us-central1", "test-workflow", dry_run=False) + + @patch('routes.imports.import_utils.get_failed_imports') + def test_rollback_ingestion_no_failed_imports(self, mock_get_failed): + mock_get_failed.return_value = [] + mock_spanner_client = MagicMock() + app.dependency_overrides[get_spanner_client] = lambda: mock_spanner_client + + payload = { + "jobId": "job-123", + "projectId": "test-project", + "location": "us-central1" + } + response = client.post("/imports/rollback", json=payload) + self.assertEqual(response.status_code, 200) + data = response.json() + self.assertEqual(data["status"], "SKIPPED") + self.assertEqual(data["revertedImports"], []) + self.assertIsNone(data["retriggeredExecution"]) + mock_spanner_client.revert_import.assert_not_called() + + @patch('routes.imports.import_utils.retrigger_workflow') + @patch('routes.imports.import_utils.get_failed_imports') + def test_rollback_ingestion_revert_failure(self, mock_get_failed, mock_retrigger): + mock_get_failed.return_value = ["imp1", "imp2"] + mock_spanner_client = MagicMock() + # First import succeeds, second fails + mock_spanner_client.revert_import.side_effect = [True, False] + app.dependency_overrides[get_spanner_client] = lambda: mock_spanner_client + + payload = { + "jobId": "job-123", + "projectId": "test-project", + "location": "us-central1" + } + response = client.post("/imports/rollback", json=payload) + self.assertEqual(response.status_code, 200) + data = response.json() + self.assertEqual(data["status"], "SKIPPED") + self.assertEqual(data["revertedImports"], ["imp1"]) + self.assertIsNone(data["retriggeredExecution"]) + mock_spanner_client.reset_retry_imports.assert_not_called() + mock_retrigger.assert_not_called() + + @patch('routes.imports.import_utils.retrigger_workflow') + @patch('routes.imports.import_utils.get_failed_imports') + def test_rollback_ingestion_dry_run(self, mock_get_failed, mock_retrigger): + mock_get_failed.return_value = ["imp1"] + mock_retrigger.return_value = None + mock_spanner_client = MagicMock() + mock_spanner_client.revert_import.return_value = True + app.dependency_overrides[get_spanner_client] = lambda: mock_spanner_client + + payload = { + "jobId": "job-123", + "projectId": "test-project", + "location": "us-central1", + "dryRun": True + } + response = client.post("/imports/rollback", json=payload) + self.assertEqual(response.status_code, 200) + data = response.json() + self.assertEqual(data["status"], "OK") + self.assertTrue(data["dryRun"]) + self.assertIn("[DRY RUN]", data["message"]) + mock_spanner_client.revert_import.assert_called_once_with("imp1", "job-123", dry_run=True) + mock_spanner_client.reset_retry_imports.assert_called_once_with(dry_run=True) + mock_retrigger.assert_called_once_with("test-project", "us-central1", "spanner-ingestion-workflow", dry_run=True) + if __name__ == '__main__': unittest.main() diff --git a/pipeline/workflow/ingestion-helper/clients/spanner.py b/pipeline/workflow/ingestion-helper/clients/spanner.py index cd3f6a18f..00de0216b 100644 --- a/pipeline/workflow/ingestion-helper/clients/spanner.py +++ b/pipeline/workflow/ingestion-helper/clients/spanner.py @@ -14,6 +14,7 @@ import logging import os +import posixpath from google.cloud import spanner from google.cloud.spanner_admin_database_v1 import DatabaseAdminClient from google.cloud.spanner_admin_database_v1.types import UpdateDatabaseDdlRequest @@ -376,6 +377,185 @@ def check_failed_imports(self) -> bool: logging.error(f'Error checking for retry imports: {e}') return True + def revert_import(self, import_name: str, job_id: str, dry_run: bool = False) -> bool | str: + """Reverts an import to its previous version in Spanner and sets it to STAGING. + + This function updates the ImportStatus and ImportVersionHistory tables in + Spanner to effectively rollback an import that failed during the Dataflow job. + + Args: + import_name: The name of the import to revert (e.g., 'foo:bar:import'). + job_id: The ID of the failed Dataflow job, used for commenting. + dry_run: Dry run mode. + + Returns: + True if the revert was successful, 'ALREADY_REVERTED' if the import + was recently reverted, or False if no version history was found or an + error occurred. + """ + short_name = import_name.split(':')[-1] + + def _revert_txn(transaction: Transaction): + # 1. Get current version from ImportStatus to know what to avoid. + status_sql = """ + SELECT LatestVersion FROM ImportStatus + WHERE ImportName = @importName + """ + status_results = transaction.execute_sql( + status_sql, + params={'importName': short_name}, + param_types={'importName': STRING}) + status_rows = list(status_results) + if not status_rows: + logging.warning(f"No ImportStatus found for '{short_name}'.") + return False + + current_latest_version_path = status_rows[0][0] + if not current_latest_version_path: + logging.warning(f"ImportStatus.LatestVersion is empty for '{short_name}'.") + return False + + current_version_name = posixpath.basename(current_latest_version_path) + + # 2. Fetch the absolute most recent version and comment from ImportVersionHistory. + sql_latest = """ + SELECT Version, Comment FROM ImportVersionHistory + WHERE ImportName = @importName + ORDER BY UpdateTimestamp DESC + LIMIT 1 + """ + results_latest = transaction.execute_sql( + sql_latest, + params={'importName': short_name}, + param_types={'importName': STRING}) + rows_latest = list(results_latest) + + if not rows_latest: + logging.warning( + f"No version history found for '{short_name}'. Cannot revert.") + return False + + latest_comment = rows_latest[0][1] if len(rows_latest[0]) > 1 else "" + + if latest_comment and "Reverted due to Dataflow failure" in latest_comment: + logging.error( + f"Import '{short_name}' was already reverted recently. Aborting to prevent infinite loop." + ) + return "ALREADY_REVERTED" + + # 3. Find the latest version in history that is different from current_version_name + sql_prev = """ + SELECT Version FROM ImportVersionHistory + WHERE ImportName = @importName AND Version != @currentVersion + ORDER BY UpdateTimestamp DESC + LIMIT 1 + """ + results_prev = transaction.execute_sql( + sql_prev, + params={'importName': short_name, 'currentVersion': current_version_name}, + param_types={'importName': STRING, 'currentVersion': STRING}) + rows_prev = list(results_prev) + + if not rows_prev: + logging.warning( + f"No previous different version found for '{short_name}'. Cannot revert.") + return False + + previous_version = rows_prev[0][0] + + # Reconstruct full path for LatestVersion + parent_dir = posixpath.dirname(current_latest_version_path) + new_latest_version_path = posixpath.join(parent_dir, previous_version) + + if dry_run: + logging.info( + f"[DRY RUN] Would revert '{short_name}' from '{current_version_name}' to last known good version: '{previous_version}' (Path: '{new_latest_version_path}')" + ) + return True + + logging.info( + f"Reverting '{short_name}' from '{current_version_name}' to last known good version: '{previous_version}' (Path: '{new_latest_version_path}')" + ) + + # 4. Update ImportStatus to point to the previous version and set state to STAGING. + param_types = { + 'version': STRING, + 'importName': STRING + } + + transaction.execute_update(""" + UPDATE ImportStatus + SET LatestVersion = @version, State = 'STAGING', StatusUpdateTimestamp = PENDING_COMMIT_TIMESTAMP() + WHERE ImportName = @importName + """, + params={ + 'version': new_latest_version_path, + 'importName': short_name + }, + param_types=param_types) + + # 5. Add an entry to ImportVersionHistory to record the revert action. + transaction.execute_update( + """ + INSERT INTO ImportVersionHistory (ImportName, Version, UpdateTimestamp, Comment) + VALUES (@importName, @version, PENDING_COMMIT_TIMESTAMP(), @comment) + """, + params={ + 'importName': short_name, + 'version': previous_version, + 'comment': f"Reverted due to Dataflow failure ({job_id})" + }, + param_types={ + 'importName': STRING, + 'version': STRING, + 'comment': STRING + }) + return True + + try: + return self.database.run_in_transaction(_revert_txn) + except Exception as e: + logging.error(f"Transaction failed for '{short_name}': {e}") + return False + + def reset_retry_imports(self, dry_run: bool = False) -> int | None: + """Resets all imports with state 'RETRY' back to 'STAGING'. + + This ensures that any imports that were stuck in a RETRY state due to + the failure are put back into the queue for processing. + + Args: + dry_run: Dry run mode. + """ + if dry_run: + try: + with self.database.snapshot() as snapshot: + results = snapshot.execute_sql( + "SELECT ImportName FROM ImportStatus WHERE State = 'RETRY'" + ) + retry_imports = [row[0] for row in results] + logging.info(f"[DRY RUN] Would reset RETRY imports to STAGING: {retry_imports}") + except Exception as e: + logging.error(f"Failed to query RETRY imports: {e}") + return None + + logging.info("Resetting all RETRY imports to STAGING...") + + def _reset_txn(transaction: Transaction): + return transaction.execute_update(""" + UPDATE ImportStatus + SET State = 'STAGING', StatusUpdateTimestamp = PENDING_COMMIT_TIMESTAMP() + WHERE State = 'RETRY' + """) + + try: + count = self.database.run_in_transaction(_reset_txn) + logging.info(f"Successfully reset {count} RETRY imports to STAGING.") + return count + except Exception as e: + logging.error(f"Failed to reset RETRY imports: {e}") + return None + def update_import_status(self, params: dict): """Updates the status for the specified import job. diff --git a/pipeline/workflow/ingestion-helper/clients/spanner_test.py b/pipeline/workflow/ingestion-helper/clients/spanner_test.py index 4e5b2e395..bac56cc47 100644 --- a/pipeline/workflow/ingestion-helper/clients/spanner_test.py +++ b/pipeline/workflow/ingestion-helper/clients/spanner_test.py @@ -372,6 +372,143 @@ def open_side_effect(file_path, mode='r', *args, **kwargs): ) self.assertEqual(statements[2].strip(), "CREATE VECTOR INDEX CustomEmbeddingIndex ON CustomEmbeddingTable(embeddings)") + @patch('google.cloud.spanner.Client') + def test_revert_import_success(self, mock_spanner_client): + mock_instance = MagicMock() + mock_db = MagicMock() + mock_spanner_client.return_value.instance.return_value = mock_instance + mock_instance.database.return_value = mock_db + + mock_txn = MagicMock() + def run_in_transaction_side_effect(callback, *args, **kwargs): + return callback(mock_txn, *args, **kwargs) + mock_db.run_in_transaction.side_effect = run_in_transaction_side_effect + + def execute_sql_side_effect(sql, params, param_types): + if "SELECT LatestVersion FROM ImportStatus" in sql: + mock_results = MagicMock() + mock_results.__iter__.return_value = [["gs://bucket/imp/v2"]] + return mock_results + elif "FROM ImportVersionHistory" in sql: + if "Version != @currentVersion" in sql: + mock_results = MagicMock() + mock_results.__iter__.return_value = [["v1"]] + return mock_results + else: + mock_results = MagicMock() + mock_results.__iter__.return_value = [["v2", "import-workflow:job-123"]] + return mock_results + return [] + + mock_txn.execute_sql.side_effect = execute_sql_side_effect + + client = SpannerClient("project", "instance", "database") + result = client.revert_import("foo:bar:imp", "job-123") + self.assertTrue(result) + self.assertEqual(mock_txn.execute_update.call_count, 2) + + @patch('google.cloud.spanner.Client') + def test_revert_import_already_reverted(self, mock_spanner_client): + mock_instance = MagicMock() + mock_db = MagicMock() + mock_spanner_client.return_value.instance.return_value = mock_instance + mock_instance.database.return_value = mock_db + + mock_txn = MagicMock() + def run_in_transaction_side_effect(callback, *args, **kwargs): + return callback(mock_txn, *args, **kwargs) + mock_db.run_in_transaction.side_effect = run_in_transaction_side_effect + + def execute_sql_side_effect(sql, params, param_types): + if "SELECT LatestVersion FROM ImportStatus" in sql: + mock_results = MagicMock() + mock_results.__iter__.return_value = [["gs://bucket/imp/v1"]] + return mock_results + elif "FROM ImportVersionHistory" in sql: + mock_results = MagicMock() + mock_results.__iter__.return_value = [["v1", "Reverted due to Dataflow failure (job-999)"]] + return mock_results + return [] + + mock_txn.execute_sql.side_effect = execute_sql_side_effect + + client = SpannerClient("project", "instance", "database") + result = client.revert_import("foo:bar:imp", "job-123") + self.assertEqual(result, "ALREADY_REVERTED") + mock_txn.execute_update.assert_not_called() + + @patch('google.cloud.spanner.Client') + def test_revert_import_dry_run(self, mock_spanner_client): + mock_instance = MagicMock() + mock_db = MagicMock() + mock_spanner_client.return_value.instance.return_value = mock_instance + mock_instance.database.return_value = mock_db + + mock_txn = MagicMock() + def run_in_transaction_side_effect(callback, *args, **kwargs): + return callback(mock_txn, *args, **kwargs) + mock_db.run_in_transaction.side_effect = run_in_transaction_side_effect + + def execute_sql_side_effect(sql, params, param_types): + if "SELECT LatestVersion FROM ImportStatus" in sql: + mock_results = MagicMock() + mock_results.__iter__.return_value = [["gs://bucket/imp/v2"]] + return mock_results + elif "FROM ImportVersionHistory" in sql: + if "Version != @currentVersion" in sql: + mock_results = MagicMock() + mock_results.__iter__.return_value = [["v1"]] + return mock_results + else: + mock_results = MagicMock() + mock_results.__iter__.return_value = [["v2", "import-workflow:job-123"]] + return mock_results + return [] + + mock_txn.execute_sql.side_effect = execute_sql_side_effect + + client = SpannerClient("project", "instance", "database") + result = client.revert_import("foo:bar:imp", "job-123", dry_run=True) + self.assertTrue(result) + mock_txn.execute_update.assert_not_called() + + @patch('google.cloud.spanner.Client') + def test_reset_retry_imports(self, mock_spanner_client): + mock_instance = MagicMock() + mock_db = MagicMock() + mock_spanner_client.return_value.instance.return_value = mock_instance + mock_instance.database.return_value = mock_db + + mock_txn = MagicMock() + mock_txn.execute_update.return_value = 3 + def run_in_transaction_side_effect(callback, *args, **kwargs): + return callback(mock_txn, *args, **kwargs) + mock_db.run_in_transaction.side_effect = run_in_transaction_side_effect + + client = SpannerClient("project", "instance", "database") + count = client.reset_retry_imports() + self.assertEqual(count, 3) + mock_txn.execute_update.assert_called_once() + + @patch('google.cloud.spanner.Client') + def test_reset_retry_imports_dry_run(self, mock_spanner_client): + mock_instance = MagicMock() + mock_db = MagicMock() + mock_spanner_client.return_value.instance.return_value = mock_instance + mock_instance.database.return_value = mock_db + + mock_snapshot = MagicMock() + mock_db.snapshot.return_value.__enter__.return_value = mock_snapshot + mock_results = MagicMock() + mock_results.__iter__.return_value = [["imp1"], ["imp2"]] + mock_snapshot.execute_sql.return_value = mock_results + + client = SpannerClient("project", "instance", "database") + count = client.reset_retry_imports(dry_run=True) + self.assertIsNone(count) + mock_snapshot.execute_sql.assert_called_once() + mock_db.run_in_transaction.assert_not_called() + if __name__ == '__main__': unittest.main() diff --git a/pipeline/workflow/ingestion-helper/routes/imports.py b/pipeline/workflow/ingestion-helper/routes/imports.py index e8fa01087..59720d987 100644 --- a/pipeline/workflow/ingestion-helper/routes/imports.py +++ b/pipeline/workflow/ingestion-helper/routes/imports.py @@ -74,6 +74,18 @@ class ImportInfoItem(BaseModel): latestVersion: str = Field(description="The latest version GCS filename") graphPath: str = Field(description="The full GCS path to the import graph files") +class RollbackRequest(BaseModel): + jobId: str = Field(description="Failed Dataflow Job ID") + projectId: Optional[str] = Field(default=None, description="GCP Project ID override") + location: Optional[str] = Field(default=None, description="GCP Location override") + workflowName: str = Field(default="spanner-ingestion-workflow", description="Workflow name to retrigger") + dryRun: bool = Field(default=False, description="Dry run mode") + +class RollbackResponse(BaseResponse): + revertedImports: List[str] = Field(default_factory=list, description="Imports reverted or that would be reverted") + retriggeredExecution: Optional[str] = Field(default=None, description="Execution ID of retriggered workflow") + dryRun: bool = Field(default=False, description="Whether this was executed in dry run mode") + router = APIRouter(prefix="/imports", tags=["imports"]) @router.post("/info", response_model=List[ImportInfoItem]) @@ -180,3 +192,60 @@ def update_import_version( status=ResponseStatus.OK, message="; ".join(updated_imports) ) + +@router.post("/rollback", response_model=RollbackResponse) +def rollback_ingestion( + req: RollbackRequest, + spanner: SpannerClient = Depends(get_spanner_client) +): + """Reverts failed imports from a Dataflow job and optionally retriggers ingestion.""" + project_id = req.projectId or config.PROJECT_ID + location = req.location or config.LOCATION + + if not project_id or not location: + raise HTTPException( + status_code=500, + detail="Project ID or Location is not configured and was not provided in request." + ) + + failed_imports = import_utils.get_failed_imports(project_id, location, req.jobId) + if not failed_imports: + return RollbackResponse( + status=ResponseStatus.SKIPPED, + message="No failed imports identified for the specified job.", + dryRun=req.dryRun + ) + + logging.info(f"Processing failed imports for rollback: {failed_imports}") + can_retrigger = True + reverted_list = [] + for imp in failed_imports: + status = spanner.revert_import(imp, req.jobId, dry_run=req.dryRun) + if status != True: + can_retrigger = False + else: + reverted_list.append(imp) + + exec_name = None + if can_retrigger: + spanner.reset_retry_imports(dry_run=req.dryRun) + exec_name = import_utils.retrigger_workflow( + project_id, location, req.workflowName, dry_run=req.dryRun + ) + msg = "Successfully rolled back failed imports and retriggered workflow." if not req.dryRun else "[DRY RUN] Would roll back failed imports and retrigger workflow." + return RollbackResponse( + status=ResponseStatus.OK, + message=msg, + revertedImports=reverted_list, + retriggeredExecution=exec_name, + dryRun=req.dryRun + ) + else: + msg = "Skipping re-trigger because at least one import rollback failed or was already reverted." + logging.warning(msg) + return RollbackResponse( + status=ResponseStatus.SKIPPED, + message=msg, + revertedImports=reverted_list, + dryRun=req.dryRun + ) diff --git a/pipeline/workflow/ingestion-helper/utils/imports.py b/pipeline/workflow/ingestion-helper/utils/imports.py index 33f9d1fac..67606b1b9 100644 --- a/pipeline/workflow/ingestion-helper/utils/imports.py +++ b/pipeline/workflow/ingestion-helper/utils/imports.py @@ -13,6 +13,7 @@ # limitations under the License. """Utility functions for the ingestion helper.""" +import json import logging import re from datetime import datetime, timezone @@ -173,3 +174,117 @@ def get_ingestion_metrics(project_id, location, job_id): 'edge_count': edge_count, 'execution_time': execution_time } + + +def get_failed_imports(project_id: str, location: str, job_id: str) -> list[str]: + """Identifies failed import names from Dataflow job messages and parameters. + + Args: + project_id: The GCP Project ID where the Dataflow job ran. + location: The GCP location of the Dataflow job (e.g., 'us-central1'). + job_id: The unique ID of the failed Dataflow job. + + Returns: + A list of strings representing the names of the imports that failed + during the Dataflow job. + """ + failed_imports = set() + try: + dataflow = build('dataflow', 'v1b3', cache_discovery=False) + logging.info(f"Fetching job details for {job_id}...") + job = dataflow.projects().locations().jobs().get( + projectId=project_id, + location=location, + jobId=job_id, + view='JOB_VIEW_ALL').execute() + + # Get all imports involved in this job from displayData. + all_import_names = set() + import_list_json = None + for item in job.get('pipelineDescription', {}).get('displayData', []): + if item.get('key') == 'importList': + import_list_json = item.get('strValue') + break + + if import_list_json: + try: + all_imports_data = json.loads(import_list_json) + all_import_names = { + i['importName'].split(':')[-1] for i in all_imports_data + } + logging.info( + f"Imports involved in this job: {all_import_names}") + except Exception as e: + logging.error(f"Error parsing importList parameter: {e}") + + logging.info(f"Fetching error messages for Dataflow job {job_id}...") + request = dataflow.projects().locations().jobs().messages().list( + projectId=project_id, + location=location, + jobId=job_id, + minimumImportance='JOB_MESSAGE_ERROR') + + while request is not None: + response = request.execute() + for msg in response.get('jobMessages', []): + text = msg.get('messageText', '') + # If we find a valid import name in the error message, add it. + for name in all_import_names: + if name in text: + failed_imports.add(name) + request = dataflow.projects().locations().jobs().messages().list_next( + previous_request=request, previous_response=response) + + # If the job failed globally and we haven't identified specific imports, + # or if we want to be safe, include all imports if the state is failed. + if job.get('currentState') in [ + 'JOB_STATE_FAILED', 'JOB_STATE_CANCELLED' + ]: + if not failed_imports: + logging.warning( + "Job failed globally. Reverting all involved imports.") + failed_imports.update(all_import_names) + else: + logging.info( + f"Job failed. Identified specific failed imports: {failed_imports}" + ) + + except HttpError as e: + logging.error(f"Error querying Dataflow API: {e}") + + return list(failed_imports) + + +def retrigger_workflow(project_id: str, location: str, workflow_name: str, dry_run: bool = False) -> str | None: + """Retriggers the Spanner ingestion workflow with a retry flag. + + Args: + project_id: The GCP Project ID where the workflow is hosted. + location: The GCP location of the workflow (e.g., 'us-central1'). + workflow_name: The name of the workflow to execute. + dry_run: Dry run mode. + + Returns: + The name of the new workflow execution, or None if dry run or error. + """ + if dry_run: + logging.info(f"[DRY RUN] Would retrigger workflow '{workflow_name}' in project '{project_id}' location '{location}' with argument: {{'is_retry': true}}") + return None + + logging.info(f"Retriggering workflow '{workflow_name}'...") + try: + service = build('workflowexecutions', 'v1', cache_discovery=False) + parent = f"projects/{project_id}/locations/{location}/workflows/{workflow_name}" + # Pass is_retry=True as a workflow argument. + execution = service.projects().locations().workflows().executions( + ).create(parent=parent, + body={ + "argument": json.dumps({"is_retry": True}) + }).execute() + exec_name = execution.get('name') + logging.info( + f"Retriggered successfully. New execution: {exec_name}") + return exec_name + except Exception as e: + logging.error(f"Failed to retrigger workflow: {e}") + return None