Skip to content

Commit 161aab7

Browse files
makinzmmakinzm
committed
feat: drop all versioned Milvus collections on teardown/update and lint
Signed-off-by: makinzm <nozomi.maki.da@gmail.com>
1 parent 175cbae commit 161aab7

2 files changed

Lines changed: 110 additions & 22 deletions

File tree

‎sdk/python/feast/infra/online_stores/milvus_online_store/milvus.py‎

Lines changed: 25 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -498,15 +498,12 @@ def update(
498498
for table in tables_to_keep:
499499
self._get_or_create_collection(config, table)
500500

501+
# Always drop the base collection plus any "_v{N}" siblings, regardless of
502+
# the current versioning flag. This handles mixed-state repos where
503+
# versioning was toggled on/off across applies and would otherwise leave
504+
# orphan collections behind in Milvus.
501505
for table in tables_to_delete:
502-
collection_name = _table_id(
503-
config.project,
504-
table,
505-
config.registry.enable_online_feature_view_versioning,
506-
)
507-
if self._collections.get(collection_name, None):
508-
self.client.drop_collection(collection_name)
509-
self._collections.pop(collection_name, None)
506+
self._drop_all_version_collections(config.project, table)
510507

511508
def plan(
512509
self, config: RepoConfig, desired_registry_proto: RegistryProto
@@ -520,15 +517,9 @@ def teardown(
520517
entities: Sequence[Entity],
521518
):
522519
self.client = self._connect(config)
520+
# See update(): drop base + all "_v{N}" siblings to handle mixed-state repos.
523521
for table in tables:
524-
collection_name = _table_id(
525-
config.project,
526-
table,
527-
config.registry.enable_online_feature_view_versioning,
528-
)
529-
if self._collections.get(collection_name, None):
530-
self.client.drop_collection(collection_name)
531-
self._collections.pop(collection_name, None)
522+
self._drop_all_version_collections(config.project, table)
532523

533524
def retrieve_online_documents_v2(
534525
self,
@@ -763,6 +754,24 @@ def retrieve_online_documents_v2(
763754
result_list.append((res_ts, entity_key_proto, res if res else None))
764755
return result_list
765756

757+
def _drop_all_version_collections(self, project: str, table: FeatureView) -> None:
758+
"""Drop the base collection and every ``_v{N}`` versioned sibling.
759+
760+
Mirrors the ``_drop_all_version_tables`` helpers in the MySQL/PostgreSQL
761+
online stores. Always called from ``update`` and ``teardown`` so a
762+
repo that toggles versioning on and off does not leave orphan
763+
collections behind in Milvus.
764+
"""
765+
base = f"{project}_{table.name}"
766+
versioned_prefix = f"{base}_v"
767+
for collection_name in self.client.list_collections():
768+
if collection_name == base or (
769+
collection_name.startswith(versioned_prefix)
770+
and collection_name[len(versioned_prefix) :].isdigit()
771+
):
772+
self.client.drop_collection(collection_name)
773+
self._collections.pop(collection_name, None)
774+
766775

767776
def _table_id(project: str, table: FeatureView, enable_versioning: bool = False) -> str:
768777
return compute_table_id(project, table, enable_versioning)

‎sdk/python/tests/unit/infra/online_store/test_milvus_versioning.py‎

Lines changed: 85 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -51,28 +51,40 @@ def test_versioning_enabled_with_version(self):
5151

5252
fv = _make_feature_view(version_number=2)
5353
config = _make_config(versioning=True)
54-
assert _table_id(config.project, fv, enable_versioning=True) == "test_project_driver_stats_v2"
54+
assert (
55+
_table_id(config.project, fv, enable_versioning=True)
56+
== "test_project_driver_stats_v2"
57+
)
5558

5659
def test_projection_version_tag_takes_priority(self):
5760
from feast.infra.online_stores.milvus_online_store.milvus import _table_id
5861

5962
fv = _make_feature_view(version_number=1, version_tag=3)
6063
config = _make_config(versioning=True)
61-
assert _table_id(config.project, fv, enable_versioning=True) == "test_project_driver_stats_v3"
64+
assert (
65+
_table_id(config.project, fv, enable_versioning=True)
66+
== "test_project_driver_stats_v3"
67+
)
6268

6369
def test_version_zero_no_suffix(self):
6470
from feast.infra.online_stores.milvus_online_store.milvus import _table_id
6571

6672
fv = _make_feature_view(version_number=0)
6773
config = _make_config(versioning=True)
68-
assert _table_id(config.project, fv, enable_versioning=True) == "test_project_driver_stats"
74+
assert (
75+
_table_id(config.project, fv, enable_versioning=True)
76+
== "test_project_driver_stats"
77+
)
6978

7079
def test_versioning_enabled_no_version_set(self):
7180
from feast.infra.online_stores.milvus_online_store.milvus import _table_id
7281

7382
fv = _make_feature_view()
7483
config = _make_config(versioning=True)
75-
assert _table_id(config.project, fv, enable_versioning=True) == "test_project_driver_stats"
84+
assert (
85+
_table_id(config.project, fv, enable_versioning=True)
86+
== "test_project_driver_stats"
87+
)
7688

7789
def test_versioning_disabled_ignores_version(self):
7890
from feast.infra.online_stores.milvus_online_store.milvus import _table_id
@@ -86,16 +98,83 @@ class TestMilvusVersionedReadSupport:
8698
"""Test that MilvusOnlineStore passes _check_versioned_read_support."""
8799

88100
def test_allowed_with_version_tag(self):
89-
from feast.infra.online_stores.milvus_online_store.milvus import MilvusOnlineStore
101+
from feast.infra.online_stores.milvus_online_store.milvus import (
102+
MilvusOnlineStore,
103+
)
90104

91105
store = MilvusOnlineStore()
92106
fv = _make_feature_view()
93107
fv.projection.version_tag = 2
94108
store._check_versioned_read_support([(fv, ["trips_today"])])
95109

96110
def test_allowed_without_version_tag(self):
97-
from feast.infra.online_stores.milvus_online_store.milvus import MilvusOnlineStore
111+
from feast.infra.online_stores.milvus_online_store.milvus import (
112+
MilvusOnlineStore,
113+
)
98114

99115
store = MilvusOnlineStore()
100116
fv = _make_feature_view()
101117
store._check_versioned_read_support([(fv, ["trips_today"])])
118+
119+
120+
class TestTeardownDropsAllVersions:
121+
"""Teardown should drop the base collection AND all versioned collections."""
122+
123+
def _build_store_with_collections(self, existing_collections):
124+
from feast.infra.online_stores.milvus_online_store.milvus import (
125+
MilvusOnlineStore,
126+
)
127+
128+
store = MilvusOnlineStore()
129+
store.client = MagicMock()
130+
store.client.list_collections.return_value = existing_collections
131+
store._connect = MagicMock(return_value=store.client)
132+
store._collections = {name: MagicMock() for name in existing_collections}
133+
return store
134+
135+
def test_teardown_drops_base_and_all_versioned_collections(self):
136+
fv = _make_feature_view()
137+
config = _make_config(versioning=True)
138+
existing = [
139+
"test_project_driver_stats",
140+
"test_project_driver_stats_v1",
141+
"test_project_driver_stats_v2",
142+
"test_project_other_view", # unrelated, must not be dropped
143+
]
144+
store = self._build_store_with_collections(existing)
145+
146+
store.teardown(config, [fv], [])
147+
148+
dropped = {call.args[0] for call in store.client.drop_collection.call_args_list}
149+
assert dropped == {
150+
"test_project_driver_stats",
151+
"test_project_driver_stats_v1",
152+
"test_project_driver_stats_v2",
153+
}
154+
assert "test_project_other_view" not in dropped
155+
156+
def test_update_drops_all_versions_for_deleted_table(self):
157+
fv = _make_feature_view()
158+
config = _make_config(versioning=True)
159+
existing = [
160+
"test_project_driver_stats",
161+
"test_project_driver_stats_v3",
162+
"test_project_driver_stats_v4",
163+
]
164+
store = self._build_store_with_collections(existing)
165+
166+
store.update(
167+
config=config,
168+
tables_to_delete=[fv],
169+
tables_to_keep=[],
170+
entities_to_delete=[],
171+
entities_to_keep=[],
172+
partial=False,
173+
)
174+
175+
dropped = {call.args[0] for call in store.client.drop_collection.call_args_list}
176+
assert dropped == {
177+
"test_project_driver_stats",
178+
"test_project_driver_stats_v3",
179+
"test_project_driver_stats_v4",
180+
}

0 commit comments

Comments
 (0)