Skip to content

Commit dbf8643

Browse files
committed
fix: Compare source options in ClickHouse and SQL Server source equality
ClickhouseSource had no __eq__ of its own, so the base comparison ignored its table and query and a moved source compared equal to the registered one. Compare both options on top of the base fields. MsSqlServerSource.__eq__ did not call the base comparison, so changes to tags, description, owner and connection_ref were ignored on re-apply. Use super().__eq__() and serialize description, tags and owner in the proto so a source read back from the registry still equals itself. Signed-off-by: NopeCharon <[email protected]>
1 parent 4d3b6a2 commit dbf8643

3 files changed

Lines changed: 57 additions & 6 deletions

File tree

‎sdk/python/feast/infra/offline_stores/contrib/clickhouse_offline_store/clickhouse_source.py‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -94,6 +94,19 @@ def __init__(
9494
connection_ref=connection_ref,
9595
)
9696

97+
def __hash__(self):
98+
return super().__hash__()
99+
100+
def __eq__(self, other):
101+
if not isinstance(other, ClickhouseSource):
102+
return False
103+
104+
return (
105+
super().__eq__(other)
106+
and self._clickhouse_options._query == other._clickhouse_options._query
107+
and self._clickhouse_options._table == other._clickhouse_options._table
108+
)
109+
97110
@staticmethod
98111
def from_proto(data_source: DataSourceProto) -> Any:
99112
assert data_source.HasField("custom_options")

‎sdk/python/feast/infra/offline_stores/contrib/mssql_offline_store/mssqlserver_source.py‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -176,7 +176,7 @@ def __eq__(self, other):
176176
return False
177177

178178
return (
179-
self.name == other.name
179+
super().__eq__(other)
180180
and self.mssqlserver_options.table_ref
181181
== other.mssqlserver_options.table_ref
182182
and self.mssqlserver_options.connection_str
@@ -226,7 +226,9 @@ def from_proto(data_source: DataSourceProto):
226226
event_timestamp_column=data_source.timestamp_field,
227227
created_timestamp_column=data_source.created_timestamp_column,
228228
date_partition_column=data_source.date_partition_column,
229+
description=data_source.description,
229230
tags=tags,
231+
owner=data_source.owner,
230232
connection_ref=ConnectionRef.from_tags(tags),
231233
)
232234

@@ -236,6 +238,9 @@ def _to_proto_impl(self) -> DataSourceProto:
236238
data_source_class_type="feast.infra.offline_stores.contrib.mssql_offline_store.mssqlserver_source.MsSqlServerSource",
237239
field_mapping=self.field_mapping,
238240
custom_options=self.mssqlserver_options.to_proto(),
241+
description=self.description,
242+
tags=self.tags,
243+
owner=self.owner,
239244
)
240245

241246
data_source_proto.timestamp_field = self.timestamp_field

‎sdk/python/tests/unit/test_data_sources.py‎

Lines changed: 38 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -568,19 +568,26 @@ def test_stream_source_without_batch_source_round_trips(stream_source):
568568
"timestamp_field": "ts",
569569
},
570570
),
571+
(
572+
"feast.infra.offline_stores.contrib.clickhouse_offline_store.clickhouse_source",
573+
"ClickhouseSource",
574+
{"table": "orders", "timestamp_field": "ts"},
575+
{"table": "orders_v2", "timestamp_field": "ts"},
576+
),
571577
],
572-
ids=["postgres", "mssql", "couchbase"],
578+
ids=["postgres", "mssql", "couchbase", "clickhouse"],
573579
)
574580
def test_data_source_eq_compares_the_source_table(
575581
module_path, class_name, location, moved_location, tmp_path
576582
):
577583
"""Pointing a source at another table makes it unequal to the old one.
578584
579585
``__eq__`` compared only the query (or, for SQL Server, the connection
580-
string) of these sources, so a source moved to another table compared
581-
equal to the registered one. The registry skips ``apply_data_source`` for
582-
an equal source, so ``feast apply`` kept the old table on the registered
583-
data source while feature views built on it moved to the new one.
586+
string, and for ClickHouse, none of its options) of these sources, so a
587+
source moved to another table compared equal to the registered one. The
588+
registry skips ``apply_data_source`` for an equal source, so ``feast
589+
apply`` kept the old table on the registered data source while feature
590+
views built on it moved to the new one.
584591
"""
585592
source_class = getattr(pytest.importorskip(module_path), class_name)
586593
source = source_class(name="src", **location)
@@ -602,3 +609,29 @@ def test_data_source_eq_compares_the_source_table(
602609
assert stored.get_table_query_string() == moved_source.get_table_query_string()
603610
finally:
604611
registry.teardown()
612+
613+
614+
def test_mssqlserver_source_eq_compares_base_fields():
615+
"""SQL Server sources compare and round-trip tags, description and owner."""
616+
mssqlserver_source = pytest.importorskip(
617+
"feast.infra.offline_stores.contrib.mssql_offline_store.mssqlserver_source"
618+
)
619+
source = mssqlserver_source.MsSqlServerSource(
620+
name="src",
621+
table_ref="orders",
622+
event_timestamp_column="ts",
623+
description="orders table",
624+
tags={"team": "a"},
625+
626+
)
627+
retagged_source = mssqlserver_source.MsSqlServerSource(
628+
name="src",
629+
table_ref="orders",
630+
event_timestamp_column="ts",
631+
description="orders table",
632+
tags={"team": "b"},
633+
634+
)
635+
636+
assert source != retagged_source
637+
assert mssqlserver_source.MsSqlServerSource.from_proto(source.to_proto()) == source

0 commit comments

Comments
 (0)