Add script to rollback failed import

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()

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.