fix: Fix list saved dataset api

ntkathole · ntkathole · commit 833696c79fa4 · 2025-06-30T19:19:35.000+05:30
Signed-off-by: ntkathole &lt;nikhilkathole2683@gmail.com&gt;
diff --git a/sdk/python/feast/api/registry/rest/saved_datasets.py b/sdk/python/feast/api/registry/rest/saved_datasets.py
@@ -32,6 +32,6 @@ def list_saved_datasets(
             tags=tags,
         )
         response = grpc_call(grpc_handler.ListSavedDatasets, req)
-        return {"saved_datasets": response.get("saved_datasets", [])}
+        return {"saved_datasets": response.get("savedDatasets", [])}
 
     return router
diff --git a/sdk/python/feast/infra/registry/proto_registry_utils.py b/sdk/python/feast/infra/registry/proto_registry_utils.py
@@ -320,7 +320,7 @@ def list_saved_datasets(
     saved_datasets = []
     for saved_dataset in registry_proto.saved_datasets:
         if saved_dataset.spec.project == project and utils.has_all_tags(
-            saved_dataset.tags, tags
+            saved_dataset.spec.tags, tags
         ):
             saved_datasets.append(SavedDataset.from_proto(saved_dataset))
     return saved_datasets
diff --git a/sdk/python/tests/unit/api/test_api_rest_registry.py b/sdk/python/tests/unit/api/test_api_rest_registry.py
@@ -7,7 +7,9 @@
 from feast import Entity, FeatureService, FeatureView, Field, FileSource
 from feast.api.registry.rest.rest_registry_server import RestRegistryServer
 from feast.feature_store import FeatureStore
+from feast.infra.offline_stores.file_source import SavedDatasetFileStorage
 from feast.repo_config import RepoConfig
+from feast.saved_dataset import SavedDataset
 from feast.types import Float64, Int64
 from feast.value_type import ValueType
 
@@ -67,8 +69,19 @@ def fastapi_test_app():
         features=[user_profile_feature_view],
     )
 
+    # Create a saved dataset for testing
+    saved_dataset_storage = SavedDatasetFileStorage(path=parquet_file_path)
+    test_saved_dataset = SavedDataset(
+        name="test_saved_dataset",
+        features=["user_profile:age", "user_profile:income"],
+        join_keys=["user_id"],
+        storage=saved_dataset_storage,
+        tags={"environment": "test", "version": "1.0"},
+    )
+
     # Apply objects
     store.apply([user_id_entity, user_profile_feature_view, user_feature_service])
+    store._registry.apply_saved_dataset(test_saved_dataset, "demo_project")
 
     # Build REST app with registered routes
     rest_server = RestRegistryServer(store)
@@ -248,3 +261,69 @@ def test_lineage_endpoint_error_handling(fastapi_test_app):
     # Test object relationships with missing parameters
     response = fastapi_test_app.get("/lineage/objects/featureView/test_fv")
     assert response.status_code == 422  # Missing required project parameter
