Skip to content

chore(warehouse-sources): split delta maintenance out of DeltaTableHelper #75978

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
estefaniarabadan merged 2 commits into from
Jul 31, 2026
Merged

chore(warehouse-sources): split delta maintenance out of DeltaTableHelper #75978

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
a0864e7
chore(warehouse-sources): split delta maintenance out of DeltaTableHe…
estefaniarabadan Jul 31, 2026
f9184ad
chore(warehouse-sources): fix test_errors module name collision
estefaniarabadan Jul 31, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
56 changes: 2 additions & 54 deletions products/warehouse_sources/backend/temporal/data_imports/pipelines/common/extract.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
)
Expand Down Expand Up @@ -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)
61 changes: 15 additions & 46 deletions products/warehouse_sources/backend/temporal/data_imports/pipelines/common/load.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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,
Expand Down Expand Up @@ -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)
Expand Down
136 changes: 0 additions & 136 deletions ...cts/warehouse_sources/backend/temporal/data_imports/pipelines/common/test/test_extract.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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

Expand Down Expand Up @@ -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(
Expand Down
Loading
Loading