diff --git a/pyproject.toml b/pyproject.toml index 1361728..367b35b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -65,6 +65,7 @@ home = ".prefect" [dependency-groups] dev = [ "aiosqlite>=0.22.1", + "freezegun>=1.5.5", "pre-commit>=4.6.0", "prefect-docker>=0.7.1", "pytest>=9.0.3", diff --git a/src/data_index/inventory_source/delta_iceberg_table.py b/src/data_index/inventory_source/delta_iceberg_table.py index 924692e..010bba2 100644 --- a/src/data_index/inventory_source/delta_iceberg_table.py +++ b/src/data_index/inventory_source/delta_iceberg_table.py @@ -1,5 +1,6 @@ from __future__ import annotations +import datetime import typing import polars @@ -8,6 +9,23 @@ from data_index.inventory_source import iceberg_table +class LookbackConfig(pydantic.BaseModel): + days: int = pydantic.Field(default=0) + hours: int = pydantic.Field(default=0) + column_name: str + + @property + def lookback_timestamp(self): + return ( + datetime.datetime.now(tz=datetime.timezone.utc) + - datetime.timedelta(days=self.days, hours=self.hours) + ).isoformat() + + @property + def lookback_filter(self): + return f"{self.column_name} >= '{self.lookback_timestamp}'" + + class DeltaIcebergTableInventorySource(pydantic.BaseModel): type: typing.Literal["delta_iceberg_table"] = pydantic.Field( default="delta_iceberg_table" @@ -15,13 +33,29 @@ class DeltaIcebergTableInventorySource(pydantic.BaseModel): source: iceberg_table.IcebergTableInventorySource sink: iceberg_table.IcebergTableInventorySource - left_on: list[str] = pydantic.Field(default=("bucket", "key", "version_id")) - right_on: list[str] = pydantic.Field(default=("bucket", "key", "version_id")) + lookback_config: LookbackConfig | None = pydantic.Field( + default=None, + description="The number of hours to look back from the current time in the given column", + ) + left_on: list[str] = pydantic.Field(default=["bucket", "key", "version_id"]) + right_on: list[str] = pydantic.Field(default=["bucket", "key", "version_id"]) def inventory( self, ) -> polars.DataFrame: + # Set the lookback configuration filters on the source + if self.lookback_config: + # Append + if self.source.table_scan_config.row_filter: + self.source.table_scan_config.row_filter = f"({self.source.table_scan_config.row_filter}) AND {self.lookback_config.lookback_filter}" + + # Write + else: + self.source.table_scan_config.row_filter = ( + self.lookback_config.lookback_filter + ) + source_df = self.source.inventory() sink_df = self.sink.inventory() diff --git a/tests/test_delta_iceberg_table_source.py b/tests/test_delta_iceberg_table_source.py index 8387be9..6424c12 100644 --- a/tests/test_delta_iceberg_table_source.py +++ b/tests/test_delta_iceberg_table_source.py @@ -1,9 +1,10 @@ +import datetime from unittest.mock import MagicMock +import freezegun import polars as pl -# Assuming your class is imported like this: -import data_index.inventory_source +import data_index.inventory_source.delta_iceberg_table def test_delta_iceberg_table_inventory_anti_join(): @@ -38,10 +39,8 @@ def test_delta_iceberg_table_inventory_anti_join(): mock_sink = MagicMock() mock_sink.inventory.return_value = sink_data - delta_source = ( - data_index.inventory_source.DeltaIcebergTableInventorySource.model_construct( - source=mock_source, sink=mock_sink - ) + delta_source = data_index.inventory_source.delta_iceberg_table.DeltaIcebergTableInventorySource.model_construct( + source=mock_source, sink=mock_sink ) result_df = delta_source.inventory() @@ -74,13 +73,11 @@ def test_delta_iceberg_table_inventory_custom_join_keys(): mock_sink = MagicMock() mock_sink.inventory.return_value = sink_data - delta_source = ( - data_index.inventory_source.DeltaIcebergTableInventorySource.model_construct( - source=mock_source, - sink=mock_sink, - left_on=["src_key"], - right_on=["snk_key"], - ) + delta_source = data_index.inventory_source.delta_iceberg_table.DeltaIcebergTableInventorySource.model_construct( + source=mock_source, + sink=mock_sink, + left_on=["src_key"], + right_on=["snk_key"], ) result_df = delta_source.inventory() @@ -97,12 +94,100 @@ def test_delta_iceberg_table_pydantic_defaults(): mock_source = MagicMock() mock_sink = MagicMock() - delta_source = ( - data_index.inventory_source.DeltaIcebergTableInventorySource.model_construct( - source=mock_source, sink=mock_sink - ) + delta_source = data_index.inventory_source.delta_iceberg_table.DeltaIcebergTableInventorySource.model_construct( + source=mock_source, sink=mock_sink ) assert delta_source.type == "delta_iceberg_table" - assert delta_source.left_on == ("bucket", "key", "version_id") - assert delta_source.right_on == ("bucket", "key", "version_id") + assert delta_source.left_on == ["bucket", "key", "version_id"] + assert delta_source.right_on == ["bucket", "key", "version_id"] + + +class MockDatetime: + # This makes MockDatetime.datetime return MockDatetime itself, + # satisfying the source code's `datetime.datetime.now()` syntax chain. + @property + def datetime(self): + return self + + def now(self, tz=None): + return datetime.datetime(2026, 7, 1, 12, 0, 0, tzinfo=datetime.timezone.utc) + + @property + def timezone(self): + return datetime.timezone + + +# 2. Patch the target module's reference to datetime.datetime +@freezegun.freeze_time(time_to_freeze="2026-07-01T12:00:00") +def test_lookback_config_time_generation(): + config = data_index.inventory_source.delta_iceberg_table.LookbackConfig( + days=1, hours=2, column_name="last_modified" + ) + + expected_iso = "2026-06-30T10:00:00+00:00" + + assert config.lookback_timestamp == expected_iso + assert config.lookback_filter == "last_modified >= '2026-06-30T10:00:00+00:00'" + + +@freezegun.freeze_time("2026-07-01T12:00:00+00:00") +def test_delta_iceberg_inventory_appends_lookback_filters(): + # Arrange Mock DataFrames and Table Sources + mock_source_df = pl.DataFrame( + {"bucket": ["b1"], "key": ["k1"], "version_id": ["v1"]}, + schema={"bucket": pl.String, "key": pl.String, "version_id": pl.String}, + ) + + # Explicitly enforce schema + mock_sink_df = pl.DataFrame( + {"bucket": [], "key": [], "version_id": []}, + schema={"bucket": pl.String, "key": pl.String, "version_id": pl.String}, + ) + + # Create Mock Iceberg Sources + mock_source = MagicMock() + mock_source.inventory.return_value = mock_source_df + # Initialize table_scan_config with no row_filter initially + mock_source.table_scan_config = MagicMock() + mock_source.table_scan_config.row_filter = None + + mock_sink = MagicMock() + mock_sink.inventory.return_value = mock_sink_df + + # Configure a 1 day, 2 hour lookback config -> target time: 2026-06-30T10:00:00+00:00 + lookback_config = data_index.inventory_source.delta_iceberg_table.LookbackConfig( + days=1, hours=2, column_name="last_modified" + ) + expected_lookback_filter = "last_modified >= '2026-06-30T10:00:00+00:00'" + + # Case A: Verify it writes directly when there is NO existing filter + inventory_source_no_filter = data_index.inventory_source.delta_iceberg_table.DeltaIcebergTableInventorySource.model_construct( + source=mock_source, + sink=mock_sink, + lookback_config=lookback_config, + ) + + # Trigger the inventory call + inventory_source_no_filter.inventory() + + assert mock_source.table_scan_config.row_filter == expected_lookback_filter + + # Case B: Verify it appends using AND when there IS an existing filter + mock_source.table_scan_config.row_filter = "status = 'active'" + + inventory_source_with_filter = data_index.inventory_source.delta_iceberg_table.DeltaIcebergTableInventorySource.model_construct( + source=mock_source, + sink=mock_sink, + lookback_config=lookback_config, + ) + + # Trigger the inventory call again + result_df = inventory_source_with_filter.inventory() + + expected_combined_filter = f"(status = 'active') AND {expected_lookback_filter}" + assert mock_source.table_scan_config.row_filter == expected_combined_filter + + # Verify the anti-join executed correctly (mock_source_df elements remain) + assert result_df.shape == (1, 3) + assert result_df.equals(mock_source_df) diff --git a/uv.lock b/uv.lock index 96be167..e8669a4 100644 --- a/uv.lock +++ b/uv.lock @@ -692,6 +692,7 @@ analysis = [ ] dev = [ { name = "aiosqlite" }, + { name = "freezegun" }, { name = "pre-commit" }, { name = "prefect-docker" }, { name = "pytest" }, @@ -738,6 +739,7 @@ analysis = [ ] dev = [ { name = "aiosqlite", specifier = ">=0.22.1" }, + { name = "freezegun", specifier = ">=1.5.5" }, { name = "pre-commit", specifier = ">=4.6.0" }, { name = "prefect-docker", specifier = ">=0.7.1" }, { name = "pytest", specifier = ">=9.0.3" }, @@ -947,6 +949,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/cf/58/8acf1b3e91c58313ce5cb67df61001fc9dcd21be4fadb76c1a2d540e09ed/fqdn-1.5.1-py3-none-any.whl", hash = "sha256:3a179af3761e4df6eb2e026ff9e1a3033d3587bf980a0b1b2e1e5d08d7358014", size = 9121, upload-time = "2021-03-11T07:16:28.351Z" }, ] +[[package]] +name = "freezegun" +version = "1.5.5" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "python-dateutil" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/95/dd/23e2f4e357f8fd3bdff613c1fe4466d21bfb00a6177f238079b17f7b1c84/freezegun-1.5.5.tar.gz", hash = "sha256:ac7742a6cc6c25a2c35e9292dfd554b897b517d2dec26891a2e8debf205cb94a", size = 35914, upload-time = "2025-08-09T10:39:08.338Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5e/2e/b41d8a1a917d6581fc27a35d05561037b048e47df50f27f8ac9c7e27a710/freezegun-1.5.5-py3-none-any.whl", hash = "sha256:cd557f4a75cf074e84bc374249b9dd491eaeacd61376b9eb3c423282211619d2", size = 19266, upload-time = "2025-08-09T10:39:06.636Z" }, +] + [[package]] name = "frozenlist" version = "1.8.0"