fix(data-imports): retry queryable-folder copy after a concurrent vacuum race

products/warehouse_sources/backend/temporal/data_imports/pipelines/common/load.py

@@ -377,6 +377,7 @@ def _get_companion_queryable_folder():
             delete_existing=True,
             existing_queryable_folder=existing_queryable_folder,
             logger=logger,
+            refresh_file_uris=delta_table_helper.get_file_uris,
         )
 
 

products/warehouse_sources/backend/temporal/data_imports/tests/test_util.py

@@ -59,6 +59,41 @@ async def test_copies_files_with_cp_file_not_copy(self):
         assert s3._cp_file.await_count == 2
         s3._copy.assert_not_awaited()
 
+    async def test_retries_with_fresh_listing_when_source_file_vanishes_mid_copy(self):
+        # A concurrent compact/vacuum pass on the same Delta table can delete a source file
+        # between get_file_uris() listing it and this copy step reading it, raising
+        # FileNotFoundError. Regression for that race: https://github.com/PostHog/posthog
+        vanished_file = "s3://bucket/job/my_table/part-0.parquet"
+        cp_file = AsyncMock(side_effect=[FileNotFoundError(vanished_file), None])
+        s3 = _fake_s3(_cp_file=cp_file)
+        refresh_file_uris = AsyncMock(return_value=["s3://bucket/job/my_table/part-1.parquet"])
+
+        with patch.object(util_module, "aget_s3_client", return_value=_FakeS3CM(s3)):
+            await prepare_s3_files_for_querying(
+                folder_path="job",
+                table_name="my_table",
+                file_uris=[vanished_file],
+                delete_existing=False,
+                refresh_file_uris=refresh_file_uris,
+            )
+
+        refresh_file_uris.assert_awaited_once()
+        assert cp_file.await_args_list[-1].args[0] == "s3://bucket/job/my_table/part-1.parquet"
+
+    async def test_propagates_vanished_source_file_without_refresh_callback(self):
+        # Callers that don't pass refresh_file_uris keep today's behavior: the race still
+        # surfaces as an error instead of being retried blindly.
+        s3 = _fake_s3(_cp_file=AsyncMock(side_effect=FileNotFoundError("gone")))
+
+        with patch.object(util_module, "aget_s3_client", return_value=_FakeS3CM(s3)):
+            with pytest.raises(FileNotFoundError):
+                await prepare_s3_files_for_querying(
+                    folder_path="job",
+                    table_name="my_table",
+                    file_uris=["s3://bucket/job/my_table/part-0.parquet"],
+                    delete_existing=False,
+                )
+
 
 @parameterized.expand(
     [

products/warehouse_sources/backend/temporal/data_imports/util.py

@@ -1,5 +1,6 @@
 import re
 import asyncio
+from collections.abc import Awaitable, Callable
 from datetime import datetime
 from typing import Literal, Optional
 from uuid import uuid4
@@ -75,6 +76,7 @@ async def prepare_s3_files_for_querying(
     preserve_table_name_casing: Optional[bool] = False,
     delete_existing: bool = True,
     logger: Optional[FilteringBoundLogger] = None,
+    refresh_file_uris: Optional[Callable[[], Awaitable[list[str]]]] = None,
 ) -> str:
     """Async version that uses s3fs native async methods for concurrent file operations."""
 
@@ -182,7 +184,18 @@ async def copy_file(file: str) -> None:
                 # S3's SlowDown rate limiting on the destination prefix.
                 await s3._cp_file(file, f"{s3_path_for_querying}/{file_name}")
 
-        await asyncio.gather(*[copy_file(file) for file in file_uris])
+        try:
+            await asyncio.gather(*[copy_file(file) for file in file_uris])
+        except FileNotFoundError as e:
+            if refresh_file_uris is None:
+                raise
+            # A concurrent compact/vacuum pass on the same Delta table (e.g. a zombie attempt
+            # from a heartbeat timeout still running) can physically delete a source file
+            # between our listing and this copy. Re-listing picks up wherever that pass left
+            # the table and retries once against the current file set.
+            await _log(f"Source file vanished mid-copy, retrying with a fresh file listing: {e}", level="error")
+            file_uris = await refresh_file_uris()
+            await asyncio.gather(*[copy_file(file) for file in file_uris])
 
         # Delete existing files after copying new ones
         if delete_existing and files_to_delete: