fix(data-imports): reuse resolved delta table in compact/vacuum maintenance

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,15 +1035,15 @@ 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.
         """
         table = await self.get_delta_table()
         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(

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: