diff --git a/Makefile b/Makefile index 4d75b5b..8919c1d 100644 --- a/Makefile +++ b/Makefile @@ -30,4 +30,7 @@ test-image: -o type=cacheonly . deploy: - uv run prefect deploy --all \ No newline at end of file + @echo "Deploying all to local Prefect..." + cat prefect.yaml + uv run prefect deploy --all + @echo "Deployed all to local Prefect!" \ No newline at end of file diff --git a/src/data_index/inventory_source/delta_iceberg_table.py b/src/data_index/inventory_source/delta_iceberg_table.py index 010bba2..ef46d6e 100644 --- a/src/data_index/inventory_source/delta_iceberg_table.py +++ b/src/data_index/inventory_source/delta_iceberg_table.py @@ -12,7 +12,7 @@ class LookbackConfig(pydantic.BaseModel): days: int = pydantic.Field(default=0) hours: int = pydantic.Field(default=0) - column_name: str + column_name: str = pydantic.Field(default="last_modified_date") @property def lookback_timestamp(self): @@ -40,22 +40,22 @@ class DeltaIcebergTableInventorySource(pydantic.BaseModel): 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 + @pydantic.model_validator(mode="after") + def apply_lookback_config(self) -> typing.Self: + """Applies the lookback configuration filters to the source table scan config.""" 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 ) + return self + + def inventory( + self, + ) -> polars.DataFrame: 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 6424c12..540f4e8 100644 --- a/tests/test_delta_iceberg_table_source.py +++ b/tests/test_delta_iceberg_table_source.py @@ -4,7 +4,9 @@ import freezegun import polars as pl +import data_index.iceberg_config import data_index.inventory_source.delta_iceberg_table +import data_index.inventory_source.iceberg_table def test_delta_iceberg_table_inventory_anti_join(): @@ -132,62 +134,96 @@ def test_lookback_config_time_generation(): @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}, +def test_delta_iceberg_inventory_appends_lookback_filters_with_real_objects(): + # ------------------------------------------------------------------------- + # Setup Shared Sink and Lookback Configurations + # ------------------------------------------------------------------------- + dummy_sink = data_index.inventory_source.iceberg_table.IcebergTableInventorySource( + type="iceberg_table", + table_config=data_index.iceberg_config.IcebergTableConfig( + catalog_config=data_index.iceberg_config.SqliteCatalogConfig( + uri="", + warehouse="", + ), + namespace="data_index", + table_name="structured_metadata_v5", + ), + table_scan_config=data_index.iceberg_config.IcebergTableScanConfig( + row_filter=None, + selected_fields=("bucket", "key", "version_id"), + case_sensitive=True, + snapshot_id=None, + limit=None, + ), ) - # 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, + # ------------------------------------------------------------------------- + source_no_filter = ( + data_index.inventory_source.iceberg_table.IcebergTableInventorySource( + type="iceberg_table", + table_config=data_index.iceberg_config.IcebergTableConfig( + catalog_config=data_index.iceberg_config.SqliteCatalogConfig( + uri="", + warehouse="", + ), + namespace="inventory", + table_name="live", + ), + table_scan_config=data_index.iceberg_config.IcebergTableScanConfig( + row_filter=None + ), # No existing filter + ) ) - # Trigger the inventory call - inventory_source_no_filter.inventory() + model_no_filter = data_index.inventory_source.delta_iceberg_table.DeltaIcebergTableInventorySource( + source=source_no_filter, + sink=dummy_sink, + lookback_config=lookback_config, + ) - assert mock_source.table_scan_config.row_filter == expected_lookback_filter + assert ( + model_no_filter.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'" + # ------------------------------------------------------------------------- + source_with_filter = ( + data_index.inventory_source.iceberg_table.IcebergTableInventorySource( + type="iceberg_table", + table_config=data_index.iceberg_config.IcebergTableConfig( + catalog_config=data_index.iceberg_config.SqliteCatalogConfig( + uri="", + warehouse="", + ), + namespace="inventory", + table_name="live", + ), + # Existing filter provided + table_scan_config=data_index.iceberg_config.IcebergTableScanConfig( + row_filter="last_modified_date >= '2026-06-21T00:00:00'", + selected_fields=("bucket", "key", "version_id", "size"), + ), + ) + ) - inventory_source_with_filter = data_index.inventory_source.delta_iceberg_table.DeltaIcebergTableInventorySource.model_construct( - source=mock_source, - sink=mock_sink, + model_with_filter = data_index.inventory_source.delta_iceberg_table.DeltaIcebergTableInventorySource( + source=source_with_filter, + sink=dummy_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) + expected_combined_filter = ( + f"(last_modified_date >= '2026-06-21T00:00:00') AND {expected_lookback_filter}" + ) + assert ( + model_with_filter.source.table_scan_config.row_filter + == expected_combined_filter + )