Add endpoint to rollback failed imports by vish-cs · Pull Request #613 · datacommonsorg/import

vish-cs · 2026-07-06T08:12:26Z

No description provided.

gemini-code-assist

Code Review

This pull request introduces a new script, retry_failed_ingestion.py, along with its corresponding unit tests, to automate the rollback and retry process for failed Dataflow ingestion jobs. Feedback on the implementation highlights a critical correctness bug where the workflow could be retriggered even if a rollback fails, and a missing pagination implementation when fetching Dataflow job messages. Additionally, it is recommended to use posixpath instead of os.path for GCS path manipulations to avoid platform-specific path separator issues, and to clean up unused imports in the test file.

codacy-production · 2026-07-06T08:15:36Z

Not up to standards ⛔

🔴 Issues 2 medium · 28 minor

Alerts:
⚠ 30 issues (≤ 0 issues of at least minor severity)

Results:
30 new issues

Category Results

Documentation 19 minor

CodeStyle 9 minor

Complexity 2 medium

View in Codacy

🟢 Metrics 74 complexity

Metric Results

Complexity 74

View in Codacy

_{NEW Get contextual insights on your PRs based on Codacy's metrics, along with PR and Jira context, without leaving GitHub. Enable AI reviewer}
_{TIP This summary will be updated as you push new changes.}

gmechali · 2026-07-07T15:01:26Z

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


This seems to only be checking dataflow for failures. What if the failures came from elsewhere in the worfklow (spanner embeddings, ingestion helper issue, bq federation).

What's the long term plan for this, should we be checking all the stages? Should the workflow be modified to propoagate errors from dataflow? (if that's even possible)

Yes, we should be taking bq failures also into account. It will require us to track aggregation status into spanner. Looking into this.

The new implementation is an endpoint within ingestion helper. It takes a workflow id as an argument and reverts all the imports for that workflow. It can be invoked manually or from the ingestion workflow in the failure handing logic for any type of failures.

gmechali · 2026-07-07T15:04:55Z

+        """
+        status_results = transaction.execute_sql(
+            status_sql,
+            params={'importName': short_name},


How would this work for a multi-provenance import?
Like if I import WHO and ILO at once?

The script checks for the dataflow logs to see if it finds a specific failed import and reverts it. Else, it falls back to reverting all the imports in that batch.

I am also thinking that it may be simpler to revert the whole batch as a v1 for the rollback flow.

The new implementation reverts the whole batch

gmechali

hey Vishal, overall I think the major concern I have is that I think the ingesiton helper is being used for too much. It shouldn't contain such major implementation logic, I think all that should live within the workflow, and the ingestion helper should just give the workflow the information it needs (like the last successful graphpath).

Happy to talk this through if you want. lmk what you think

gmechali · 2026-07-16T13:57:45Z

+            logging.error(f"Error querying IngestionHistory for workflow '{workflow_id}': {e}")
+        return []
+
+    def revert_import(self,


Can you move most of this logic in a separate file in util/rollback_helper.py or something?

And just keep the spanner code in here, so get_imports_for_workflow, fetch_import_version_history, set_rollback_status etc..

I think this spanner client should just be simple helper functions to execute spanner queries. The actual logic for how these helpers are called belong elsewhere.

gmechali · 2026-07-16T13:58:24Z

            logging.error(f'Error checking for retry imports: {e}')
            return True

+    def get_imports_for_workflow(self, workflow_id: str) -> list[str]:


ultra nit: "from" instead of for?

gmechali · 2026-07-16T14:04:42Z

+
+            if not rows_history:
+                logging.warning(
+                    f"No version history found in ImportVersionHistory for '{short_name}'. Cannot revert.")


So what should happen here to the database?

do we leave the run in FAILED state, the staleness timestamp should read before this fialed run start. Ideally they get alerted of this failuyre and have to manually intervene, is that right?

gmechali · 2026-07-16T14:06:37Z

+            latest_version = rows_history[0][0]
+            latest_comment = rows_history[0][1] if len(rows_history[0]) > 1 else ""
+
+            if latest_comment and ("Reverted due to Dataflow failure" in latest_comment or "Reverted batch workflow" in latest_comment):


Using these string comments seems super dangerous to me. Shouldn't we use the Status instead?

It's really not transparent that if you modify the comment, you ll modify actual workflow behavior. You know this now, but whoever is working on this in 9 months won't. Can't we just fetch the Status on IngestionHistory row, add a new rollback status and use that to determine the state?

gmechali · 2026-07-16T14:07:56Z

+                logging.error(
+                    f"Import '{short_name}' was already reverted recently. Aborting to prevent infinite loop."
+                )
+                return "ALREADY_REVERTED", None


What exactly is this function supposed to return? I see you returning a True/False or now thing "ALREADY_REVERTED" string.

gmechali · 2026-07-16T14:15:15Z

+
+class RollbackResponse(BaseResponse):
+    revertedImports: Dict[str, str] = Field(default_factory=dict, description="Map of import names to their reverted versions")
+    dryRun: bool = Field(default=False, description="Whether this was executed in dry run mode")


Why do we need the dry run in the response?

gmechali · 2026-07-16T14:22:42Z

+                )
+                return "ALREADY_REVERTED", None
+
+            previous_version = None


I kind of disagree with the architecture of this. I think the rollback shouldn't be triggered by the ingestion helper. I think the ingestion helper should be called to give the worfklow the information it needs to do the rollback.

I think we can use a subworkflow to package everything from updating the ingesiton_history, to dataflow, to aggregations.

Then the worfklow, on failure, can aks the ingestion Helper for the graphPath of the last successful ingestions for these imports. Then the worfklow can re-update the ingestionHistory/Importversionhistory statuses, re-laucnh dataflow with new params, relaunch postprocessing etc...

Otherwise, with this implementation, we're just complete obfuscating the process for the rollback.

gmechali · 2026-07-16T14:24:12Z

+
+    logging.info(f"Retriggering workflow '{workflow_name}'...")
+    try:
+        service = build('workflowexecutions', 'v1', cache_discovery=False)


I dont think we should make this a separate workflow, but a subworkflow.

The first workflow acquires the lock, I believe it should NOT be allowed to release the lock until the database is back in a good state post rollback. That's why I think it should also be a subworkflow.

gemini-code-assist Bot reviewed Jul 6, 2026

View reviewed changes

vish-cs force-pushed the rollback branch 2 times, most recently from a3696c4 to 5bb4e0a Compare July 6, 2026 08:47

vish-cs requested a review from gmechali July 6, 2026 08:48

gmechali reviewed Jul 7, 2026

View reviewed changes

vish-cs force-pushed the rollback branch from 5bb4e0a to a2ce24c Compare July 8, 2026 05:32

vish-cs force-pushed the rollback branch from a2ce24c to 64b5599 Compare July 15, 2026 17:10

vish-cs changed the title ~~Add script to rollback failed import~~ Add endpoint to rollback failed imports Jul 15, 2026

vish-cs force-pushed the rollback branch from 64b5599 to 9407517 Compare July 16, 2026 05:34

Add script to rollback failed import

02e405c

vish-cs force-pushed the rollback branch from 9407517 to 02e405c Compare July 16, 2026 05:40

vish-cs requested a review from gmechali July 16, 2026 05:40

gmechali reviewed Jul 16, 2026

View reviewed changes

Uh oh!

Conversation

vish-cs commented Jul 6, 2026

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

codacy-production Bot commented Jul 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Not up to standards ⛔

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gmechali left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

codacy-production Bot commented Jul 6, 2026 •

edited

Loading