From 077fc7829b2be1ca4df950b1641eb258ad6c316b Mon Sep 17 00:00:00 2001 From: estefania Date: Fri, 31 Jul 2026 21:32:47 +0200 Subject: [PATCH] chore(warehouse-sources): extract DeltaWriter from DeltaTableHelper Phase 3 (PR 3) of the warehouse-sources untangling, after #75978 and #76067: the core write/merge path and commit-metadata idempotency move to core/delta/writer.py, leaving DeltaTableHelper as a pure table access and lifecycle object (~270 lines: cached open, corruption detection, reset, file listing, first-sync state). DeltaWriter is a stateless wrapper over the helper, mirroring DeltaMaintenance and Scd2DeltaWriter: write_to_deltalake becomes DeltaWriter.write, together with the deltalite canary (moved verbatim, mid-rollout), batch dedupe, schema-mismatch fallback, and the idempotency pair has_commit_with_metadata / has_batch_been_committed. Tag-write and tag-read stay together deliberately: only the terminal commit of a multi-commit write may carry the (run_uuid, batch_index) tag, and that invariant spans both halves. _commit_matches becomes the module-level commit_matches, shared with the SCD2 tagging test. Callers updated: the v2 pipeline chunk write, the v3 load processor, the CDC companion seeding in common/load.py, and the redelivery fallback in load/idempotency.py (which now wraps the helper in a DeltaWriter). Debug log prefixes change from 'write_to_deltalake:' to 'write:' with the rename; no behavior changes otherwise. --- .../activities/materialize_view.py | 2 +- .../data_imports/pipelines/common/load.py | 3 +- .../pipelines/core/delta/test/test_scd2.py | 6 +- .../pipelines/core/delta/test/test_writer.py | 831 ++++++++++++++++++ .../pipelines/core/delta/writer.py | 528 +++++++++++ .../pipelines/core/delta_table_helper.py | 498 +---------- .../pipelines/core/deltalite_write.py | 2 +- .../core/test/test_delta_table_helper.py | 812 +---------------- .../pipelines/pipeline_v2/pipeline.py | 3 +- .../pipelines/pipeline_v3/load/idempotency.py | 5 +- .../pipelines/pipeline_v3/load/processor.py | 5 +- .../pipeline_v3/load/test_idempotency.py | 12 +- .../pipeline_v3/load/test_processor.py | 9 +- .../data_imports/sources/mysql/mysql.py | 2 +- .../data_imports/tests/e2e/test_end_to_end.py | 2 +- 15 files changed, 1395 insertions(+), 1325 deletions(-) create mode 100644 products/warehouse_sources/backend/temporal/data_imports/pipelines/core/delta/test/test_writer.py create mode 100644 products/warehouse_sources/backend/temporal/data_imports/pipelines/core/delta/writer.py diff --git a/posthog/temporal/data_modeling/activities/materialize_view.py b/posthog/temporal/data_modeling/activities/materialize_view.py index 5b8a463d8193..0aa88db7fb5c 100644 --- a/posthog/temporal/data_modeling/activities/materialize_view.py +++ b/posthog/temporal/data_modeling/activities/materialize_view.py @@ -505,7 +505,7 @@ async def materialize_view_activity(inputs: MaterializeViewInputs) -> Materializ pa_schema: pa.Schema | None = None # write each batch as its own delta commit, imitating the data_imports pipeline - # (DeltaTableHelper.write_to_deltalake): the first batch overwrites — creating the + # (DeltaWriter.write): the first batch overwrites — creating the # table from the exact arrow schema, which pins column case like `personId` — and # later batches append with schema_mode="merge". this keeps peak memory at ~one # batch (hogql_table yields ~100MB combined batches) and, because each write is a 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 5c27244171ce..5dd3bfaabaff 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 @@ -168,6 +168,7 @@ async def _seed_cdc_companion_from_snapshot( SCD2_VALID_FROM_COLUMN, SCD2_VALID_TO_COLUMN, ) + from products.warehouse_sources.backend.temporal.data_imports.pipelines.core.delta.writer import DeltaWriter from products.warehouse_sources.backend.temporal.data_imports.pipelines.core.delta_table_helper import ( DeltaTableHelper, ) @@ -242,7 +243,7 @@ def _read_next_batch(r: pa.RecordBatchReader) -> pa.RecordBatch | None: # Plain append — the companion table is freshly reset so there are no existing # rows to close, making SCD2 merge unnecessary. - await companion_helper.write_to_deltalake( + await DeltaWriter(companion_helper).write( data=batch_table, write_type="append", should_overwrite_table=False, diff --git a/products/warehouse_sources/backend/temporal/data_imports/pipelines/core/delta/test/test_scd2.py b/products/warehouse_sources/backend/temporal/data_imports/pipelines/core/delta/test/test_scd2.py index 1441932481e6..466743a58d6e 100644 --- a/products/warehouse_sources/backend/temporal/data_imports/pipelines/core/delta/test/test_scd2.py +++ b/products/warehouse_sources/backend/temporal/data_imports/pipelines/core/delta/test/test_scd2.py @@ -9,7 +9,7 @@ import pyarrow.compute as pc from products.warehouse_sources.backend.temporal.data_imports.pipelines.core.delta.scd2 import Scd2DeltaWriter -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.writer import commit_matches from products.warehouse_sources.backend.temporal.data_imports.pipelines.core.test.test_delta_table_helper import ( _decimal_array, _make_local_helper, @@ -95,8 +95,8 @@ async def test_only_terminal_append_commit_carries_metadata(self, tmp_path: Path await _make_writer(delta_path).write(data=batch, primary_keys=["id"], commit_metadata=metadata) history = deltalake.DeltaTable(delta_path).history() - # Newest first: [append (WRITE), close (MERGE), seed (WRITE)]. _commit_matches is the + # Newest first: [append (WRITE), close (MERGE), seed (WRITE)]. commit_matches is the # same layout-agnostic check has_batch_been_committed uses for redelivery dedup. - tagged = [c["operation"] for c in history if DeltaTableHelper._commit_matches(c, metadata)] + tagged = [c["operation"] for c in history if commit_matches(c, metadata)] assert tagged == ["WRITE"] assert any(c["operation"] == "MERGE" for c in history) diff --git a/products/warehouse_sources/backend/temporal/data_imports/pipelines/core/delta/test/test_writer.py b/products/warehouse_sources/backend/temporal/data_imports/pipelines/core/delta/test/test_writer.py new file mode 100644 index 000000000000..43a51ca2a65a --- /dev/null +++ b/products/warehouse_sources/backend/temporal/data_imports/pipelines/core/delta/test/test_writer.py @@ -0,0 +1,831 @@ +import json +from decimal import Decimal +from pathlib import Path +from types import SimpleNamespace +from typing import Any, cast + +import pytest +from unittest.mock import AsyncMock, MagicMock, patch + +import pyarrow as pa +import deltalake +import pyarrow.compute as pc +from parameterized import parameterized + +from products.warehouse_sources.backend.temporal.data_imports.pipelines.core.arrow_utils import ( + SchemaColumnTypeChangedException, + evolve_pyarrow_schema, + first_per_pk_table, +) +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.writer import ( + DeltaWriter, + _deltalite_write_stats, + _merge_predicate_ops, +) +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.test.test_delta_table_helper import ( + _decimal_array, + _make_local_helper, + _make_logger, + _table_is_misaligned, +) + +_WRITER_MODULE = "products.warehouse_sources.backend.temporal.data_imports.pipelines.core.delta.writer" + + +_COMMIT_LAYOUT_CASES: list[tuple[str, list[dict], dict, bool]] = [ + # nested dict layout (older delta-rs / fallback form) + ( + "nested_dict_exact_match", + [{"userMetadata": {"run_uuid": "abc", "batch_index": "0"}}], + {"run_uuid": "abc", "batch_index": "0"}, + True, + ), + # delta-rs 1.x flat layout: custom_metadata entries inlined onto the commit dict + ( + "flat_inlined_exact_match", + [{"operation": "WRITE", "timestamp": 1, "run_uuid": "abc", "batch_index": "0", "version": 1}], + {"run_uuid": "abc", "batch_index": "0"}, + True, + ), + ( + "flat_missing_one_required_key", + [{"operation": "WRITE", "run_uuid": "abc", "version": 1}], + {"run_uuid": "abc", "batch_index": "0"}, + False, + ), + # nested JSON-string layout (some delta-rs versions serialize userMetadata as JSON) + ( + "nested_json_string_exact_match", + [{"userMetadata": json.dumps({"run_uuid": "abc", "batch_index": "0"})}], + {"run_uuid": "abc", "batch_index": "0"}, + True, + ), + # match is a subset of the metadata — should still match + ( + "match_is_subset", + [{"userMetadata": {"run_uuid": "abc", "batch_index": "0", "extra": "field"}}], + {"run_uuid": "abc"}, + True, + ), + # multiple commits, none matching + ( + "no_match_in_history", + [ + {"userMetadata": {"run_uuid": "other", "batch_index": "9"}}, + {"userMetadata": {"run_uuid": "abc", "batch_index": "1"}}, + ], + {"run_uuid": "abc", "batch_index": "0"}, + False, + ), + # commits without any custom metadata at all + ( + "no_metadata_on_any_commit", + [{"operation": "WRITE"}, {}], + {"run_uuid": "abc"}, + False, + ), + # one commit has invalid JSON userMetadata, the next is a valid match — still found + ( + "invalid_json_string_skipped_then_match", + [ + {"userMetadata": "not-valid-json{"}, + {"userMetadata": {"run_uuid": "abc"}}, + ], + {"run_uuid": "abc"}, + True, + ), +] + + +def _make_writer() -> DeltaWriter: + return DeltaWriter(DeltaTableHelper(resource_name="t", job=MagicMock(), logger=_make_logger())) + + +class TestHasCommitWithMetadata: + @pytest.mark.asyncio + async def test_returns_false_when_no_delta_table(self): + writer = _make_writer() + with patch.object(writer._table, "get_delta_table", AsyncMock(return_value=None)): + assert await writer.has_commit_with_metadata({"run_uuid": "abc", "batch_index": "0"}) is False + + @parameterized.expand( + [(name, history, match, expected) for (name, history, match, expected) in _COMMIT_LAYOUT_CASES] + ) + @pytest.mark.asyncio + async def test_layout(self, _name: str, history: list[dict], match: dict, expected: bool): + writer = _make_writer() + mock_delta = MagicMock() + mock_delta.history = MagicMock(return_value=history) + + with patch.object(writer._table, "get_delta_table", AsyncMock(return_value=mock_delta)): + assert await writer.has_commit_with_metadata(match) is expected + + @pytest.mark.asyncio + async def test_scan_limit_passed_to_history(self): + writer = _make_writer() + mock_delta = MagicMock() + mock_delta.history = MagicMock(return_value=[]) + + with patch.object(writer._table, "get_delta_table", AsyncMock(return_value=mock_delta)): + await writer.has_commit_with_metadata({"k": "v"}, scan_limit=123) + + mock_delta.history.assert_called_once_with(limit=123) + + +class TestHasBatchBeenCommitted: + @parameterized.expand( + [ + ("string_run_uuid_int_batch", "run-123", 5, True), + ("zero_batch_index", "run-1", 0, False), + ] + ) + @pytest.mark.asyncio + async def test_wraps_has_commit_with_metadata( + self, _name: str, run_uuid: str, batch_index: int, mocked_return: bool + ): + writer = _make_writer() + with patch.object(writer, "has_commit_with_metadata", AsyncMock(return_value=mocked_return)) as m: + result = await writer.has_batch_been_committed(run_uuid, batch_index) + + assert result is mocked_return + m.assert_called_once_with({"run_uuid": run_uuid, "batch_index": str(batch_index)}) + + +class TestWriteToDeltalakeCommitMetadataPassThrough: + """Covers that commit_metadata is forwarded to deltalake.write_deltalake as CommitProperties.""" + + @parameterized.expand( + [ + ("no_metadata", None, None), + ("with_metadata", {"run_uuid": "abc", "batch_index": "2"}, {"run_uuid": "abc", "batch_index": "2"}), + ] + ) + @pytest.mark.asyncio + async def test_full_refresh_passes_commit_properties( + self, + _name: str, + commit_metadata: dict[str, str] | None, + expected_custom_metadata: dict[str, str] | None, + ): + import pyarrow as pa + + helper = DeltaTableHelper(resource_name="t", job=MagicMock(), logger=_make_logger()) + data = pa.table({"id": [1, 2, 3]}) + mock_delta = MagicMock() + mock_delta.schema = MagicMock(return_value=MagicMock(to_arrow=MagicMock(return_value=data.schema))) + + with ( + patch.object(helper, "get_delta_table", AsyncMock(return_value=mock_delta)), + patch(f"{_WRITER_MODULE}.evolve_delta_schema", AsyncMock(return_value=mock_delta)), + patch("deltalake.write_deltalake") as mock_write, + ): + await DeltaWriter(helper).write( + data=data, + write_type="full_refresh", + should_overwrite_table=False, + primary_keys=None, + commit_metadata=commit_metadata, + ) + + assert mock_write.called + _, kwargs = mock_write.call_args + commit_properties = kwargs["commit_properties"] + if expected_custom_metadata is None: + assert commit_properties is None + else: + assert isinstance(commit_properties, deltalake.CommitProperties) + assert commit_properties.custom_metadata == expected_custom_metadata + + +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.""" + fields: list[pa.Field] = [ + pa.field("id", pa.int64()), + pa.field("name", pa.string()), + pa.field("_dlt_id", pa.string(), nullable=False), + pa.field("_dlt_load_id", pa.string(), nullable=False), + ] + if partitioned: + fields.append(pa.field(PARTITION_KEY, pa.string())) + + data_dict: dict[str, Any] = { + "id": pa.array([1, 2]), + "name": pa.array(["a", "b"]), + "_dlt_id": pa.array(["id1", "id2"]), + "_dlt_load_id": pa.array(["load1", "load1"]), + } + if partitioned: + data_dict[PARTITION_KEY] = pa.array(["p0", "p0"]) + + table = pa.table(data_dict, schema=pa.schema(fields)) + deltalake.write_deltalake(path, table, partition_by=PARTITION_KEY if partitioned else None) + return deltalake.DeltaTable(path) + + +def _v3_batch(*, partitioned: bool = False) -> pa.Table: + """Build an incoming batch the way pipeline_v3 does: no _dlt_* columns.""" + data_dict: dict[str, Any] = {"id": pa.array([3, 4]), "name": pa.array(["c", "d"])} + if partitioned: + data_dict[PARTITION_KEY] = pa.array(["p0", "p0"]) + return pa.table(data_dict) + + +class TestLegacyDltTableReconciliation: + """Pipeline_v3 must handle dlt-created Delta tables with NOT NULL _dlt_* columns.""" + + def test_raw_merge_rejects_missing_non_nullable_columns(self, tmp_path: Path) -> None: + """Baseline: proves delta-rs rejects merges when non-nullable columns are absent + from the source batch. This is the root cause of the production failures.""" + delta_path = str(tmp_path / "table") + _create_legacy_delta_table(delta_path) + batch = _v3_batch() + dt = deltalake.DeltaTable(delta_path) + + with pytest.raises(Exception, match="(?i)(invalid data|non-nullable|validation|not found)"): + dt.merge( + source=batch, + source_alias="source", + target_alias="target", + predicate="source.id = target.id", + ).when_matched_update_all().when_not_matched_insert_all().execute() + + @pytest.mark.parametrize("partitioned", [False, True], ids=["flat", "partitioned"]) + @pytest.mark.asyncio + async def test_incremental_merge_into_legacy_table(self, partitioned: bool, tmp_path: Path) -> None: + delta_path = str(tmp_path / "table") + dt = _create_legacy_delta_table(delta_path, partitioned=partitioned) + + helper = _make_local_helper(delta_path) + batch = evolve_pyarrow_schema(_v3_batch(partitioned=partitioned), dt.schema()) + + result = await DeltaWriter(helper).write( + data=batch, + write_type="incremental", + should_overwrite_table=False, + primary_keys=["id"], + ) + + final = result.to_pyarrow_table() + assert final.num_rows == 4 + assert set(final.column("id").to_pylist()) == {1, 2, 3, 4} + + new_rows = final.filter(pc.is_in(final.column("id"), value_set=pa.array([3, 4]))) + assert all(v == "" for v in new_rows.column("_dlt_id").to_pylist()) + assert all(v == "" for v in new_rows.column("_dlt_load_id").to_pylist()) + + @pytest.mark.asyncio + async def test_append_to_legacy_table(self, tmp_path: Path) -> None: + delta_path = str(tmp_path / "table") + dt = _create_legacy_delta_table(delta_path) + + helper = _make_local_helper(delta_path) + batch = evolve_pyarrow_schema(_v3_batch(), dt.schema()) + + result = await DeltaWriter(helper).write( + data=batch, + write_type="append", + should_overwrite_table=False, + primary_keys=None, + ) + + final = result.to_pyarrow_table() + assert final.num_rows == 4 + assert set(final.column("id").to_pylist()) == {1, 2, 3, 4} + + @pytest.mark.asyncio + async def test_full_refresh_overwrite_on_legacy_table(self, tmp_path: Path) -> None: + delta_path = str(tmp_path / "table") + dt = _create_legacy_delta_table(delta_path) + + helper = _make_local_helper(delta_path) + batch = evolve_pyarrow_schema(_v3_batch(), dt.schema()) + + result = await DeltaWriter(helper).write( + data=batch, + write_type="full_refresh", + should_overwrite_table=True, + primary_keys=None, + ) + + final = result.to_pyarrow_table() + assert final.num_rows == 2 + assert all(v == "" for v in final.column("_dlt_id").to_pylist()) + assert all(v == "" for v in final.column("_dlt_load_id").to_pylist()) + + @pytest.mark.asyncio + async def test_v3_native_table_still_merges(self, tmp_path: Path) -> None: + delta_path = str(tmp_path / "table") + fields: list[pa.Field] = [pa.field("id", pa.int64()), pa.field("name", pa.string())] + schema = pa.schema(fields) + deltalake.write_deltalake(delta_path, pa.table({"id": [1, 2], "name": ["a", "b"]}, schema=schema)) + + helper = _make_local_helper(delta_path) + batch = pa.table({"id": [3], "name": ["c"]}) + + result = await DeltaWriter(helper).write( + data=batch, + write_type="incremental", + should_overwrite_table=False, + primary_keys=["id"], + ) + + final = result.to_pyarrow_table() + assert final.num_rows == 3 + assert set(final.column("id").to_pylist()) == {1, 2, 3} + + +class TestAppendDecimalReconciliation: + """Appending a decimal column that outgrew decimal128 must reconcile to the stored type. + + A batch whose numeric column exceeds decimal128 is promoted to decimal256, which + `evolve_pyarrow_schema` renders to text for the Delta write. Arrow emits scientific + notation for scale-heavy zeros (e.g. '0E-18'), which delta-rs can't parse back into + the stored decimal — an opaque, infinitely-retrying DeltaError on the append path. + """ + + def _seed_decimal_table(self, delta_path: str) -> deltalake.DeltaTable: + table = pa.table( + {"id": pa.array([1], type=pa.int64()), "amount": pa.array([Decimal("1.5")], type=pa.decimal128(38, 10))} + ) + deltalake.write_deltalake(delta_path, table) + return deltalake.DeltaTable(delta_path) + + @pytest.mark.asyncio + async def test_scale_heavy_batch_is_rounded_to_stored_type(self, tmp_path: Path) -> None: + delta_path = str(tmp_path / "table") + dt = self._seed_decimal_table(delta_path) + helper = _make_local_helper(delta_path) + + # Values fit decimal128's integer budget but carry more scale than the stored column, + # so they land as decimal256 and evolve renders them to text (the zero as '0E-18'). + batch = evolve_pyarrow_schema( + pa.table( + { + "id": pa.array([2, 3], type=pa.int64()), + "amount": pa.array( + [Decimal("0.12345678901234567890"), Decimal("0E-18")], type=pa.decimal256(76, 20) + ), + } + ), + dt.schema(), + ) + assert pa.types.is_string(batch.schema.field("amount").type) + + result = await DeltaWriter(helper).write( + data=batch, write_type="append", should_overwrite_table=False, primary_keys=None + ) + + final = result.to_pyarrow_table() + assert final.schema.field("amount").type == pa.decimal128(38, 10) + assert set(final.column("id").to_pylist()) == {1, 2, 3} + assert Decimal("0") in final.column("amount").to_pylist() + + @pytest.mark.asyncio + async def test_integer_overflow_batch_raises_clean_non_retryable(self, tmp_path: Path) -> None: + delta_path = str(tmp_path / "table") + dt = self._seed_decimal_table(delta_path) + helper = _make_local_helper(delta_path) + + batch = evolve_pyarrow_schema( + pa.table( + { + "id": pa.array([2, 3], type=pa.int64()), + "amount": pa.array([Decimal("1" + "0" * 35 + ".5"), Decimal("0E-18")], type=pa.decimal256(76, 18)), + } + ), + dt.schema(), + ) + + with pytest.raises(SchemaColumnTypeChangedException): + await DeltaWriter(helper).write( + data=batch, write_type="append", should_overwrite_table=False, primary_keys=None + ) + + +class TestSchemaEvolutionNullability: + """A column added mid-table-lifetime always predates its own addition: every file the + table already holds was written without it, so `optimize.compact()` must be able to + treat those rows as null for that column. If schema evolution adds the column as NOT + NULL — which happens whenever the batch that introduces it has no nulls, since delta-rs + takes the new field's nullability straight from the incoming Arrow field — compaction + later fails with "Non-nullable column '' is missing from the physical schema".""" + + @pytest.mark.asyncio + async def test_compact_survives_a_column_added_by_an_all_non_null_batch(self, tmp_path: Path) -> None: + delta_path = str(tmp_path / "table") + deltalake.write_deltalake(delta_path, pa.table({"id": pa.array([1, 2], type=pa.int64())})) + + helper = _make_local_helper(delta_path) + + # The incoming field is non-nullable because every value in *this* batch is + # non-null — exactly how upstream Arrow construction infers it, unrelated to + # whether the column can appear in prior or future batches. + fields: list[pa.Field] = [pa.field("id", pa.int64()), pa.field("status", pa.string(), nullable=False)] + batch_schema = pa.schema(fields) + batch = pa.table( + {"id": pa.array([3, 4], type=pa.int64()), "status": pa.array(["ok", "ok"])}, schema=batch_schema + ) + + result = await DeltaWriter(helper).write( + data=batch, write_type="append", should_overwrite_table=False, primary_keys=None + ) + status_field = next(f for f in result.schema().fields if f.name == "status") + assert status_field.nullable is True + + await DeltaMaintenance(helper).compact_table() + + final = result.to_pyarrow_table() + by_id = dict(zip(final.column("id").to_pylist(), final.column("status").to_pylist())) + assert by_id == {1: None, 2: None, 3: "ok", 4: "ok"} + + +class TestIncrementalBatchDeduplication: + """Duplicate PKs in a source batch must never reach the Delta write. + + `when_not_matched_insert_all` inserts every unmatched source row, so a batch with a + repeated PK seeds duplicate rows in the table; every later merge then multi-matches + those rows and the join blows up (the OOM loop seen with sources whose primary keys + aren't actually unique). + """ + + @parameterized.expand( + [ + ("keep_first", "first", ["a1", "b1"]), + ("keep_last", "last", ["a2", "b1"]), + ] + ) + def test_first_per_pk_table_keep_modes(self, _name, keep, expected_names): + table = pa.table({"id": [1, 1, 2], "name": ["a1", "a2", "b1"]}) + + result = first_per_pk_table(table, ["id"], keep=keep).sort_by("id") + + assert result.column("id").to_pylist() == [1, 2] + assert result.column("name").to_pylist() == expected_names + + @pytest.mark.asyncio + async def test_incremental_merge_dedupes_duplicate_source_rows(self, tmp_path: Path) -> None: + delta_path = str(tmp_path / "table") + deltalake.write_deltalake(delta_path, pa.table({"id": [1], "name": ["old"]})) + + helper = _make_local_helper(delta_path) + # id=2 appears twice in one batch — without dedup both copies get inserted. + batch = pa.table({"id": [1, 2, 2], "name": ["updated", "first_copy", "second_copy"]}) + + result = await DeltaWriter(helper).write( + data=batch, + write_type="incremental", + should_overwrite_table=False, + primary_keys=["id"], + ) + + final = result.to_pyarrow_table().sort_by("id") + assert final.column("id").to_pylist() == [1, 2] + # The last occurrence of a duplicated key carries the freshest data. + assert final.column("name").to_pylist() == ["updated", "second_copy"] + cast(AsyncMock, helper._logger.awarning).assert_awaited_once() + + @pytest.mark.asyncio + async def test_first_sync_append_dedupes_duplicate_source_rows(self, tmp_path: Path) -> None: + delta_path = str(tmp_path / "table") + + helper = _make_local_helper(delta_path) + batch = pa.table({"id": [1, 1], "name": ["first_copy", "second_copy"]}) + + result = await DeltaWriter(helper).write( + data=batch, + write_type="incremental", + should_overwrite_table=False, + primary_keys=["id"], + ) + + final = result.to_pyarrow_table() + assert final.column("id").to_pylist() == [1] + assert final.column("name").to_pylist() == ["second_copy"] + + +class TestUnpartitionedTableWithPartitionKeyColumn: + """A Delta table can carry `_ph_partition_key` in its schema while its + partition_columns metadata is empty `[]` — e.g. the SchemaMismatchError fallback in + DeltaWriter.write rewrites with partition_by=None while the column is still in the + data, or evolve_pyarrow_schema re-adds the column to a batch headed for an + unpartitioned table. DeltaWriter.write derives partitioning from column *presence*, + so it then passes partition_by=_ph_partition_key against a table delta-rs considers + unpartitioned and raises: + "Specified table partitioning does not match table partitioning: expected: [], got: [_ph_partition_key]" + """ + + def _seed_unpartitioned_table_with_partition_column(self, delta_path: str) -> None: + # _ph_partition_key is a plain column; the table is NOT partitioned by it. + deltalake.write_deltalake( + delta_path, + pa.table({"id": pa.array([1, 2]), PARTITION_KEY: pa.array(["p0", "p0"])}), + partition_by=None, + ) + dt = deltalake.DeltaTable(delta_path) + assert dt.metadata().partition_columns == [] + assert PARTITION_KEY in dt.schema().to_arrow().names + + @pytest.mark.parametrize( + "write_type,primary_keys,should_overwrite,expected_ids", + [ + # append/incremental keep the existing rows; full_refresh overwrites them. Each + # routes through a distinct write branch, all of which previously raised against + # the unpartitioned-but-column-present table. + ("append", None, False, {1, 2, 3, 4}), + ("incremental", ["id"], False, {1, 2, 3, 4}), + ("full_refresh", None, True, {2, 3, 4}), + ], + ids=["append", "incremental_merge", "full_refresh_overwrite"], + ) + @pytest.mark.asyncio + async def test_write_does_not_partition_unpartitioned_table( + self, + write_type: str, + primary_keys: list[str] | None, + should_overwrite: bool, + expected_ids: set[int], + tmp_path: Path, + ) -> None: + delta_path = str(tmp_path / "table") + self._seed_unpartitioned_table_with_partition_column(delta_path) + + helper = _make_local_helper(delta_path) + # id=2 already exists (merge updates it); id=3,4 are new. + batch = pa.table({"id": pa.array([2, 3, 4]), PARTITION_KEY: pa.array(["p0", "p0", "p0"])}) + + result = await DeltaWriter(helper).write( + data=batch, + write_type=write_type, # type: ignore[arg-type] + should_overwrite_table=should_overwrite, + primary_keys=primary_keys, + ) + + final = result.to_pyarrow_table() + assert set(final.column("id").to_pylist()) == expected_ids + # The table stays unpartitioned — we don't fight its existing layout. + assert result.metadata().partition_columns == [] + + +class TestWriteMisalignedDecimalEndToEnd: + """Writes a misaligned-decimal batch through the real delta-rs write path. Without the + realignment guard, delta-rs would abort the process; with it, the write succeeds.""" + + @pytest.mark.parametrize( + "write_type,should_overwrite", + [("full_refresh", True), ("append", False), ("incremental", False)], + ) + @pytest.mark.asyncio + async def test_write_misaligned_decimal_to_local_delta( + self, write_type: str, should_overwrite: bool, tmp_path: Path + ) -> None: + delta_path = str(tmp_path / "table") + # Seed the table so incremental/append have an existing target to write into. + deltalake.write_deltalake( + delta_path, + pa.table({"id": pa.array([1, 2]), "amount": _decimal_array([5, 6], misaligned=False)}), + ) + + helper = _make_local_helper(delta_path) + batch = pa.table({"id": pa.array([3, 4]), "amount": _decimal_array([7, 8], misaligned=True)}) + assert _table_is_misaligned(batch) is True + + result = await DeltaWriter(helper).write( + data=batch, + write_type=write_type, # type: ignore[arg-type] + should_overwrite_table=should_overwrite, + primary_keys=["id"] if write_type == "incremental" else None, + ) + + final = result.to_pyarrow_table() + amounts = set(final.column("amount").to_pylist()) + if should_overwrite: + assert set(final.column("id").to_pylist()) == {3, 4} + else: + assert {3, 4}.issubset(set(final.column("id").to_pylist())) + assert {7, 8}.issubset(amounts) + + +class TestNullSafeMergePredicate: + """The incremental-merge match must be NULL-safe. + + Regression for the duplicate-accumulation bug found by the deltalite shadow canary: composite + keys with nullable columns (e.g. GoogleAds report resources keyed on `segments.*`) matched with + bare `source.c = target.c` never match on NULL (`NULL = NULL` is NULL), so the row is re-inserted + on every incremental sync and the table silently grows. + """ + + def test_predicate_ops_are_null_safe(self): + assert _merge_predicate_ops(["id", "seg"]) == [ + "(source.id IS NOT DISTINCT FROM target.id)", + "(source.seg IS NOT DISTINCT FROM target.seg)", + ] + + @staticmethod + def _seed_then_merge(path: Path, predicate_ops: list[str]) -> pa.Table: + # Seed one row whose composite key has a NULL component, then merge the same key with a new value. + seed = pa.table( + { + "id": pa.array([1], pa.int64()), + "seg": pa.array([None], pa.string()), + "val": pa.array(["a"], pa.string()), + } + ) + deltalake.write_deltalake(str(path), seed, mode="overwrite") + source = pa.table( + { + "id": pa.array([1], pa.int64()), + "seg": pa.array([None], pa.string()), + "val": pa.array(["b"], pa.string()), + } + ) + ( + deltalake.DeltaTable(str(path)) + .merge(source=source, source_alias="source", target_alias="target", predicate=" AND ".join(predicate_ops)) + .when_matched_update_all() + .when_not_matched_insert_all() + .execute() + ) + return deltalake.DeltaTable(str(path)).to_pyarrow_table() + + def test_null_composite_key_row_matches_instead_of_duplicating(self, tmp_path): + result = self._seed_then_merge(tmp_path / "safe", _merge_predicate_ops(["id", "seg"])) + assert result.num_rows == 1 + assert result.column("val").to_pylist() == ["b"] + + def test_non_null_key_still_matches(self, tmp_path): + # The null-safe form must not change behaviour for ordinary (non-NULL) keys. + seed = pa.table({"id": pa.array([1], pa.int64()), "seg": pa.array(["MOBILE"]), "val": pa.array(["a"])}) + deltalake.write_deltalake(str(tmp_path / "nn"), seed, mode="overwrite") + source = pa.table({"id": pa.array([1], pa.int64()), "seg": pa.array(["MOBILE"]), "val": pa.array(["b"])}) + ( + deltalake.DeltaTable(str(tmp_path / "nn")) + .merge( + source=source, + source_alias="source", + target_alias="target", + predicate=" AND ".join(_merge_predicate_ops(["id", "seg"])), + ) + .when_matched_update_all() + .when_not_matched_insert_all() + .execute() + ) + result = deltalake.DeltaTable(str(tmp_path / "nn")).to_pyarrow_table() + assert result.num_rows == 1 + assert result.column("val").to_pylist() == ["b"] + + def test_bare_equality_duplicates_null_key_row(self, tmp_path): + # Documents the pre-fix behaviour the null-safe predicate corrects. + result = self._seed_then_merge(tmp_path / "unsafe", ["source.id = target.id", "source.seg = target.seg"]) + assert result.num_rows == 2 + + +class TestDeltaliteWritePath: + """Phase 2: deltalite performs the real incremental merge, gated solely by a per-schema flag, with + a hard fallback to the delta-rs MERGE so a deltalite failure can never fail a sync.""" + + _FLAG = ( + "products.warehouse_sources.backend.temporal.data_imports.pipelines.core." + "deltalite_write.is_deltalite_write_enabled" + ) + + @pytest.fixture(autouse=True) + def _preload_write_metrics(self): + # _write_via_deltalite lazily imports the pipeline_v3 metrics module. These tests fake the + # `deltalite` module via patch.dict(sys.modules, ...), which on exit restores the enter-time + # snapshot and thus evicts any module first imported *inside* the block. If the metrics module + # were first loaded there, the next test would re-execute it and hit "Duplicated timeseries" in + # the global Prometheus registry. Loading it here (at setup, before any patch.dict) keeps it in + # the snapshot so it survives. Imported at runtime — not module top — to keep the heavy + # pipeline_v3 chain off collection. + from products.warehouse_sources.backend.temporal.data_imports.pipelines.pipeline_v3.load import ( # noqa: F401 + metrics, + ) + + def _helper(self) -> DeltaTableHelper: + return DeltaTableHelper(resource_name="t", job=MagicMock(team_id=2, schema_id="sch-1"), logger=_make_logger()) + + async def _call(self, helper: DeltaTableHelper) -> bool: + return await DeltaWriter(helper)._write_via_deltalite( + existing_delta_table=MagicMock(), + data=pa.table({"id": pa.array([1], pa.int64())}), + normalized_primary_keys=["id"], + use_partitioning=False, + commit_metadata={"run_uuid": "abc"}, + ) + + @pytest.mark.asyncio + async def test_skips_without_primary_keys(self): + # No primary keys => nothing to key an upsert on; fall back without even evaluating the flag. + with patch(self._FLAG) as flag: + wrote = await DeltaWriter(self._helper())._write_via_deltalite( + existing_delta_table=MagicMock(), + data=pa.table({"id": pa.array([1], pa.int64())}), + normalized_primary_keys=[], + use_partitioning=False, + commit_metadata=None, + ) + assert wrote is False + flag.assert_not_called() + + def test_write_stats_flattens_scalar_getters_only(self): + # Enumerates scalar attributes (so future crate fields flow through) and drops methods/non-scalars. + stats = SimpleNamespace(version=7, rows_inserted=2, files_added=1, _private=9) + stats.helper = lambda: None # callable attribute must be ignored + assert _deltalite_write_stats(stats) == {"version": 7, "rows_inserted": 2, "files_added": 1} + + @pytest.mark.asyncio + async def test_falls_back_when_flag_disabled(self): + with patch(self._FLAG, return_value=False): + assert await self._call(self._helper()) is False + + @pytest.mark.asyncio + async def test_writes_via_deltalite_when_enabled(self): + logger = _make_logger() # captured so we can inspect the structured log without hitting the typed attr + helper = DeltaTableHelper(resource_name="t", job=MagicMock(team_id=2, schema_id="sch-1"), logger=logger) + existing = MagicMock() + # SimpleNamespace stands in for the pyo3 UpsertStats: predictable scalar getters for the structured log. + fake_stats = SimpleNamespace( + version=5, partitions_touched=1, rows_inserted=3, rows_updated=2, rows_copied=10, null_pk_rows=0 + ) + fake_table = MagicMock() + fake_table.upsert.return_value = fake_stats + fake_deltalite = MagicMock() + fake_deltalite.DeltaLiteTable.open.return_value = fake_table + with ( + patch(self._FLAG, return_value=True), + patch.dict("sys.modules", {"deltalite": fake_deltalite}), + patch.object(helper, "_get_delta_table_uri", AsyncMock(return_value="s3://b/t")), + patch.object(helper, "_get_credentials", return_value={"AWS_REGION": "us-east-1"}), + ): + wrote = await DeltaWriter(helper)._write_via_deltalite( + existing_delta_table=existing, + data=pa.table({"id": pa.array([1], pa.int64())}), + normalized_primary_keys=["id"], + use_partitioning=True, + commit_metadata={"run_uuid": "abc"}, + ) + assert wrote is True + fake_deltalite.DeltaLiteTable.open.assert_called_once_with("s3://b/t", {"AWS_REGION": "us-east-1"}) + fake_table.upsert.assert_called_once() + # PARTITION_KEY is passed as the partition arg when the table is partitioned. + assert fake_table.upsert.call_args.args[2] == PARTITION_KEY + existing.update_incremental.assert_called_once() + # The commit is logged with the UpsertStats fields as structured keys + a duration, so it's parseable. + logger.ainfo.assert_called_once() + log_kwargs = logger.ainfo.call_args.kwargs + assert log_kwargs["version"] == 5 + assert log_kwargs["rows_inserted"] == 3 + assert log_kwargs["partitions_touched"] == 1 + assert "duration_ms" in log_kwargs + + @pytest.mark.asyncio + async def test_falls_back_when_deltalite_raises(self): + helper = self._helper() + fake_table = MagicMock() + fake_table.upsert.side_effect = RuntimeError("commit conflict, retries exhausted") + fake_deltalite = MagicMock() + fake_deltalite.DeltaLiteTable.open.return_value = fake_table + with ( + patch(self._FLAG, return_value=True), + patch.dict("sys.modules", {"deltalite": fake_deltalite}), + patch.object(helper, "_get_delta_table_uri", AsyncMock(return_value="s3://b/t")), + patch.object(helper, "_get_credentials", return_value={}), + ): + wrote = await self._call(helper) + assert wrote is False # deltalite blew up -> caller falls through to the delta-rs MERGE + + @parameterized.expand([("refresh",), ("log",)]) + @pytest.mark.asyncio + async def test_post_commit_failure_does_not_fall_back(self, failing_step: str): + # Once the upsert commits, NO post-commit step (handle refresh, log, metric) may raise into the + # caller — that would return False / bubble up and re-run the MERGE on top of deltalite's commit. + logger = _make_logger() # set the side effect on the mock before it becomes the typed _logger attr + existing = MagicMock() + if failing_step == "refresh": + existing.update_incremental.side_effect = RuntimeError("post-commit refresh boom") + else: + logger.ainfo.side_effect = RuntimeError("post-commit log boom") + helper = DeltaTableHelper(resource_name="t", job=MagicMock(team_id=2, schema_id="sch-1"), logger=logger) + fake_table = MagicMock() + fake_table.upsert.return_value = MagicMock(version=5, rows_inserted=1, rows_updated=0, rows_copied=0) + fake_deltalite = MagicMock() + fake_deltalite.DeltaLiteTable.open.return_value = fake_table + with ( + patch(self._FLAG, return_value=True), + patch.dict("sys.modules", {"deltalite": fake_deltalite}), + patch.object(helper, "_get_delta_table_uri", AsyncMock(return_value="s3://b/t")), + patch.object(helper, "_get_credentials", return_value={}), + ): + wrote = await DeltaWriter(helper)._write_via_deltalite( + existing_delta_table=existing, + data=pa.table({"id": pa.array([1], pa.int64())}), + normalized_primary_keys=["id"], + use_partitioning=False, + commit_metadata=None, + ) + assert wrote is True # committed; the post-commit failure is swallowed + fake_table.upsert.assert_called_once() diff --git a/products/warehouse_sources/backend/temporal/data_imports/pipelines/core/delta/writer.py b/products/warehouse_sources/backend/temporal/data_imports/pipelines/core/delta/writer.py new file mode 100644 index 000000000000..00332b5e6d6b --- /dev/null +++ b/products/warehouse_sources/backend/temporal/data_imports/pipelines/core/delta/writer.py @@ -0,0 +1,528 @@ +import json +import time +import asyncio +import contextlib +from collections.abc import Callable, Sequence +from typing import TYPE_CHECKING, Any, Literal + +import pyarrow as pa +import deltalake +import pyarrow.compute as pc +import deltalake.exceptions + +from posthog.exceptions_capture import capture_exception +from posthog.sync import database_sync_to_async_pool + +from products.warehouse_sources.backend.temporal.data_imports.pipelines.core.arrow_utils import ( + align_incoming_decimals_to_delta, + first_per_pk_table, + normalize_column_name, + realign_decimal_buffers, +) +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.evolution import evolve_delta_schema +from products.warehouse_sources.backend.temporal.data_imports.pipelines.core.delta.ops import ( + delta_merge_spill_kwargs, + execute_with_conflict_retry, +) + +if TYPE_CHECKING: + from products.warehouse_sources.backend.temporal.data_imports.pipelines.core.delta_table_helper import ( + DeltaTableHelper, + ) + + +def _write_deltalake( + table_or_uri: str | deltalake.DeltaTable, + table_data: pa.Table, + partition_by: str | None, + mode: Literal["error", "append", "overwrite", "ignore"], + schema_mode: Literal["merge", "overwrite"] | None, + commit_properties: deltalake.CommitProperties | None = None, +) -> None: + deltalake.write_deltalake( + table_or_uri=table_or_uri, + data=table_data, + partition_by=partition_by, + mode=mode, + schema_mode=schema_mode, + commit_properties=commit_properties, + ) + + +def _merge_predicate_ops(normalized_primary_keys: list[str]) -> list[str]: + """Per-key merge match conditions, using NULL-safe equality. + + delta-rs matches source↔target with plain `source.c = target.c`, which is NULL-*un*safe: + `NULL = NULL` evaluates to NULL (not true). Composite keys with nullable columns — e.g. the + GoogleAds report resources keyed on `segments.ad_network_type` / `segments.click_type` / + `segments.device`, which are frequently NULL — therefore never match their existing target row, + so `when_not_matched_insert_all` re-inserts them on *every* incremental sync and the table + silently accumulates a duplicate per NULL-keyed row. `IS NOT DISTINCT FROM` treats NULL == NULL, + matching the source dedup (`first_per_pk_table` groups NULLs together) and stopping the drift. + + Each term is parenthesised: delta-rs's predicate parser (1.6.1) mis-associates a bare + `a IS NOT DISTINCT FROM b AND c IS NOT DISTINCT FROM d` (it groups `b AND c`), so the parens are + required for it to plan. + """ + return [f"(source.{c} IS NOT DISTINCT FROM target.{c})" for c in normalized_primary_keys] + + +def _deltalite_write_stats(stats: Any) -> dict[str, int | float | str | bool]: + """Flatten a deltalite ``UpsertStats`` into scalar log fields for structured, parseable output. + + Enumerates the object's public scalar attributes (the pyo3 ``#[pyo3(get)]`` getters — version, + partitions_touched, files_added/removed/carried_over/probed, rows_updated/inserted/copied, + source_rows, null_pk_rows, …) rather than a fixed list, so fields added crate-side later (e.g. + per-phase timings) surface automatically. Best-effort — a stats change must never break the write. + """ + fields: dict[str, int | float | str | bool] = {} + for name in dir(stats): + if name.startswith("_"): + continue + try: + value = getattr(stats, name) + except Exception: # noqa: BLE001 - a flaky getter must not break logging a committed write + continue + if isinstance(value, bool | int | float | str): + fields[name] = value + return fields + + +def commit_matches(commit: dict[str, Any], match: dict[str, str]) -> bool: + """Return True iff every (k, v) in `match` is present in this commit's metadata. + + Handles both the flat layout (delta-rs 1.x inlines custom_metadata onto the + top-level commit dict) and a nested `userMetadata` key (older/other layouts). + """ + if all(commit.get(k) == v for k, v in match.items()): + return True + + raw = commit.get("userMetadata") + if raw is None: + return False + + if isinstance(raw, str): + try: + nested = json.loads(raw) + except (json.JSONDecodeError, ValueError): + return False + elif isinstance(raw, dict): + nested = raw + else: + return False + + return all(nested.get(k) == v for k, v in match.items()) + + +class DeltaWriter: + """The core write/merge path for one schema's Delta table, plus commit-metadata idempotency. + + Stateless over a `DeltaTableHelper`, which holds the cached table handle and the first-sync + flag — construct one at the call site. Tagging commits with `commit_metadata` and reading the + tags back (`has_batch_been_committed`) live together because they are two halves of one + contract: only the terminal commit of a multi-commit write may carry the tag, or a redelivery + after a mid-write crash would treat the batch as done and lose data. + """ + + def __init__(self, table: "DeltaTableHelper") -> None: + self._table = table + self._logger = table.logger + + async def _dedupe_incremental_batch( + self, data: pa.Table, primary_keys: Sequence[Any], use_partitioning: bool + ) -> pa.Table: + """Drop all but the last occurrence of each PK (+ partition) tuple in a batch.""" + dedupe_keys = [n for x in primary_keys if (n := normalize_column_name(x)) in data.column_names] + if not dedupe_keys: + return data + if use_partitioning: + dedupe_keys.append(PARTITION_KEY) + + deduped = first_per_pk_table(data, dedupe_keys, keep="last") + dropped = data.num_rows - deduped.num_rows + if dropped > 0: + await self._logger.awarning( + f"write: dropped {dropped} duplicate primary-key rows " + f"(keys={dedupe_keys}) from a batch of {data.num_rows} before writing" + ) + return deduped + + async def _write_via_deltalite( + self, + *, + existing_delta_table: deltalake.DeltaTable, + data: pa.Table, + normalized_primary_keys: list[str], + use_partitioning: bool, + commit_metadata: dict[str, str] | None, + ) -> bool: + """Phase 2: perform the incremental merge via deltalite instead of the delta-rs MERGE. + + Returns True if deltalite committed the write (caller then skips the delta-rs MERGE), or False + to fall back to the MERGE. Falls back on *anything* — flag off, import failure, deltalite error / + commit conflict / refusal — so switching a schema to deltalite can only change which engine + writes, never whether the sync succeeds; the worst case is today's behaviour. Controlled solely + by the per-schema ``data-warehouse-deltalite-write`` feature flag (no env switch), so it can be + ramped / killed entirely from the flag UI without a deploy. + """ + if not normalized_primary_keys: + return False + + # The flag check is a rollout gate, not part of the write: a flag miss (off) or any error here + # (including the import) must fall back to the delta-rs MERGE *silently*. + try: + from products.warehouse_sources.backend.temporal.data_imports.pipelines.core.deltalite_write import ( + is_deltalite_write_enabled, + ) + + enabled = await database_sync_to_async_pool(is_deltalite_write_enabled)( + self._table.job.team_id, str(self._table.job.schema_id), None + ) + except Exception: # noqa: BLE001 - a flag-eval / import error just means "don't use deltalite" + return False + if not enabled: + return False + + # deltalite is enabled. Only the upsert *commit* gates the fallback: a pre-commit failure means + # nothing was written, so we re-run the delta-rs MERGE. Anything AFTER the commit is best-effort + # bookkeeping and must NOT return False — otherwise the MERGE would re-run on top of deltalite's + # already-committed write. (Lazy metrics import keeps the heavy pipeline_v3 chain off the module + # import path — circular.) + try: + import deltalite + + from products.warehouse_sources.backend.temporal.data_imports.pipelines.pipeline_v3.load.metrics import ( + DELTALITE_WRITE_DURATION_SECONDS, + DELTALITE_WRITE_TOTAL, + ) + + uri = await self._table.get_table_uri() + storage_options = self._table.get_storage_options() + partition_key = PARTITION_KEY if use_partitioning else None + + def _upsert() -> Any: + table = deltalite.DeltaLiteTable.open(uri, storage_options) + return table.upsert( + data, + list(normalized_primary_keys), + partition_key, + commit_metadata=commit_metadata, + ) + + started = time.perf_counter() + stats = await asyncio.to_thread(_upsert) + duration_s = time.perf_counter() - started + except Exception as e: # noqa: BLE001 - pre-commit failure: nothing committed, fall back to MERGE + await self._logger.awarning( + f"deltalite write failed; falling back to delta-rs MERGE (sync unaffected): {e}" + ) + try: + DELTALITE_WRITE_TOTAL.labels(outcome="fallback").inc() + except Exception: # noqa: BLE001 - the metrics import itself failed; the warning is enough + pass + return False + + # Committed — the real table is now deltalite's output. NOTHING past this point may raise into + # the caller: an exception here would leave `deltalite_wrote` unset and either fail/retry the + # sync or re-run the delta-rs MERGE on top of deltalite's already-committed write. So every + # post-commit step (handle refresh, log, metric) is wrapped best-effort and we always return True. + try: + # Refresh the in-memory delta-rs handle to deltalite's new version so the table returned by + # write (and any subsequent reads) reflects the real state. + await asyncio.to_thread(existing_delta_table.update_incremental) + # Structured, parseable stats (parity with the old `Delta Merge Stats: {json}` line): every + # UpsertStats field becomes its own log key, plus the wall-clock duration. `_deltalite_write_stats` + # enumerates the pyo3 getters, so fields added crate-side later (e.g. per-phase timings) flow + # through here without a code change. + await self._logger.ainfo( + "deltalite write: committed", + duration_ms=round(duration_s * 1000), + **_deltalite_write_stats(stats), + ) + DELTALITE_WRITE_TOTAL.labels(outcome="written").inc() + DELTALITE_WRITE_DURATION_SECONDS.observe(duration_s) + except Exception as e: # noqa: BLE001 - the write is committed; bookkeeping must never raise + with contextlib.suppress(Exception): + await self._logger.awarning(f"deltalite write committed but post-commit bookkeeping failed: {e}") + return True + + async def write( + self, + data: pa.Table, + write_type: Literal["incremental", "full_refresh", "append"], + should_overwrite_table: bool, + primary_keys: Sequence[Any] | None, + progress_callback: Callable[[], None] | None = None, + commit_metadata: dict[str, str] | None = None, + ) -> deltalake.DeltaTable: + # Guard against delta-rs aborting the worker on misaligned decimal buffers (see + # realign_decimal_buffers). Sub-tables derived below via filter()/take() are + # freshly allocated by pyarrow and so inherit safe alignment. + data = realign_decimal_buffers(data) + + delta_table = await self._table.get_delta_table() + + if delta_table: + delta_table = await evolve_delta_schema(delta_table, data.schema) + + is_first_sync = self._table.is_first_sync + await self._logger.adebug( + f"write: is_first_sync = {is_first_sync}. should_overwrite_table = {should_overwrite_table}" + ) + + use_partitioning = False + if PARTITION_KEY in data.column_names: + use_partitioning = True + await self._logger.adebug(f"Using partitioning on {PARTITION_KEY}") + + # The column can exist without the table being partitioned by it; defer to the + # table's real partition_columns or delta-rs rejects the write as a mismatch. + if use_partitioning and delta_table is not None: + existing_partition_columns = getattr(delta_table.metadata(), "partition_columns", None) or [] + if PARTITION_KEY not in existing_partition_columns: + use_partitioning = False + await self._logger.adebug( + f"Existing table is not partitioned by {PARTITION_KEY}; skipping partitioning to match its layout" + ) + + commit_properties: deltalake.CommitProperties | None = ( + deltalake.CommitProperties(custom_metadata=commit_metadata) if commit_metadata else None + ) + + if write_type == "incremental" and primary_keys: + # Sources can emit the same key twice in one batch (re-listed parents, retried + # pages, genuinely non-unique upstream ids). The merge treats PK (+ partition) + # as row identity, and duplicates on the source side either error the merge or + # get double-inserted by `when_not_matched_insert_all` — after which every later + # merge multi-matches those rows and blows up. Keep only the last occurrence. + data = await self._dedupe_incremental_batch(data, primary_keys, use_partitioning) + + if write_type == "incremental" and delta_table is not None and not is_first_sync: + if not primary_keys or len(primary_keys) == 0: + raise Exception("Primary key required for incremental syncs") + + # The merge casts every source column to its stored column type; a scale-heavy decimal + # column (e.g. decimal128(38, 32)) overflows that cast on larger values. Align to the + # stored types up front so the merge cast is a no-op, or raise a clean reset signal. + data = align_incoming_decimals_to_delta(data, delta_table.schema()) + + existing_delta_table = delta_table + + await self._logger.adebug(f"write: merging...") + + # Normalize keys and check the keys actually exist in the dataset + py_table_column_names = data.column_names + normalized_primary_keys: list[str] = [] + for x in primary_keys: + n = normalize_column_name(x) + if n in py_table_column_names: + normalized_primary_keys.append(n) + + predicate_ops = _merge_predicate_ops(normalized_primary_keys) + + # Phase 2 canary: try deltalite for the real merge. On success the delta-rs MERGE (and the + # now-redundant forward shadow) below are skipped; on any failure this returns False and we + # fall through to the MERGE, so a deltalite issue can never fail the sync. + deltalite_wrote = await self._write_via_deltalite( + existing_delta_table=existing_delta_table, + data=data, + normalized_primary_keys=normalized_primary_keys, + use_partitioning=use_partitioning, + commit_metadata=commit_metadata, + ) + + if not deltalite_wrote and use_partitioning: + predicate_ops.append(f"source.{PARTITION_KEY} = target.{PARTITION_KEY}") + + # Group the table by the partition key and merge multiple times with streamed_exec=True for optimised merging + unique_partitions = list(pc.unique(data[PARTITION_KEY])) + + await self._logger.adebug(f"Running {len(unique_partitions)} optimised merges") + + # Only tag the FINAL partition merge with `commit_properties`. Intermediate + # merges must remain untagged so a crash mid-loop doesn't leave behind a + # tagged commit that would cause `has_batch_been_committed` to skip the + # remaining partitions on Kafka redelivery (which would lose data). + last_partition_index = len(unique_partitions) - 1 + for i, partition in enumerate(unique_partitions): + partition_predicate_ops = predicate_ops.copy() + partition_predicate_ops.append(f"target.{PARTITION_KEY} = '{partition}'") + predicate = " AND ".join(partition_predicate_ops) + + filtered_table = data.filter(pc.equal(data[PARTITION_KEY], partition)) + + await self._logger.adebug(f"Merging partition={partition} with predicate={predicate}") + + merge_commit_properties = commit_properties if i == last_partition_index else None + + # Bind the current loop values as defaults so a conflict retry (which re-calls + # this closure) can't accidentally pick up a later iteration's values. + def _do_merge( + filtered_table: pa.Table = filtered_table, + predicate: str = predicate, + merge_commit_properties: deltalake.CommitProperties | None = merge_commit_properties, + ) -> dict: + return ( + existing_delta_table.merge( + source=filtered_table, + source_alias="source", + target_alias="target", + predicate=predicate, + streamed_exec=True, + commit_properties=merge_commit_properties, + **delta_merge_spill_kwargs(), + ) + .when_matched_update_all() + .when_not_matched_insert_all() + .execute() + ) + + merge_stats = await execute_with_conflict_retry( + existing_delta_table, _do_merge, "write: merge", self._logger + ) + + await self._logger.adebug(f"Delta Merge Stats: {json.dumps(merge_stats)}") + + if progress_callback: + progress_callback() + elif not deltalite_wrote: + # Single merge call → safe to tag directly; this is the terminal commit. + def _do_merge_unpartitioned(data: pa.Table, predicate_ops: list[str]): + return ( + existing_delta_table.merge( + source=data, + source_alias="source", + target_alias="target", + predicate=" AND ".join(predicate_ops), + streamed_exec=False, + commit_properties=commit_properties, + **delta_merge_spill_kwargs(), + ) + .when_matched_update_all() + .when_not_matched_insert_all() + .execute() + ) + + merge_stats = await execute_with_conflict_retry( + existing_delta_table, + lambda: _do_merge_unpartitioned(data, predicate_ops), + "write: merge", + self._logger, + ) + await self._logger.adebug(f"Delta Merge Stats: {json.dumps(merge_stats)}") + elif ( + write_type == "full_refresh" + or (write_type == "incremental" and delta_table is None) + or (write_type == "incremental" and is_first_sync) + ): + mode: Literal["error", "append", "overwrite", "ignore"] = "append" + schema_mode: Literal["merge", "overwrite"] | None = "merge" + if should_overwrite_table or delta_table is None: + mode = "overwrite" + schema_mode = "overwrite" + + await self._logger.adebug(f"write: mode = {mode}") + + if delta_table is None: + storage_options = self._table.get_storage_options() + delta_uri = await self._table.get_table_uri() + delta_table = await asyncio.to_thread( + deltalake.DeltaTable.create, + table_uri=delta_uri, + schema=data.schema, + storage_options=storage_options, + partition_by=PARTITION_KEY if use_partitioning else None, + ) + + try: + await asyncio.to_thread( + _write_deltalake, + delta_table, + data, + partition_by=PARTITION_KEY if use_partitioning else None, + mode=mode, + schema_mode=schema_mode, + commit_properties=commit_properties, + ) + except deltalake.exceptions.SchemaMismatchError as e: + await self._logger.adebug("SchemaMismatchError: attempting to overwrite schema instead", exc_info=e) + capture_exception(e) + + await asyncio.to_thread( + _write_deltalake, + delta_table, + data, + partition_by=None, + mode=mode, + schema_mode="overwrite", + commit_properties=commit_properties, + ) + elif write_type == "append": + if delta_table is None: + storage_options = self._table.get_storage_options() + delta_uri = await self._table.get_table_uri() + delta_table = await asyncio.to_thread( + deltalake.DeltaTable.create, + table_uri=delta_uri, + schema=data.schema, + storage_options=storage_options, + partition_by=PARTITION_KEY if use_partitioning else None, + ) + else: + # An append re-casts each source column to its stored type, same as a merge. A decimal + # column that outgrew decimal128 arrives here as text (decimal256 renders to string), + # and delta-rs can't parse the scientific notation arrow emits for scale-heavy zeros + # (e.g. '0E-18') back into the stored decimal — an opaque DeltaError that retries + # forever. Align to the stored decimal types up front, exactly as the merge path does. + data = align_incoming_decimals_to_delta(data, delta_table.schema()) + + await self._logger.adebug(f"write: write_type = append") + + await asyncio.to_thread( + _write_deltalake, + delta_table, + data, + partition_by=PARTITION_KEY if use_partitioning else None, + mode="append", + schema_mode="merge", + commit_properties=commit_properties, + ) + + delta_table = await self._table.get_delta_table() + assert delta_table is not None + + return delta_table + + async def has_commit_with_metadata(self, match: dict[str, str], *, scan_limit: int = 50) -> bool: + """Check whether any recent delta commit has custom metadata matching all entries in `match`. + + Used to detect that a given (run_uuid, batch_index) has already been written + even when a faster external dedup cache (e.g. Redis) is missing the marker — + the canonical case is a writer crash between a successful `write` + and the subsequent cache update. + + delta-rs `history()` returns commits where `CommitProperties.custom_metadata` + entries are flattened directly into the commit dict alongside `operation`, + `timestamp`, etc. Older versions nested them under a `userMetadata` key, so + we accept both layouts for forward compatibility. + """ + delta_table = await self._table.get_delta_table() + if delta_table is None: + return False + + history = await asyncio.to_thread(delta_table.history, limit=scan_limit) + + for commit in history: + if commit_matches(commit, match): + return True + + return False + + async def has_batch_been_committed(self, run_uuid: str, batch_index: int) -> bool: + """Check whether a specific (run_uuid, batch_index) has already been committed to delta. + + Thin wrapper around `has_commit_with_metadata` so callers don't need to know + the metadata schema used for idempotency tagging. + """ + return await self.has_commit_with_metadata({"run_uuid": run_uuid, "batch_index": str(batch_index)}) 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 ebcf00292fe0..d248265ad53f 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 @@ -1,15 +1,9 @@ -import json -import time import asyncio -import contextlib -from collections.abc import Callable, Sequence -from typing import Any, Literal +from typing import Any from django.conf import settings -import pyarrow as pa import deltalake as deltalake -import pyarrow.compute as pc import deltalake.exceptions from structlog.types import FilteringBoundLogger @@ -20,21 +14,11 @@ from products.warehouse_sources.backend.models.external_data_job import ExternalDataJob from products.warehouse_sources.backend.temporal.data_imports.naming_convention import NamingConvention from products.warehouse_sources.backend.temporal.data_imports.pipelines.core.arrow_utils import ( - align_incoming_decimals_to_delta, conditional_lru_cache_async, - first_per_pk_table, - normalize_column_name, - realign_decimal_buffers, ) -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.errors import ( is_transient_object_store_error, ) -from products.warehouse_sources.backend.temporal.data_imports.pipelines.core.delta.evolution import evolve_delta_schema -from products.warehouse_sources.backend.temporal.data_imports.pipelines.core.delta.ops import ( - delta_merge_spill_kwargs, - execute_with_conflict_retry, -) # _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. @@ -82,63 +66,6 @@ async def _purge_s3_prefix_once(s3: Any, uri: str) -> None: await s3._rm(uri, recursive=True) -def _write_deltalake( - table_or_uri: str | deltalake.DeltaTable, - table_data: pa.Table, - partition_by: str | None, - mode: Literal["error", "append", "overwrite", "ignore"], - schema_mode: Literal["merge", "overwrite"] | None, - commit_properties: deltalake.CommitProperties | None = None, -) -> None: - deltalake.write_deltalake( - table_or_uri=table_or_uri, - data=table_data, - partition_by=partition_by, - mode=mode, - schema_mode=schema_mode, - commit_properties=commit_properties, - ) - - -def _merge_predicate_ops(normalized_primary_keys: list[str]) -> list[str]: - """Per-key merge match conditions, using NULL-safe equality. - - delta-rs matches source↔target with plain `source.c = target.c`, which is NULL-*un*safe: - `NULL = NULL` evaluates to NULL (not true). Composite keys with nullable columns — e.g. the - GoogleAds report resources keyed on `segments.ad_network_type` / `segments.click_type` / - `segments.device`, which are frequently NULL — therefore never match their existing target row, - so `when_not_matched_insert_all` re-inserts them on *every* incremental sync and the table - silently accumulates a duplicate per NULL-keyed row. `IS NOT DISTINCT FROM` treats NULL == NULL, - matching the source dedup (`first_per_pk_table` groups NULLs together) and stopping the drift. - - Each term is parenthesised: delta-rs's predicate parser (1.6.1) mis-associates a bare - `a IS NOT DISTINCT FROM b AND c IS NOT DISTINCT FROM d` (it groups `b AND c`), so the parens are - required for it to plan. - """ - return [f"(source.{c} IS NOT DISTINCT FROM target.{c})" for c in normalized_primary_keys] - - -def _deltalite_write_stats(stats: Any) -> dict[str, int | float | str | bool]: - """Flatten a deltalite ``UpsertStats`` into scalar log fields for structured, parseable output. - - Enumerates the object's public scalar attributes (the pyo3 ``#[pyo3(get)]`` getters — version, - partitions_touched, files_added/removed/carried_over/probed, rows_updated/inserted/copied, - source_rows, null_pk_rows, …) rather than a fixed list, so fields added crate-side later (e.g. - per-phase timings) surface automatically. Best-effort — a stats change must never break the write. - """ - fields: dict[str, int | float | str | bool] = {} - for name in dir(stats): - if name.startswith("_"): - continue - try: - value = getattr(stats, name) - except Exception: # noqa: BLE001 - a flaky getter must not break logging a committed write - continue - if isinstance(value, bool | int | float | str): - fields[name] = value - return fields - - def delta_storage_options() -> dict[str, str]: """delta-rs storage options for the data-warehouse bucket, independent of any import job — so a read path (e.g. the person-property backfill) can open a Delta table without constructing a full @@ -325,426 +252,3 @@ async def get_file_uris(self) -> list[str]: return [] return await asyncio.to_thread(delta_table.file_uris) - - async def _dedupe_incremental_batch( - self, data: pa.Table, primary_keys: Sequence[Any], use_partitioning: bool - ) -> pa.Table: - """Drop all but the last occurrence of each PK (+ partition) tuple in a batch.""" - dedupe_keys = [n for x in primary_keys if (n := normalize_column_name(x)) in data.column_names] - if not dedupe_keys: - return data - if use_partitioning: - dedupe_keys.append(PARTITION_KEY) - - deduped = first_per_pk_table(data, dedupe_keys, keep="last") - dropped = data.num_rows - deduped.num_rows - if dropped > 0: - await self._logger.awarning( - f"write_to_deltalake: dropped {dropped} duplicate primary-key rows " - f"(keys={dedupe_keys}) from a batch of {data.num_rows} before writing" - ) - return deduped - - async def _write_via_deltalite( - self, - *, - existing_delta_table: deltalake.DeltaTable, - data: pa.Table, - normalized_primary_keys: list[str], - use_partitioning: bool, - commit_metadata: dict[str, str] | None, - ) -> bool: - """Phase 2: perform the incremental merge via deltalite instead of the delta-rs MERGE. - - Returns True if deltalite committed the write (caller then skips the delta-rs MERGE), or False - to fall back to the MERGE. Falls back on *anything* — flag off, import failure, deltalite error / - commit conflict / refusal — so switching a schema to deltalite can only change which engine - writes, never whether the sync succeeds; the worst case is today's behaviour. Controlled solely - by the per-schema ``data-warehouse-deltalite-write`` feature flag (no env switch), so it can be - ramped / killed entirely from the flag UI without a deploy. - """ - if not normalized_primary_keys: - return False - - # The flag check is a rollout gate, not part of the write: a flag miss (off) or any error here - # (including the import) must fall back to the delta-rs MERGE *silently*. - try: - from products.warehouse_sources.backend.temporal.data_imports.pipelines.core.deltalite_write import ( - is_deltalite_write_enabled, - ) - - enabled = await database_sync_to_async_pool(is_deltalite_write_enabled)( - self._job.team_id, str(self._job.schema_id), None - ) - except Exception: # noqa: BLE001 - a flag-eval / import error just means "don't use deltalite" - return False - if not enabled: - return False - - # deltalite is enabled. Only the upsert *commit* gates the fallback: a pre-commit failure means - # nothing was written, so we re-run the delta-rs MERGE. Anything AFTER the commit is best-effort - # bookkeeping and must NOT return False — otherwise the MERGE would re-run on top of deltalite's - # already-committed write. (Lazy metrics import keeps the heavy pipeline_v3 chain off the module - # import path — circular.) - try: - import deltalite - - from products.warehouse_sources.backend.temporal.data_imports.pipelines.pipeline_v3.load.metrics import ( - DELTALITE_WRITE_DURATION_SECONDS, - DELTALITE_WRITE_TOTAL, - ) - - uri = await self._get_delta_table_uri() - storage_options = self._get_credentials() - partition_key = PARTITION_KEY if use_partitioning else None - - def _upsert() -> Any: - table = deltalite.DeltaLiteTable.open(uri, storage_options) - return table.upsert( - data, - list(normalized_primary_keys), - partition_key, - commit_metadata=commit_metadata, - ) - - started = time.perf_counter() - stats = await asyncio.to_thread(_upsert) - duration_s = time.perf_counter() - started - except Exception as e: # noqa: BLE001 - pre-commit failure: nothing committed, fall back to MERGE - await self._logger.awarning( - f"deltalite write failed; falling back to delta-rs MERGE (sync unaffected): {e}" - ) - try: - DELTALITE_WRITE_TOTAL.labels(outcome="fallback").inc() - except Exception: # noqa: BLE001 - the metrics import itself failed; the warning is enough - pass - return False - - # Committed — the real table is now deltalite's output. NOTHING past this point may raise into - # the caller: an exception here would leave `deltalite_wrote` unset and either fail/retry the - # sync or re-run the delta-rs MERGE on top of deltalite's already-committed write. So every - # post-commit step (handle refresh, log, metric) is wrapped best-effort and we always return True. - try: - # Refresh the in-memory delta-rs handle to deltalite's new version so the table returned by - # write_to_deltalake (and any subsequent reads) reflects the real state. - await asyncio.to_thread(existing_delta_table.update_incremental) - # Structured, parseable stats (parity with the old `Delta Merge Stats: {json}` line): every - # UpsertStats field becomes its own log key, plus the wall-clock duration. `_deltalite_write_stats` - # enumerates the pyo3 getters, so fields added crate-side later (e.g. per-phase timings) flow - # through here without a code change. - await self._logger.ainfo( - "deltalite write: committed", - duration_ms=round(duration_s * 1000), - **_deltalite_write_stats(stats), - ) - DELTALITE_WRITE_TOTAL.labels(outcome="written").inc() - DELTALITE_WRITE_DURATION_SECONDS.observe(duration_s) - except Exception as e: # noqa: BLE001 - the write is committed; bookkeeping must never raise - with contextlib.suppress(Exception): - await self._logger.awarning(f"deltalite write committed but post-commit bookkeeping failed: {e}") - return True - - async def write_to_deltalake( - self, - data: pa.Table, - write_type: Literal["incremental", "full_refresh", "append"], - should_overwrite_table: bool, - primary_keys: Sequence[Any] | None, - progress_callback: Callable[[], None] | None = None, - commit_metadata: dict[str, str] | None = None, - ) -> deltalake.DeltaTable: - # Guard against delta-rs aborting the worker on misaligned decimal buffers (see - # realign_decimal_buffers). Sub-tables derived below via filter()/take() are - # freshly allocated by pyarrow and so inherit safe alignment. - data = realign_decimal_buffers(data) - - delta_table = await self.get_delta_table() - - if delta_table: - delta_table = await evolve_delta_schema(delta_table, data.schema) - - await self._logger.adebug( - f"write_to_deltalake: _is_first_sync = {self._is_first_sync}. should_overwrite_table = {should_overwrite_table}" - ) - - use_partitioning = False - if PARTITION_KEY in data.column_names: - use_partitioning = True - await self._logger.adebug(f"Using partitioning on {PARTITION_KEY}") - - # The column can exist without the table being partitioned by it; defer to the - # table's real partition_columns or delta-rs rejects the write as a mismatch. - if use_partitioning and delta_table is not None: - existing_partition_columns = getattr(delta_table.metadata(), "partition_columns", None) or [] - if PARTITION_KEY not in existing_partition_columns: - use_partitioning = False - await self._logger.adebug( - f"Existing table is not partitioned by {PARTITION_KEY}; skipping partitioning to match its layout" - ) - - commit_properties: deltalake.CommitProperties | None = ( - deltalake.CommitProperties(custom_metadata=commit_metadata) if commit_metadata else None - ) - - if write_type == "incremental" and primary_keys: - # Sources can emit the same key twice in one batch (re-listed parents, retried - # pages, genuinely non-unique upstream ids). The merge treats PK (+ partition) - # as row identity, and duplicates on the source side either error the merge or - # get double-inserted by `when_not_matched_insert_all` — after which every later - # merge multi-matches those rows and blows up. Keep only the last occurrence. - data = await self._dedupe_incremental_batch(data, primary_keys, use_partitioning) - - if write_type == "incremental" and delta_table is not None and not self._is_first_sync: - if not primary_keys or len(primary_keys) == 0: - raise Exception("Primary key required for incremental syncs") - - # The merge casts every source column to its stored column type; a scale-heavy decimal - # column (e.g. decimal128(38, 32)) overflows that cast on larger values. Align to the - # stored types up front so the merge cast is a no-op, or raise a clean reset signal. - data = align_incoming_decimals_to_delta(data, delta_table.schema()) - - existing_delta_table = delta_table - - await self._logger.adebug(f"write_to_deltalake: merging...") - - # Normalize keys and check the keys actually exist in the dataset - py_table_column_names = data.column_names - normalized_primary_keys: list[str] = [] - for x in primary_keys: - n = normalize_column_name(x) - if n in py_table_column_names: - normalized_primary_keys.append(n) - - predicate_ops = _merge_predicate_ops(normalized_primary_keys) - - # Phase 2 canary: try deltalite for the real merge. On success the delta-rs MERGE (and the - # now-redundant forward shadow) below are skipped; on any failure this returns False and we - # fall through to the MERGE, so a deltalite issue can never fail the sync. - deltalite_wrote = await self._write_via_deltalite( - existing_delta_table=existing_delta_table, - data=data, - normalized_primary_keys=normalized_primary_keys, - use_partitioning=use_partitioning, - commit_metadata=commit_metadata, - ) - - if not deltalite_wrote and use_partitioning: - predicate_ops.append(f"source.{PARTITION_KEY} = target.{PARTITION_KEY}") - - # Group the table by the partition key and merge multiple times with streamed_exec=True for optimised merging - unique_partitions = list(pc.unique(data[PARTITION_KEY])) - - await self._logger.adebug(f"Running {len(unique_partitions)} optimised merges") - - # Only tag the FINAL partition merge with `commit_properties`. Intermediate - # merges must remain untagged so a crash mid-loop doesn't leave behind a - # tagged commit that would cause `has_batch_been_committed` to skip the - # remaining partitions on Kafka redelivery (which would lose data). - last_partition_index = len(unique_partitions) - 1 - for i, partition in enumerate(unique_partitions): - partition_predicate_ops = predicate_ops.copy() - partition_predicate_ops.append(f"target.{PARTITION_KEY} = '{partition}'") - predicate = " AND ".join(partition_predicate_ops) - - filtered_table = data.filter(pc.equal(data[PARTITION_KEY], partition)) - - await self._logger.adebug(f"Merging partition={partition} with predicate={predicate}") - - merge_commit_properties = commit_properties if i == last_partition_index else None - - # Bind the current loop values as defaults so a conflict retry (which re-calls - # this closure) can't accidentally pick up a later iteration's values. - def _do_merge( - filtered_table: pa.Table = filtered_table, - predicate: str = predicate, - merge_commit_properties: deltalake.CommitProperties | None = merge_commit_properties, - ) -> dict: - return ( - existing_delta_table.merge( - source=filtered_table, - source_alias="source", - target_alias="target", - predicate=predicate, - streamed_exec=True, - commit_properties=merge_commit_properties, - **delta_merge_spill_kwargs(), - ) - .when_matched_update_all() - .when_not_matched_insert_all() - .execute() - ) - - 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)}") - - if progress_callback: - progress_callback() - elif not deltalite_wrote: - # Single merge call → safe to tag directly; this is the terminal commit. - def _do_merge_unpartitioned(data: pa.Table, predicate_ops: list[str]): - return ( - existing_delta_table.merge( - source=data, - source_alias="source", - target_alias="target", - predicate=" AND ".join(predicate_ops), - streamed_exec=False, - commit_properties=commit_properties, - **delta_merge_spill_kwargs(), - ) - .when_matched_update_all() - .when_not_matched_insert_all() - .execute() - ) - - 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 ( - write_type == "full_refresh" - or (write_type == "incremental" and delta_table is None) - or (write_type == "incremental" and self._is_first_sync) - ): - mode: Literal["error", "append", "overwrite", "ignore"] = "append" - schema_mode: Literal["merge", "overwrite"] | None = "merge" - if should_overwrite_table or delta_table is None: - mode = "overwrite" - schema_mode = "overwrite" - - await self._logger.adebug(f"write_to_deltalake: mode = {mode}") - - if delta_table is None: - storage_options = self._get_credentials() - delta_uri = await self._get_delta_table_uri() - delta_table = await asyncio.to_thread( - deltalake.DeltaTable.create, - table_uri=delta_uri, - schema=data.schema, - storage_options=storage_options, - partition_by=PARTITION_KEY if use_partitioning else None, - ) - - try: - await asyncio.to_thread( - _write_deltalake, - delta_table, - data, - partition_by=PARTITION_KEY if use_partitioning else None, - mode=mode, - schema_mode=schema_mode, - commit_properties=commit_properties, - ) - except deltalake.exceptions.SchemaMismatchError as e: - await self._logger.adebug("SchemaMismatchError: attempting to overwrite schema instead", exc_info=e) - capture_exception(e) - - await asyncio.to_thread( - _write_deltalake, - delta_table, - data, - partition_by=None, - mode=mode, - schema_mode="overwrite", - commit_properties=commit_properties, - ) - elif write_type == "append": - if delta_table is None: - storage_options = self._get_credentials() - delta_uri = await self._get_delta_table_uri() - delta_table = await asyncio.to_thread( - deltalake.DeltaTable.create, - table_uri=delta_uri, - schema=data.schema, - storage_options=storage_options, - partition_by=PARTITION_KEY if use_partitioning else None, - ) - else: - # An append re-casts each source column to its stored type, same as a merge. A decimal - # column that outgrew decimal128 arrives here as text (decimal256 renders to string), - # and delta-rs can't parse the scientific notation arrow emits for scale-heavy zeros - # (e.g. '0E-18') back into the stored decimal — an opaque DeltaError that retries - # forever. Align to the stored decimal types up front, exactly as the merge path does. - data = align_incoming_decimals_to_delta(data, delta_table.schema()) - - await self._logger.adebug(f"write_to_deltalake: write_type = append") - - await asyncio.to_thread( - _write_deltalake, - delta_table, - data, - partition_by=PARTITION_KEY if use_partitioning else None, - mode="append", - schema_mode="merge", - commit_properties=commit_properties, - ) - - delta_table = await self.get_delta_table() - assert delta_table is not None - - return delta_table - - async def has_commit_with_metadata(self, match: dict[str, str], *, scan_limit: int = 50) -> bool: - """Check whether any recent delta commit has custom metadata matching all entries in `match`. - - Used to detect that a given (run_uuid, batch_index) has already been written - even when a faster external dedup cache (e.g. Redis) is missing the marker — - the canonical case is a writer crash between a successful `write_to_deltalake` - and the subsequent cache update. - - delta-rs `history()` returns commits where `CommitProperties.custom_metadata` - entries are flattened directly into the commit dict alongside `operation`, - `timestamp`, etc. Older versions nested them under a `userMetadata` key, so - we accept both layouts for forward compatibility. - """ - delta_table = await self.get_delta_table() - if delta_table is None: - return False - - history = await asyncio.to_thread(delta_table.history, limit=scan_limit) - - for commit in history: - if self._commit_matches(commit, match): - return True - - return False - - @staticmethod - def _commit_matches(commit: dict[str, Any], match: dict[str, str]) -> bool: - """Return True iff every (k, v) in `match` is present in this commit's metadata. - - Handles both the flat layout (delta-rs 1.x inlines custom_metadata onto the - top-level commit dict) and a nested `userMetadata` key (older/other layouts). - """ - if all(commit.get(k) == v for k, v in match.items()): - return True - - raw = commit.get("userMetadata") - if raw is None: - return False - - if isinstance(raw, str): - try: - nested = json.loads(raw) - except (json.JSONDecodeError, ValueError): - return False - elif isinstance(raw, dict): - nested = raw - else: - return False - - return all(nested.get(k) == v for k, v in match.items()) - - async def has_batch_been_committed(self, run_uuid: str, batch_index: int) -> bool: - """Check whether a specific (run_uuid, batch_index) has already been committed to delta. - - Thin wrapper around `has_commit_with_metadata` so callers don't need to know - the metadata schema used for idempotency tagging. - """ - return await self.has_commit_with_metadata({"run_uuid": run_uuid, "batch_index": str(batch_index)}) diff --git a/products/warehouse_sources/backend/temporal/data_imports/pipelines/core/deltalite_write.py b/products/warehouse_sources/backend/temporal/data_imports/pipelines/core/deltalite_write.py index 8b73b71d578a..176b8f873214 100644 --- a/products/warehouse_sources/backend/temporal/data_imports/pipelines/core/deltalite_write.py +++ b/products/warehouse_sources/backend/temporal/data_imports/pipelines/core/deltalite_write.py @@ -2,7 +2,7 @@ When the ``data-warehouse-deltalite-write`` flag matches a schema, deltalite performs the incremental merge (``DeltaLiteTable.upsert``) instead of the delta-rs ``MERGE`` — see -``DeltaTableHelper._write_via_deltalite``. A deltalite failure falls back to the MERGE, so the flag +``DeltaWriter._write_via_deltalite``. A deltalite failure falls back to the MERGE, so the flag only ever changes *which engine* writes, never whether the sync succeeds. """ 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 bba2d6fcd681..44f99557195b 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 @@ -1,8 +1,4 @@ -import json -from decimal import Decimal -from pathlib import Path -from types import SimpleNamespace -from typing import Any, cast +from typing import cast import pytest from unittest.mock import AsyncMock, MagicMock, patch @@ -11,22 +7,10 @@ import pyarrow as pa import deltalake -import pyarrow.compute as pc from parameterized import parameterized -from products.warehouse_sources.backend.temporal.data_imports.pipelines.core.arrow_utils import ( - SchemaColumnTypeChangedException, - evolve_pyarrow_schema, - first_per_pk_table, - realign_decimal_buffers, -) -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, - _deltalite_write_stats, - _merge_predicate_ops, -) +from products.warehouse_sources.backend.temporal.data_imports.pipelines.core.arrow_utils import realign_decimal_buffers +from products.warehouse_sources.backend.temporal.data_imports.pipelines.core.delta_table_helper import DeltaTableHelper _HELPER_MODULE = "products.warehouse_sources.backend.temporal.data_imports.pipelines.core.delta_table_helper" @@ -76,71 +60,6 @@ def helper(): return DeltaTableHelper(resource_name="test_resource", job=MagicMock(), logger=_make_logger()) -_COMMIT_LAYOUT_CASES: list[tuple[str, list[dict], dict, bool]] = [ - # nested dict layout (older delta-rs / fallback form) - ( - "nested_dict_exact_match", - [{"userMetadata": {"run_uuid": "abc", "batch_index": "0"}}], - {"run_uuid": "abc", "batch_index": "0"}, - True, - ), - # delta-rs 1.x flat layout: custom_metadata entries inlined onto the commit dict - ( - "flat_inlined_exact_match", - [{"operation": "WRITE", "timestamp": 1, "run_uuid": "abc", "batch_index": "0", "version": 1}], - {"run_uuid": "abc", "batch_index": "0"}, - True, - ), - ( - "flat_missing_one_required_key", - [{"operation": "WRITE", "run_uuid": "abc", "version": 1}], - {"run_uuid": "abc", "batch_index": "0"}, - False, - ), - # nested JSON-string layout (some delta-rs versions serialize userMetadata as JSON) - ( - "nested_json_string_exact_match", - [{"userMetadata": json.dumps({"run_uuid": "abc", "batch_index": "0"})}], - {"run_uuid": "abc", "batch_index": "0"}, - True, - ), - # match is a subset of the metadata — should still match - ( - "match_is_subset", - [{"userMetadata": {"run_uuid": "abc", "batch_index": "0", "extra": "field"}}], - {"run_uuid": "abc"}, - True, - ), - # multiple commits, none matching - ( - "no_match_in_history", - [ - {"userMetadata": {"run_uuid": "other", "batch_index": "9"}}, - {"userMetadata": {"run_uuid": "abc", "batch_index": "1"}}, - ], - {"run_uuid": "abc", "batch_index": "0"}, - False, - ), - # commits without any custom metadata at all - ( - "no_metadata_on_any_commit", - [{"operation": "WRITE"}, {}], - {"run_uuid": "abc"}, - False, - ), - # one commit has invalid JSON userMetadata, the next is a valid match — still found - ( - "invalid_json_string_skipped_then_match", - [ - {"userMetadata": "not-valid-json{"}, - {"userMetadata": {"run_uuid": "abc"}}, - ], - {"run_uuid": "abc"}, - True, - ), -] - - class TestStorageOptionsCommitSafety: # Re-adding AWS_S3_ALLOW_UNSAFE_RENAME unconditionally would silently restore # the legacy rename backend, which has no commit-conflict detection. @@ -171,54 +90,6 @@ def test_conditional_put_on_unsafe_rename_gated( assert ("AWS_S3_ALLOW_UNSAFE_RENAME" in options) is allow_unsafe -class TestHasCommitWithMetadata: - @pytest.mark.asyncio - async def test_returns_false_when_no_delta_table(self, helper: DeltaTableHelper): - with patch.object(helper, "get_delta_table", AsyncMock(return_value=None)): - assert await helper.has_commit_with_metadata({"run_uuid": "abc", "batch_index": "0"}) is False - - @parameterized.expand( - [(name, history, match, expected) for (name, history, match, expected) in _COMMIT_LAYOUT_CASES] - ) - @pytest.mark.asyncio - async def test_layout(self, _name: str, history: list[dict], match: dict, expected: bool): - helper = DeltaTableHelper(resource_name="t", job=MagicMock(), logger=_make_logger()) - mock_delta = MagicMock() - mock_delta.history = MagicMock(return_value=history) - - with patch.object(helper, "get_delta_table", AsyncMock(return_value=mock_delta)): - assert await helper.has_commit_with_metadata(match) is expected - - @pytest.mark.asyncio - async def test_scan_limit_passed_to_history(self, helper: DeltaTableHelper): - mock_delta = MagicMock() - mock_delta.history = MagicMock(return_value=[]) - - with patch.object(helper, "get_delta_table", AsyncMock(return_value=mock_delta)): - await helper.has_commit_with_metadata({"k": "v"}, scan_limit=123) - - mock_delta.history.assert_called_once_with(limit=123) - - -class TestHasBatchBeenCommitted: - @parameterized.expand( - [ - ("string_run_uuid_int_batch", "run-123", 5, True), - ("zero_batch_index", "run-1", 0, False), - ] - ) - @pytest.mark.asyncio - async def test_wraps_has_commit_with_metadata( - self, _name: str, run_uuid: str, batch_index: int, mocked_return: bool - ): - helper = DeltaTableHelper(resource_name="t", job=MagicMock(), logger=_make_logger()) - with patch.object(helper, "has_commit_with_metadata", AsyncMock(return_value=mocked_return)) as m: - result = await helper.has_batch_been_committed(run_uuid, batch_index) - - assert result is mocked_return - m.assert_called_once_with({"run_uuid": run_uuid, "batch_index": str(batch_index)}) - - 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]] = [ @@ -314,86 +185,6 @@ async def test_is_deltatable_transient_error_is_not_captured_but_still_reraised( assert helper.is_first_sync is False -class TestWriteToDeltalakeCommitMetadataPassThrough: - """Covers that commit_metadata is forwarded to deltalake.write_deltalake as CommitProperties.""" - - @parameterized.expand( - [ - ("no_metadata", None, None), - ("with_metadata", {"run_uuid": "abc", "batch_index": "2"}, {"run_uuid": "abc", "batch_index": "2"}), - ] - ) - @pytest.mark.asyncio - async def test_full_refresh_passes_commit_properties( - self, - _name: str, - commit_metadata: dict[str, str] | None, - expected_custom_metadata: dict[str, str] | None, - ): - import pyarrow as pa - - helper = DeltaTableHelper(resource_name="t", job=MagicMock(), logger=_make_logger()) - data = pa.table({"id": [1, 2, 3]}) - mock_delta = MagicMock() - mock_delta.schema = MagicMock(return_value=MagicMock(to_arrow=MagicMock(return_value=data.schema))) - - with ( - patch.object(helper, "get_delta_table", AsyncMock(return_value=mock_delta)), - patch(f"{_HELPER_MODULE}.evolve_delta_schema", AsyncMock(return_value=mock_delta)), - patch("deltalake.write_deltalake") as mock_write, - ): - await helper.write_to_deltalake( - data=data, - write_type="full_refresh", - should_overwrite_table=False, - primary_keys=None, - commit_metadata=commit_metadata, - ) - - assert mock_write.called - _, kwargs = mock_write.call_args - commit_properties = kwargs["commit_properties"] - if expected_custom_metadata is None: - assert commit_properties is None - else: - assert isinstance(commit_properties, deltalake.CommitProperties) - assert commit_properties.custom_metadata == expected_custom_metadata - - -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.""" - fields: list[pa.Field] = [ - pa.field("id", pa.int64()), - pa.field("name", pa.string()), - pa.field("_dlt_id", pa.string(), nullable=False), - pa.field("_dlt_load_id", pa.string(), nullable=False), - ] - if partitioned: - fields.append(pa.field(PARTITION_KEY, pa.string())) - - data_dict: dict[str, Any] = { - "id": pa.array([1, 2]), - "name": pa.array(["a", "b"]), - "_dlt_id": pa.array(["id1", "id2"]), - "_dlt_load_id": pa.array(["load1", "load1"]), - } - if partitioned: - data_dict[PARTITION_KEY] = pa.array(["p0", "p0"]) - - table = pa.table(data_dict, schema=pa.schema(fields)) - deltalake.write_deltalake(path, table, partition_by=PARTITION_KEY if partitioned else None) - return deltalake.DeltaTable(path) - - -def _v3_batch(*, partitioned: bool = False) -> pa.Table: - """Build an incoming batch the way pipeline_v3 does: no _dlt_* columns.""" - data_dict: dict[str, Any] = {"id": pa.array([3, 4]), "name": pa.array(["c", "d"])} - if partitioned: - data_dict[PARTITION_KEY] = pa.array(["p0", "p0"]) - return pa.table(data_dict) - - def _make_local_helper(delta_uri: str) -> DeltaTableHelper: """DeltaTableHelper that reads/writes a local filesystem path instead of S3.""" helper = DeltaTableHelper(resource_name="test", job=MagicMock(), logger=_make_logger()) @@ -403,342 +194,6 @@ def _make_local_helper(delta_uri: str) -> DeltaTableHelper: return helper -class TestLegacyDltTableReconciliation: - """Pipeline_v3 must handle dlt-created Delta tables with NOT NULL _dlt_* columns.""" - - def test_raw_merge_rejects_missing_non_nullable_columns(self, tmp_path: Path) -> None: - """Baseline: proves delta-rs rejects merges when non-nullable columns are absent - from the source batch. This is the root cause of the production failures.""" - delta_path = str(tmp_path / "table") - _create_legacy_delta_table(delta_path) - batch = _v3_batch() - dt = deltalake.DeltaTable(delta_path) - - with pytest.raises(Exception, match="(?i)(invalid data|non-nullable|validation|not found)"): - dt.merge( - source=batch, - source_alias="source", - target_alias="target", - predicate="source.id = target.id", - ).when_matched_update_all().when_not_matched_insert_all().execute() - - @pytest.mark.parametrize("partitioned", [False, True], ids=["flat", "partitioned"]) - @pytest.mark.asyncio - async def test_incremental_merge_into_legacy_table(self, partitioned: bool, tmp_path: Path) -> None: - delta_path = str(tmp_path / "table") - dt = _create_legacy_delta_table(delta_path, partitioned=partitioned) - - helper = _make_local_helper(delta_path) - batch = evolve_pyarrow_schema(_v3_batch(partitioned=partitioned), dt.schema()) - - result = await helper.write_to_deltalake( - data=batch, - write_type="incremental", - should_overwrite_table=False, - primary_keys=["id"], - ) - - final = result.to_pyarrow_table() - assert final.num_rows == 4 - assert set(final.column("id").to_pylist()) == {1, 2, 3, 4} - - new_rows = final.filter(pc.is_in(final.column("id"), value_set=pa.array([3, 4]))) - assert all(v == "" for v in new_rows.column("_dlt_id").to_pylist()) - assert all(v == "" for v in new_rows.column("_dlt_load_id").to_pylist()) - - @pytest.mark.asyncio - async def test_append_to_legacy_table(self, tmp_path: Path) -> None: - delta_path = str(tmp_path / "table") - dt = _create_legacy_delta_table(delta_path) - - helper = _make_local_helper(delta_path) - batch = evolve_pyarrow_schema(_v3_batch(), dt.schema()) - - result = await helper.write_to_deltalake( - data=batch, - write_type="append", - should_overwrite_table=False, - primary_keys=None, - ) - - final = result.to_pyarrow_table() - assert final.num_rows == 4 - assert set(final.column("id").to_pylist()) == {1, 2, 3, 4} - - @pytest.mark.asyncio - async def test_full_refresh_overwrite_on_legacy_table(self, tmp_path: Path) -> None: - delta_path = str(tmp_path / "table") - dt = _create_legacy_delta_table(delta_path) - - helper = _make_local_helper(delta_path) - batch = evolve_pyarrow_schema(_v3_batch(), dt.schema()) - - result = await helper.write_to_deltalake( - data=batch, - write_type="full_refresh", - should_overwrite_table=True, - primary_keys=None, - ) - - final = result.to_pyarrow_table() - assert final.num_rows == 2 - assert all(v == "" for v in final.column("_dlt_id").to_pylist()) - assert all(v == "" for v in final.column("_dlt_load_id").to_pylist()) - - @pytest.mark.asyncio - async def test_v3_native_table_still_merges(self, tmp_path: Path) -> None: - delta_path = str(tmp_path / "table") - fields: list[pa.Field] = [pa.field("id", pa.int64()), pa.field("name", pa.string())] - schema = pa.schema(fields) - deltalake.write_deltalake(delta_path, pa.table({"id": [1, 2], "name": ["a", "b"]}, schema=schema)) - - helper = _make_local_helper(delta_path) - batch = pa.table({"id": [3], "name": ["c"]}) - - result = await helper.write_to_deltalake( - data=batch, - write_type="incremental", - should_overwrite_table=False, - primary_keys=["id"], - ) - - final = result.to_pyarrow_table() - assert final.num_rows == 3 - assert set(final.column("id").to_pylist()) == {1, 2, 3} - - -class TestAppendDecimalReconciliation: - """Appending a decimal column that outgrew decimal128 must reconcile to the stored type. - - A batch whose numeric column exceeds decimal128 is promoted to decimal256, which - `evolve_pyarrow_schema` renders to text for the Delta write. Arrow emits scientific - notation for scale-heavy zeros (e.g. '0E-18'), which delta-rs can't parse back into - the stored decimal — an opaque, infinitely-retrying DeltaError on the append path. - """ - - def _seed_decimal_table(self, delta_path: str) -> deltalake.DeltaTable: - table = pa.table( - {"id": pa.array([1], type=pa.int64()), "amount": pa.array([Decimal("1.5")], type=pa.decimal128(38, 10))} - ) - deltalake.write_deltalake(delta_path, table) - return deltalake.DeltaTable(delta_path) - - @pytest.mark.asyncio - async def test_scale_heavy_batch_is_rounded_to_stored_type(self, tmp_path: Path) -> None: - delta_path = str(tmp_path / "table") - dt = self._seed_decimal_table(delta_path) - helper = _make_local_helper(delta_path) - - # Values fit decimal128's integer budget but carry more scale than the stored column, - # so they land as decimal256 and evolve renders them to text (the zero as '0E-18'). - batch = evolve_pyarrow_schema( - pa.table( - { - "id": pa.array([2, 3], type=pa.int64()), - "amount": pa.array( - [Decimal("0.12345678901234567890"), Decimal("0E-18")], type=pa.decimal256(76, 20) - ), - } - ), - dt.schema(), - ) - assert pa.types.is_string(batch.schema.field("amount").type) - - result = await helper.write_to_deltalake( - data=batch, write_type="append", should_overwrite_table=False, primary_keys=None - ) - - final = result.to_pyarrow_table() - assert final.schema.field("amount").type == pa.decimal128(38, 10) - assert set(final.column("id").to_pylist()) == {1, 2, 3} - assert Decimal("0") in final.column("amount").to_pylist() - - @pytest.mark.asyncio - async def test_integer_overflow_batch_raises_clean_non_retryable(self, tmp_path: Path) -> None: - delta_path = str(tmp_path / "table") - dt = self._seed_decimal_table(delta_path) - helper = _make_local_helper(delta_path) - - batch = evolve_pyarrow_schema( - pa.table( - { - "id": pa.array([2, 3], type=pa.int64()), - "amount": pa.array([Decimal("1" + "0" * 35 + ".5"), Decimal("0E-18")], type=pa.decimal256(76, 18)), - } - ), - dt.schema(), - ) - - with pytest.raises(SchemaColumnTypeChangedException): - await helper.write_to_deltalake( - data=batch, write_type="append", should_overwrite_table=False, primary_keys=None - ) - - -class TestSchemaEvolutionNullability: - """A column added mid-table-lifetime always predates its own addition: every file the - table already holds was written without it, so `optimize.compact()` must be able to - treat those rows as null for that column. If schema evolution adds the column as NOT - NULL — which happens whenever the batch that introduces it has no nulls, since delta-rs - takes the new field's nullability straight from the incoming Arrow field — compaction - later fails with "Non-nullable column '' is missing from the physical schema".""" - - @pytest.mark.asyncio - async def test_compact_survives_a_column_added_by_an_all_non_null_batch(self, tmp_path: Path) -> None: - delta_path = str(tmp_path / "table") - deltalake.write_deltalake(delta_path, pa.table({"id": pa.array([1, 2], type=pa.int64())})) - - helper = _make_local_helper(delta_path) - - # The incoming field is non-nullable because every value in *this* batch is - # non-null — exactly how upstream Arrow construction infers it, unrelated to - # whether the column can appear in prior or future batches. - fields: list[pa.Field] = [pa.field("id", pa.int64()), pa.field("status", pa.string(), nullable=False)] - batch_schema = pa.schema(fields) - batch = pa.table( - {"id": pa.array([3, 4], type=pa.int64()), "status": pa.array(["ok", "ok"])}, schema=batch_schema - ) - - result = await helper.write_to_deltalake( - data=batch, write_type="append", should_overwrite_table=False, primary_keys=None - ) - status_field = next(f for f in result.schema().fields if f.name == "status") - assert status_field.nullable is True - - await DeltaMaintenance(helper).compact_table() - - final = result.to_pyarrow_table() - by_id = dict(zip(final.column("id").to_pylist(), final.column("status").to_pylist())) - assert by_id == {1: None, 2: None, 3: "ok", 4: "ok"} - - -class TestIncrementalBatchDeduplication: - """Duplicate PKs in a source batch must never reach the Delta write. - - `when_not_matched_insert_all` inserts every unmatched source row, so a batch with a - repeated PK seeds duplicate rows in the table; every later merge then multi-matches - those rows and the join blows up (the OOM loop seen with sources whose primary keys - aren't actually unique). - """ - - @parameterized.expand( - [ - ("keep_first", "first", ["a1", "b1"]), - ("keep_last", "last", ["a2", "b1"]), - ] - ) - def test_first_per_pk_table_keep_modes(self, _name, keep, expected_names): - table = pa.table({"id": [1, 1, 2], "name": ["a1", "a2", "b1"]}) - - result = first_per_pk_table(table, ["id"], keep=keep).sort_by("id") - - assert result.column("id").to_pylist() == [1, 2] - assert result.column("name").to_pylist() == expected_names - - @pytest.mark.asyncio - async def test_incremental_merge_dedupes_duplicate_source_rows(self, tmp_path: Path) -> None: - delta_path = str(tmp_path / "table") - deltalake.write_deltalake(delta_path, pa.table({"id": [1], "name": ["old"]})) - - helper = _make_local_helper(delta_path) - # id=2 appears twice in one batch — without dedup both copies get inserted. - batch = pa.table({"id": [1, 2, 2], "name": ["updated", "first_copy", "second_copy"]}) - - result = await helper.write_to_deltalake( - data=batch, - write_type="incremental", - should_overwrite_table=False, - primary_keys=["id"], - ) - - final = result.to_pyarrow_table().sort_by("id") - assert final.column("id").to_pylist() == [1, 2] - # The last occurrence of a duplicated key carries the freshest data. - assert final.column("name").to_pylist() == ["updated", "second_copy"] - cast(AsyncMock, helper._logger.awarning).assert_awaited_once() - - @pytest.mark.asyncio - async def test_first_sync_append_dedupes_duplicate_source_rows(self, tmp_path: Path) -> None: - delta_path = str(tmp_path / "table") - - helper = _make_local_helper(delta_path) - batch = pa.table({"id": [1, 1], "name": ["first_copy", "second_copy"]}) - - result = await helper.write_to_deltalake( - data=batch, - write_type="incremental", - should_overwrite_table=False, - primary_keys=["id"], - ) - - final = result.to_pyarrow_table() - assert final.column("id").to_pylist() == [1] - assert final.column("name").to_pylist() == ["second_copy"] - - -class TestUnpartitionedTableWithPartitionKeyColumn: - """A Delta table can carry `_ph_partition_key` in its schema while its - partition_columns metadata is empty `[]` — e.g. the SchemaMismatchError fallback in - write_to_deltalake rewrites with partition_by=None while the column is still in the - data, or evolve_pyarrow_schema re-adds the column to a batch headed for an - unpartitioned table. write_to_deltalake derives partitioning from column *presence*, - so it then passes partition_by=_ph_partition_key against a table delta-rs considers - unpartitioned and raises: - "Specified table partitioning does not match table partitioning: expected: [], got: [_ph_partition_key]" - """ - - def _seed_unpartitioned_table_with_partition_column(self, delta_path: str) -> None: - # _ph_partition_key is a plain column; the table is NOT partitioned by it. - deltalake.write_deltalake( - delta_path, - pa.table({"id": pa.array([1, 2]), PARTITION_KEY: pa.array(["p0", "p0"])}), - partition_by=None, - ) - dt = deltalake.DeltaTable(delta_path) - assert dt.metadata().partition_columns == [] - assert PARTITION_KEY in dt.schema().to_arrow().names - - @pytest.mark.parametrize( - "write_type,primary_keys,should_overwrite,expected_ids", - [ - # append/incremental keep the existing rows; full_refresh overwrites them. Each - # routes through a distinct write branch, all of which previously raised against - # the unpartitioned-but-column-present table. - ("append", None, False, {1, 2, 3, 4}), - ("incremental", ["id"], False, {1, 2, 3, 4}), - ("full_refresh", None, True, {2, 3, 4}), - ], - ids=["append", "incremental_merge", "full_refresh_overwrite"], - ) - @pytest.mark.asyncio - async def test_write_does_not_partition_unpartitioned_table( - self, - write_type: str, - primary_keys: list[str] | None, - should_overwrite: bool, - expected_ids: set[int], - tmp_path: Path, - ) -> None: - delta_path = str(tmp_path / "table") - self._seed_unpartitioned_table_with_partition_column(delta_path) - - helper = _make_local_helper(delta_path) - # id=2 already exists (merge updates it); id=3,4 are new. - batch = pa.table({"id": pa.array([2, 3, 4]), PARTITION_KEY: pa.array(["p0", "p0", "p0"])}) - - result = await helper.write_to_deltalake( - data=batch, - write_type=write_type, # type: ignore[arg-type] - should_overwrite_table=should_overwrite, - primary_keys=primary_keys, - ) - - final = result.to_pyarrow_table() - assert set(final.column("id").to_pylist()) == expected_ids - # The table stays unpartitioned — we don't fight its existing layout. - assert result.metadata().partition_columns == [] - - class TestRealignDecimalBuffers: """delta-rs aborts the worker on 8-byte-aligned Decimal128 buffers; we realign them to pyarrow's 64-byte allocator before any Delta write. See delta-io/delta-rs#3884.""" @@ -810,45 +265,6 @@ def test_empty_decimal_table(self) -> None: assert result.schema == table.schema -class TestWriteMisalignedDecimalEndToEnd: - """Writes a misaligned-decimal batch through the real delta-rs write path. Without the - realignment guard, delta-rs would abort the process; with it, the write succeeds.""" - - @pytest.mark.parametrize( - "write_type,should_overwrite", - [("full_refresh", True), ("append", False), ("incremental", False)], - ) - @pytest.mark.asyncio - async def test_write_misaligned_decimal_to_local_delta( - self, write_type: str, should_overwrite: bool, tmp_path: Path - ) -> None: - delta_path = str(tmp_path / "table") - # Seed the table so incremental/append have an existing target to write into. - deltalake.write_deltalake( - delta_path, - pa.table({"id": pa.array([1, 2]), "amount": _decimal_array([5, 6], misaligned=False)}), - ) - - helper = _make_local_helper(delta_path) - batch = pa.table({"id": pa.array([3, 4]), "amount": _decimal_array([7, 8], misaligned=True)}) - assert _table_is_misaligned(batch) is True - - result = await helper.write_to_deltalake( - data=batch, - write_type=write_type, # type: ignore[arg-type] - should_overwrite_table=should_overwrite, - primary_keys=["id"] if write_type == "incremental" else None, - ) - - final = result.to_pyarrow_table() - amounts = set(final.column("amount").to_pylist()) - if should_overwrite: - assert set(final.column("id").to_pylist()) == {3, 4} - else: - assert {3, 4}.issubset(set(final.column("id").to_pylist())) - assert {7, 8}.issubset(amounts) - - class TestIsTableCorrupted: _MODULE = "products.warehouse_sources.backend.temporal.data_imports.pipelines.core.delta_table_helper" @@ -881,225 +297,3 @@ async def test_is_table_corrupted(self, _name: str, is_delta: bool, open_exc: Ex result = await helper.is_table_corrupted() assert result is expected - - -class TestNullSafeMergePredicate: - """The incremental-merge match must be NULL-safe. - - Regression for the duplicate-accumulation bug found by the deltalite shadow canary: composite - keys with nullable columns (e.g. GoogleAds report resources keyed on `segments.*`) matched with - bare `source.c = target.c` never match on NULL (`NULL = NULL` is NULL), so the row is re-inserted - on every incremental sync and the table silently grows. - """ - - def test_predicate_ops_are_null_safe(self): - assert _merge_predicate_ops(["id", "seg"]) == [ - "(source.id IS NOT DISTINCT FROM target.id)", - "(source.seg IS NOT DISTINCT FROM target.seg)", - ] - - @staticmethod - def _seed_then_merge(path: Path, predicate_ops: list[str]) -> pa.Table: - # Seed one row whose composite key has a NULL component, then merge the same key with a new value. - seed = pa.table( - { - "id": pa.array([1], pa.int64()), - "seg": pa.array([None], pa.string()), - "val": pa.array(["a"], pa.string()), - } - ) - deltalake.write_deltalake(str(path), seed, mode="overwrite") - source = pa.table( - { - "id": pa.array([1], pa.int64()), - "seg": pa.array([None], pa.string()), - "val": pa.array(["b"], pa.string()), - } - ) - ( - deltalake.DeltaTable(str(path)) - .merge(source=source, source_alias="source", target_alias="target", predicate=" AND ".join(predicate_ops)) - .when_matched_update_all() - .when_not_matched_insert_all() - .execute() - ) - return deltalake.DeltaTable(str(path)).to_pyarrow_table() - - def test_null_composite_key_row_matches_instead_of_duplicating(self, tmp_path): - result = self._seed_then_merge(tmp_path / "safe", _merge_predicate_ops(["id", "seg"])) - assert result.num_rows == 1 - assert result.column("val").to_pylist() == ["b"] - - def test_non_null_key_still_matches(self, tmp_path): - # The null-safe form must not change behaviour for ordinary (non-NULL) keys. - seed = pa.table({"id": pa.array([1], pa.int64()), "seg": pa.array(["MOBILE"]), "val": pa.array(["a"])}) - deltalake.write_deltalake(str(tmp_path / "nn"), seed, mode="overwrite") - source = pa.table({"id": pa.array([1], pa.int64()), "seg": pa.array(["MOBILE"]), "val": pa.array(["b"])}) - ( - deltalake.DeltaTable(str(tmp_path / "nn")) - .merge( - source=source, - source_alias="source", - target_alias="target", - predicate=" AND ".join(_merge_predicate_ops(["id", "seg"])), - ) - .when_matched_update_all() - .when_not_matched_insert_all() - .execute() - ) - result = deltalake.DeltaTable(str(tmp_path / "nn")).to_pyarrow_table() - assert result.num_rows == 1 - assert result.column("val").to_pylist() == ["b"] - - def test_bare_equality_duplicates_null_key_row(self, tmp_path): - # Documents the pre-fix behaviour the null-safe predicate corrects. - result = self._seed_then_merge(tmp_path / "unsafe", ["source.id = target.id", "source.seg = target.seg"]) - assert result.num_rows == 2 - - -class TestDeltaliteWritePath: - """Phase 2: deltalite performs the real incremental merge, gated solely by a per-schema flag, with - a hard fallback to the delta-rs MERGE so a deltalite failure can never fail a sync.""" - - _FLAG = ( - "products.warehouse_sources.backend.temporal.data_imports.pipelines.core." - "deltalite_write.is_deltalite_write_enabled" - ) - - @pytest.fixture(autouse=True) - def _preload_write_metrics(self): - # _write_via_deltalite lazily imports the pipeline_v3 metrics module. These tests fake the - # `deltalite` module via patch.dict(sys.modules, ...), which on exit restores the enter-time - # snapshot and thus evicts any module first imported *inside* the block. If the metrics module - # were first loaded there, the next test would re-execute it and hit "Duplicated timeseries" in - # the global Prometheus registry. Loading it here (at setup, before any patch.dict) keeps it in - # the snapshot so it survives. Imported at runtime — not module top — to keep the heavy - # pipeline_v3 chain off collection. - from products.warehouse_sources.backend.temporal.data_imports.pipelines.pipeline_v3.load import ( # noqa: F401 - metrics, - ) - - def _helper(self) -> DeltaTableHelper: - return DeltaTableHelper(resource_name="t", job=MagicMock(team_id=2, schema_id="sch-1"), logger=_make_logger()) - - async def _call(self, helper: DeltaTableHelper) -> bool: - return await helper._write_via_deltalite( - existing_delta_table=MagicMock(), - data=pa.table({"id": pa.array([1], pa.int64())}), - normalized_primary_keys=["id"], - use_partitioning=False, - commit_metadata={"run_uuid": "abc"}, - ) - - @pytest.mark.asyncio - async def test_skips_without_primary_keys(self): - # No primary keys => nothing to key an upsert on; fall back without even evaluating the flag. - with patch(self._FLAG) as flag: - wrote = await self._helper()._write_via_deltalite( - existing_delta_table=MagicMock(), - data=pa.table({"id": pa.array([1], pa.int64())}), - normalized_primary_keys=[], - use_partitioning=False, - commit_metadata=None, - ) - assert wrote is False - flag.assert_not_called() - - def test_write_stats_flattens_scalar_getters_only(self): - # Enumerates scalar attributes (so future crate fields flow through) and drops methods/non-scalars. - stats = SimpleNamespace(version=7, rows_inserted=2, files_added=1, _private=9) - stats.helper = lambda: None # callable attribute must be ignored - assert _deltalite_write_stats(stats) == {"version": 7, "rows_inserted": 2, "files_added": 1} - - @pytest.mark.asyncio - async def test_falls_back_when_flag_disabled(self): - with patch(self._FLAG, return_value=False): - assert await self._call(self._helper()) is False - - @pytest.mark.asyncio - async def test_writes_via_deltalite_when_enabled(self): - logger = _make_logger() # captured so we can inspect the structured log without hitting the typed attr - helper = DeltaTableHelper(resource_name="t", job=MagicMock(team_id=2, schema_id="sch-1"), logger=logger) - existing = MagicMock() - # SimpleNamespace stands in for the pyo3 UpsertStats: predictable scalar getters for the structured log. - fake_stats = SimpleNamespace( - version=5, partitions_touched=1, rows_inserted=3, rows_updated=2, rows_copied=10, null_pk_rows=0 - ) - fake_table = MagicMock() - fake_table.upsert.return_value = fake_stats - fake_deltalite = MagicMock() - fake_deltalite.DeltaLiteTable.open.return_value = fake_table - with ( - patch(self._FLAG, return_value=True), - patch.dict("sys.modules", {"deltalite": fake_deltalite}), - patch.object(helper, "_get_delta_table_uri", AsyncMock(return_value="s3://b/t")), - patch.object(helper, "_get_credentials", return_value={"AWS_REGION": "us-east-1"}), - ): - wrote = await helper._write_via_deltalite( - existing_delta_table=existing, - data=pa.table({"id": pa.array([1], pa.int64())}), - normalized_primary_keys=["id"], - use_partitioning=True, - commit_metadata={"run_uuid": "abc"}, - ) - assert wrote is True - fake_deltalite.DeltaLiteTable.open.assert_called_once_with("s3://b/t", {"AWS_REGION": "us-east-1"}) - fake_table.upsert.assert_called_once() - # PARTITION_KEY is passed as the partition arg when the table is partitioned. - assert fake_table.upsert.call_args.args[2] == PARTITION_KEY - existing.update_incremental.assert_called_once() - # The commit is logged with the UpsertStats fields as structured keys + a duration, so it's parseable. - logger.ainfo.assert_called_once() - log_kwargs = logger.ainfo.call_args.kwargs - assert log_kwargs["version"] == 5 - assert log_kwargs["rows_inserted"] == 3 - assert log_kwargs["partitions_touched"] == 1 - assert "duration_ms" in log_kwargs - - @pytest.mark.asyncio - async def test_falls_back_when_deltalite_raises(self): - helper = self._helper() - fake_table = MagicMock() - fake_table.upsert.side_effect = RuntimeError("commit conflict, retries exhausted") - fake_deltalite = MagicMock() - fake_deltalite.DeltaLiteTable.open.return_value = fake_table - with ( - patch(self._FLAG, return_value=True), - patch.dict("sys.modules", {"deltalite": fake_deltalite}), - patch.object(helper, "_get_delta_table_uri", AsyncMock(return_value="s3://b/t")), - patch.object(helper, "_get_credentials", return_value={}), - ): - wrote = await self._call(helper) - assert wrote is False # deltalite blew up -> caller falls through to the delta-rs MERGE - - @parameterized.expand([("refresh",), ("log",)]) - @pytest.mark.asyncio - async def test_post_commit_failure_does_not_fall_back(self, failing_step: str): - # Once the upsert commits, NO post-commit step (handle refresh, log, metric) may raise into the - # caller — that would return False / bubble up and re-run the MERGE on top of deltalite's commit. - logger = _make_logger() # set the side effect on the mock before it becomes the typed _logger attr - existing = MagicMock() - if failing_step == "refresh": - existing.update_incremental.side_effect = RuntimeError("post-commit refresh boom") - else: - logger.ainfo.side_effect = RuntimeError("post-commit log boom") - helper = DeltaTableHelper(resource_name="t", job=MagicMock(team_id=2, schema_id="sch-1"), logger=logger) - fake_table = MagicMock() - fake_table.upsert.return_value = MagicMock(version=5, rows_inserted=1, rows_updated=0, rows_copied=0) - fake_deltalite = MagicMock() - fake_deltalite.DeltaLiteTable.open.return_value = fake_table - with ( - patch(self._FLAG, return_value=True), - patch.dict("sys.modules", {"deltalite": fake_deltalite}), - patch.object(helper, "_get_delta_table_uri", AsyncMock(return_value="s3://b/t")), - patch.object(helper, "_get_credentials", return_value={}), - ): - wrote = await helper._write_via_deltalite( - existing_delta_table=existing, - data=pa.table({"id": pa.array([1], pa.int64())}), - normalized_primary_keys=["id"], - use_partitioning=False, - commit_metadata=None, - ) - assert wrote is True # committed; the post-commit failure is swallowed - fake_table.upsert.assert_called_once() 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 e263ce5620bd..bf4abe986ad5 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 @@ -53,6 +53,7 @@ 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.writer import DeltaWriter 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 @@ -343,7 +344,7 @@ async def _process_pa_table( should_overwrite_table = index == 0 and not resuming_sync - delta_table = await self._delta_table_helper.write_to_deltalake( + delta_table = await DeltaWriter(self._delta_table_helper).write( pa_table, write_type, should_overwrite_table=should_overwrite_table, diff --git a/products/warehouse_sources/backend/temporal/data_imports/pipelines/pipeline_v3/load/idempotency.py b/products/warehouse_sources/backend/temporal/data_imports/pipelines/pipeline_v3/load/idempotency.py index 1e9d810e7caa..38af50dd19cc 100644 --- a/products/warehouse_sources/backend/temporal/data_imports/pipelines/pipeline_v3/load/idempotency.py +++ b/products/warehouse_sources/backend/temporal/data_imports/pipelines/pipeline_v3/load/idempotency.py @@ -8,6 +8,7 @@ from posthog.exceptions_capture import capture_exception from posthog.redis import get_client +from products.warehouse_sources.backend.temporal.data_imports.pipelines.core.delta.writer import DeltaWriter from products.warehouse_sources.backend.temporal.data_imports.pipelines.core.delta_table_helper import DeltaTableHelper logger = structlog.get_logger(__name__) @@ -62,7 +63,7 @@ def is_batch_already_processed( Slow path (post-crash recovery): if Redis has no flag and a `DeltaTableHelper` is provided, scan recent delta commits for a commit whose userMetadata matches this (run_uuid, batch_index). This catches the narrow writer-crash window - between `write_to_deltalake` committing and `mark_batch_as_processed` running — + between `DeltaWriter.write` committing and `mark_batch_as_processed` running — on Kafka redelivery we'd otherwise re-write the same batch and produce duplicate rows. """ @@ -76,7 +77,7 @@ def is_batch_already_processed( return False try: - return async_to_sync(delta_table_helper.has_batch_been_committed)(run_uuid, batch_index) + return async_to_sync(DeltaWriter(delta_table_helper).has_batch_been_committed)(run_uuid, batch_index) except Exception as e: # Failing open here would re-enable the duplicate-write race we're fixing, # so we log and surface the error to the caller (which will retry the message). diff --git a/products/warehouse_sources/backend/temporal/data_imports/pipelines/pipeline_v3/load/processor.py b/products/warehouse_sources/backend/temporal/data_imports/pipelines/pipeline_v3/load/processor.py index 77ac295dcf20..c2f77f5aa55c 100644 --- a/products/warehouse_sources/backend/temporal/data_imports/pipelines/pipeline_v3/load/processor.py +++ b/products/warehouse_sources/backend/temporal/data_imports/pipelines/pipeline_v3/load/processor.py @@ -40,6 +40,7 @@ ) 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.scd2 import Scd2DeltaWriter +from products.warehouse_sources.backend.temporal.data_imports.pipelines.core.delta.writer import DeltaWriter 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 ( @@ -664,7 +665,7 @@ def process_message( # Build the helper early so the idempotency check can use it as a # delta-history fallback when the Redis dedup flag is missing — the case - # where the writer crashed between `write_to_deltalake` committing and + # where the writer crashed between `DeltaWriter.write` committing and # `mark_batch_as_processed` being called. job = ExternalDataJob.objects.prefetch_related("schema", "schema__source", "schema__table").get( id=export_signal.job_id @@ -813,7 +814,7 @@ def process_message( with DELTA_WRITE_DURATION_SECONDS.labels( team_id=team_id_str, schema_id=schema_id_str, write_type=write_type ).time(): - delta_table = async_to_sync(delta_table_helper.write_to_deltalake)( + delta_table = async_to_sync(DeltaWriter(delta_table_helper).write)( data=pa_table, write_type=write_type, should_overwrite_table=should_overwrite_table, diff --git a/products/warehouse_sources/backend/temporal/data_imports/pipelines/pipeline_v3/load/test_idempotency.py b/products/warehouse_sources/backend/temporal/data_imports/pipelines/pipeline_v3/load/test_idempotency.py index 0fc4921cd2d2..f0857e6b84e8 100644 --- a/products/warehouse_sources/backend/temporal/data_imports/pipelines/pipeline_v3/load/test_idempotency.py +++ b/products/warehouse_sources/backend/temporal/data_imports/pipelines/pipeline_v3/load/test_idempotency.py @@ -9,9 +9,8 @@ mark_batch_as_processed, ) -REDIS_CLIENT_PATH = ( - "products.warehouse_sources.backend.temporal.data_imports.pipelines.pipeline_v3.load.idempotency.get_redis_client" -) +_IDEMPOTENCY_MODULE = "products.warehouse_sources.backend.temporal.data_imports.pipelines.pipeline_v3.load.idempotency" +REDIS_CLIENT_PATH = f"{_IDEMPOTENCY_MODULE}.get_redis_client" def _redis_client(exists_value: int | None) -> MagicMock | None: @@ -52,6 +51,13 @@ def test_key_format(self): class TestIsBatchAlreadyProcessed: + @pytest.fixture(autouse=True) + def _writer_wraps_helper(self): + # The slow path wraps the helper in DeltaWriter(helper).has_batch_been_committed; an identity + # stand-in keeps the helper-shaped mocks below driving the same decision matrix. + with patch(f"{_IDEMPOTENCY_MODULE}.DeltaWriter", side_effect=lambda helper: helper): + yield + @parameterized.expand( [ # (name, redis_exists, helper_state, expected_result) diff --git a/products/warehouse_sources/backend/temporal/data_imports/pipelines/pipeline_v3/load/test_processor.py b/products/warehouse_sources/backend/temporal/data_imports/pipelines/pipeline_v3/load/test_processor.py index 333bf9d0164d..81c2c5bb9452 100644 --- a/products/warehouse_sources/backend/temporal/data_imports/pipelines/pipeline_v3/load/test_processor.py +++ b/products/warehouse_sources/backend/temporal/data_imports/pipelines/pipeline_v3/load/test_processor.py @@ -231,6 +231,7 @@ class TestProcessMessageOwnershipGate: @patch(f"{_PROCESSOR}.posthoganalytics") @patch(f"{_PROCESSOR}.read_parquet", return_value=pa.table({"id": [1]})) @patch(f"{_PROCESSOR}.is_batch_already_processed", return_value=False) + @patch(f"{_PROCESSOR}.DeltaWriter") @patch(f"{_PROCESSOR}.Scd2DeltaWriter") @patch(f"{_PROCESSOR}.DeltaTableHelper") @patch(f"{_PROCESSOR}.ExternalDataJob") @@ -243,13 +244,13 @@ def test_lost_ownership_blocks_delta_write( mock_job_model: MagicMock, mock_helper_cls: MagicMock, mock_scd2_cls: MagicMock, + mock_writer_cls: MagicMock, _already: MagicMock, _read: MagicMock, _analytics: MagicMock, ) -> None: helper = mock_helper_cls.return_value helper.get_delta_table = AsyncMock(return_value=None) - helper.write_to_deltalake = AsyncMock() mock_job_model.objects.prefetch_related.return_value.get.return_value = MagicMock() def verify_ownership() -> None: @@ -258,7 +259,7 @@ def verify_ownership() -> None: with pytest.raises(_LeaseLost): process_message(_message(), verify_ownership=verify_ownership) - helper.write_to_deltalake.assert_not_called() + mock_writer_cls.return_value.write.assert_not_called() mock_scd2_cls.return_value.write.assert_not_called() @patch(f"{_PROCESSOR}.posthoganalytics") @@ -479,6 +480,7 @@ class TestPostImportTrigger: @patch(f"{_PROCESSOR}.run_post_load_operations", new_callable=AsyncMock, return_value="folder") @patch(f"{_PROCESSOR}.read_parquet", return_value=pa.table({"id": [1]})) @patch(f"{_PROCESSOR}.is_batch_already_processed", return_value=False) + @patch(f"{_PROCESSOR}.DeltaWriter") @patch(f"{_PROCESSOR}.DeltaTableHelper") @patch(f"{_PROCESSOR}.ExternalDataJob") @patch(f"{_PROCESSOR}.s3fs") @@ -487,6 +489,7 @@ def test_final_batch_triggers_post_import_once( _s3fs: MagicMock, mock_job_model: MagicMock, mock_helper_cls: MagicMock, + mock_writer_cls: MagicMock, _already: MagicMock, _read: MagicMock, _post_load: AsyncMock, @@ -501,7 +504,7 @@ def test_final_batch_triggers_post_import_once( delta_table.file_uris.return_value = [] helper = mock_helper_cls.return_value helper.get_delta_table = AsyncMock(return_value=None) - helper.write_to_deltalake = AsyncMock(return_value=delta_table) + mock_writer_cls.return_value.write = AsyncMock(return_value=delta_table) mock_job_model.objects.prefetch_related.return_value.get.return_value = MagicMock() process_message(_message(is_final_batch=True)) diff --git a/products/warehouse_sources/backend/temporal/data_imports/sources/mysql/mysql.py b/products/warehouse_sources/backend/temporal/data_imports/sources/mysql/mysql.py index abe1d4facade..c0e29ec0ad2d 100644 --- a/products/warehouse_sources/backend/temporal/data_imports/sources/mysql/mysql.py +++ b/products/warehouse_sources/backend/temporal/data_imports/sources/mysql/mysql.py @@ -1468,7 +1468,7 @@ def get_rows() -> Iterator[Any]: # the retry path can't safely restart from the original # cursor: the delta merge only dedupes rows for `incremental` # writes into an existing table (see - # `delta_table_helper.write_to_deltalake`), so full-refresh + # `DeltaWriter.write`), so full-refresh # and first-ever-sync scenarios would get silent duplicates # on replay. The observed bad-plan failure fails before any # rows stream, so this guard is defensive — it enforces the 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 a92f76373db5..4bb08e0fd833 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 @@ -3402,7 +3402,7 @@ async def test_v3_delta_commit_metadata_and_idempotency_fallback(team, stripe_cu the Redis idempotency flag is missing. This exercises the writer-side idempotency gap: if the writer crashes between - `write_to_deltalake` committing and `mark_batch_as_processed` running, Kafka redelivery + `DeltaWriter.write` committing and `mark_batch_as_processed` running, Kafka redelivery would otherwise re-write the same batch and produce duplicate rows. The delta-history fallback closes that gap. """