diff --git a/products/warehouse_sources/backend/temporal/data_imports/pipelines/common/load.py b/products/warehouse_sources/backend/temporal/data_imports/pipelines/common/load.py index 70ee1b491ebb..65299e549f61 100644 --- a/products/warehouse_sources/backend/temporal/data_imports/pipelines/common/load.py +++ b/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, ) diff --git a/products/warehouse_sources/backend/temporal/data_imports/tests/test_util.py b/products/warehouse_sources/backend/temporal/data_imports/tests/test_util.py index 78186c307d37..03286464762d 100644 --- a/products/warehouse_sources/backend/temporal/data_imports/tests/test_util.py +++ b/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( [ diff --git a/products/warehouse_sources/backend/temporal/data_imports/util.py b/products/warehouse_sources/backend/temporal/data_imports/util.py index 6e5f9a9f3027..9d9e3fbb76ac 100644 --- a/products/warehouse_sources/backend/temporal/data_imports/util.py +++ b/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: