retry until incoming tables are available

Author: cgsheeh

src/lando/utils/management/commands/etl.py

@@ -12,7 +12,9 @@
 
 from django.core.management.base import BaseCommand, CommandError, CommandParser
 from django.db.models import Model, Q, QuerySet
+from google.api_core.exceptions import NotFound
 from google.cloud import bigquery
+from google.cloud.bigquery import Table
 from more_itertools import chunked
 
 from lando.main.models import BaseModel
@@ -271,6 +273,48 @@ def setup(self, transformers: list[ModelTransformer]) -> None:
             )
             logger.info("Created incoming table for %s.", transformer.name)
 
+        self.wait_for_incoming_tables()
+
+    def wait_for_incoming_tables(
+        self, max_retries: int = 5, retry_base_delay_s: float = 1.0
+    ) -> None:
+        """Wait for all incoming tables to be visible to the BigQuery API.
+
+        BigQuery's streaming insert API is eventually consistent with table
+        creation, so a newly created table may not be immediately available.
+        This method polls `get_table` for each incoming table until all are
+        confirmed visible, preventing `NotFound` errors on the first insert.
+        """
+        pending: list[Table] = list(self.incoming_tables.values())
+
+        for attempt in range(max_retries):
+            still_pending: list[Table] = []
+            for table in pending:
+                try:
+                    self.bq_client.get_table(sql_table_id(table))
+                except NotFound:
+                    still_pending.append(table)
+
+            if not still_pending:
+                logger.debug("All incoming tables are ready.")
+                return
+
+            delay = retry_base_delay_s * (2**attempt)
+            logger.warning(
+                "Waiting %.1fs for %d incoming table(s) to become visible.",
+                delay,
+                len(still_pending),
+            )
+            time.sleep(delay)
+            pending = still_pending
+
+        table_ids = [sql_table_id(table) for table in pending]
+        self.cleanup_incoming_tables()
+        raise CommandError(
+            f"Incoming tables not visible after {max_retries} retries: "
+            f"{', '.join(table_ids)}"
+        )
+
     def load(
         self, transformer: ModelTransformer, queryset: QuerySet, chunk_size: int = 500
     ) -> int:

src/lando/utils/tests/test_etl.py

@@ -8,6 +8,7 @@
 from django.contrib.auth.models import User
 from django.core.management import call_command
 from django.core.management.base import CommandError
+from google.api_core.exceptions import NotFound
 
 from lando.main.models.uplift import (
     RevisionUpliftJob,
@@ -17,6 +18,7 @@
     UpliftSubmission,
 )
 from lando.utils.management.commands.etl import (
+    BigQueryLoader,
     Command,
     JsonLinesLoader,
     RepoTransformer,
@@ -545,3 +547,59 @@ def test_extract_filters_out_records_before_cutoff():
     assert (
         old_assessment.id not in result_ids
     ), "Records before the cutoff should be excluded."
+
+
+def test_wait_for_incoming_tables_succeeds_immediately():
+    """Should return immediately when all tables are visible."""
+    mock_client = MagicMock()
+    loader = BigQueryLoader(mock_client)
+
+    incoming_table = MagicMock()
+    incoming_table.project = "proj"
+    incoming_table.dataset_id = "ds"
+    incoming_table.table_id = "tbl_incoming"
+    loader.incoming_tables["key"] = incoming_table
+
+    loader.wait_for_incoming_tables(retry_base_delay_s=0)
+
+    assert (
+        mock_client.get_table.call_count == 1
+    ), "Should call `get_table` once per incoming table."
+
+
+def test_wait_for_incoming_tables_retries_on_not_found():
+    """Should retry when `get_table` raises `NotFound`, then succeed."""
+    mock_client = MagicMock()
+    mock_client.get_table.side_effect = [
+        NotFound("Table not found."),
+        MagicMock(),
+    ]
+    loader = BigQueryLoader(mock_client)
+
+    incoming_table = MagicMock()
+    incoming_table.project = "proj"
+    incoming_table.dataset_id = "ds"
+    incoming_table.table_id = "tbl_incoming"
+    loader.incoming_tables["key"] = incoming_table
+
+    loader.wait_for_incoming_tables(retry_base_delay_s=0)
+
+    assert (
+        mock_client.get_table.call_count == 2
+    ), "Should have retried after `NotFound`."
+
+
+def test_wait_for_incoming_tables_raises_after_max_retries():
+    """Should raise `CommandError` when tables never become visible."""
+    mock_client = MagicMock()
+    mock_client.get_table.side_effect = NotFound("Table not found.")
+    loader = BigQueryLoader(mock_client)
+
+    incoming_table = MagicMock()
+    incoming_table.project = "proj"
+    incoming_table.dataset_id = "ds"
+    incoming_table.table_id = "tbl_incoming"
+    loader.incoming_tables["key"] = incoming_table
+
+    with pytest.raises(CommandError, match="Incoming tables not visible"):
+        loader.wait_for_incoming_tables(max_retries=2, retry_base_delay_s=0)