diff --git a/products/warehouse_sources/backend/temporal/data_imports/pipelines/common/extract.py b/products/warehouse_sources/backend/temporal/data_imports/pipelines/common/extract.py index 7b3a07b2e2c7..5a032a65a49f 100644 --- a/products/warehouse_sources/backend/temporal/data_imports/pipelines/common/extract.py +++ b/products/warehouse_sources/backend/temporal/data_imports/pipelines/common/extract.py @@ -19,11 +19,10 @@ DuplicatePrimaryKeysException, ) from products.warehouse_sources.backend.temporal.data_imports.pipelines.core.cdp_producer import CDPProducer -from products.warehouse_sources.backend.temporal.data_imports.pipelines.core.delta_table_helper import ( - DeltaTableHelper, - is_transient_maintenance_error, +from products.warehouse_sources.backend.temporal.data_imports.pipelines.core.delta.errors import ( is_transient_object_store_error, ) +from products.warehouse_sources.backend.temporal.data_imports.pipelines.core.delta_table_helper import DeltaTableHelper from products.warehouse_sources.backend.temporal.data_imports.pipelines.core.person_property_row_sink import ( PersonPropertyRowSink, ) @@ -697,54 +696,3 @@ async def person_property_sink_clear_chunks(sink: PersonPropertyRowSink): async def stage_chunk_for_person_property_sink(sink: PersonPropertyRowSink, index: int, pa_table: pa.Table): if await sink.should_stage(): await sink.stage_chunk(chunk=index, table=pa_table) - - -async def run_pre_write_defensive_compact( - delta_table_helper: DeltaTableHelper, - schema: "ExternalDataSchema", - resource: SourceResponse, - logger: FilteringBoundLogger, -) -> None: - """Best-effort pre-write compact + vacuum at the start of a sync run. - - Delegates to `DeltaTableHelper.run_maintenance`, which compacts a fragmented Delta - target (a sync that arrived fragmented because earlier attempts failed before - reaching `_post_run_operations` — keeping the subsequent per-partition merge scans - cheap) and otherwise vacuums on a commit-count cadence so a table that OOMs its merge - every run and never reaches post-load compaction still sheds tombstones (the - ~99%-dead-file tables). The helper returns the single vacuum watermark to persist; - the CDC post-load path in `common/load.py` writes the same watermark, and both merge - via `update_sync_type_config_keys` under a row lock. Wrapped in try/except so a - maintenance failure never blocks the actual sync; the original error path is unaffected. - A transient infra error (see `is_transient_maintenance_error`) — an object-store hiccup, a racy - concurrent-maintenance DeltaError, or an app-DB connection blip — is logged at warning instead of - captured. The next sync's maintenance pass retries it from scratch, so it isn't a bug in this - function, just a temporary blip talking to our own S3 bucket, delta table, or app DB. - - Used by both `PipelineNonDLT.run` (v2) and `PipelineV3.run` to keep the behaviour - identical across pipelines without each having to know how to look up `partition_count` - or how to swallow maintenance errors. - """ - try: - from products.warehouse_sources.backend.models.external_data_schema import ( # noqa: PLC0415 — Django model import kept off this activity module's load path - update_sync_type_config_keys, - ) - - partition_count_for_compact = schema.partition_count or resource.partition_count - last_vacuum_version = schema.last_vacuum_version - commit_threshold = settings.DATA_WAREHOUSE_VACUUM_COMMIT_THRESHOLD - new_version = await delta_table_helper.run_maintenance( - partition_count=partition_count_for_compact, - last_vacuum_version=last_vacuum_version, - commit_threshold=commit_threshold, - ) - if new_version is not None and new_version != last_vacuum_version: - await database_sync_to_async_pool(update_sync_type_config_keys)( - schema.id, schema.team_id, updates={"last_vacuum_version": new_version} - ) - except Exception as e: - if is_transient_maintenance_error(e): - await logger.awarning(f"Pre-write maintenance skipped: transient infra error: {e}") - return - capture_exception(e) - await logger.aexception(f"Pre-write maintenance failed: {e}", exc_info=e) 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 75720465779e..5c27244171ce 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 @@ -1,6 +1,5 @@ from typing import TYPE_CHECKING, Any, Literal, Optional, Protocol -from django.conf import settings from django.db.models import F import pyarrow as pa @@ -13,11 +12,7 @@ from posthog.temporal.common.logger import get_logger from products.warehouse_sources.backend.models.external_data_job import ExternalDataJob -from products.warehouse_sources.backend.models.external_data_schema import ( - ExternalDataSchema, - process_incremental_value, - update_sync_type_config_keys, -) +from products.warehouse_sources.backend.models.external_data_schema import ExternalDataSchema, process_incremental_value from products.warehouse_sources.backend.models.table import DataWarehouseTable from products.warehouse_sources.backend.temporal.data_imports.pipelines.common.db_retry import ( retry_on_operational_error, @@ -278,58 +273,32 @@ async def _run_delta_maintenance( is_cdc_companion: bool, logger: FilteringBoundLogger, ) -> None: - from products.warehouse_sources.backend.temporal.data_imports.pipelines.core.delta_table_helper import ( # noqa: PLC0415 — keeps the heavy deltalake dep off this module's top-level import path - is_transient_object_store_error, + from products.warehouse_sources.backend.temporal.data_imports.pipelines.core.delta.errors import ( # noqa: PLC0415 — keeps the heavy deltalake dep off this module's top-level import path + is_transient_maintenance_error, + ) + from products.warehouse_sources.backend.temporal.data_imports.pipelines.core.delta.maintenance import ( # noqa: PLC0415 — keeps the heavy deltalake dep off this module's top-level import path + DeltaMaintenance, ) + maintenance = DeltaMaintenance(delta_table_helper) if schema.is_cdc: # CDC finals land once per tick per changed schema, so unconditional compaction would run # near-continuously after mostly-tiny merges. Use threshold/cadence maintenance instead: # compact when fragmented, otherwise vacuum once enough commits have accrued. logger.debug("Running threshold-based delta maintenance") - try: - with POST_LOAD_DURATION_SECONDS.labels(operation="maintenance").time(): - # Only md5 partitioning persists a partition_count; datetime/numerical modes leave it - # None, and the companion is a different table than the one schema.partition_count - # describes. Without a count the threshold math treats the table as one partition and - # any >200-file table compacts every tick, so derive it from the table's actual layout - # (one directory per partition value in the delta log's file paths). - partition_count = None if is_cdc_companion else schema.partition_count - if partition_count is None: - file_uris = await delta_table_helper.get_file_uris() - partition_count = len({uri.rsplit("/", 1)[0] for uri in file_uris}) or None - - # One schema can back two delta tables (snapshot + _cdc companion) whose delta - # versions are unrelated numbers, so each table's vacuum cadence gets its own - # watermark key — sharing one would corrupt both cadences. - watermark_key = "last_vacuum_version_cdc" if is_cdc_companion else "last_vacuum_version" - last_vacuum_version = schema.last_vacuum_version_cdc if is_cdc_companion else schema.last_vacuum_version - commit_threshold = settings.DATA_WAREHOUSE_VACUUM_COMMIT_THRESHOLD - new_version = await delta_table_helper.run_maintenance( - partition_count=partition_count, - last_vacuum_version=last_vacuum_version, - commit_threshold=commit_threshold, - ) - if new_version is not None and new_version != last_vacuum_version: - await database_sync_to_async_pool(update_sync_type_config_keys)( - schema.id, schema.team_id, updates={watermark_key: new_version} - ) - except Exception as e: - if is_transient_object_store_error(e): - # A rate-limited or connectivity blip talking to our own S3 bucket isn't a bug - the - # next tick's maintenance pass retries the same idempotent cleanup. - logger.warning(f"Delta maintenance skipped: transient object-store error: {e}") - else: - capture_exception(e) - logger.exception(f"Delta maintenance failed: {e}", exc_info=e) + with POST_LOAD_DURATION_SECONDS.labels(operation="maintenance").time(): + await maintenance.run_scheduled(schema, is_cdc_companion=is_cdc_companion) else: logger.debug("Triggering compaction and vacuuming on delta table") try: with POST_LOAD_DURATION_SECONDS.labels(operation="compact").time(): - await delta_table_helper.compact_table() + await maintenance.compact_table() except Exception as e: - if is_transient_object_store_error(e): - logger.warning(f"Compaction skipped: transient object-store error: {e}") + if is_transient_maintenance_error(e): + # A rate-limited or connectivity blip talking to our own S3 bucket (or a concurrent + # maintenance pass losing a file race) isn't a bug - the next sync's maintenance pass + # retries the same idempotent cleanup. + logger.warning(f"Compaction skipped: transient infra error: {e}") else: capture_exception(e) logger.exception(f"Compaction failed: {e}", exc_info=e) diff --git a/products/warehouse_sources/backend/temporal/data_imports/pipelines/common/test/test_extract.py b/products/warehouse_sources/backend/temporal/data_imports/pipelines/common/test/test_extract.py index c23f61401707..09b92b4825ab 100644 --- a/products/warehouse_sources/backend/temporal/data_imports/pipelines/common/test/test_extract.py +++ b/products/warehouse_sources/backend/temporal/data_imports/pipelines/common/test/test_extract.py @@ -5,9 +5,6 @@ from posthog.test.base import BaseTest from unittest.mock import AsyncMock, MagicMock, patch -from django.db import InterfaceError, OperationalError - -import deltalake.exceptions from asgiref.sync import async_to_sync from parameterized import parameterized @@ -21,7 +18,6 @@ persist_primary_keys, report_heartbeat_timeout, resolve_primary_keys, - run_pre_write_defensive_compact, ) from products.warehouse_sources.backend.temporal.data_imports.util import NonRetryableException @@ -130,138 +126,6 @@ async def _call(*args, **kwargs): logger.aexception.assert_awaited_once() -class TestRunPreWriteDefensiveCompact: - @parameterized.expand( - [ - # (schema_partition_count, resource_partition_count, expected_passed_to_run_maintenance) - ("schema_value_wins", 10, 72, 10), - ("falls_back_to_resource", None, 72, 72), - ("both_none_passes_none", None, None, None), - ] - ) - @pytest.mark.asyncio - async def test_resolves_partition_count_schema_over_resource( - self, _name: str, schema_count: int | None, resource_count: int | None, expected: int | None - ): - run_maintenance = AsyncMock(return_value=None) - helper = MagicMock(run_maintenance=run_maintenance) - - await run_pre_write_defensive_compact( - helper, - MagicMock(partition_count=schema_count, sync_type_config={}), - MagicMock(partition_count=resource_count), - MagicMock(aexception=AsyncMock()), - ) - - assert run_maintenance.await_args is not None - assert run_maintenance.await_args.kwargs["partition_count"] == expected - - @pytest.mark.asyncio - async def test_swallows_maintenance_failure(self): - # The whole point of the wrapper: a maintenance error must never propagate and - # block the sync — it's captured and logged instead. - helper = MagicMock(run_maintenance=AsyncMock(side_effect=RuntimeError("maintenance blew up"))) - logger = MagicMock(aexception=AsyncMock()) - - schema = MagicMock(partition_count=5, sync_type_config={}) - with patch(f"{_EXTRACT_MODULE}.capture_exception") as mock_capture: - await run_pre_write_defensive_compact(helper, schema, MagicMock(partition_count=None), logger) - - mock_capture.assert_called_once() - logger.aexception.assert_awaited_once() - - @parameterized.expand( - [ - ( - "credentials_loading", - OSError, - "Operation not supported: an error occurred while loading credentials: dispatch failure: timeout", - ), - ( - "credential_provider_not_enabled", - OSError, - "Operation not supported: the credential provider was not enabled: no providers in chain provided credentials", - ), - ( - "generic_s3_error", - OSError, - "Generic S3 error: Error getting list response body: operation timed out", - ), - ( - # table.vacuum()/optimize.compact() surface the identical object-store error text - # wrapped in DeltaError instead of OSError (unlike is_deltatable()'s OSError) — the - # exact shape of the issue this test guards against. - "generic_s3_error_as_delta_error", - deltalake.exceptions.DeltaError, - "Generic error: Kernel error: Error interacting with object store: Generic S3 error: " - "Server returned non-2xx status code: 503 Service Unavailable: SlowDown", - ), - ] - ) - @pytest.mark.asyncio - async def test_logs_transient_object_store_error_without_capturing( - self, _name: str, error_cls: type[Exception], error_message: str - ): - # A transient blip talking to our own delta S3 bucket (credential-provider or connectivity - # errors from delta-rs) isn't a bug in this function — it shouldn't flood error tracking the - # way an actual maintenance bug does (see test_swallows_maintenance_failure above). - helper = MagicMock(run_maintenance=AsyncMock(side_effect=error_cls(error_message))) - logger = MagicMock(aexception=AsyncMock(), awarning=AsyncMock()) - - schema = MagicMock(partition_count=5, sync_type_config={}) - with patch(f"{_EXTRACT_MODULE}.capture_exception") as mock_capture: - await run_pre_write_defensive_compact(helper, schema, MagicMock(partition_count=None), logger) - - mock_capture.assert_not_called() - logger.awarning.assert_awaited_once() - logger.aexception.assert_not_awaited() - - @pytest.mark.asyncio - async def test_logs_transient_delta_maintenance_race_without_capturing(self): - # Regression: a concurrent optimize/vacuum pass on the same table (e.g. a zombie Temporal - # attempt racing its own retry) can have `optimize.compact` scan a file the other attempt - # already vacuumed away. Nothing gets committed when the scan fails, so the table isn't - # corrupted — this must be treated the same as the object-store blips above, not captured. - error = deltalake.exceptions.DeltaError( - "Failed to parse parquet: Optimize selected-file scan failed while scanning data: " - "Object at location .../part-0.parquet not found: 404 Not Found" - ) - helper = MagicMock(run_maintenance=AsyncMock(side_effect=error)) - logger = MagicMock(aexception=AsyncMock(), awarning=AsyncMock()) - - schema = MagicMock(partition_count=5, sync_type_config={}) - with patch(f"{_EXTRACT_MODULE}.capture_exception") as mock_capture: - await run_pre_write_defensive_compact(helper, schema, MagicMock(partition_count=None), logger) - - mock_capture.assert_not_called() - logger.awarning.assert_awaited_once() - logger.aexception.assert_not_awaited() - - @parameterized.expand( - [ - ("dns_resolution_failure", OperationalError, "[Errno -2] Name or service not known"), - ("pooler_dropped_connection", InterfaceError, "connection already closed"), - ] - ) - @pytest.mark.asyncio - async def test_logs_transient_db_connection_error_without_capturing( - self, _name: str, error_cls: type[Exception], error_message: str - ): - # A DNS/pooler blip hit while resolving `job.folder_path()` on a pooled app-DB connection - # (e.g. inside `_get_delta_table_uri`) isn't a maintenance bug either — same treatment as - # the object-store blips above, so a self-healing retry doesn't flood error tracking. - helper = MagicMock(run_maintenance=AsyncMock(side_effect=error_cls(error_message))) - logger = MagicMock(aexception=AsyncMock(), awarning=AsyncMock()) - - schema = MagicMock(partition_count=5, sync_type_config={}) - with patch(f"{_EXTRACT_MODULE}.capture_exception") as mock_capture: - await run_pre_write_defensive_compact(helper, schema, MagicMock(partition_count=None), logger) - - mock_capture.assert_not_called() - logger.awarning.assert_awaited_once() - logger.aexception.assert_not_awaited() - - class TestReportHeartbeatTimeoutRecording(BaseTest): def _schema(self) -> ExternalDataSchema: source = ExternalDataSource.objects.create( diff --git a/products/warehouse_sources/backend/temporal/data_imports/pipelines/common/test/test_load.py b/products/warehouse_sources/backend/temporal/data_imports/pipelines/common/test/test_load.py index 13c9817cb36e..65218debcb3f 100644 --- a/products/warehouse_sources/backend/temporal/data_imports/pipelines/common/test/test_load.py +++ b/products/warehouse_sources/backend/temporal/data_imports/pipelines/common/test/test_load.py @@ -16,6 +16,7 @@ run_post_load_operations, update_job_row_count, ) +from products.warehouse_sources.backend.temporal.data_imports.pipelines.core.delta.maintenance import DeltaMaintenance _LOAD_MODULE = "products.warehouse_sources.backend.temporal.data_imports.pipelines.common.load" _DB_RETRY_MODULE = "products.warehouse_sources.backend.temporal.data_imports.pipelines.common.db_retry" @@ -39,12 +40,10 @@ def _make_schema(*, is_cdc: bool, sync_type_config: dict | None = None, partitio return schema -def _make_helper(*, run_maintenance_returns: int | None = None, file_uris: list[str] | None = None) -> MagicMock: +def _make_helper(*, file_uris: list[str] | None = None) -> MagicMock: return MagicMock( get_delta_table=AsyncMock(return_value=MagicMock()), get_file_uris=AsyncMock(return_value=file_uris or []), - compact_table=AsyncMock(), - run_maintenance=AsyncMock(return_value=run_maintenance_returns), ) @@ -53,19 +52,23 @@ async def _run_post_load( helper: MagicMock, *, cdc_write_mode: str | None = None, -) -> tuple[MagicMock, AsyncMock]: + compact_error: Exception | None = None, +) -> tuple[AsyncMock, AsyncMock, AsyncMock]: job = MagicMock() job.id = uuid.uuid4() job.team_id = schema.team_id logger = MagicMock(adebug=AsyncMock(), ainfo=AsyncMock()) prepare_s3 = AsyncMock(return_value="orders__query_1") + run_scheduled = AsyncMock() + compact_table = AsyncMock(side_effect=compact_error) with ( patch(f"{_LOAD_MODULE}.prepare_s3_files_for_querying", prepare_s3), patch(f"{_LOAD_MODULE}.notify_revenue_analytics_that_sync_has_completed", AsyncMock()), patch(f"{_LOAD_MODULE}.sync_revenue_analytics_views", MagicMock()), - patch(f"{_LOAD_MODULE}.update_sync_type_config_keys", MagicMock()) as update_config, patch(f"{_LOAD_MODULE}.DataWarehouseTable", MagicMock()), + patch.object(DeltaMaintenance, "run_scheduled", run_scheduled), + patch.object(DeltaMaintenance, "compact_table", compact_table), patch(f"{_PIPELINE_SYNC_MODULE}.update_last_synced_at", AsyncMock()), patch(f"{_PIPELINE_SYNC_MODULE}.validate_schema_and_update_table", AsyncMock()), patch(f"{_PIPELINE_SYNC_MODULE}.register_cdc_companion_table", AsyncMock()), @@ -82,100 +85,43 @@ async def _run_post_load( logger=logger, cdc_write_mode=cdc_write_mode, ) - return update_config, prepare_s3 + return run_scheduled, compact_table, prepare_s3 class TestRunPostLoadDeltaMaintenance: + """Post-load picks the right maintenance flavor per schema kind; the threshold/watermark + mechanics themselves are covered in core/delta/test/test_maintenance.py.""" + @pytest.mark.asyncio async def test_cdc_schema_uses_threshold_maintenance_not_unconditional_compact(self): # The incident behavior this guards: CDC finals land every tick, so an unconditional # compact_table here means hundreds of compact+vacuum cycles per hour on a busy source. schema = _make_schema(is_cdc=True, sync_type_config={"last_vacuum_version": 41}) - helper = _make_helper() - - await _run_post_load(schema, helper, cdc_write_mode="incremental") - - helper.compact_table.assert_not_awaited() - assert helper.run_maintenance.await_args is not None - assert helper.run_maintenance.await_args.kwargs == { - "partition_count": 7, - "last_vacuum_version": 41, - "commit_threshold": 100, - } - - @pytest.mark.asyncio - async def test_missing_partition_count_is_derived_from_table_layout(self): - # datetime/numerical-partitioned schemas persist no partition_count. Passing None through - # makes the threshold math treat the table as one partition, so any >200-file table would - # compact every tick again — the exact behavior this change removes. - schema = _make_schema(is_cdc=True, partition_count=None) - helper = _make_helper( - file_uris=[ - "s3://bucket/orders/_ph_partition_key=2026-01/a.parquet", - "s3://bucket/orders/_ph_partition_key=2026-01/b.parquet", - "s3://bucket/orders/_ph_partition_key=2026-02/c.parquet", - ] - ) - await _run_post_load(schema, helper, cdc_write_mode="incremental") + run_scheduled, compact_table, _ = await _run_post_load(schema, _make_helper(), cdc_write_mode="incremental") - assert helper.run_maintenance.await_args is not None - assert helper.run_maintenance.await_args.kwargs["partition_count"] == 2 + compact_table.assert_not_awaited() + run_scheduled.assert_awaited_once_with(schema, is_cdc_companion=False) @pytest.mark.asyncio async def test_non_cdc_schema_keeps_unconditional_compact(self): schema = _make_schema(is_cdc=False) - helper = _make_helper() - update_config, _ = await _run_post_load(schema, helper) + run_scheduled, compact_table, _ = await _run_post_load(schema, _make_helper()) - helper.compact_table.assert_awaited_once() - helper.run_maintenance.assert_not_awaited() - update_config.assert_not_called() + compact_table.assert_awaited_once() + run_scheduled.assert_not_awaited() @pytest.mark.asyncio - async def test_cdc_companion_uses_its_own_watermark_key(self): - # The snapshot and _cdc companion are different delta tables with unrelated versions, so - # the companion must run cadence maintenance against last_vacuum_version_cdc — reading or - # writing the snapshot's last_vacuum_version would corrupt both cadences, and skipping - # cadence maintenance entirely would let companion tombstones accumulate until the - # file-count thresholds happen to trip. Partition count is derived from its own layout — - # schema.partition_count describes the snapshot table. + async def test_cdc_companion_write_runs_companion_maintenance(self): + # The snapshot and _cdc companion are different delta tables, so a companion (scd2_append) + # write must run maintenance in companion mode — run_scheduled then uses the companion's own + # watermark key and layout instead of the snapshot's (see test_maintenance.TestRunScheduled). schema = _make_schema(is_cdc=True, sync_type_config={"last_vacuum_version": 41, "last_vacuum_version_cdc": 7}) - helper = _make_helper(run_maintenance_returns=9, file_uris=["s3://bucket/orders_cdc/a.parquet"]) - - update_config, _ = await _run_post_load(schema, helper, cdc_write_mode="scd2_append") - - assert helper.run_maintenance.await_args is not None - assert helper.run_maintenance.await_args.kwargs == { - "partition_count": 1, - "last_vacuum_version": 7, - "commit_threshold": 100, - } - update_config.assert_called_once_with(schema.id, schema.team_id, updates={"last_vacuum_version_cdc": 9}) - @parameterized.expand( - [ - # run_maintenance returning a version must persist it — a lost watermark means - # vacuum_if_stale re-seeds forever and the table never vacuums. - ("new_version_persists", 55, True), - ("no_change_skips_write", None, False), - ("same_version_skips_write", 41, False), - ] - ) - @pytest.mark.asyncio - async def test_watermark_persistence(self, _name: str, returned_version: int | None, expect_write: bool): - schema = _make_schema(is_cdc=True, sync_type_config={"last_vacuum_version": 41}) - helper = _make_helper(run_maintenance_returns=returned_version) + run_scheduled, _, _ = await _run_post_load(schema, _make_helper(), cdc_write_mode="scd2_append") - update_config, _ = await _run_post_load(schema, helper, cdc_write_mode="incremental") - - if expect_write: - update_config.assert_called_once_with( - schema.id, schema.team_id, updates={"last_vacuum_version": returned_version} - ) - else: - update_config.assert_not_called() + run_scheduled.assert_awaited_once_with(schema, is_cdc_companion=True) @parameterized.expand([("non_cdc", False), ("cdc", True)]) @pytest.mark.asyncio @@ -187,7 +133,7 @@ async def test_prepares_s3_files_with_post_maintenance_file_list(self, _name: st post_maintenance_uris = ["s3://bucket/orders/compacted.parquet"] helper = _make_helper(file_uris=post_maintenance_uris) - _, prepare_s3 = await _run_post_load(schema, helper, cdc_write_mode="incremental" if is_cdc else None) + _, _, prepare_s3 = await _run_post_load(schema, helper, cdc_write_mode="incremental" if is_cdc else None) prepare_s3.assert_awaited_once() assert prepare_s3.await_args is not None @@ -195,43 +141,22 @@ async def test_prepares_s3_files_with_post_maintenance_file_list(self, _name: st @parameterized.expand( [ - # A genuine maintenance bug must still be captured for visibility. - ("genuine_bug", RuntimeError("maintenance blew up"), True), + # A genuine compaction bug must still be captured for visibility. + ("genuine_bug", RuntimeError("compaction blew up"), True), # A transient S3 rate-limit/connectivity blip is already non-fatal here (the next - # tick's maintenance retries the same idempotent cleanup) and must not be promoted + # sync's maintenance retries the same idempotent cleanup) and must not be promoted # into a fresh error-tracking issue — the regression this guards. ("transient_s3_slowdown", OSError("Generic S3 error: Please reduce your request rate."), False), ] ) @pytest.mark.asyncio - async def test_maintenance_failure_handling(self, _name: str, error: Exception, expect_capture: bool): - # A maintenance hiccup must not fail the final batch — the rest of post-load - # (queryable folder prep, table registration) still has to run or the job wedges. - schema = _make_schema(is_cdc=True) - helper = _make_helper() - helper.run_maintenance = AsyncMock(side_effect=error) - - with patch(f"{_LOAD_MODULE}.capture_exception") as mock_capture: - _, prepare_s3 = await _run_post_load(schema, helper, cdc_write_mode="incremental") - - assert mock_capture.called is expect_capture - prepare_s3.assert_awaited_once() - - @parameterized.expand( - [ - ("genuine_bug", RuntimeError("compaction blew up"), True), - ("transient_s3_slowdown", OSError("Generic S3 error: Please reduce your request rate."), False), - ] - ) - @pytest.mark.asyncio async def test_compact_failure_handling(self, _name: str, error: Exception, expect_capture: bool): - # Same non-fatal handling as maintenance, for the non-CDC unconditional compact_table path. + # A compaction hiccup must not fail the final batch — the rest of post-load + # (queryable folder prep, table registration) still has to run or the job wedges. schema = _make_schema(is_cdc=False) - helper = _make_helper() - helper.compact_table = AsyncMock(side_effect=error) with patch(f"{_LOAD_MODULE}.capture_exception") as mock_capture: - _, prepare_s3 = await _run_post_load(schema, helper) + _, _, prepare_s3 = await _run_post_load(schema, _make_helper(), compact_error=error) assert mock_capture.called is expect_capture prepare_s3.assert_awaited_once() diff --git a/products/warehouse_sources/backend/temporal/data_imports/pipelines/core/delta/errors.py b/products/warehouse_sources/backend/temporal/data_imports/pipelines/core/delta/errors.py new file mode 100644 index 000000000000..af5744c26f57 --- /dev/null +++ b/products/warehouse_sources/backend/temporal/data_imports/pipelines/core/delta/errors.py @@ -0,0 +1,69 @@ +from django.db import InterfaceError, OperationalError + +import botocore.exceptions +import deltalake.exceptions + +# Substrings of the object-store errors raised talking to our own S3-backed data-warehouse bucket +# that are transient and self-recovering, not a bug in our code or a customer credential problem: +# - the first three come from delta-rs's Rust `object_store` crate inside `DeltaTable.is_deltatable()` +# (IMDS/STS blips, dispatch timeouts) +# - "Please reduce your request rate" is S3's SlowDown throttling response, surfaced by s3fs/aiobotocore +# when a bulk operation (e.g. `_purge_s3_prefix`'s list-then-delete) outruns the bucket's request-rate limit +# A retry (of the same idempotent operation) clears these, so they shouldn't be treated the same as a +# bug in our logic. +TRANSIENT_OBJECT_STORE_ERRORS = ( + "an error occurred while loading credentials", + "the credential provider was not enabled", + "Generic S3 error", + "Please reduce your request rate", +) + + +def is_transient_object_store_error(error: BaseException) -> bool: + """True for a transient object-store error, however it happened to surface. + + `DeltaTable.is_deltatable()` raises these as a plain `OSError`, but table-level operations + (e.g. `vacuum()`, `optimize.compact()`) wrap the identical underlying object-store error text in + `deltalake.exceptions.DeltaError` instead — same blip, different exception type depending on + which delta-rs entry point hit it. `_purge_s3_prefix`'s s3fs/aiobotocore calls can also raise a + bare `botocore.exceptions.NoCredentialsError` unwrapped — the same IMDS/STS credential-provider + blip, just surfaced by aiobotocore's own credential resolution instead of delta-rs's Rust + `object_store` crate. `NoCredentialsError`'s message is a fixed, generic string (no needle to + match), but hitting our own instance-role-authenticated bucket always means the same transient + resolution hiccup, so it's recognized by type rather than by message. + """ + if isinstance(error, botocore.exceptions.NoCredentialsError): + return True + return isinstance(error, OSError | deltalake.exceptions.DeltaError) and any( + needle in str(error) for needle in TRANSIENT_OBJECT_STORE_ERRORS + ) + + +# `optimize.compact` plans its rewrite against the file list at the start of its scan, then reads +# those files. A concurrent maintenance pass on the same table (e.g. a Temporal activity attempt that +# heartbeat-timed-out but keeps running as a zombie — see this package's README on the equivalent +# unfenced race for repartition) can vacuum one of those files out from under the scan before it gets +# read, which delta-rs surfaces as this DeltaError. The scan failing here means the optimize aborted +# before committing anything — the table is left exactly as it was, just still fragmented — so this is +# safe to skip and retry on the next maintenance pass, not a bug in our logic. +TRANSIENT_DELTA_MAINTENANCE_ERRORS = ("Optimize selected-file scan failed",) + + +def is_transient_delta_maintenance_error(error: BaseException) -> bool: + return isinstance(error, deltalake.exceptions.DeltaError) and any( + needle in str(error) for needle in TRANSIENT_DELTA_MAINTENANCE_ERRORS + ) + + +def is_transient_maintenance_error(error: BaseException) -> bool: + """Infra blips seen during delta maintenance that aren't a maintenance bug. + + Covers S3/object-store hiccups reaching our own data-warehouse bucket (see + `is_transient_object_store_error` above), racy concurrent-maintenance DeltaErrors (see + `is_transient_delta_maintenance_error` above), and app-DB connection blips (DNS, pooler drops) hit + while resolving `job.folder_path()` on a pooled connection — the same `OperationalError`/`InterfaceError` + classification used for this failure class in `repartition_table.py`'s `_is_transient_infra_error`. + """ + if isinstance(error, OperationalError | InterfaceError): + return True + return is_transient_object_store_error(error) or is_transient_delta_maintenance_error(error) diff --git a/products/warehouse_sources/backend/temporal/data_imports/pipelines/core/delta/maintenance.py b/products/warehouse_sources/backend/temporal/data_imports/pipelines/core/delta/maintenance.py new file mode 100644 index 000000000000..e79a14c4618c --- /dev/null +++ b/products/warehouse_sources/backend/temporal/data_imports/pipelines/core/delta/maintenance.py @@ -0,0 +1,274 @@ +import json +import asyncio +from typing import TYPE_CHECKING + +from django.conf import settings + +import deltalake +import posthoganalytics + +from posthog.exceptions_capture import capture_exception +from posthog.sync import database_sync_to_async_pool +from posthog.utils import get_machine_id + +from products.warehouse_sources.backend.models.external_data_schema import update_sync_type_config_keys +from products.warehouse_sources.backend.temporal.data_imports.pipelines.core.delta.errors import ( + is_transient_maintenance_error, +) +from products.warehouse_sources.backend.temporal.data_imports.pipelines.core.delta.ops import ( + execute_with_conflict_retry, +) + +if TYPE_CHECKING: + from products.warehouse_sources.backend.models.external_data_schema import ExternalDataSchema + from products.warehouse_sources.backend.temporal.data_imports.pipelines.core.delta_table_helper import ( + DeltaTableHelper, + ) + +# A defensive compact fires when EITHER threshold is exceeded. +# +# Calibrated against the production file-count distribution (delta merge stats across +# all teams): total files per table sit at p50≈60, p90≈470, p95≈850, with a long tail +# (p99≈12k → ~14s merges; an observed pathological case hit ~82k files → ~45s merges). +# Merge planning time tracks TOTAL files, not files-per-partition — delta still +# enumerates every file's metadata even when partition pruning skips reading them — so +# we gate on both: +# +# - files-per-partition: bounds per-partition fragmentation and rescues partitioned +# (esp. md5) tables, where a merge touches every partition. 200 sits well above the +# healthy steady state (compaction runs at the end of each successful sync) yet +# triggers long before a table reaches the slow tail. +# - total files: a partition-count-independent backstop so a table with a high +# partition_count can't accumulate tens of thousands of files (each adding to merge +# planning time) while staying under the per-partition bar. 5,000 is above p95 (~850) +# — so healthy tables never trip it — and well below the p99/pathological tail. +# +# Tune further once the admin fragmentation view gives per-customer distributions. +DEFAULT_COMPACT_FILES_PER_PARTITION_THRESHOLD = 200 +DEFAULT_COMPACT_TOTAL_FILES_THRESHOLD = 5000 + + +class DeltaMaintenance: + """Compaction, vacuuming, and the vacuum-watermark cadence for one schema's Delta table. + + Stateless over a `DeltaTableHelper`, which holds the cached table handle — construct one at the + call site whenever maintenance is needed. `run_scheduled` is the policy entry point shared by + the pre-write defensive pass (both pipelines, so a sync that arrived at a fragmented table + cleans up before adding to the pile) and the CDC post-load pass; `compact_table` is the + unconditional post-load compaction for non-CDC syncs. + """ + + def __init__(self, table: "DeltaTableHelper") -> None: + self._table = table + self._logger = table.logger + + 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(self, table: deltalake.DeltaTable) -> None: + await self._logger.adebug("Compacting table...") + compact_stats = await execute_with_conflict_retry( + table, lambda: table.optimize.compact(), "compact_table", self._logger + ) + await self._logger.adebug(json.dumps(compact_stats)) + + async def compact_table(self) -> None: + table = await self._table.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: + """Vacuum tombstoned files once enough commits have accrued since the last vacuum. + + Decoupled from merge success (called pre-write) so a table that OOMs its merge every run still + gets cleaned — the post-load compaction never runs for it, which is how tables reach ~99% dead + files. Vacuum only deletes dead files (an S3 LIST + delete), so unlike `compact_table`'s + `optimize.compact` (which rewrites partitions) it is memory-safe even on an oversized table. + + Uses the delta version (commit count) as a cheap proxy for tombstone accumulation — no S3 LIST to + decide. Returns the current version to persist as the new watermark when it vacuumed, on first + encounter, or when the table was recreated (both reseed the watermark without vacuuming); + None when nothing changed. + """ + table = await self._table.get_delta_table() + if table is None: + return None + + version = await asyncio.to_thread(table.version) + if last_vacuum_version is None or version < last_vacuum_version: + # First encounter: seed the watermark without vacuuming so existing tables clean up gradually + # over the next `commit_threshold` commits rather than all vacuuming at once on deploy. + # A version below the watermark means the table was reset/recreated (delta versions are + # monotonic within one incarnation) and no reset path clears the persisted watermark — + # left alone it would block the cadence until the new table out-versioned the old one. + return version + + commits_since = version - last_vacuum_version + if commits_since < commit_threshold: + await self._logger.adebug( + f"vacuum_if_stale: skipping, {commits_since} commits since last vacuum (< {commit_threshold})" + ) + return None + + await self._logger.ainfo( + f"vacuum_if_stale: {commits_since} commits since last vacuum (>= {commit_threshold}), vacuuming" + ) + 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. + posthoganalytics.capture( + distinct_id=get_machine_id(), + event="warehouse_delta_vacuumed", + properties={ + "team_id": self._table.job.team_id, + "schema_id": str(self._table.job.schema_id), + "source_id": str(self._table.job.pipeline_id), + "resource_name": self._table.resource_name, + "commits_since_last_vacuum": commits_since, + "delta_version": version, + }, + ) + except Exception as e: + capture_exception(e) + return version + + async def compact_if_fragmented( + self, + partition_count: int | None, + threshold: int = DEFAULT_COMPACT_FILES_PER_PARTITION_THRESHOLD, + total_threshold: int = DEFAULT_COMPACT_TOTAL_FILES_THRESHOLD, + ) -> bool: + """Run compact + vacuum if the table is fragmented past either threshold. + + Fragmented = files-per-partition > `threshold` OR total files > `total_threshold`. + The total-files backstop matters because delta enumerates every file's metadata + during a merge even when partition pruning skips reading them, so merge planning + time tracks total files — a high partition_count must not let a table accumulate + tens of thousands of files while staying under the per-partition bar. + + When `partition_count` is None it is derived from the table's actual layout (the + distinct file directories in the delta log, no extra I/O) — only md5 partitioning + persists a count on the schema, so datetime/numerical-partitioned tables always + arrive here with None. + + Returns True if compaction ran, False if it was skipped. Cheap when the table is + 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-load compaction) cleans up before adding to the pile. + """ + table = await self._table.get_delta_table() + if table is None: + return False + + 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 + # table root. Without this, a partitioned table with no persisted count reads as + # one giant partition and trips the per-partition threshold on every run. + partition_count = len({uri.rsplit("/", 1)[0] for uri in file_uris}) + # Treat unpartitioned tables as one "partition" for the threshold math. + effective_partitions = max(partition_count or 1, 1) + files_per_partition = total_files / effective_partitions + + fragmented = files_per_partition > threshold or total_files > total_threshold + stats = ( + f"total_files={total_files}, partitions={effective_partitions}, " + f"files_per_partition={files_per_partition:.1f}, threshold={threshold}, " + f"total_threshold={total_threshold}" + ) + if not fragmented: + await self._logger.adebug(f"compact_if_fragmented: skipping ({stats})") + return False + + await self._logger.ainfo(f"compact_if_fragmented: triggering compact ({stats})") + await self._compact(table) + await self._vacuum(table) + return True + + async def run_maintenance( + self, + partition_count: int | None, + last_vacuum_version: int | None, + commit_threshold: int, + ) -> int | None: + """Single threshold-maintenance step: compact if fragmented, else vacuum on commit cadence. + + The two triggers are orthogonal — fragmentation (active file count) vs. commit cadence (tombstone + accrual) — but they share one outcome, the vacuum watermark. `compact_if_fragmented` already + vacuums as part of compaction, so when it runs it supersedes the cadence vacuum (no double vacuum + in one run) and the watermark advances to the post-compaction version. When nothing was fragmented, + fall through to `vacuum_if_stale`. Returns the single delta version to persist as the new + `last_vacuum_version` watermark, or None when nothing changed — `run_scheduled` persists it. + """ + compacted = await self.compact_if_fragmented(partition_count=partition_count) + if compacted: + table = await self._table.get_delta_table() + if table is None: + return None + # Compaction (which vacuumed) added a commit, advancing the version; reset the cadence + # watermark to it so the next vacuum is measured from this cleanup, not the old baseline. + return await asyncio.to_thread(table.version) + return await self.vacuum_if_stale(last_vacuum_version, commit_threshold) + + async def run_scheduled( + self, + schema: "ExternalDataSchema", + *, + is_cdc_companion: bool = False, + partition_count_fallback: int | None = None, + ) -> None: + """Best-effort threshold maintenance owning the vacuum-watermark lifecycle for `schema`. + + Reads the right watermark, runs `run_maintenance`, and persists the returned watermark via + `update_sync_type_config_keys` (row-locked merge) — both call sites (the pre-write defensive + pass and the CDC post-load pass) share this, so the watermark can't drift between them. + + One schema can back two delta tables (snapshot + `_cdc` companion) whose delta versions are + unrelated numbers, so each table's vacuum cadence gets its own watermark key — sharing one + would corrupt both cadences. The companion also ignores `schema.partition_count` (it + describes the snapshot table); `compact_if_fragmented` derives the companion's count from + its actual layout instead. + + Never raises: a maintenance failure must not block the sync, and the next scheduled pass + retries the same idempotent cleanup. A transient infra error (see + `is_transient_maintenance_error`) — an object-store hiccup, a racy concurrent-maintenance + DeltaError, or an app-DB connection blip — is logged at warning instead of captured. + """ + try: + if is_cdc_companion: + partition_count = None + watermark_key = "last_vacuum_version_cdc" + last_vacuum_version = schema.last_vacuum_version_cdc + else: + partition_count = schema.partition_count or partition_count_fallback + watermark_key = "last_vacuum_version" + last_vacuum_version = schema.last_vacuum_version + + new_version = await self.run_maintenance( + partition_count=partition_count, + last_vacuum_version=last_vacuum_version, + commit_threshold=settings.DATA_WAREHOUSE_VACUUM_COMMIT_THRESHOLD, + ) + if new_version is not None and new_version != last_vacuum_version: + await database_sync_to_async_pool(update_sync_type_config_keys)( + schema.id, schema.team_id, updates={watermark_key: new_version} + ) + except Exception as e: + if is_transient_maintenance_error(e): + await self._logger.awarning(f"Delta maintenance skipped: transient infra error: {e}") + return + capture_exception(e) + await self._logger.aexception(f"Delta maintenance failed: {e}", exc_info=e) diff --git a/products/warehouse_sources/backend/temporal/data_imports/pipelines/core/delta/ops.py b/products/warehouse_sources/backend/temporal/data_imports/pipelines/core/delta/ops.py new file mode 100644 index 000000000000..30f3950fca18 --- /dev/null +++ b/products/warehouse_sources/backend/temporal/data_imports/pipelines/core/delta/ops.py @@ -0,0 +1,40 @@ +import asyncio +from collections.abc import Callable + +import deltalake +import deltalake.exceptions +from structlog.types import FilteringBoundLogger + +# Delta's conflict checker raises CommitFailedError the moment a concurrent commit invalidates what +# a committing operation read — a merge predicate, or optimize.compact's file-rewrite plan — unlike a +# plain version-bump race, delta-rs does not consume max_commit_retries or retry this itself (see +# delta-rs kernel/transaction/conflict_checker.rs), because resolving it safely requires re-reading +# the table and re-running the operation, which is exactly what its "must be rerun" error message +# asks the caller to do. +DELTA_MERGE_CONFLICT_RETRIES = 3 + + +async def execute_with_conflict_retry( + table: deltalake.DeltaTable, + operation_fn: Callable[[], dict], + operation_name: str, + logger: FilteringBoundLogger, +) -> dict: + """Run a Delta operation that commits (merge, optimize.compact, ...), refreshing the table + and re-running it on a commit conflict. + + See DELTA_MERGE_CONFLICT_RETRIES for why this can't rely on delta-rs's own retry budget. + """ + attempt = 0 + while True: + try: + return await asyncio.to_thread(operation_fn) + except deltalake.exceptions.CommitFailedError: + if attempt >= DELTA_MERGE_CONFLICT_RETRIES: + raise + attempt += 1 + await logger.awarning( + f"{operation_name}: commit conflict, retrying with refreshed table " + f"(attempt {attempt}/{DELTA_MERGE_CONFLICT_RETRIES})" + ) + await asyncio.to_thread(table.update_incremental) diff --git a/products/warehouse_sources/backend/temporal/data_imports/pipelines/core/delta/test/test_delta_errors.py b/products/warehouse_sources/backend/temporal/data_imports/pipelines/core/delta/test/test_delta_errors.py new file mode 100644 index 000000000000..161c6307c5b9 --- /dev/null +++ b/products/warehouse_sources/backend/temporal/data_imports/pipelines/core/delta/test/test_delta_errors.py @@ -0,0 +1,86 @@ +from django.db import InterfaceError, OperationalError + +import deltalake +import botocore.exceptions +from parameterized import parameterized + +from products.warehouse_sources.backend.temporal.data_imports.pipelines.core.delta.errors import ( + is_transient_delta_maintenance_error, + is_transient_maintenance_error, + is_transient_object_store_error, +) + + +class TestIsTransientObjectStoreError: + @parameterized.expand( + [ + ( + "credential_provider_not_enabled_os_error", + OSError( + "Operation not supported: the credential provider was not enabled: " + "no providers in chain provided credentials" + ), + True, + ), + ("unrelated_os_error", OSError("Permission denied: bucket policy forbids this operation"), False), + ( + # s3fs/aiobotocore's own credential resolution (distinct from delta-rs's Rust + # object_store crate) can raise this bare, unwrapped — same IMDS/STS blip + # hitting our own instance-role-authenticated bucket, different client library. + "bare_no_credentials_error", + botocore.exceptions.NoCredentialsError(), + True, + ), + ("unrelated_exception_type", ValueError("some other unrelated failure"), False), + ] + ) + def test_classifies_transient_errors(self, _name: str, error: Exception, expected: bool): + assert is_transient_object_store_error(error) is expected + + +class TestIsTransientDeltaMaintenanceError: + @parameterized.expand( + [ + # A concurrent optimize/vacuum losing the race on a file scan: safe to skip and retry. + ( + "optimize_scan_file_not_found", + deltalake.exceptions.DeltaError( + "Failed to parse parquet: Optimize selected-file scan failed while scanning data: " + "Object at location .../part-0.parquet not found: 404 Not Found" + ), + True, + ), + # Other DeltaErrors are real failures (e.g. a genuinely corrupt log) and must still be captured. + ("unrelated_delta_error", deltalake.exceptions.DeltaError("no protocol found in delta log"), False), + # Same message shape but not the DeltaError type delta-rs actually raises for it. + ("wrong_exception_type", RuntimeError("Optimize selected-file scan failed"), False), + ] + ) + def test_matches_only_the_racy_optimize_scan_signature(self, _name: str, error: Exception, expected: bool): + assert is_transient_delta_maintenance_error(error) is expected + + +class TestIsTransientMaintenanceError: + @parameterized.expand( + [ + # A DNS/pooler blip hit while resolving `job.folder_path()` on a pooled app-DB + # connection (e.g. inside `_get_delta_table_uri`) isn't a maintenance bug. + ("dns_resolution_failure", OperationalError("[Errno -2] Name or service not known"), True), + ("pooler_dropped_connection", InterfaceError("connection already closed"), True), + # The object-store and concurrent-maintenance classifiers must stay folded in — vacuum() + # and optimize.compact() surface both shapes through this single maintenance check. + ( + "object_store_blip", + OSError("Generic S3 error: Error getting list response body: operation timed out"), + True, + ), + ( + "concurrent_maintenance_race", + deltalake.exceptions.DeltaError("Optimize selected-file scan failed while scanning data"), + True, + ), + ("genuine_bug", RuntimeError("maintenance blew up"), False), + ] + ) + def test_classifies_transient_errors(self, _name: str, error: Exception, expected: bool): + assert is_transient_maintenance_error(error) is expected diff --git a/products/warehouse_sources/backend/temporal/data_imports/pipelines/core/delta/test/test_maintenance.py b/products/warehouse_sources/backend/temporal/data_imports/pipelines/core/delta/test/test_maintenance.py new file mode 100644 index 000000000000..a97cbbb19daf --- /dev/null +++ b/products/warehouse_sources/backend/temporal/data_imports/pipelines/core/delta/test/test_maintenance.py @@ -0,0 +1,372 @@ +import pytest +from unittest.mock import AsyncMock, MagicMock, patch + +import deltalake +from parameterized import parameterized + +from products.warehouse_sources.backend.temporal.data_imports.pipelines.core.delta.maintenance import DeltaMaintenance + +_MAINTENANCE_MODULE = "products.warehouse_sources.backend.temporal.data_imports.pipelines.core.delta.maintenance" + + +def _make_logger(): + return MagicMock(adebug=AsyncMock(), ainfo=AsyncMock(), awarning=AsyncMock(), aexception=AsyncMock()) + + +def _make_maintenance(delta_table: MagicMock | None) -> DeltaMaintenance: + table_ref = MagicMock() + table_ref.logger = _make_logger() + table_ref.get_delta_table = AsyncMock(return_value=delta_table) + return DeltaMaintenance(table_ref) + + +def _passthrough_pool(fn): + async def _call(*args, **kwargs): + return fn(*args, **kwargs) + + return _call + + +class TestCompactIfFragmented: + """Defensive compaction fires on files-per-partition OR total-files threshold.""" + + @pytest.mark.asyncio + async def test_skips_when_no_delta_table(self): + ran = await _make_maintenance(None).compact_if_fragmented(partition_count=10) + assert ran is False + + # (case_name, file_count, partition_count, threshold_kw, expected_ran) + # threshold_kw=None means "use default threshold" — exercises the prod path. + _THRESHOLD_CASES: list[tuple[str, int, int | None, int | None, bool]] = [ + # 100 / 10 = 10 fpp, well below default 200 -> skip + ("below_default_threshold", 100, 10, None, False), + # 5,000 / 10 = 500 fpp, well above default 200 -> fire + ("above_default_threshold", 5_000, 10, None, True), + # partition_count=None on an unpartitioned layout derives 1 partition; 250 fpp >> 200 -> fire + ("unpartitioned_above_default", 250, None, None, True), + # Custom threshold: 100 / 10 = 10 fpp, threshold=5 -> fire + ("custom_threshold_fires", 100, 10, 5, True), + # Boundary: exactly at threshold -> `>` not `>=`, so skip + ("exactly_at_default_threshold", 2_000, 10, None, False), + # Total-files backstop: 6,000 / 100 = 60 fpp (under the per-partition bar) but + # total 6,000 > 5,000 default total threshold -> fire. Guards high-partition tables. + ("total_cap_fires_under_per_partition", 6_000, 100, None, True), + # Under both bars: 4,000 / 100 = 40 fpp and total 4,000 < 5,000 -> skip. + ("below_both_thresholds", 4_000, 100, None, False), + ] + + @parameterized.expand(_THRESHOLD_CASES) + @pytest.mark.asyncio + async def test_threshold( + self, + _name: str, + file_count: int, + partition_count: int | None, + threshold_kw: int | None, + expected_ran: bool, + ): + 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) + maintenance = _make_maintenance(mock_delta) + with ( + patch.object(maintenance, "_compact", AsyncMock()) as mock_compact, + patch.object(maintenance, "_vacuum", AsyncMock()) as mock_vacuum, + ): + kwargs: dict = {"partition_count": partition_count} + if threshold_kw is not None: + kwargs["threshold"] = threshold_kw + ran = await maintenance.compact_if_fragmented(**kwargs) + + assert ran is expected_ran + if expected_ran: + 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]] = [ + # 300 files / 3 derived partitions = 100 fpp < 200 and total < 5,000 -> skip. + # Before derivation, None meant 1 partition (300 fpp) and this healthy table + # compacted on every run. + ("healthy_partitioned_table_skips", 100, 3, False), + # Genuinely fragmented per partition: 750/3 = 250 fpp > 200 -> still fires. + ("fragmented_partitioned_table_fires", 250, 3, True), + ] + + @parameterized.expand(_DERIVATION_CASES) + @pytest.mark.asyncio + async def test_partition_count_derived_from_layout( + self, _name: str, files_per_dir: int, dir_count: int, expected_ran: bool + ): + # Only md5 partitioning persists a partition_count; datetime/numerical schemas pass + # None. The count must come from the layout or every >200-file partitioned table + # would defensively compact at the start of every sync run. + file_uris = [ + f"s3://bucket/table/_ph_partition_key={d}/f{i}.parquet" + 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) + maintenance = _make_maintenance(mock_delta) + with ( + patch.object(maintenance, "_compact", AsyncMock()) as mock_compact, + patch.object(maintenance, "_vacuum", AsyncMock()) as mock_vacuum, + ): + ran = await maintenance.compact_if_fragmented(partition_count=None) + + assert ran is expected_ran + if expected_ran: + 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): + # 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=[]) + maintenance = _make_maintenance(mock_delta) + + await maintenance.compact_table() + + maintenance._table.get_delta_table.assert_called_once() + mock_delta.optimize.compact.assert_called_once() + mock_delta.vacuum.assert_called_once() + + @pytest.mark.asyncio + async def test_retries_compact_on_commit_conflict_then_succeeds(self): + # compact_table's optimize.compact() commits a REMOVE+ADD when rewriting fragmented files — + # the same commit-conflict shape as a merge (see test_ops.TestExecuteWithConflictRetry). + # Regression coverage for a CommitFailedError propagating straight out of compact_table on + # the first conflict instead of retrying with a refreshed table, like the write merges do. + mock_delta = MagicMock() + mock_delta.optimize.compact = MagicMock( + side_effect=[ + deltalake.exceptions.CommitFailedError( + "Commit failed: a concurrent transaction deleted data this operation read." + ), + {"numFilesAdded": 1}, + ] + ) + mock_delta.vacuum = MagicMock(return_value=[]) + + await _make_maintenance(mock_delta).compact_table() + + assert mock_delta.optimize.compact.call_count == 2 + mock_delta.update_incremental.assert_called_once() + + +class TestVacuumIfStale: + @parameterized.expand( + [ + # (last_vacuum_version, expect_vacuum, expected_return) — current version=150, threshold=100. + # First encounter must seed the watermark WITHOUT vacuuming (else every existing table vacuums + # at once on deploy); below threshold must skip (else vacuum runs every sync); at/above threshold + # must vacuum (else tombstones accumulate forever on tables that never reach post-load compaction). + ("first_encounter_seeds_no_vacuum", None, False, 150), + ("below_threshold_skips", 100, False, None), + ("at_threshold_vacuums", 50, True, 150), + ("above_threshold_vacuums", 40, True, 150), + # A watermark above the current version means the table was reset/recreated (delta + # versions are monotonic within one incarnation) and no reset path clears the persisted + # watermark — it must reseed, not block the cadence until the version catches up. + ("stale_watermark_from_recreated_table_reseeds", 999, False, 150), + ] + ) + @pytest.mark.asyncio + async def test_vacuum_cadence( + self, _name: str, last_version: int | None, expect_vacuum: bool, expected_return: int | None + ): + table = MagicMock() + table.version = MagicMock(return_value=150) + maintenance = _make_maintenance(table) + with ( + patch.object(maintenance, "_vacuum", new=AsyncMock()) as vacuum, + patch(f"{_MAINTENANCE_MODULE}.posthoganalytics") as ph, + ): + result = await maintenance.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: + assert ph.capture.call_args.kwargs["event"] == "warehouse_delta_vacuumed" + + +class TestRunMaintenance: + """run_maintenance is the single threshold-maintenance step: compaction supersedes the cadence vacuum.""" + + @pytest.mark.asyncio + async def test_compaction_supersedes_vacuum_and_advances_watermark(self): + # Fragmented table: compact runs (and vacuums as part of it), so the cadence vacuum is skipped — + # no double vacuum in one run — and the watermark advances to the post-compaction version. + table = MagicMock(version=MagicMock(return_value=200)) + maintenance = _make_maintenance(table) + with ( + patch.object(maintenance, "compact_if_fragmented", new=AsyncMock(return_value=True)), + patch.object(maintenance, "vacuum_if_stale", new=AsyncMock()) as vacuum_if_stale, + ): + result = await maintenance.run_maintenance(partition_count=10, last_vacuum_version=50, commit_threshold=100) + + assert result == 200 + vacuum_if_stale.assert_not_awaited() + + @pytest.mark.asyncio + async def test_falls_through_to_vacuum_when_not_fragmented(self): + # Not fragmented → no compaction; fall through to the commit-cadence vacuum and return its watermark. + maintenance = _make_maintenance(MagicMock()) + with ( + patch.object(maintenance, "compact_if_fragmented", new=AsyncMock(return_value=False)), + patch.object(maintenance, "vacuum_if_stale", new=AsyncMock(return_value=150)) as vacuum_if_stale, + ): + result = await maintenance.run_maintenance(partition_count=10, last_vacuum_version=40, commit_threshold=100) + + assert result == 150 + vacuum_if_stale.assert_awaited_once_with(40, 100) + + +class TestRunScheduled: + """run_scheduled owns the vacuum-watermark lifecycle for both call sites (pre-write defensive + pass and CDC post-load), so watermark-key selection, persistence gating, and the never-raise + contract all live here.""" + + def _schema(self) -> MagicMock: + schema = MagicMock() + schema.partition_count = 10 + schema.last_vacuum_version = 41 + schema.last_vacuum_version_cdc = 7 + return schema + + async def _run( + self, + maintenance: DeltaMaintenance, + schema: MagicMock, + *, + run_maintenance_result: int | None | Exception = None, + is_cdc_companion: bool = False, + partition_count_fallback: int | None = None, + ) -> tuple[AsyncMock, MagicMock, MagicMock]: + run_maintenance = ( + AsyncMock(side_effect=run_maintenance_result) + if isinstance(run_maintenance_result, Exception) + else AsyncMock(return_value=run_maintenance_result) + ) + with ( + patch.object(maintenance, "run_maintenance", run_maintenance), + patch(f"{_MAINTENANCE_MODULE}.database_sync_to_async_pool", _passthrough_pool), + patch(f"{_MAINTENANCE_MODULE}.update_sync_type_config_keys") as update_config, + patch(f"{_MAINTENANCE_MODULE}.capture_exception") as capture, + ): + await maintenance.run_scheduled( + schema, is_cdc_companion=is_cdc_companion, partition_count_fallback=partition_count_fallback + ) + return run_maintenance, update_config, capture + + @parameterized.expand( + [ + # (name, is_cdc_companion, schema_partition_count, fallback, expected_count, expected_last, expected_key) + # The schema's persisted count wins over the source's fallback. + ("main_schema_count_wins", False, 10, 72, 10, 41, "last_vacuum_version"), + # md5-less schemas persist no count; the source-provided fallback applies. + ("main_falls_back_to_source_count", False, None, 72, 72, 41, "last_vacuum_version"), + ("main_both_none_derives_downstream", False, None, None, None, 41, "last_vacuum_version"), + # The snapshot and _cdc companion are different delta tables with unrelated versions, so + # the companion must use last_vacuum_version_cdc — sharing a key corrupts both cadences — + # and must ignore schema.partition_count, which describes the snapshot table's layout. + ("companion_own_key_and_layout", True, 10, 72, None, 7, "last_vacuum_version_cdc"), + ] + ) + @pytest.mark.asyncio + async def test_partition_count_and_watermark_key_selection( + self, + _name: str, + is_cdc_companion: bool, + schema_count: int | None, + fallback: int | None, + expected_count: int | None, + expected_last: int, + expected_key: str, + ): + schema = self._schema() + schema.partition_count = schema_count + run_maintenance, update_config, _ = await self._run( + _make_maintenance(MagicMock()), + schema, + run_maintenance_result=99, + is_cdc_companion=is_cdc_companion, + partition_count_fallback=fallback, + ) + + assert run_maintenance.await_args is not None + assert run_maintenance.await_args.kwargs["partition_count"] == expected_count + assert run_maintenance.await_args.kwargs["last_vacuum_version"] == expected_last + update_config.assert_called_once_with(schema.id, schema.team_id, updates={expected_key: 99}) + + @parameterized.expand( + [ + # run_maintenance returning a version must persist it — a lost watermark means + # vacuum_if_stale re-seeds forever and the table never vacuums. + ("new_version_persists", 55, True), + ("no_change_skips_write", None, False), + ("same_version_skips_write", 41, False), + ] + ) + @pytest.mark.asyncio + async def test_watermark_persistence_gating(self, _name: str, returned_version: int | None, expect_write: bool): + schema = self._schema() + _, update_config, _ = await self._run( + _make_maintenance(MagicMock()), schema, run_maintenance_result=returned_version + ) + + if expect_write: + update_config.assert_called_once_with( + schema.id, schema.team_id, updates={"last_vacuum_version": returned_version} + ) + else: + update_config.assert_not_called() + + @parameterized.expand( + [ + # A genuine maintenance bug must be captured for visibility but never raise — the sync + # itself must proceed either way. The full transient-vs-genuine classification matrix + # lives in test_errors.py; this covers the two handling paths. + ("genuine_bug_captured", RuntimeError("maintenance blew up"), True), + # A transient infra blip self-heals on the next scheduled pass and must not be promoted + # into a fresh error-tracking issue. + ("transient_blip_warned_only", OSError("Generic S3 error: Please reduce your request rate."), False), + ] + ) + @pytest.mark.asyncio + async def test_never_raises(self, _name: str, error: Exception, expect_capture: bool): + logger = _make_logger() + table_ref = MagicMock() + table_ref.logger = logger + table_ref.get_delta_table = AsyncMock(return_value=MagicMock()) + maintenance = DeltaMaintenance(table_ref) + + _, update_config, capture = await self._run(maintenance, self._schema(), run_maintenance_result=error) + + assert capture.called is expect_capture + update_config.assert_not_called() + if expect_capture: + logger.aexception.assert_awaited_once() + else: + logger.awarning.assert_awaited_once() diff --git a/products/warehouse_sources/backend/temporal/data_imports/pipelines/core/delta/test/test_ops.py b/products/warehouse_sources/backend/temporal/data_imports/pipelines/core/delta/test/test_ops.py new file mode 100644 index 000000000000..3aff814569ae --- /dev/null +++ b/products/warehouse_sources/backend/temporal/data_imports/pipelines/core/delta/test/test_ops.py @@ -0,0 +1,74 @@ +import pytest +from unittest.mock import AsyncMock, MagicMock + +import deltalake + +from products.warehouse_sources.backend.temporal.data_imports.pipelines.core.delta.ops import ( + DELTA_MERGE_CONFLICT_RETRIES, + execute_with_conflict_retry, +) + + +def _make_logger(): + return MagicMock(adebug=AsyncMock(), ainfo=AsyncMock(), awarning=AsyncMock(), aerror=AsyncMock()) + + +class TestExecuteWithConflictRetry: + """A committing operation's CommitFailedError means delta-rs's conflict checker rejected the + commit outright, without spending any of its own internal retry budget (see the comment on + DELTA_MERGE_CONFLICT_RETRIES). Regression coverage for the sync dying on the first such + conflict instead of refreshing the table and re-running the operation, as the error's own + "must be rerun" message calls for. Shared by merges and `compact_table`'s optimize.compact.""" + + @pytest.mark.asyncio + async def test_succeeds_without_retry(self): + table = MagicMock() + operation_fn = MagicMock(return_value={"num_output_rows": 1}) + + result = await execute_with_conflict_retry(table, operation_fn, "op", _make_logger()) + + assert result == {"num_output_rows": 1} + operation_fn.assert_called_once() + table.update_incremental.assert_not_called() + + @pytest.mark.asyncio + async def test_retries_on_conflict_then_succeeds(self): + table = MagicMock() + operation_fn = MagicMock( + side_effect=[ + deltalake.exceptions.CommitFailedError("Commit failed: a concurrent transactions added new data."), + {"num_output_rows": 1}, + ] + ) + + result = await execute_with_conflict_retry(table, operation_fn, "op", _make_logger()) + + assert result == {"num_output_rows": 1} + assert operation_fn.call_count == 2 + table.update_incremental.assert_called_once() + + @pytest.mark.asyncio + async def test_gives_up_after_exhausting_retries(self): + table = MagicMock() + operation_fn = MagicMock( + side_effect=deltalake.exceptions.CommitFailedError( + "Commit failed: a concurrent transactions added new data." + ) + ) + + with pytest.raises(deltalake.exceptions.CommitFailedError): + await execute_with_conflict_retry(table, operation_fn, "op", _make_logger()) + + assert operation_fn.call_count == DELTA_MERGE_CONFLICT_RETRIES + 1 + assert table.update_incremental.call_count == DELTA_MERGE_CONFLICT_RETRIES + + @pytest.mark.asyncio + async def test_other_errors_propagate_without_retry(self): + table = MagicMock() + operation_fn = MagicMock(side_effect=ValueError("not a commit conflict")) + + with pytest.raises(ValueError): + await execute_with_conflict_retry(table, operation_fn, "op", _make_logger()) + + operation_fn.assert_called_once() + table.update_incremental.assert_not_called() 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 ea8bd5a14371..309755fa7317 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 @@ -5,21 +5,17 @@ from typing import Any, Literal from django.conf import settings -from django.db import InterfaceError, OperationalError import numpy as np import pyarrow as pa import deltalake as deltalake import pyarrow.compute as pc -import posthoganalytics -import botocore.exceptions import deltalake.exceptions from dlt.common.libs.deltalake import ensure_delta_compatible_arrow_schema from structlog.types import FilteringBoundLogger from posthog.exceptions_capture import capture_exception from posthog.sync import database_sync_to_async_pool -from posthog.utils import get_machine_id from products.data_warehouse.backend.facade.api import aget_s3_client, ensure_bucket_exists from products.warehouse_sources.backend.models.external_data_job import ExternalDataJob @@ -31,108 +27,18 @@ pyarrow_schema_from_arrow_exportable, ) from products.warehouse_sources.backend.temporal.data_imports.pipelines.core.consts import PARTITION_KEY - -# A pre-write defensive compact fires when EITHER threshold is exceeded. -# -# Calibrated against the production file-count distribution (delta merge stats across -# all teams): total files per table sit at p50≈60, p90≈470, p95≈850, with a long tail -# (p99≈12k → ~14s merges; an observed pathological case hit ~82k files → ~45s merges). -# Merge planning time tracks TOTAL files, not files-per-partition — delta still -# enumerates every file's metadata even when partition pruning skips reading them — so -# we gate on both: -# -# - files-per-partition: bounds per-partition fragmentation and rescues partitioned -# (esp. md5) tables, where a merge touches every partition. 200 sits well above the -# healthy steady state (compaction runs at the end of each successful sync) yet -# triggers long before a table reaches the slow tail. -# - total files: a partition-count-independent backstop so a table with a high -# partition_count can't accumulate tens of thousands of files (each adding to merge -# planning time) while staying under the per-partition bar. 5,000 is above p95 (~850) -# — so healthy tables never trip it — and well below the p99/pathological tail. -# -# Tune further once the admin fragmentation view gives per-customer distributions. -DEFAULT_COMPACT_FILES_PER_PARTITION_THRESHOLD = 200 -DEFAULT_COMPACT_TOTAL_FILES_THRESHOLD = 5000 - -# Substrings of the object-store errors raised talking to our own S3-backed data-warehouse bucket -# that are transient and self-recovering, not a bug in our code or a customer credential problem: -# - the first three come from delta-rs's Rust `object_store` crate inside `DeltaTable.is_deltatable()` -# (IMDS/STS blips, dispatch timeouts) -# - "Please reduce your request rate" is S3's SlowDown throttling response, surfaced by s3fs/aiobotocore -# when a bulk operation (e.g. `_purge_s3_prefix`'s list-then-delete) outruns the bucket's request-rate limit -# A retry (of the same idempotent operation) clears these, so they shouldn't be treated the same as a -# bug in our logic. -TRANSIENT_OBJECT_STORE_ERRORS = ( - "an error occurred while loading credentials", - "the credential provider was not enabled", - "Generic S3 error", - "Please reduce your request rate", +from products.warehouse_sources.backend.temporal.data_imports.pipelines.core.delta.errors import ( + is_transient_object_store_error, +) +from products.warehouse_sources.backend.temporal.data_imports.pipelines.core.delta.ops import ( + execute_with_conflict_retry, ) - - -def is_transient_object_store_error(error: BaseException) -> bool: - """True for a transient object-store error, however it happened to surface. - - `DeltaTable.is_deltatable()` raises these as a plain `OSError`, but table-level operations - (e.g. `vacuum()`, `optimize.compact()`) wrap the identical underlying object-store error text in - `deltalake.exceptions.DeltaError` instead — same blip, different exception type depending on - which delta-rs entry point hit it. `_purge_s3_prefix`'s s3fs/aiobotocore calls can also raise a - bare `botocore.exceptions.NoCredentialsError` unwrapped — the same IMDS/STS credential-provider - blip, just surfaced by aiobotocore's own credential resolution instead of delta-rs's Rust - `object_store` crate. `NoCredentialsError`'s message is a fixed, generic string (no needle to - match), but hitting our own instance-role-authenticated bucket always means the same transient - resolution hiccup, so it's recognized by type rather than by message. - """ - if isinstance(error, botocore.exceptions.NoCredentialsError): - return True - return isinstance(error, OSError | deltalake.exceptions.DeltaError) and any( - needle in str(error) for needle in TRANSIENT_OBJECT_STORE_ERRORS - ) - - -# Delta's conflict checker raises CommitFailedError the moment a concurrent commit invalidates what -# a committing operation read — a merge predicate, or optimize.compact's file-rewrite plan — unlike a -# plain version-bump race, delta-rs does not consume max_commit_retries or retry this itself (see -# delta-rs kernel/transaction/conflict_checker.rs), because resolving it safely requires re-reading -# the table and re-running the operation, which is exactly what its "must be rerun" error message -# asks the caller to do. -DELTA_MERGE_CONFLICT_RETRIES = 3 - -# `optimize.compact` plans its rewrite against the file list at the start of its scan, then reads -# those files. A concurrent maintenance pass on the same table (e.g. a Temporal activity attempt that -# heartbeat-timed-out but keeps running as a zombie — see this package's README on the equivalent -# unfenced race for repartition) can vacuum one of those files out from under the scan before it gets -# read, which delta-rs surfaces as this DeltaError. The scan failing here means the optimize aborted -# before committing anything — the table is left exactly as it was, just still fragmented — so this is -# safe to skip and retry on the next maintenance pass, not a bug in our logic. -TRANSIENT_DELTA_MAINTENANCE_ERRORS = ("Optimize selected-file scan failed",) - - -def is_transient_delta_maintenance_error(error: BaseException) -> bool: - return isinstance(error, deltalake.exceptions.DeltaError) and any( - needle in str(error) for needle in TRANSIENT_DELTA_MAINTENANCE_ERRORS - ) - # _purge_s3_prefix is idempotent (every step is existence-gated), so retrying it whole after a brief # backoff is as safe as retrying a single failed call, and simpler. _PURGE_S3_PREFIX_MAX_ATTEMPTS = 4 -def is_transient_maintenance_error(error: BaseException) -> bool: - """Infra blips seen during pre-write maintenance that aren't a maintenance bug. - - Covers S3/object-store hiccups reaching our own data-warehouse bucket (see - `is_transient_object_store_error` above), racy concurrent-maintenance DeltaErrors (see - `is_transient_delta_maintenance_error` above), and app-DB connection blips (DNS, pooler drops) hit - while resolving `job.folder_path()` on a pooled connection — the same `OperationalError`/`InterfaceError` - classification used for this failure class in `repartition_table.py`'s `_is_transient_infra_error`. - """ - if isinstance(error, OperationalError | InterfaceError): - return True - return is_transient_object_store_error(error) or is_transient_delta_maintenance_error(error) - - def _delta_merge_spill_kwargs() -> dict[str, int]: """delta-rs `merge` kwargs that let DataFusion spill to disk instead of OOMing on large merges. @@ -361,6 +267,18 @@ def __init__( def is_first_sync(self) -> bool: return self._is_first_sync + @property + def job(self) -> ExternalDataJob: + return self._job + + @property + def resource_name(self) -> str: + return self._resource_name + + @property + def logger(self) -> FilteringBoundLogger: + return self._logger + def _get_credentials(self): return delta_storage_options() @@ -501,28 +419,6 @@ async def get_file_uris(self) -> list[str]: return await asyncio.to_thread(delta_table.file_uris) - async def _execute_with_conflict_retry( - self, table: deltalake.DeltaTable, operation_fn: Callable[[], dict], operation_name: str - ) -> dict: - """Run a Delta operation that commits (merge, optimize.compact, ...), refreshing the table - and re-running it on a commit conflict. - - See DELTA_MERGE_CONFLICT_RETRIES for why this can't rely on delta-rs's own retry budget. - """ - attempt = 0 - while True: - try: - return await asyncio.to_thread(operation_fn) - except deltalake.exceptions.CommitFailedError: - if attempt >= DELTA_MERGE_CONFLICT_RETRIES: - raise - attempt += 1 - await self._logger.awarning( - f"{operation_name}: commit conflict, retrying with refreshed table " - f"(attempt {attempt}/{DELTA_MERGE_CONFLICT_RETRIES})" - ) - await asyncio.to_thread(table.update_incremental) - async def _dedupe_incremental_batch( self, data: pa.Table, primary_keys: Sequence[Any], use_partitioning: bool ) -> pa.Table: @@ -762,8 +658,8 @@ def _do_merge( .execute() ) - merge_stats = await self._execute_with_conflict_retry( - existing_delta_table, _do_merge, "write_to_deltalake: merge" + merge_stats = await execute_with_conflict_retry( + existing_delta_table, _do_merge, "write_to_deltalake: merge", self._logger ) await self._logger.adebug(f"Delta Merge Stats: {json.dumps(merge_stats)}") @@ -788,10 +684,11 @@ def _do_merge_unpartitioned(data: pa.Table, predicate_ops: list[str]): .execute() ) - merge_stats = await self._execute_with_conflict_retry( + merge_stats = await execute_with_conflict_retry( existing_delta_table, lambda: _do_merge_unpartitioned(data, predicate_ops), "write_to_deltalake: merge", + self._logger, ) await self._logger.adebug(f"Delta Merge Stats: {json.dumps(merge_stats)}") elif ( @@ -945,10 +842,11 @@ def _do_scd2_close(first_per_pk: pa.Table, predicate: str) -> dict: .execute() ) - close_stats = await self._execute_with_conflict_retry( + close_stats = await execute_with_conflict_retry( existing_delta_table, lambda: _do_scd2_close(first_per_pk, predicate), "write_scd2_to_deltalake: close merge", + self._logger, ) await self._logger.adebug(f"SCD2 close stats: {json.dumps(close_stats)}") @@ -1034,173 +932,3 @@ async def has_batch_been_committed(self, run_uuid: str, batch_index: int) -> boo the metadata schema used for idempotency tagging. """ return await self.has_commit_with_metadata({"run_uuid": run_uuid, "batch_index": str(batch_index)}) - - 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(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)) - - 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: - """Vacuum tombstoned files once enough commits have accrued since the last vacuum. - - Decoupled from merge success (called pre-write) so a table that OOMs its merge every run still - gets cleaned — the post-load compaction never runs for it, which is how tables reach ~99% dead - files. Vacuum only deletes dead files (an S3 LIST + delete), so unlike `compact_table`'s - `optimize.compact` (which rewrites partitions) it is memory-safe even on an oversized table. - - Uses the delta version (commit count) as a cheap proxy for tombstone accumulation — no S3 LIST to - decide. Returns the current version to persist as the new watermark when it vacuumed, on first - encounter, or when the table was recreated (both reseed the watermark without vacuuming); - None when nothing changed. - """ - table = await self.get_delta_table() - if table is None: - return None - - version = await asyncio.to_thread(table.version) - if last_vacuum_version is None or version < last_vacuum_version: - # First encounter: seed the watermark without vacuuming so existing tables clean up gradually - # over the next `commit_threshold` commits rather than all vacuuming at once on deploy. - # A version below the watermark means the table was reset/recreated (delta versions are - # monotonic within one incarnation) and no reset path clears the persisted watermark — - # left alone it would block the cadence until the new table out-versioned the old one. - return version - - commits_since = version - last_vacuum_version - if commits_since < commit_threshold: - await self._logger.adebug( - f"vacuum_if_stale: skipping, {commits_since} commits since last vacuum (< {commit_threshold})" - ) - return None - - await self._logger.ainfo( - f"vacuum_if_stale: {commits_since} commits since last vacuum (>= {commit_threshold}), vacuuming" - ) - 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. - posthoganalytics.capture( - distinct_id=get_machine_id(), - event="warehouse_delta_vacuumed", - properties={ - "team_id": self._job.team_id, - "schema_id": str(self._job.schema_id), - "source_id": str(self._job.pipeline_id), - "resource_name": self._resource_name, - "commits_since_last_vacuum": commits_since, - "delta_version": version, - }, - ) - except Exception as e: - capture_exception(e) - return version - - async def compact_if_fragmented( - self, - partition_count: int | None, - threshold: int = DEFAULT_COMPACT_FILES_PER_PARTITION_THRESHOLD, - total_threshold: int = DEFAULT_COMPACT_TOTAL_FILES_THRESHOLD, - ) -> bool: - """Run compact + vacuum if the table is fragmented past either threshold. - - Fragmented = files-per-partition > `threshold` OR total files > `total_threshold`. - The total-files backstop matters because delta enumerates every file's metadata - during a merge even when partition pruning skips reading them, so merge planning - time tracks total files — a high partition_count must not let a table accumulate - tens of thousands of files while staying under the per-partition bar. - - When `partition_count` is None it is derived from the table's actual layout (the - distinct file directories in the delta log, no extra I/O) — only md5 partitioning - persists a count on the schema, so datetime/numerical-partitioned tables always - arrive here with None. - - Returns True if compaction ran, False if it was skipped. Cheap when the table is - 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 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 - # table root. Without this, a partitioned table with no persisted count reads as - # one giant partition and trips the per-partition threshold on every run. - partition_count = len({uri.rsplit("/", 1)[0] for uri in file_uris}) - # Treat unpartitioned tables as one "partition" for the threshold math. - effective_partitions = max(partition_count or 1, 1) - files_per_partition = total_files / effective_partitions - - fragmented = files_per_partition > threshold or total_files > total_threshold - stats = ( - f"total_files={total_files}, partitions={effective_partitions}, " - f"files_per_partition={files_per_partition:.1f}, threshold={threshold}, " - f"total_threshold={total_threshold}" - ) - if not fragmented: - await self._logger.adebug(f"compact_if_fragmented: skipping ({stats})") - return False - - await self._logger.ainfo(f"compact_if_fragmented: triggering compact ({stats})") - await self._compact(table) - await self._vacuum(table) - return True - - async def run_maintenance( - self, - partition_count: int | None, - last_vacuum_version: int | None, - commit_threshold: int, - ) -> int | None: - """Single pre-write maintenance entry point: compact if fragmented, else vacuum on commit cadence. - - The two triggers are orthogonal — fragmentation (active file count) vs. commit cadence (tombstone - accrual) — but they share one outcome, the vacuum watermark. `compact_if_fragmented` already - vacuums as part of compaction, so when it runs it supersedes the cadence vacuum (no double vacuum - in one run) and the watermark advances to the post-compaction version. When nothing was fragmented, - fall through to `vacuum_if_stale`. Returns the single delta version to persist as the new - `last_vacuum_version` watermark, or None when nothing changed. Callers persist the watermark - via `update_sync_type_config_keys` (row-locked merge): the pre-write defensive path in - `common/extract.py` and the CDC post-load path in `common/load.py`. - """ - compacted = await self.compact_if_fragmented(partition_count=partition_count) - if compacted: - table = await self.get_delta_table() - if table is None: - return None - # Compaction (which vacuumed) added a commit, advancing the version; reset the cadence - # watermark to it so the next vacuum is measured from this cleanup, not the old baseline. - return await asyncio.to_thread(table.version) - return await self.vacuum_if_stale(last_vacuum_version, commit_threshold) 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 c456a396dc7b..879922227138 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 @@ -12,7 +12,6 @@ import pyarrow as pa import deltalake import pyarrow.compute as pc -import botocore.exceptions from parameterized import parameterized from products.warehouse_sources.backend.temporal.data_imports.pipelines.core.arrow_utils import ( @@ -20,15 +19,13 @@ evolve_pyarrow_schema, ) from products.warehouse_sources.backend.temporal.data_imports.pipelines.core.consts import PARTITION_KEY +from products.warehouse_sources.backend.temporal.data_imports.pipelines.core.delta.maintenance import DeltaMaintenance from products.warehouse_sources.backend.temporal.data_imports.pipelines.core.delta_table_helper import ( - DELTA_MERGE_CONFLICT_RETRIES, DeltaTableHelper, _delta_merge_spill_kwargs, _first_per_pk_table, _merge_predicate_ops, _realign_decimal_buffers, - is_transient_delta_maintenance_error, - is_transient_object_store_error, ) @@ -246,130 +243,6 @@ async def test_wraps_has_commit_with_metadata( m.assert_called_once_with({"run_uuid": run_uuid, "batch_index": str(batch_index)}) -class TestCompactIfFragmented: - """Pre-write defensive compaction fires on files-per-partition OR total-files threshold.""" - - @pytest.mark.asyncio - async def test_skips_when_no_delta_table(self, helper: DeltaTableHelper): - with patch.object(helper, "get_delta_table", AsyncMock(return_value=None)): - ran = await helper.compact_if_fragmented(partition_count=10) - assert ran is False - - # (case_name, file_count, partition_count, threshold_kw, expected_ran) - # threshold_kw=None means "use default threshold" — exercises the prod path. - _THRESHOLD_CASES: list[tuple[str, int, int | None, int | None, bool]] = [ - # 100 / 10 = 10 fpp, well below default 200 -> skip - ("below_default_threshold", 100, 10, None, False), - # 5,000 / 10 = 500 fpp, well above default 200 -> fire - ("above_default_threshold", 5_000, 10, None, True), - # partition_count=None on an unpartitioned layout derives 1 partition; 250 fpp >> 200 -> fire - ("unpartitioned_above_default", 250, None, None, True), - # Custom threshold: 100 / 10 = 10 fpp, threshold=5 -> fire - ("custom_threshold_fires", 100, 10, 5, True), - # Boundary: exactly at threshold -> `>` not `>=`, so skip - ("exactly_at_default_threshold", 2_000, 10, None, False), - # Total-files backstop: 6,000 / 100 = 60 fpp (under the per-partition bar) but - # total 6,000 > 5,000 default total threshold -> fire. Guards high-partition tables. - ("total_cap_fires_under_per_partition", 6_000, 100, None, True), - # Under both bars: 4,000 / 100 = 40 fpp and total 4,000 < 5,000 -> skip. - ("below_both_thresholds", 4_000, 100, None, False), - ] - - @parameterized.expand(_THRESHOLD_CASES) - @pytest.mark.asyncio - async def test_threshold( - self, - _name: str, - file_count: int, - partition_count: int | None, - threshold_kw: int | None, - expected_ran: bool, - ): - helper = DeltaTableHelper(resource_name="t", job=MagicMock(), logger=_make_logger()) - 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, "_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: - kwargs["threshold"] = threshold_kw - ran = await helper.compact_if_fragmented(**kwargs) - - assert ran is expected_ran - if expected_ran: - 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]] = [ - # 300 files / 3 derived partitions = 100 fpp < 200 and total < 5,000 -> skip. - # Before derivation, None meant 1 partition (300 fpp) and this healthy table - # compacted on every run. - ("healthy_partitioned_table_skips", 100, 3, False), - # Genuinely fragmented per partition: 750/3 = 250 fpp > 200 -> still fires. - ("fragmented_partitioned_table_fires", 250, 3, True), - ] - - @parameterized.expand(_DERIVATION_CASES) - @pytest.mark.asyncio - async def test_partition_count_derived_from_layout( - self, _name: str, files_per_dir: int, dir_count: int, expected_ran: bool - ): - # Only md5 partitioning persists a partition_count; datetime/numerical schemas pass - # None. The count must come from the layout or every >200-file partitioned table - # would defensively compact at the start of every sync run. - helper = DeltaTableHelper(resource_name="t", job=MagicMock(), logger=_make_logger()) - file_uris = [ - f"s3://bucket/table/_ph_partition_key={d}/f{i}.parquet" - 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=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_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: # (case_name, error_message, expect_heal) — heal = wipe the table and fall back to first-sync mode _ERROR_CASES: list[tuple[str, str, bool]] = [ @@ -511,100 +384,6 @@ async def test_full_refresh_passes_commit_properties( assert commit_properties.custom_metadata == expected_custom_metadata -class TestExecuteWithConflictRetry: - """A committing operation's CommitFailedError means delta-rs's conflict checker rejected the - commit outright, without spending any of its own internal retry budget (see the comment on - DELTA_MERGE_CONFLICT_RETRIES). Regression coverage for the sync dying on the first such - conflict instead of refreshing the table and re-running the operation, as the error's own - "must be rerun" message calls for. Shared by merges and `compact_table`'s optimize.compact.""" - - def _helper(self) -> DeltaTableHelper: - return DeltaTableHelper(resource_name="t", job=MagicMock(), logger=_make_logger()) - - @pytest.mark.asyncio - async def test_succeeds_without_retry(self): - helper = self._helper() - table = MagicMock() - operation_fn = MagicMock(return_value={"num_output_rows": 1}) - - result = await helper._execute_with_conflict_retry(table, operation_fn, "op") - - assert result == {"num_output_rows": 1} - operation_fn.assert_called_once() - table.update_incremental.assert_not_called() - - @pytest.mark.asyncio - async def test_retries_on_conflict_then_succeeds(self): - helper = self._helper() - table = MagicMock() - operation_fn = MagicMock( - side_effect=[ - deltalake.exceptions.CommitFailedError("Commit failed: a concurrent transactions added new data."), - {"num_output_rows": 1}, - ] - ) - - result = await helper._execute_with_conflict_retry(table, operation_fn, "op") - - assert result == {"num_output_rows": 1} - assert operation_fn.call_count == 2 - table.update_incremental.assert_called_once() - - @pytest.mark.asyncio - async def test_gives_up_after_exhausting_retries(self): - helper = self._helper() - table = MagicMock() - operation_fn = MagicMock( - side_effect=deltalake.exceptions.CommitFailedError( - "Commit failed: a concurrent transactions added new data." - ) - ) - - with pytest.raises(deltalake.exceptions.CommitFailedError): - await helper._execute_with_conflict_retry(table, operation_fn, "op") - - assert operation_fn.call_count == DELTA_MERGE_CONFLICT_RETRIES + 1 - assert table.update_incremental.call_count == DELTA_MERGE_CONFLICT_RETRIES - - @pytest.mark.asyncio - async def test_other_errors_propagate_without_retry(self): - helper = self._helper() - table = MagicMock() - operation_fn = MagicMock(side_effect=ValueError("not a commit conflict")) - - with pytest.raises(ValueError): - await helper._execute_with_conflict_retry(table, operation_fn, "op") - - operation_fn.assert_called_once() - table.update_incremental.assert_not_called() - - -class TestCompactTableConflictRetry: - """compact_table's optimize.compact() commits a REMOVE+ADD when rewriting fragmented files — - the same commit-conflict shape as a merge (see TestExecuteWithConflictRetry). Regression - coverage for a CommitFailedError propagating straight out of compact_table on the first - conflict instead of retrying with a refreshed table, like write_to_deltalake's merges do.""" - - @pytest.mark.asyncio - async def test_retries_compact_on_commit_conflict_then_succeeds(self, helper: DeltaTableHelper): - mock_delta = MagicMock() - mock_delta.optimize.compact = MagicMock( - side_effect=[ - deltalake.exceptions.CommitFailedError( - "Commit failed: a concurrent transaction deleted data this operation read." - ), - {"numFilesAdded": 1}, - ] - ) - mock_delta.vacuum = MagicMock(return_value=[]) - - with patch.object(helper, "get_delta_table", AsyncMock(return_value=mock_delta)): - await helper.compact_table() - - assert mock_delta.optimize.compact.call_count == 2 - mock_delta.update_incremental.assert_called_once() - - def _create_legacy_delta_table(path: str, *, partitioned: bool = False) -> deltalake.DeltaTable: """Seed a Delta table that mimics what the old dlt pipeline created: business columns plus NOT NULL _dlt_id and _dlt_load_id.""" @@ -850,7 +629,7 @@ async def test_compact_survives_a_column_added_by_an_all_non_null_batch(self, tm status_field = next(f for f in result.schema().fields if f.name == "status") assert status_field.nullable is True - await helper.compact_table() + await DeltaMaintenance(helper).compact_table() final = result.to_pyarrow_table() by_id = dict(zip(final.column("id").to_pylist(), final.column("status").to_pylist())) @@ -1135,87 +914,6 @@ async def test_write_scd2_misaligned_decimal_to_local_delta(self, tmp_path: Path assert closed.column("valid_to").to_pylist() == [ts2] -class TestVacuumIfStale: - def _helper(self) -> DeltaTableHelper: - return DeltaTableHelper("t", MagicMock(), MagicMock(adebug=AsyncMock(), ainfo=AsyncMock()), False) - - @parameterized.expand( - [ - # (last_vacuum_version, expect_vacuum, expected_return) — current version=150, threshold=100. - # First encounter must seed the watermark WITHOUT vacuuming (else every existing table vacuums - # at once on deploy); below threshold must skip (else vacuum runs every sync); at/above threshold - # must vacuum (else tombstones accumulate forever on tables that never reach post-load compaction). - ("first_encounter_seeds_no_vacuum", None, False, 150), - ("below_threshold_skips", 100, False, None), - ("at_threshold_vacuums", 50, True, 150), - ("above_threshold_vacuums", 40, True, 150), - # A watermark above the current version means the table was reset/recreated (delta - # versions are monotonic within one incarnation) and no reset path clears the persisted - # watermark — it must reseed, not block the cadence until the version catches up. - ("stale_watermark_from_recreated_table_reseeds", 999, False, 150), - ] - ) - @pytest.mark.asyncio - async def test_vacuum_cadence( - self, _name: str, last_version: int | None, expect_vacuum: bool, expected_return: int | None - ): - module = "products.warehouse_sources.backend.temporal.data_imports.pipelines.core.delta_table_helper" - helper = self._helper() - table = MagicMock() - table.version = MagicMock(return_value=150) - with ( - patch.object(helper, "get_delta_table", new=AsyncMock(return_value=table)), - 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: - assert ph.capture.call_args.kwargs["event"] == "warehouse_delta_vacuumed" - - -class TestRunMaintenance: - """run_maintenance is the single pre-write entry point: compaction supersedes the cadence vacuum.""" - - def _helper(self) -> DeltaTableHelper: - return DeltaTableHelper("t", MagicMock(), MagicMock(adebug=AsyncMock(), ainfo=AsyncMock()), False) - - @pytest.mark.asyncio - async def test_compaction_supersedes_vacuum_and_advances_watermark(self): - # Fragmented table: compact runs (and vacuums as part of it), so the cadence vacuum is skipped — - # no double vacuum in one run — and the watermark advances to the post-compaction version. - helper = self._helper() - table = MagicMock(version=MagicMock(return_value=200)) - with ( - patch.object(helper, "compact_if_fragmented", new=AsyncMock(return_value=True)), - patch.object(helper, "get_delta_table", new=AsyncMock(return_value=table)), - patch.object(helper, "vacuum_if_stale", new=AsyncMock()) as vacuum_if_stale, - ): - result = await helper.run_maintenance(partition_count=10, last_vacuum_version=50, commit_threshold=100) - - assert result == 200 - vacuum_if_stale.assert_not_awaited() - - @pytest.mark.asyncio - async def test_falls_through_to_vacuum_when_not_fragmented(self): - # Not fragmented → no compaction; fall through to the commit-cadence vacuum and return its watermark. - helper = self._helper() - with ( - patch.object(helper, "compact_if_fragmented", new=AsyncMock(return_value=False)), - patch.object(helper, "vacuum_if_stale", new=AsyncMock(return_value=150)) as vacuum_if_stale, - ): - result = await helper.run_maintenance(partition_count=10, last_vacuum_version=40, commit_threshold=100) - - assert result == 150 - vacuum_if_stale.assert_awaited_once_with(40, 100) - - class TestIsTableCorrupted: _MODULE = "products.warehouse_sources.backend.temporal.data_imports.pipelines.core.delta_table_helper" @@ -1250,55 +948,6 @@ async def test_is_table_corrupted(self, _name: str, is_delta: bool, open_exc: Ex assert result is expected -class TestIsTransientDeltaMaintenanceError: - @parameterized.expand( - [ - # A concurrent optimize/vacuum losing the race on a file scan: safe to skip and retry. - ( - "optimize_scan_file_not_found", - deltalake.exceptions.DeltaError( - "Failed to parse parquet: Optimize selected-file scan failed while scanning data: " - "Object at location .../part-0.parquet not found: 404 Not Found" - ), - True, - ), - # Other DeltaErrors are real failures (e.g. a genuinely corrupt log) and must still be captured. - ("unrelated_delta_error", deltalake.exceptions.DeltaError("no protocol found in delta log"), False), - # Same message shape but not the DeltaError type delta-rs actually raises for it. - ("wrong_exception_type", RuntimeError("Optimize selected-file scan failed"), False), - ] - ) - def test_matches_only_the_racy_optimize_scan_signature(self, _name: str, error: Exception, expected: bool): - assert is_transient_delta_maintenance_error(error) is expected - - -class TestIsTransientObjectStoreError: - @parameterized.expand( - [ - ( - "credential_provider_not_enabled_os_error", - OSError( - "Operation not supported: the credential provider was not enabled: " - "no providers in chain provided credentials" - ), - True, - ), - ("unrelated_os_error", OSError("Permission denied: bucket policy forbids this operation"), False), - ( - # s3fs/aiobotocore's own credential resolution (distinct from delta-rs's Rust - # object_store crate) can raise this bare, unwrapped — same IMDS/STS blip - # hitting our own instance-role-authenticated bucket, different client library. - "bare_no_credentials_error", - botocore.exceptions.NoCredentialsError(), - True, - ), - ("unrelated_exception_type", ValueError("some other unrelated failure"), False), - ] - ) - def test_classifies_transient_errors(self, _name: str, error: Exception, expected: bool): - assert is_transient_object_store_error(error) is expected - - class TestNullSafeMergePredicate: """The incremental-merge match must be NULL-safe. diff --git a/products/warehouse_sources/backend/temporal/data_imports/pipelines/pipeline_v2/pipeline.py b/products/warehouse_sources/backend/temporal/data_imports/pipelines/pipeline_v2/pipeline.py index 532cc8993735..e263ce5620bd 100644 --- a/products/warehouse_sources/backend/temporal/data_imports/pipelines/pipeline_v2/pipeline.py +++ b/products/warehouse_sources/backend/temporal/data_imports/pipelines/pipeline_v2/pipeline.py @@ -28,7 +28,6 @@ person_property_sink_clear_chunks, reset_rows_synced_if_needed, resolve_primary_keys, - run_pre_write_defensive_compact, setup_row_tracking_with_billing_check, should_check_shutdown, stage_chunk_for_person_property_sink, @@ -53,6 +52,7 @@ from products.warehouse_sources.backend.temporal.data_imports.pipelines.core.async_iterate import async_iterate from products.warehouse_sources.backend.temporal.data_imports.pipelines.core.batcher import Batcher from products.warehouse_sources.backend.temporal.data_imports.pipelines.core.cdp_producer import CDPProducer +from products.warehouse_sources.backend.temporal.data_imports.pipelines.core.delta.maintenance import DeltaMaintenance from products.warehouse_sources.backend.temporal.data_imports.pipelines.core.delta_table_helper import DeltaTableHelper from products.warehouse_sources.backend.temporal.data_imports.pipelines.core.hogql_schema import HogQLSchema from products.warehouse_sources.backend.temporal.data_imports.pipelines.core.partitioning import setup_partitioning @@ -206,11 +206,10 @@ async def run(self) -> PipelineResult: # Defensive pre-write compaction so a sync that arrived at a fragmented # Delta target (e.g. earlier attempts that failed before reaching # `_post_run_operations`) cleans up before adding more small files. Skipped - # cheaply when the table is healthy. Shared implementation lives in - # `extract.run_pre_write_defensive_compact` so the v3 pipeline matches. + # cheaply when the table is healthy; see DeltaMaintenance.run_scheduled. if not is_first_ever_sync: - await run_pre_write_defensive_compact( - self._delta_table_helper, self._schema, self._resource, self._logger + await DeltaMaintenance(self._delta_table_helper).run_scheduled( + self._schema, partition_count_fallback=self._resource.partition_count ) async for item in async_iterate(self._resource.items()): diff --git a/products/warehouse_sources/backend/temporal/data_imports/pipelines/pipeline_v3/pipeline.py b/products/warehouse_sources/backend/temporal/data_imports/pipelines/pipeline_v3/pipeline.py index a32fde212d3f..6d3073288474 100644 --- a/products/warehouse_sources/backend/temporal/data_imports/pipelines/pipeline_v3/pipeline.py +++ b/products/warehouse_sources/backend/temporal/data_imports/pipelines/pipeline_v3/pipeline.py @@ -36,7 +36,6 @@ person_property_sink_clear_chunks, reset_rows_synced_if_needed, resolve_primary_keys, - run_pre_write_defensive_compact, setup_row_tracking_with_billing_check, should_check_shutdown, stage_chunk_for_person_property_sink, @@ -57,6 +56,7 @@ from products.warehouse_sources.backend.temporal.data_imports.pipelines.core.async_iterate import async_iterate from products.warehouse_sources.backend.temporal.data_imports.pipelines.core.batcher import Batcher from products.warehouse_sources.backend.temporal.data_imports.pipelines.core.cdp_producer import CDPProducer +from products.warehouse_sources.backend.temporal.data_imports.pipelines.core.delta.maintenance import DeltaMaintenance from products.warehouse_sources.backend.temporal.data_imports.pipelines.core.delta_table_helper import DeltaTableHelper from products.warehouse_sources.backend.temporal.data_imports.pipelines.core.hogql_schema import HogQLSchema from products.warehouse_sources.backend.temporal.data_imports.pipelines.core.person_property_row_sink import ( @@ -311,12 +311,11 @@ async def run(self) -> PipelineResult: if is_fresh_sync: self._pg_producer.is_first_ever_sync = True - # Defensive pre-write compaction. See `extract.run_pre_write_defensive_compact` - # for rationale; shared with the v2 pipeline so the threshold + error handling - # stay in lockstep. + # Defensive pre-write compaction so a sync that arrived at a fragmented Delta + # target cleans up before adding more small files; see DeltaMaintenance.run_scheduled. if not is_fresh_sync: - await run_pre_write_defensive_compact( - self._delta_table_helper, self._schema, self._resource, self._logger + await DeltaMaintenance(self._delta_table_helper).run_scheduled( + self._schema, partition_count_fallback=self._resource.partition_count ) async for item in async_iterate(self._resource.items()): diff --git a/products/warehouse_sources/backend/temporal/data_imports/tests/e2e/conftest.py b/products/warehouse_sources/backend/temporal/data_imports/tests/e2e/conftest.py index 7d8ea9dc34c8..6f9cdc708129 100644 --- a/products/warehouse_sources/backend/temporal/data_imports/tests/e2e/conftest.py +++ b/products/warehouse_sources/backend/temporal/data_imports/tests/e2e/conftest.py @@ -31,7 +31,7 @@ from products.warehouse_sources.backend.facade.models import ExternalDataJob, get_latest_run_if_exists from products.warehouse_sources.backend.models.external_table_definitions import external_tables from products.warehouse_sources.backend.temporal.data_imports.external_data_job import ExternalDataJobWorkflow -from products.warehouse_sources.backend.temporal.data_imports.pipelines.core.delta_table_helper import DeltaTableHelper +from products.warehouse_sources.backend.temporal.data_imports.pipelines.core.delta.maintenance import DeltaMaintenance from products.warehouse_sources.backend.temporal.data_imports.pipelines.pipeline_v3.postgres_queue.jobs_db import ( BATCH_TABLE, STATUS_TABLE, @@ -214,7 +214,7 @@ async def run_external_data_job_workflow( DATAWAREHOUSE_LOCAL_BUCKET_REGION="us-east-1", DATAWAREHOUSE_BUCKET_DOMAIN="objectstorage:19000", ), - mock.patch.object(DeltaTableHelper, "compact_table") as mock_compact_table, + mock.patch.object(DeltaMaintenance, "compact_table") as mock_compact_table, mock.patch( "products.warehouse_sources.backend.temporal.data_imports.external_data_job.get_data_import_finished_metric" ) as mock_get_data_import_finished_metric, diff --git a/products/warehouse_sources/backend/temporal/data_imports/tests/e2e/test_end_to_end.py b/products/warehouse_sources/backend/temporal/data_imports/tests/e2e/test_end_to_end.py index e5d12b058fdd..a92f76373db5 100644 --- a/products/warehouse_sources/backend/temporal/data_imports/tests/e2e/test_end_to_end.py +++ b/products/warehouse_sources/backend/temporal/data_imports/tests/e2e/test_end_to_end.py @@ -71,6 +71,7 @@ from products.warehouse_sources.backend.temporal.data_imports.cdp_producer_job import CDPProducerJobWorkflow from products.warehouse_sources.backend.temporal.data_imports.external_data_job import ExternalDataJobWorkflow from products.warehouse_sources.backend.temporal.data_imports.pipelines.core.consts import PARTITION_KEY +from products.warehouse_sources.backend.temporal.data_imports.pipelines.core.delta.maintenance import DeltaMaintenance from products.warehouse_sources.backend.temporal.data_imports.pipelines.core.delta_table_helper import DeltaTableHelper from products.warehouse_sources.backend.temporal.data_imports.pipelines.pipeline_v2.pipeline import PipelineNonDLT from products.warehouse_sources.backend.temporal.data_imports.pipelines.pipeline_v3.load.processor import ( @@ -389,7 +390,7 @@ async def _run( ) with ( - mock.patch.object(DeltaTableHelper, "compact_table") as mock_compact_table, + mock.patch.object(DeltaMaintenance, "compact_table") as mock_compact_table, mock.patch( "products.warehouse_sources.backend.temporal.data_imports.external_data_job.get_data_import_finished_metric" ) as mock_get_data_import_finished_metric, @@ -2787,7 +2788,7 @@ async def test_append_only_table(team, mock_stripe_client): sync_type_config={"incremental_field": "created", "incremental_field_type": "integer"}, ) - with mock.patch.object(DeltaTableHelper, "compact_table"): + with mock.patch.object(DeltaMaintenance, "compact_table"): await _execute_run(str(uuid.uuid4()), inputs, []) run_for_replay = await sync_to_async( @@ -3867,7 +3868,7 @@ async def test_stripe_webhook_s3_charges(team, stripe_charge, mock_stripe_client assert len(files.get("Contents", [])) == 1 # Run the pipeline again to ingest the webhook parquet - with mock.patch.object(DeltaTableHelper, "compact_table"): + with mock.patch.object(DeltaMaintenance, "compact_table"): workflow_id = str(uuid.uuid4()) await _execute_run(workflow_id, inputs, stripe_charge["data"]) @@ -4058,7 +4059,7 @@ async def test_stripe_webhook_consumer_e2e(team, stripe_charge, mock_stripe_clie consumer._consumer.commit.assert_called_once_with(asynchronous=False) # 6. Run the import pipeline to ingest the parquet - with mock.patch.object(DeltaTableHelper, "compact_table"): + with mock.patch.object(DeltaMaintenance, "compact_table"): workflow_id = str(uuid.uuid4()) await _execute_run(workflow_id, inputs, stripe_charge["data"]) diff --git a/products/warehouse_sources/backend/temporal/data_imports/workflow_activities/import_data_sync.py b/products/warehouse_sources/backend/temporal/data_imports/workflow_activities/import_data_sync.py index 12b93b34bbcc..1048a46d6be3 100644 --- a/products/warehouse_sources/backend/temporal/data_imports/workflow_activities/import_data_sync.py +++ b/products/warehouse_sources/backend/temporal/data_imports/workflow_activities/import_data_sync.py @@ -37,7 +37,7 @@ from products.warehouse_sources.backend.temporal.data_imports.pipelines.core.arrow_utils import ( SchemaColumnTypeChangedException, ) -from products.warehouse_sources.backend.temporal.data_imports.pipelines.core.delta_table_helper import ( +from products.warehouse_sources.backend.temporal.data_imports.pipelines.core.delta.errors import ( is_transient_object_store_error, ) from products.warehouse_sources.backend.temporal.data_imports.pipelines.core.typings import PipelineResult diff --git a/products/warehouse_sources/backend/tests/api/test_external_data_source_end_to_end.py b/products/warehouse_sources/backend/tests/api/test_external_data_source_end_to_end.py index 260412a33844..a20ef47a82ac 100644 --- a/products/warehouse_sources/backend/tests/api/test_external_data_source_end_to_end.py +++ b/products/warehouse_sources/backend/tests/api/test_external_data_source_end_to_end.py @@ -25,7 +25,7 @@ ) from products.warehouse_sources.backend.facade.types import DataWarehouseManagedViewSetKind from products.warehouse_sources.backend.temporal.data_imports.external_data_job import ExternalDataJobWorkflow -from products.warehouse_sources.backend.temporal.data_imports.pipelines.core.delta_table_helper import DeltaTableHelper +from products.warehouse_sources.backend.temporal.data_imports.pipelines.core.delta.maintenance import DeltaMaintenance from products.warehouse_sources.backend.temporal.data_imports.settings import ACTIVITIES from products.warehouse_sources.backend.temporal.data_imports.sources.stripe.constants import ( SUBSCRIPTION_RESOURCE_NAME as STRIPE_SUBSCRIPTION_RESOURCE_NAME, @@ -92,7 +92,7 @@ async def _run(team, source: ExternalDataSource, schema: ExternalDataSchema): ) with ( - mock.patch.object(DeltaTableHelper, "compact_table"), + mock.patch.object(DeltaMaintenance, "compact_table"), mock.patch( "products.warehouse_sources.backend.temporal.data_imports.external_data_job.get_data_import_finished_metric" ),