+
+
+def test_saved_datasets_via_rest(fastapi_test_app):
+    # Test list saved datasets endpoint
+    response = fastapi_test_app.get("/saved_datasets?project=demo_project")
+    assert response.status_code == 200
+    response_data = response.json()
+    assert "saved_datasets" in response_data
+    assert isinstance(response_data["saved_datasets"], list)
+    assert len(response_data["saved_datasets"]) == 1
+
+    saved_dataset = response_data["saved_datasets"][0]
+    assert saved_dataset["spec"]["name"] == "test_saved_dataset"
+    assert "user_profile:age" in saved_dataset["spec"]["features"]
+    assert "user_profile:income" in saved_dataset["spec"]["features"]
+    assert "user_id" in saved_dataset["spec"]["joinKeys"]
+    assert saved_dataset["spec"]["tags"]["environment"] == "test"
+    assert saved_dataset["spec"]["tags"]["version"] == "1.0"
+
+    # Test get specific saved dataset endpoint
+    response = fastapi_test_app.get(
+        "/saved_datasets/test_saved_dataset?project=demo_project"
+    )
+    assert response.status_code == 200
+    response_data = response.json()
+    assert response_data["spec"]["name"] == "test_saved_dataset"
+    assert "user_profile:age" in response_data["spec"]["features"]
+    assert "user_profile:income" in response_data["spec"]["features"]
+
+    # Test with allow_cache parameter
+    response = fastapi_test_app.get(
+        "/saved_datasets/test_saved_dataset?project=demo_project&allow_cache=false"
+    )
+    assert response.status_code == 200
+    assert response.json()["spec"]["name"] == "test_saved_dataset"
+
+    # Test with tags filter
+    response = fastapi_test_app.get(
+        "/saved_datasets?project=demo_project&tags=environment:test"
+    )
+    assert response.status_code == 200
+    assert len(response.json()["saved_datasets"]) == 1
+
+    # Test with non-matching tags filter
+    response = fastapi_test_app.get(
+        "/saved_datasets?project=demo_project&tags=environment:production"
+    )
+    assert response.status_code == 200
+    assert len(response.json()["saved_datasets"]) == 0
+
+    # Test with multiple tags filter
+    response = fastapi_test_app.get(
+        "/saved_datasets?project=demo_project&tags=environment:test&tags=version:1.0"
+    )
+    assert response.status_code == 200
+    assert len(response.json()["saved_datasets"]) == 1
+
+    # Test non-existent saved dataset
+    response = fastapi_test_app.get("/saved_datasets/non_existent?project=demo_project")
+    assert response.status_code == 404
+
+    # Test missing project parameter
+    response = fastapi_test_app.get("/saved_datasets/test_saved_dataset")
+    assert (
+        response.status_code == 422
+    )  # Unprocessable Entity for missing required query param
diff --git a/sdk/python/tests/unit/infra/registry/test_proto_registry_utils.py b/sdk/python/tests/unit/infra/registry/test_proto_registry_utils.py
@@ -0,0 +1,56 @@
+from feast.infra.registry.proto_registry_utils import list_saved_datasets
+from feast.protos.feast.core.Registry_pb2 import Registry as RegistryProto
+from feast.protos.feast.core.SavedDataset_pb2 import SavedDataset as SavedDatasetProto
+from feast.protos.feast.core.SavedDataset_pb2 import (
+    SavedDatasetSpec,
+    SavedDatasetStorage,
+)
+from feast.saved_dataset import SavedDataset
+
+
+class TestRegistryProto:
+    """Test class for proto_registry_utils functions"""
+
+    def test_list_saved_datasets_uses_spec_tags(self):
+        """Test that list_saved_datasets correctly uses saved_dataset.spec.tags for tag filtering"""
+        registry = RegistryProto()
+        registry.version_id = "test_version"
+        saved_dataset = SavedDatasetProto()
+        spec = SavedDatasetSpec()
+        spec.name = "test_dataset"
+        spec.project = "test_project"
+        spec.features.extend(["feature1", "feature2"])
+        spec.join_keys.extend(["entity1"])
+        spec.full_feature_names = False
+        spec.tags["environment"] = "production"
+        spec.tags["team"] = "ml-team"
+        storage = SavedDatasetStorage()
+        file_options = storage.file_storage
+        file_options.uri = "test_path.parquet"
+        spec.storage.CopyFrom(storage)
+
+        saved_dataset.spec.CopyFrom(spec)
+        registry.saved_datasets.append(saved_dataset)
+
+        # Test that filtering by tags works correctly using spec.tags
+        result = list_saved_datasets(
+            registry, "test_project", {"environment": "production"}
+        )
+
+        assert len(result) == 1
+        assert isinstance(result[0], SavedDataset)
+        assert result[0].name == "test_dataset"
+        assert result[0].tags == {"environment": "production", "team": "ml-team"}
+
+        # Test that non-matching tags filter correctly
+        result = list_saved_datasets(
+            registry, "test_project", {"environment": "staging"}
+        )
+        assert len(result) == 0
+
+        # Test that multiple tag filtering works
+        result = list_saved_datasets(
+            registry, "test_project", {"environment": "production", "team": "ml-team"}
+        )
+        assert len(result) == 1
+        assert result[0].name == "test_dataset"

Original file line number	Diff line number	Diff line change
`@@ -32,6 +32,6 @@ def list_saved_datasets(`
`32`	`32`	`tags=tags,`
`33`	`33`	`)`
`34`	`34`	`response = grpc_call(grpc_handler.ListSavedDatasets, req)`
`35`		`- return {"saved_datasets": response.get("saved_datasets", [])}`
	`35`	`+ return {"saved_datasets": response.get("savedDatasets", [])}`
`36`	`36`
`37`	`37`	`return router`