Skip to content

Commit 0d02614

Browse files
nquinn408nquinn408
authored
feat: Implement RegistryServer.Proto RPC with RBAC-filtered response (#6558) (#6552)
1 parent 0de9196 commit 0d02614

2 files changed

Lines changed: 225 additions & 0 deletions

File tree

‎sdk/python/feast/registry_server.py‎

Lines changed: 93 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@
3535
str_to_auth_manager_type,
3636
)
3737
from feast.project import Project
38+
from feast.protos.feast.core.Registry_pb2 import Registry as RegistryProto
3839
from feast.protos.feast.registry import RegistryServer_pb2, RegistryServer_pb2_grpc
3940
from feast.protos.feast.registry.RegistryServer_pb2 import Feature, ListFeaturesResponse
4041
from feast.saved_dataset import SavedDataset, ValidationReference
@@ -181,6 +182,98 @@ def __init__(self, registry: BaseRegistry, store=None) -> None:
181182
self.proxied_registry = registry
182183
self.store = store
183184

185+
def Proto(self, request: Empty, context) -> RegistryProto:
186+
"""Build a RegistryProto from individually RBAC-filtered list calls.
187+
188+
The ``RegistryServer.Proto`` RPC must honor the same permission checks as the
189+
other RPCs rather than returning ``proxied_registry.proto()`` directly, which
190+
would bypass RBAC and expose every object (entities, feature views, data
191+
sources, permissions, projects, etc.) regardless of authorization.
192+
193+
Each object type is filtered with ``permitted_resources(..., DESCRIBE)``: under
194+
``NoAuthConfig`` this is a no-op (the full registry is returned, so remote
195+
registries keep working), while with auth enabled the caller only sees the
196+
objects they are permitted to ``DESCRIBE``.
197+
"""
198+
199+
def describable(resources: list) -> list:
200+
return permitted_resources(
201+
resources=cast(list[FeastObject], resources),
202+
actions=AuthzedAction.DESCRIBE,
203+
)
204+
205+
registry_proto = RegistryProto()
206+
207+
for project in describable(self.proxied_registry.list_projects()):
208+
registry_proto.projects.append(project.to_proto())
209+
project_name = project.name
210+
211+
for entity in describable(
212+
self.proxied_registry.list_entities(project=project_name)
213+
):
214+
registry_proto.entities.append(entity.to_proto())
215+
216+
for data_source in describable(
217+
self.proxied_registry.list_data_sources(project=project_name)
218+
):
219+
registry_proto.data_sources.append(data_source.to_proto())
220+
221+
for feature_view in describable(
222+
self.proxied_registry.list_feature_views(project=project_name)
223+
):
224+
registry_proto.feature_views.append(feature_view.to_proto())
225+
226+
for stream_feature_view in describable(
227+
self.proxied_registry.list_stream_feature_views(project=project_name)
228+
):
229+
registry_proto.stream_feature_views.append(
230+
stream_feature_view.to_proto()
231+
)
232+
233+
for on_demand_feature_view in describable(
234+
self.proxied_registry.list_on_demand_feature_views(project=project_name)
235+
):
236+
registry_proto.on_demand_feature_views.append(
237+
on_demand_feature_view.to_proto()
238+
)
239+
240+
for label_view in describable(
241+
self.proxied_registry.list_label_views(project=project_name)
242+
):
243+
registry_proto.label_views.append(label_view.to_proto())
244+
245+
for feature_service in describable(
246+
self.proxied_registry.list_feature_services(project=project_name)
247+
):
248+
registry_proto.feature_services.append(feature_service.to_proto())
249+
250+
for saved_dataset in describable(
251+
self.proxied_registry.list_saved_datasets(project=project_name)
252+
):
253+
registry_proto.saved_datasets.append(saved_dataset.to_proto())
254+
255+
for validation_reference in describable(
256+
self.proxied_registry.list_validation_references(project=project_name)
257+
):
258+
registry_proto.validation_references.append(
259+
validation_reference.to_proto()
260+
)
261+
262+
for permission in describable(
263+
self.proxied_registry.list_permissions(project=project_name)
264+
):
265+
registry_proto.permissions.append(permission.to_proto())
266+
267+
# Carry the registry's real last_updated/version_id rather than stamping "now":
268+
# this proto is rebuilt from individual list calls (for RBAC filtering), but it must
269+
# not look like a fresh commit on every call — clients such as the remote feature
270+
# server key cache freshness off this metadata. Reading these two scalar fields from
271+
# the source proto leaks nothing RBAC-protected (no objects are copied from it).
272+
source_proto = self.proxied_registry.proto()
273+
registry_proto.last_updated.CopyFrom(source_proto.last_updated)
274+
registry_proto.version_id = source_proto.version_id
275+
return registry_proto
276+
184277
def ApplyEntity(self, request: RegistryServer_pb2.ApplyEntityRequest, context):
185278
entity = cast(
186279
Entity,

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

Lines changed: 132 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,132 @@
1+
"""Unit tests for the ``RegistryServer.Proto`` RPC (issue #6558).
2+
3+
The RPC must build the ``RegistryProto`` from individually RBAC-filtered list calls
4+
rather than returning ``proxied_registry.proto()`` directly (which would bypass
5+
permissions). Under ``NoAuthConfig`` filtering is a no-op, so the full registry is
6+
returned; with auth enabled only ``DESCRIBE``-permitted objects are included.
7+
"""
8+
9+
from datetime import datetime, timezone
10+
from unittest.mock import patch
11+
12+
from google.protobuf.empty_pb2 import Empty
13+
14+
from feast.data_source import DataSource
15+
from feast.entity import Entity
16+
from feast.feast_object import FeastObject
17+
from feast.project import Project
18+
from feast.protos.feast.core.Registry_pb2 import Registry as RegistryProto
19+
from feast.registry_server import RegistryServer
20+
from feast.value_type import ValueType
21+
22+
# The registry's authentic metadata, returned by _FakeRegistry.proto(). Proto() must carry these
23+
# through rather than stamping "now" (issue #6558 review feedback).
24+
_REGISTRY_LAST_UPDATED = datetime(2024, 1, 2, 3, 4, 5, tzinfo=timezone.utc)
25+
_REGISTRY_VERSION_ID = "test-version-id"
26+
27+
28+
class _FakeRegistry:
29+
"""Minimal BaseRegistry stand-in exposing only the list_* calls Proto uses."""
30+
31+
def __init__(self, projects, entities_by_project, data_sources_by_project):
32+
self._projects = projects
33+
self._entities = entities_by_project
34+
self._data_sources = data_sources_by_project
35+
36+
def proto(self) -> RegistryProto:
37+
# Source of the authentic last_updated / version_id metadata. Proto() reads only these
38+
# scalar fields from here (no objects), so RBAC filtering is unaffected.
39+
proto = RegistryProto()
40+
proto.version_id = _REGISTRY_VERSION_ID
41+
proto.last_updated.FromDatetime(_REGISTRY_LAST_UPDATED)
42+
return proto
43+
44+
def list_projects(self, allow_cache: bool = False, tags=None):
45+
return self._projects
46+
47+
def list_entities(self, project: str, allow_cache: bool = False, tags=None):
48+
return self._entities.get(project, [])
49+
50+
def list_data_sources(self, project: str, allow_cache: bool = False, tags=None):
51+
return self._data_sources.get(project, [])
52+
53+
# Every other object type is empty for this fixture.
54+
def _empty(self, *args, **kwargs):
55+
return []
56+
57+
list_feature_views = _empty
58+
list_stream_feature_views = _empty
59+
list_on_demand_feature_views = _empty
60+
list_label_views = _empty
61+
list_feature_services = _empty
62+
list_saved_datasets = _empty
63+
list_validation_references = _empty
64+
list_permissions = _empty
65+
66+
67+
def _entity(name: str) -> Entity:
68+
return Entity(name=name, value_type=ValueType.STRING)
69+
70+
71+
def _data_source(name: str) -> DataSource:
72+
from feast.infra.offline_stores.file_source import FileSource
73+
74+
return FileSource(name=name, path=f"/tmp/{name}.parquet", timestamp_field="ts")
75+
76+
77+
def _build_server() -> tuple[RegistryServer, _FakeRegistry]:
78+
registry = _FakeRegistry(
79+
projects=[Project(name="proj_a"), Project(name="proj_b")],
80+
entities_by_project={
81+
"proj_a": [_entity("driver"), _entity("customer")],
82+
"proj_b": [_entity("merchant")],
83+
},
84+
data_sources_by_project={"proj_a": [_data_source("src_a")]},
85+
)
86+
return RegistryServer(registry), registry # type: ignore[arg-type]
87+
88+
89+
def test_proto_returns_full_registry_when_no_auth():
90+
"""NoAuthConfig (no security manager) -> every object across all projects."""
91+
server, _ = _build_server()
92+
93+
result = server.Proto(Empty(), context=None)
94+
95+
assert {p.spec.name for p in result.projects} == {"proj_a", "proj_b"}
96+
assert {e.spec.name for e in result.entities} == {"driver", "customer", "merchant"}
97+
assert {d.name for d in result.data_sources} == {"src_a"}
98+
# last_updated / version_id are carried from the registry's real proto (not stamped "now"),
99+
# so cache consumers see the registry's authentic freshness metadata.
100+
assert result.version_id == _REGISTRY_VERSION_ID
101+
assert result.last_updated.ToDatetime(tzinfo=timezone.utc) == _REGISTRY_LAST_UPDATED
102+
103+
104+
def test_proto_filters_by_describe_permission():
105+
"""With RBAC, only DESCRIBE-permitted objects are included."""
106+
server, _ = _build_server()
107+
108+
# Simulate a security manager that permits everything except the "customer"
109+
# entity, regardless of object type (filters by DESCRIBE).
110+
def fake_permitted(resources: list[FeastObject], actions):
111+
return [r for r in resources if getattr(r, "name", None) != "customer"]
112+
113+
with patch(
114+
"feast.registry_server.permitted_resources", side_effect=fake_permitted
115+
) as mocked:
116+
result = server.Proto(Empty(), context=None)
117+
118+
assert mocked.called
119+
# "customer" is filtered out; everything else survives.
120+
assert {e.spec.name for e in result.entities} == {"driver", "merchant"}
121+
assert {p.spec.name for p in result.projects} == {"proj_a", "proj_b"}
122+
assert {d.name for d in result.data_sources} == {"src_a"}
123+
124+
125+
def test_proto_empty_registry():
126+
"""No projects -> empty (but valid) RegistryProto, not an error."""
127+
server = RegistryServer(_FakeRegistry([], {}, {})) # type: ignore[arg-type]
128+
129+
result = server.Proto(Empty(), context=None)
130+
131+
assert len(result.projects) == 0
132+
assert len(result.entities) == 0

0 commit comments

Comments
 (0)