diff --git a/products/warehouse_sources/backend/temporal/data_imports/pipelines/core/delta_table_helper.py b/products/warehouse_sources/backend/temporal/data_imports/pipelines/core/delta_table_helper.py index ea946a7d16e4..119b8a9ab7de 100644 --- a/products/warehouse_sources/backend/temporal/data_imports/pipelines/core/delta_table_helper.py +++ b/products/warehouse_sources/backend/temporal/data_imports/pipelines/core/delta_table_helper.py @@ -926,29 +926,37 @@ async def has_batch_been_committed(self, run_uuid: str, batch_index: int) -> boo """ return await self.has_commit_with_metadata({"run_uuid": run_uuid, "batch_index": str(batch_index)}) - async def vacuum_table(self) -> None: - table = await self.get_delta_table() - if table is None: - raise Exception("Deltatable not found") - + async def _vacuum(self, table: deltalake.DeltaTable) -> None: await self._logger.adebug("Vacuuming table...") vacuum_stats = await asyncio.to_thread( table.vacuum, retention_hours=24, enforce_retention_duration=False, dry_run=False ) await self._logger.adebug(json.dumps(vacuum_stats)) - async def compact_table(self) -> None: - table = await self.get_delta_table() - if table is None: - raise Exception("Deltatable not found") - + async def _compact(self, table: deltalake.DeltaTable) -> None: await self._logger.adebug("Compacting table...") compact_stats = await self._execute_with_conflict_retry( table, lambda: table.optimize.compact(), "compact_table" ) await self._logger.adebug(json.dumps(compact_stats)) - await self.vacuum_table() + async def vacuum_table(self) -> None: + table = await self.get_delta_table() + if table is None: + raise Exception("Deltatable not found") + + await self._vacuum(table) + + async def compact_table(self) -> None: + table = await self.get_delta_table() + if table is None: + raise Exception("Deltatable not found") + + await self._compact(table) + # Reuse the table already resolved above instead of re-fetching it: `get_delta_table` + # is cached only opportunistically, so a re-fetch here can race a concurrent sync of a + # different table evicting this table's cache entry and spuriously report it missing. + await self._vacuum(table) await self._logger.adebug("Compacting and vacuuming complete") async def vacuum_if_stale(self, last_vacuum_version: int | None, commit_threshold: int) -> int | None: @@ -987,7 +995,7 @@ async def vacuum_if_stale(self, last_vacuum_version: int | None, commit_threshol await self._logger.ainfo( f"vacuum_if_stale: {commits_since} commits since last vacuum (>= {commit_threshold}), vacuuming" ) - await self.vacuum_table() + await self._vacuum(table) try: # Observability for the maintenance path — how often tables vacuum and how much log churn # accrued between vacuums. Best-effort: telemetry must never break the sync. @@ -1027,7 +1035,7 @@ async def compact_if_fragmented( arrive here with None. Returns True if compaction ran, False if it was skipped. Cheap when the table is - healthy: one S3 LIST via `get_file_uris`. Intended for pre-write defensive cleanup + healthy: one S3 LIST via `table.file_uris`. Intended for pre-write defensive cleanup so a sync that arrived at a fragmented state (e.g. an earlier attempt that failed before reaching `_post_run_operations`) cleans up before adding to the pile. """ @@ -1035,7 +1043,7 @@ async def compact_if_fragmented( if table is None: return False - file_uris = await self.get_file_uris() + file_uris = await asyncio.to_thread(table.file_uris) total_files = len(file_uris) if partition_count is None: # One directory per partition value; unpartitioned tables collapse to the single @@ -1057,7 +1065,8 @@ async def compact_if_fragmented( return False await self._logger.ainfo(f"compact_if_fragmented: triggering compact ({stats})") - await self.compact_table() + await self._compact(table) + await self._vacuum(table) return True async def run_maintenance( diff --git a/products/warehouse_sources/backend/temporal/data_imports/pipelines/core/test/test_delta_table_helper.py b/products/warehouse_sources/backend/temporal/data_imports/pipelines/core/test/test_delta_table_helper.py index 644c6f62feb9..a1f58e4af4bd 100644 --- a/products/warehouse_sources/backend/temporal/data_imports/pipelines/core/test/test_delta_table_helper.py +++ b/products/warehouse_sources/backend/temporal/data_imports/pipelines/core/test/test_delta_table_helper.py @@ -283,12 +283,13 @@ async def test_threshold( expected_ran: bool, ): helper = DeltaTableHelper(resource_name="t", job=MagicMock(), logger=_make_logger()) - mock_delta = MagicMock() file_uris = [f"s3://bucket/table/f{i}.parquet" for i in range(file_count)] + mock_delta = MagicMock() + mock_delta.file_uris = MagicMock(return_value=file_uris) with ( patch.object(helper, "get_delta_table", AsyncMock(return_value=mock_delta)), - patch.object(helper, "get_file_uris", AsyncMock(return_value=file_uris)), - patch.object(helper, "compact_table", AsyncMock()) as mock_compact, + patch.object(helper, "_compact", AsyncMock()) as mock_compact, + patch.object(helper, "_vacuum", AsyncMock()) as mock_vacuum, ): kwargs: dict = {"partition_count": partition_count} if threshold_kw is not None: @@ -297,9 +298,11 @@ async def test_threshold( assert ran is expected_ran if expected_ran: - mock_compact.assert_called_once() + mock_compact.assert_called_once_with(mock_delta) + mock_vacuum.assert_called_once_with(mock_delta) else: mock_compact.assert_not_called() + mock_vacuum.assert_not_called() # (case_name, files_per_dir, dir_count, expected_ran) _DERIVATION_CASES: list[tuple[str, int, int, bool]] = [ @@ -325,18 +328,43 @@ async def test_partition_count_derived_from_layout( for d in range(dir_count) for i in range(files_per_dir) ] + mock_delta = MagicMock() + mock_delta.file_uris = MagicMock(return_value=file_uris) with ( - patch.object(helper, "get_delta_table", AsyncMock(return_value=MagicMock())), - patch.object(helper, "get_file_uris", AsyncMock(return_value=file_uris)), - patch.object(helper, "compact_table", AsyncMock()) as mock_compact, + patch.object(helper, "get_delta_table", AsyncMock(return_value=mock_delta)), + patch.object(helper, "_compact", AsyncMock()) as mock_compact, + patch.object(helper, "_vacuum", AsyncMock()) as mock_vacuum, ): ran = await helper.compact_if_fragmented(partition_count=None) assert ran is expected_ran if expected_ran: - mock_compact.assert_called_once() + mock_compact.assert_called_once_with(mock_delta) + mock_vacuum.assert_called_once_with(mock_delta) else: mock_compact.assert_not_called() + mock_vacuum.assert_not_called() + + +class TestCompactTable: + @pytest.mark.asyncio + async def test_does_not_refetch_table_for_the_vacuum_step(self, helper: DeltaTableHelper): + # Regression: compact_table used to finish its own compact, then call vacuum_table(), + # which called get_delta_table() again instead of reusing the table already in hand. + # get_delta_table() is cached only opportunistically (a concurrent sync of a different + # table can evict this table's cache entry), so that second call could come back None + # and raise "Deltatable not found" right after a successful compact. Asserting a single + # get_delta_table() call locks in that the vacuum step reuses the resolved table instead + # of re-deriving it. + mock_delta = MagicMock() + mock_delta.optimize.compact = MagicMock(return_value={}) + mock_delta.vacuum = MagicMock(return_value=[]) + with patch.object(helper, "get_delta_table", AsyncMock(return_value=mock_delta)) as mock_get: + await helper.compact_table() + + mock_get.assert_called_once() + mock_delta.optimize.compact.assert_called_once() + mock_delta.vacuum.assert_called_once() class TestGetDeltaTableUnrecoverableErrors: @@ -1134,13 +1162,15 @@ async def test_vacuum_cadence( table.version = MagicMock(return_value=150) with ( patch.object(helper, "get_delta_table", new=AsyncMock(return_value=table)), - patch.object(helper, "vacuum_table", new=AsyncMock()) as vacuum, + patch.object(helper, "_vacuum", new=AsyncMock()) as vacuum, patch(f"{module}.posthoganalytics") as ph, ): result = await helper.vacuum_if_stale(last_version, 100) assert result == expected_return assert vacuum.await_count == (1 if expect_vacuum else 0) + if expect_vacuum: + vacuum.assert_awaited_once_with(table) # The observability event fires exactly when a vacuum runs — not on seed/skip — so the cadence is measurable. assert ph.capture.call_count == (1 if expect_vacuum else 0) if expect_vacuum: