Skip to content

Commit 934d341

Browse files
xKomilntkathole
authored andcommitted
fix: Call ApplySavedDataset RPC instead of ApplyFeatureService in RemoteRegistry.apply_saved_dataset()
RemoteRegistry.apply_saved_dataset() builds an ApplySavedDatasetRequest but sends it via the ApplyFeatureService RPC. The server has a dedicated ApplySavedDataset RPC that expects exactly this request type; calling ApplyFeatureService instead means the server tries to interpret an ApplySavedDatasetRequest payload as an ApplyFeatureServiceRequest, so registering a saved dataset through a remote registry never actually persists it as intended. Adds a regression test asserting the correct RPC is invoked. Signed-off-by: xKomil <[email protected]>
1 parent fd4d49d commit 934d341

2 files changed

Lines changed: 29 additions & 1 deletion

File tree

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -489,7 +489,7 @@ def apply_saved_dataset(
489489
request = RegistryServer_pb2.ApplySavedDatasetRequest(
490490
saved_dataset=saved_dataset.to_proto(), project=project, commit=commit
491491
)
492-
self.stub.ApplyFeatureService(request)
492+
self.stub.ApplySavedDataset(request)
493493

494494
def delete_saved_dataset(self, name: str, project: str, commit: bool = True):
495495
request = RegistryServer_pb2.DeleteSavedDatasetRequest(

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

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,8 +18,10 @@
1818
import pytest
1919
from google.protobuf.timestamp_pb2 import Timestamp
2020

21+
from feast.infra.offline_stores.file_source import SavedDatasetFileStorage
2122
from feast.infra.registry.remote import RemoteRegistry
2223
from feast.protos.feast.registry import RegistryServer_pb2
24+
from feast.saved_dataset import SavedDataset
2325

2426

2527
@pytest.fixture
@@ -83,3 +85,29 @@ def test_updated_since_none(remote_registry):
8385
remote_registry.stub.ListAllFeatureViews.call_args[0][0]
8486
)
8587
assert not request.HasField("updated_since")
88+
89+
90+
def test_apply_saved_dataset_calls_apply_saved_dataset_rpc(remote_registry):
91+
"""apply_saved_dataset() must invoke the ApplySavedDataset RPC, not
92+
ApplyFeatureService -- calling the wrong RPC means the server tries to
93+
deserialize an ApplySavedDatasetRequest as an ApplyFeatureServiceRequest,
94+
which the server cannot handle correctly."""
95+
saved_dataset = SavedDataset(
96+
name="test_saved_dataset",
97+
features=["feature_view:feature"],
98+
join_keys=["entity_id"],
99+
storage=SavedDatasetFileStorage(path="/tmp/does-not-need-to-exist.parquet"),
100+
tags={},
101+
)
102+
103+
remote_registry.apply_saved_dataset(saved_dataset, "project", commit=True)
104+
105+
remote_registry.stub.ApplySavedDataset.assert_called_once()
106+
remote_registry.stub.ApplyFeatureService.assert_not_called()
107+
108+
request: RegistryServer_pb2.ApplySavedDatasetRequest = (
109+
remote_registry.stub.ApplySavedDataset.call_args[0][0]
110+
)
111+
assert request.saved_dataset.spec.name == "test_saved_dataset"
112+
assert request.project == "project"
113+
assert request.commit is True

0 commit comments

Comments
 (0)