Skip to content

add filter on model validate #60

New issue

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

Sign up for GitHub

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

Already on GitHub? Sign in to your account

Merged
thommodin merged 1 commit into from
Jul 1, 2026
Merged

add filter on model validate #60

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 4 additions & 1 deletion Makefile
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,4 +30,7 @@ test-image:
-o type=cacheonly .

deploy:
uv run prefect deploy --all
@echo "Deploying all to local Prefect..."
cat prefect.yaml
uv run prefect deploy --all
@echo "Deployed all to local Prefect!"
18 changes: 9 additions & 9 deletions src/data_index/inventory_source/delta_iceberg_table.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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()

Expand Down
120 changes: 78 additions & 42 deletions tests/test_delta_iceberg_table_source.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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():
Expand Down Expand Up @@ -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
)
Loading