Skip to content

Commit 81e1546

Browse files
Arnavsharma2ntkathole
authored andcommitted
fix: Snapshot tag filters in registry cache keys
Signed-off-by: Arnav Sharma <[email protected]>
1 parent 2deea25 commit 81e1546

3 files changed

Lines changed: 136 additions & 3 deletions

File tree

‎sdk/python/feast/infra/registry/proto_registry_utils.py‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -63,8 +63,14 @@ def wrapper(
6363
nonlocal cache_key, cache_value
6464

6565
kwargs_key = tuple(sorted(kwargs.items())) if kwargs else ()
66-
key = tuple(
67-
[id(registry_proto), registry_proto.version_id, project, tags, kwargs_key]
66+
# Snapshot the filter so caller mutations cannot change a cached key.
67+
tags_key = frozenset(tags.items()) if tags is not None else None
68+
key = (
69+
id(registry_proto),
70+
registry_proto.version_id,
71+
project,
72+
tags_key,
73+
kwargs_key,
6874
)
6975

7076
if key == cache_key:

‎sdk/python/tests/unit/infra/registry/test_proto_registry_utils.py‎

Lines changed: 28 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,38 @@
1-
from feast.infra.registry.proto_registry_utils import list_saved_datasets
1+
from unittest.mock import patch
2+
3+
from feast import Entity
4+
from feast.infra.registry.proto_registry_utils import list_entities, list_saved_datasets
25
from feast.protos.feast.core.Registry_pb2 import Registry as RegistryProto
36
from feast.protos.feast.core.SavedDataset_pb2 import SavedDataset as SavedDatasetProto
47
from feast.protos.feast.core.SavedDataset_pb2 import (
58
SavedDatasetSpec,
69
SavedDatasetStorage,
710
)
811
from feast.saved_dataset import SavedDataset
12+
from feast.value_type import ValueType
13+
14+
15+
def test_equivalent_tag_filters_reuse_cached_entities() -> None:
16+
registry = RegistryProto(version_id="test_version")
17+
entity = Entity(
18+
name="driver",
19+
value_type=ValueType.INT64,
20+
tags={"environment": "production", "team": "ml"},
21+
).to_proto()
22+
entity.spec.project = "test_project"
23+
registry.entities.append(entity)
24+
25+
with patch.object(Entity, "from_proto", wraps=Entity.from_proto) as from_proto:
26+
first = list_entities(
27+
registry, "test_project", {"environment": "production", "team": "ml"}
28+
)
29+
second = list_entities(
30+
registry, "test_project", {"team": "ml", "environment": "production"}
31+
)
32+
33+
assert [obj.name for obj in first] == ["driver"]
34+
assert [obj.name for obj in second] == ["driver"]
35+
from_proto.assert_called_once()
936

1037

1138
class TestRegistryProto:
Lines changed: 100 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,100 @@
1+
from __future__ import annotations
2+
3+
from pathlib import Path
4+
5+
import pytest
6+
7+
from feast import Entity, FeatureStore, FeatureView, Field
8+
from feast.infra.offline_stores.file_source import FileSource
9+
from feast.infra.registry.registry import Registry
10+
from feast.infra.registry.sql import SqlRegistry
11+
from feast.repo_config import RepoConfig
12+
from feast.types import Int64
13+
from feast.value_type import ValueType
14+
15+
16+
@pytest.mark.parametrize("registry_type", ["file", "sql"])
17+
@pytest.mark.parametrize("resource_type", ["entities", "feature_views"])
18+
@pytest.mark.parametrize(
19+
"initial_tags,updated_tags,initial_names,updated_names",
20+
[
21+
(
22+
{"environment": "production"},
23+
{"environment": "staging"},
24+
{"production"},
25+
{"staging"},
26+
),
27+
({}, {"environment": "production"}, {"production", "staging"}, {"production"}),
28+
({"environment": "production"}, {}, {"production"}, {"production", "staging"}),
29+
],
30+
ids=["change-value", "add-tag", "remove-tag"],
31+
)
32+
def test_cached_list_uses_current_tag_filter(
33+
tmp_path: Path,
34+
registry_type: str,
35+
resource_type: str,
36+
initial_tags: dict[str, str],
37+
updated_tags: dict[str, str],
38+
initial_names: set[str],
39+
updated_names: set[str],
40+
) -> None:
41+
"""Reusing a tag dictionary must not return results for its previous contents."""
42+
registry_path = tmp_path / "registry.db"
43+
registry_config = (
44+
{"path": str(registry_path), "cache_ttl_seconds": 0}
45+
if registry_type == "file"
46+
else {
47+
"registry_type": "sql",
48+
"path": f"sqlite:///{registry_path}",
49+
"cache_ttl_seconds": 0,
50+
}
51+
)
52+
store = FeatureStore(
53+
config=RepoConfig(
54+
project="test_project", registry=registry_config, provider="local"
55+
)
56+
)
57+
try:
58+
for environment in ("production", "staging"):
59+
tags = {"environment": environment}
60+
entity = Entity(name=environment, value_type=ValueType.INT64, tags=tags)
61+
store.registry.apply_entity(entity, store.project)
62+
store.registry.apply_feature_view(
63+
FeatureView(
64+
name=environment,
65+
entities=[entity],
66+
schema=[
67+
Field(name=environment, dtype=Int64),
68+
Field(name="value", dtype=Int64),
69+
],
70+
source=FileSource(
71+
path="unused.parquet", timestamp_field="event_timestamp"
72+
),
73+
tags=tags,
74+
),
75+
store.project,
76+
)
77+
store.refresh_registry()
78+
79+
def list_names(tags: dict[str, str]) -> set[str]:
80+
if resource_type == "entities":
81+
return {
82+
obj.name for obj in store.list_entities(allow_cache=True, tags=tags)
83+
}
84+
return {
85+
obj.name
86+
for obj in store.list_feature_views(allow_cache=True, tags=tags)
87+
}
88+
89+
query_tags = dict(initial_tags)
90+
assert list_names(query_tags) == initial_names
91+
92+
query_tags.clear()
93+
query_tags.update(updated_tags)
94+
assert list_names(query_tags) == updated_names
95+
# A separate caller with an equivalent filter must also see correct results.
96+
assert list_names(dict(updated_tags)) == updated_names
97+
finally:
98+
registry = store.registry
99+
assert isinstance(registry, (Registry, SqlRegistry))
100+
registry.teardown()

0 commit comments

Comments
 (0)