feat: Added relationship support to all API endpoints (#5496)

Signed-off-by: ntkathole <[email protected]>

Author: ntkathole

docs/reference/feature-servers/registry-server.md

@@ -6,7 +6,9 @@
 from feast.api.registry.rest.rest_utils import (
     create_grpc_pagination_params,
     create_grpc_sorting_params,
+    get_object_relationships,
     get_pagination_params,
+    get_relationships_for_objects,
     get_sorting_params,
     grpc_call,
     parse_tags,
@@ -22,6 +24,9 @@ def get_data_source_router(grpc_handler) -> APIRouter:
     @router.get("/data_sources")
     def list_data_sources(
         project: str = Query(...),
+        include_relationships: bool = Query(
+            False, description="Include relationships for each data source"
+        ),
         allow_cache: bool = Query(default=True),
         tags: Dict[str, str] = Depends(parse_tags),
         pagination_params: dict = Depends(get_pagination_params),
@@ -34,24 +39,46 @@ def list_data_sources(
             pagination=create_grpc_pagination_params(pagination_params),
             sorting=create_grpc_sorting_params(sorting_params),
         )
-
         response = grpc_call(grpc_handler.ListDataSources, req)
-        return {
-            "data_sources": response.get("dataSources", []),
+        data_sources = response.get("dataSources", [])
+
+        result = {
+            "data_sources": data_sources,
             "pagination": response.get("pagination", {}),
         }
 
+        if include_relationships:
+            relationships = get_relationships_for_objects(
+                grpc_handler, data_sources, "dataSource", project, allow_cache
+            )
+            result["relationships"] = relationships
+
+        return result
+
     @router.get("/data_sources/{name}")
     def get_data_source(
         name: str,
         project: str = Query(...),
+        include_relationships: bool = Query(
+            False, description="Include relationships for this data source"
+        ),
         allow_cache: bool = Query(default=True),
     ):
         req = RegistryServer_pb2.GetDataSourceRequest(
             name=name,
             project=project,
             allow_cache=allow_cache,
         )
-        return grpc_call(grpc_handler.GetDataSource, req)
+        data_source = grpc_call(grpc_handler.GetDataSource, req)
+
+        result = data_source
+
+        if include_relationships:
+            relationships = get_object_relationships(
+                grpc_handler, "dataSource", name, project, allow_cache
+            )
+            result["relationships"] = relationships
+
+        return result
 
     return router

sdk/python/feast/api/registry/rest/data_sources.py

@@ -5,7 +5,9 @@
 from feast.api.registry.rest.rest_utils import (
     create_grpc_pagination_params,
     create_grpc_sorting_params,
+    get_object_relationships,
     get_pagination_params,
+    get_relationships_for_objects,
     get_sorting_params,
     grpc_call,
 )
@@ -21,6 +23,9 @@ def get_entity_router(grpc_handler) -> APIRouter:
     def list_entities(
         project: str = Query(...),
         allow_cache: bool = Query(default=True),
+        include_relationships: bool = Query(
+            False, description="Include relationships for each entity"
+        ),
         pagination_params: dict = Depends(get_pagination_params),
         sorting_params: dict = Depends(get_sorting_params),
     ):
@@ -30,20 +35,46 @@ def list_entities(
             pagination=create_grpc_pagination_params(pagination_params),
             sorting=create_grpc_sorting_params(sorting_params),
         )
+        response = grpc_call(grpc_handler.ListEntities, req)
+        entities = response.get("entities", [])
 
-        return grpc_call(grpc_handler.ListEntities, req)
+        result = {
+            "entities": entities,
+            "pagination": response.get("pagination", {}),
+        }
+
+        if include_relationships:
+            relationships = get_relationships_for_objects(
+                grpc_handler, entities, "entity", project, allow_cache
+            )
+            result["relationships"] = relationships
+
+        return result
 
     @router.get("/entities/{name}")
     def get_entity(
         name: str,
         project: str = Query(...),
+        include_relationships: bool = Query(
+            False, description="Include relationships for this entity"
+        ),
         allow_cache: bool = Query(default=True),
     ):
         req = RegistryServer_pb2.GetEntityRequest(
             name=name,
             project=project,
             allow_cache=allow_cache,
         )
-        return grpc_call(grpc_handler.GetEntity, req)
+        entity = grpc_call(grpc_handler.GetEntity, req)
+
+        result = entity
+
+        if include_relationships:
+            relationships = get_object_relationships(
+                grpc_handler, "entity", name, project, allow_cache
+            )
+            result["relationships"] = relationships
+
+        return result
 
     return router

sdk/python/feast/api/registry/rest/entities.py

@@ -5,7 +5,9 @@
 from feast.api.registry.rest.rest_utils import (
     create_grpc_pagination_params,
     create_grpc_sorting_params,
+    get_object_relationships,
     get_pagination_params,
+    get_relationships_for_objects,
     get_sorting_params,
     grpc_call,
     parse_tags,
@@ -19,6 +21,9 @@ def get_feature_service_router(grpc_handler) -> APIRouter:
     @router.get("/feature_services")
     def list_feature_services(
         project: str = Query(...),
+        include_relationships: bool = Query(
+            False, description="Include relationships for each feature service"
+        ),
         allow_cache: bool = Query(default=True),
         tags: Dict[str, str] = Depends(parse_tags),
         pagination_params: dict = Depends(get_pagination_params),
@@ -31,19 +36,46 @@ def list_feature_services(
             pagination=create_grpc_pagination_params(pagination_params),
             sorting=create_grpc_sorting_params(sorting_params),
         )
-        return grpc_call(grpc_handler.ListFeatureServices, req)
+        response = grpc_call(grpc_handler.ListFeatureServices, req)
+        feature_services = response.get("featureServices", [])
+
+        result = {
+            "featureServices": feature_services,
+            "pagination": response.get("pagination", {}),
+        }
+
+        if include_relationships:
+            relationships = get_relationships_for_objects(
+                grpc_handler, feature_services, "featureService", project, allow_cache
+            )
+            result["relationships"] = relationships
+
+        return result
 
     @router.get("/feature_services/{name}")
     def get_feature_service(
         name: str,
         project: str = Query(...),
+        include_relationships: bool = Query(
+            False, description="Include relationships for this feature service"
+        ),
         allow_cache: bool = Query(default=True),
     ):
         req = RegistryServer_pb2.GetFeatureServiceRequest(
             name=name,
             project=project,
             allow_cache=allow_cache,
         )
-        return grpc_call(grpc_handler.GetFeatureService, req)
+        feature_service = grpc_call(grpc_handler.GetFeatureService, req)
+
+        result = feature_service
+
+        if include_relationships:
+            relationships = get_object_relationships(
+                grpc_handler, "featureService", name, project, allow_cache
+            )
+            result["relationships"] = relationships
+
+        return result
 
     return router

sdk/python/feast/api/registry/rest/feature_services.py

@@ -5,21 +5,42 @@
 from feast.api.registry.rest.rest_utils import (
     create_grpc_pagination_params,
     create_grpc_sorting_params,
+    get_object_relationships,
     get_pagination_params,
+    get_relationships_for_objects,
     get_sorting_params,
     grpc_call,
     parse_tags,
 )
 from feast.registry_server import RegistryServer_pb2
 
 
+def _extract_feature_view_from_any(any_feature_view: dict) -> dict:
+    """Extract the specific feature view type and data from an AnyFeatureView object.
+
+    Args:
+        any_feature_view: Dictionary containing the AnyFeatureView data
+
+    Returns:
+        Dictionary with 'type' and feature view data, or empty dict if no valid type found
+    """
+    for key, value in any_feature_view.items():
+        if value:
+            return {"type": key, **value}
+
+    return {}
+
+
 def get_feature_view_router(grpc_handler) -> APIRouter:
     router = APIRouter()
 
     @router.get("/feature_views/{name}")
     def get_any_feature_view(
         name: str,
         project: str = Query(...),
+        include_relationships: bool = Query(
+            False, description="Include relationships for this feature view"
+        ),
         allow_cache: bool = Query(True),
     ):
         req = RegistryServer_pb2.GetAnyFeatureViewRequest(
@@ -29,18 +50,24 @@ def get_any_feature_view(
         )
         response = grpc_call(grpc_handler.GetAnyFeatureView, req)
         any_feature_view = response.get("anyFeatureView", {})
-        feature_view = (
-            any_feature_view.get("featureView")
-            or any_feature_view.get("onDemandFeatureView")
-            or any_feature_view.get("streamFeatureView")
-            or {}
-        )
-        return {"featureView": feature_view}
+
+        result = _extract_feature_view_from_any(any_feature_view)
+
+        if include_relationships:
+            relationships = get_object_relationships(
+                grpc_handler, "featureView", name, project, allow_cache
+            )
+            result["relationships"] = relationships
+
+        return result
 
     @router.get("/feature_views")
     def list_all_feature_views(
         project: str = Query(...),
         allow_cache: bool = Query(default=True),
+        include_relationships: bool = Query(
+            False, description="Include relationships for each feature view"
+        ),
         tags: Dict[str, str] = Depends(parse_tags),
         pagination_params: dict = Depends(get_pagination_params),
         sorting_params: dict = Depends(get_sorting_params),
@@ -52,6 +79,27 @@ def list_all_feature_views(
             pagination=create_grpc_pagination_params(pagination_params),
             sorting=create_grpc_sorting_params(sorting_params),
         )
-        return grpc_call(grpc_handler.ListAllFeatureViews, req)
+        response = grpc_call(grpc_handler.ListAllFeatureViews, req)
+        any_feature_views = response.get("featureViews", [])
+
+        # Extract the specific type of feature view from each AnyFeatureView
+        feature_views = []
+        for any_feature_view in any_feature_views:
+            feature_view = _extract_feature_view_from_any(any_feature_view)
+            if feature_view:
+                feature_views.append(feature_view)
+
+        result = {
+            "featureViews": feature_views,
+            "pagination": response.get("pagination", {}),
+        }
+
+        if include_relationships:
+            relationships = get_relationships_for_objects(
+                grpc_handler, feature_views, "featureView", project, allow_cache
+            )
+            result["relationships"] = relationships
+
+        return result
 
     return router

sdk/python/feast/api/registry/rest/feature_views.py

@@ -3,7 +3,9 @@
 from feast.api.registry.rest.rest_utils import (
     create_grpc_pagination_params,
     create_grpc_sorting_params,
+    get_object_relationships,
     get_pagination_params,
+    get_relationships_for_objects,
     get_sorting_params,
     grpc_call,
 )
@@ -17,19 +19,37 @@ def get_permission_router(grpc_handler) -> APIRouter:
     def get_permission(
         name: str,
         project: str = Query(...),
+        include_relationships: bool = Query(
+            False, description="Include relationships for this permission"
+        ),
         allow_cache: bool = Query(True),
     ):
         req = RegistryServer_pb2.GetPermissionRequest(
             name=name,
             project=project,
             allow_cache=allow_cache,
         )
-        return {"permission": grpc_call(grpc_handler.GetPermission, req)}
+        permission = grpc_call(grpc_handler.GetPermission, req)
+
+        result = permission
+
+        # Note: permissions may not have relationships in the traditional sense
+        # but we include the functionality for consistency
+        if include_relationships:
+            relationships = get_object_relationships(
+                grpc_handler, "permission", name, project, allow_cache
+            )
+            result["relationships"] = relationships
+
+        return result
 
     @router.get("/permissions")
     def list_permissions(
         project: str = Query(...),
         allow_cache: bool = Query(default=True),
+        include_relationships: bool = Query(
+            False, description="Include relationships for each permission"
+        ),
         pagination_params: dict = Depends(get_pagination_params),
         sorting_params: dict = Depends(get_sorting_params),
     ):
@@ -39,6 +59,20 @@ def list_permissions(
             pagination=create_grpc_pagination_params(pagination_params),
             sorting=create_grpc_sorting_params(sorting_params),
         )
-        return grpc_call(grpc_handler.ListPermissions, req)
+        response = grpc_call(grpc_handler.ListPermissions, req)
+        permissions = response.get("permissions", [])
+
+        result = {
+            "permissions": permissions,
+            "pagination": response.get("pagination", {}),
+        }
+
+        if include_relationships:
+            relationships = get_relationships_for_objects(
+                grpc_handler, permissions, "permission", project, allow_cache
+            )
+            result["relationships"] = relationships
+
+        return result
 
     return router