zenml-nightly 0.75.0.dev20250311__py3-none-any.whl → 0.75.0.dev20250313__py3-none-any.whl

zenml/zen_server/routers/auth_endpoints.py

@@ -66,7 +66,10 @@ from zenml.zen_server.auth import (
 from zenml.zen_server.exceptions import error_response
 from zenml.zen_server.rate_limit import rate_limit_requests
 from zenml.zen_server.rbac.models import Action, ResourceType
-from zenml.zen_server.rbac.utils import verify_permission
+from zenml.zen_server.rbac.utils import (
+    verify_permission,
+    verify_permission_for_model,
+)
 from zenml.zen_server.utils import (
     get_ip_location,
     handle_exceptions,
@@ -544,10 +547,6 @@ def api_token(
             response=None,
         ).access_token
 
-    verify_permission(
-        resource_type=ResourceType.PIPELINE_RUN, action=Action.CREATE
-    )
-
     schedule_id = schedule_id or token.schedule_id
     pipeline_run_id = pipeline_run_id or token.pipeline_run_id
     step_run_id = step_run_id or token.step_run_id
@@ -583,15 +582,18 @@ def api_token(
             f"step run {token.step_run_id}."
         )
 
+    workspace_id: Optional[UUID] = None
+
     if schedule_id:
         # The schedule must exist
         try:
-            schedule = zen_store().get_schedule(schedule_id, hydrate=False)
+            schedule = zen_store().get_schedule(schedule_id, hydrate=True)
         except KeyError:
             raise ValueError(
                 f"Schedule {schedule_id} does not exist and API tokens cannot "
                 "be generated for non-existent schedules for security reasons."
             )
+        workspace_id = schedule.workspace.id
 
         if not schedule.active:
             raise ValueError(
@@ -602,7 +604,7 @@ def api_token(
     if pipeline_run_id:
         # The pipeline run must exist and the run must not be concluded
         try:
-            pipeline_run = zen_store().get_run(pipeline_run_id, hydrate=False)
+            pipeline_run = zen_store().get_run(pipeline_run_id, hydrate=True)
         except KeyError:
             raise ValueError(
                 f"Pipeline run {pipeline_run_id} does not exist and API tokens "
@@ -610,6 +612,10 @@ def api_token(
                 "security reasons."
             )
 
+        verify_permission_for_model(model=pipeline_run, action=Action.READ)
+
+        workspace_id = pipeline_run.workspace.id
+
         if pipeline_run.status.is_finished:
             raise ValueError(
                 f"The execution of pipeline run {pipeline_run_id} has already "
@@ -627,6 +633,8 @@ def api_token(
                 "be generated for non-existent step runs for security reasons."
             )
 
+        workspace_id = step_run.workspace.id
+
         if step_run.status.is_finished:
             raise ValueError(
                 f"The execution of step run {step_run_id} has already "
@@ -634,6 +642,13 @@ def api_token(
                 "for security reasons."
             )
 
+    assert workspace_id is not None
+    verify_permission(
+        resource_type=ResourceType.PIPELINE_RUN,
+        action=Action.CREATE,
+        workspace_id=workspace_id,
+    )
+
     return generate_access_token(
         user_id=token.user_id,
         # Keep the original API key and device token scopes

zenml/zen_server/routers/code_repositories_endpoints.py

@@ -13,6 +13,7 @@
 #  permissions and limitations under the License.
 """Endpoint definitions for code repositories."""
 
+from typing import Optional, Union
 from uuid import UUID
 
 from fastapi import APIRouter, Depends, Security
@@ -20,6 +21,7 @@ from fastapi import APIRouter, Depends, Security
 from zenml.constants import API, CODE_REPOSITORIES, VERSION_1
 from zenml.models import (
     CodeRepositoryFilter,
+    CodeRepositoryRequest,
     CodeRepositoryResponse,
     CodeRepositoryUpdate,
     Page,
@@ -27,12 +29,16 @@ from zenml.models import (
 from zenml.zen_server.auth import AuthContext, authorize
 from zenml.zen_server.exceptions import error_response
 from zenml.zen_server.rbac.endpoint_utils import (
+    verify_permissions_and_create_entity,
     verify_permissions_and_delete_entity,
     verify_permissions_and_get_entity,
     verify_permissions_and_list_entities,
     verify_permissions_and_update_entity,
 )
 from zenml.zen_server.rbac.models import ResourceType
+from zenml.zen_server.routers.workspaces_endpoints import (
+    router as workspace_router,
+)
 from zenml.zen_server.utils import (
     handle_exceptions,
     make_dependable,
@@ -46,16 +52,61 @@ router = APIRouter(
 )
 
 
+@router.post(
+    "",
+    responses={401: error_response, 409: error_response, 422: error_response},
+)
+# TODO: the workspace scoped endpoint is only kept for dashboard compatibility
+# and can be removed after the migration
+@workspace_router.post(
+    "/{workspace_name_or_id}" + CODE_REPOSITORIES,
+    responses={401: error_response, 409: error_response, 422: error_response},
+    deprecated=True,
+    tags=["code_repositories"],
+)
+@handle_exceptions
+def create_code_repository(
+    code_repository: CodeRepositoryRequest,
+    workspace_name_or_id: Optional[Union[str, UUID]] = None,
+    _: AuthContext = Security(authorize),
+) -> CodeRepositoryResponse:
+    """Creates a code repository.
+
+    Args:
+        code_repository: Code repository to create.
+        workspace_name_or_id: Optional name or ID of the workspace.
+
+    Returns:
+        The created code repository.
+    """
+    if workspace_name_or_id:
+        workspace = zen_store().get_workspace(workspace_name_or_id)
+        code_repository.workspace = workspace.id
+
+    return verify_permissions_and_create_entity(
+        request_model=code_repository,
+        create_method=zen_store().create_code_repository,
+    )
+
+
 @router.get(
     "",
-    response_model=Page[CodeRepositoryResponse],
     responses={401: error_response, 404: error_response, 422: error_response},
 )
+# TODO: the workspace scoped endpoint is only kept for dashboard compatibility
+# and can be removed after the migration
+@workspace_router.get(
+    "/{workspace_name_or_id}" + CODE_REPOSITORIES,
+    responses={401: error_response, 404: error_response, 422: error_response},
+    deprecated=True,
+    tags=["code_repositories"],
+)
 @handle_exceptions
 def list_code_repositories(
     filter_model: CodeRepositoryFilter = Depends(
         make_dependable(CodeRepositoryFilter)
     ),
+    workspace_name_or_id: Optional[Union[str, UUID]] = None,
     hydrate: bool = False,
     _: AuthContext = Security(authorize),
 ) -> Page[CodeRepositoryResponse]:
@@ -64,12 +115,16 @@ def list_code_repositories(
     Args:
         filter_model: Filter model used for pagination, sorting,
             filtering.
+        workspace_name_or_id: Optional name or ID of the workspace.
         hydrate: Flag deciding whether to hydrate the output model(s)
             by including metadata fields in the response.
 
     Returns:
         Page of code repository objects.
     """
+    if workspace_name_or_id:
+        filter_model.workspace = workspace_name_or_id
+
     return verify_permissions_and_list_entities(
         filter_model=filter_model,
         resource_type=ResourceType.CODE_REPOSITORY,
@@ -80,7 +135,6 @@ def list_code_repositories(
 
 @router.get(
     "/{code_repository_id}",
-    response_model=CodeRepositoryResponse,
     responses={401: error_response, 404: error_response, 422: error_response},
 )
 @handle_exceptions
@@ -108,7 +162,6 @@ def get_code_repository(
 
 @router.put(
     "/{code_repository_id}",
-    response_model=CodeRepositoryResponse,
     responses={401: error_response, 404: error_response, 422: error_response},
 )
 @handle_exceptions

zenml/zen_server/routers/devices_endpoints.py

@@ -54,7 +54,6 @@ router = APIRouter(
 
 @router.get(
     "",
-    response_model=Page[OAuthDeviceResponse],
     responses={401: error_response, 404: error_response, 422: error_response},
 )
 @handle_exceptions
@@ -85,7 +84,6 @@ def list_authorized_devices(
 
 @router.get(
     "/{device_id}",
-    response_model=OAuthDeviceResponse,
     responses={401: error_response, 404: error_response, 422: error_response},
 )
 @handle_exceptions
@@ -136,7 +134,6 @@ def get_authorization_device(
 
 @router.put(
     "/{device_id}",
-    response_model=OAuthDeviceResponse,
     responses={401: error_response, 404: error_response, 422: error_response},
 )
 @handle_exceptions
@@ -173,7 +170,6 @@ def update_authorized_device(
 
 @router.put(
     "/{device_id}" + DEVICE_VERIFY,
-    response_model=OAuthDeviceResponse,
     responses={401: error_response, 404: error_response, 422: error_response},
 )
 @handle_exceptions

zenml/zen_server/routers/event_source_endpoints.py

@@ -57,7 +57,6 @@ event_source_router = APIRouter(
 
 @event_source_router.get(
     "",
-    response_model=Page[EventSourceResponse],
     responses={401: error_response, 404: error_response, 422: error_response},
 )
 @handle_exceptions
@@ -132,7 +131,6 @@ def list_event_sources(
 
 @event_source_router.get(
     "/{event_source_id}",
-    response_model=EventSourceResponse,
     responses={401: error_response, 404: error_response, 422: error_response},
 )
 @handle_exceptions
@@ -185,7 +183,6 @@ def get_event_source(
 
 @event_source_router.post(
     "",
-    response_model=EventSourceResponse,
     responses={401: error_response, 409: error_response, 422: error_response},
 )
 @handle_exceptions
@@ -222,14 +219,12 @@ def create_event_source(
 
     return verify_permissions_and_create_entity(
         request_model=event_source,
-        resource_type=ResourceType.EVENT_SOURCE,
         create_method=event_source_handler.create_event_source,
     )
 
 
 @event_source_router.put(
     "/{event_source_id}",
-    response_model=EventSourceResponse,
     responses={401: error_response, 404: error_response, 422: error_response},
 )
 @handle_exceptions

zenml/zen_server/routers/flavors_endpoints.py

@@ -51,7 +51,6 @@ router = APIRouter(
 
 @router.get(
     "",
-    response_model=Page[FlavorResponse],
     responses={401: error_response, 404: error_response, 422: error_response},
 )
 @handle_exceptions
@@ -81,7 +80,6 @@ def list_flavors(
 
 @router.get(
     "/{flavor_id}",
-    response_model=FlavorResponse,
     responses={401: error_response, 404: error_response, 422: error_response},
 )
 @handle_exceptions
@@ -107,7 +105,6 @@ def get_flavor(
 
 @router.post(
     "",
-    response_model=FlavorResponse,
     responses={401: error_response, 409: error_response, 422: error_response},
 )
 @handle_exceptions
@@ -125,14 +122,12 @@ def create_flavor(
     """
     return verify_permissions_and_create_entity(
         request_model=flavor,
-        resource_type=ResourceType.FLAVOR,
         create_method=zen_store().create_flavor,
     )
 
 
 @router.put(
     "/{flavor_id}",
-    response_model=FlavorResponse,
     responses={401: error_response, 409: error_response, 422: error_response},
 )
 @handle_exceptions

zenml/zen_server/routers/logs_endpoints.py

@@ -42,7 +42,6 @@ router = APIRouter(
 
 @router.get(
     "/{logs_id}",
-    response_model=LogsResponse,
     responses={401: error_response, 404: error_response, 422: error_response},
 )
 @handle_exceptions

zenml/zen_server/routers/model_versions_endpoints.py

@@ -13,7 +13,7 @@
 #  permissions and limitations under the License.
 """Endpoint definitions for models."""
 
-from typing import Union
+from typing import Optional, Union
 from uuid import UUID
 
 from fastapi import APIRouter, Depends, Security
@@ -24,6 +24,7 @@ from zenml.constants import (
     MODEL_VERSION_ARTIFACTS,
     MODEL_VERSION_PIPELINE_RUNS,
     MODEL_VERSIONS,
+    MODELS,
     RUNS,
     VERSION_1,
 )
@@ -35,22 +36,33 @@ from zenml.models import (
     ModelVersionPipelineRunFilter,
     ModelVersionPipelineRunRequest,
     ModelVersionPipelineRunResponse,
+    ModelVersionRequest,
     ModelVersionResponse,
     ModelVersionUpdate,
 )
 from zenml.models.v2.base.page import Page
 from zenml.zen_server.auth import AuthContext, authorize
 from zenml.zen_server.exceptions import error_response
-from zenml.zen_server.rbac.models import Action, ResourceType
+from zenml.zen_server.rbac.endpoint_utils import (
+    verify_permissions_and_create_entity,
+    verify_permissions_and_delete_entity,
+)
+from zenml.zen_server.rbac.models import Action
 from zenml.zen_server.rbac.utils import (
     dehydrate_page,
     dehydrate_response_model,
-    get_allowed_resource_ids,
     verify_permission_for_model,
 )
+from zenml.zen_server.routers.models_endpoints import (
+    router as model_router,
+)
+from zenml.zen_server.routers.workspaces_endpoints import (
+    router as workspace_router,
+)
 from zenml.zen_server.utils import (
     handle_exceptions,
     make_dependable,
+    set_filter_workspace_scope,
     zen_store,
 )
 
@@ -66,9 +78,54 @@ router = APIRouter(
 )
 
 
+@router.post(
+    "",
+    responses={401: error_response, 409: error_response, 422: error_response},
+)
+# TODO: the workspace scoped endpoint is only kept for dashboard compatibility
+# and can be removed after the migration
+@workspace_router.post(
+    "/{workspace_name_or_id}" + MODELS + "/{model_id}" + MODEL_VERSIONS,
+    responses={401: error_response, 409: error_response, 422: error_response},
+    deprecated=True,
+    tags=["model_versions"],
+)
+@handle_exceptions
+def create_model_version(
+    model_version: ModelVersionRequest,
+    model_id: Optional[UUID] = None,
+    workspace_name_or_id: Optional[Union[str, UUID]] = None,
+    _: AuthContext = Security(authorize),
+) -> ModelVersionResponse:
+    """Creates a model version.
+
+    Args:
+        model_version: Model version to create.
+        model_id: Optional ID of the model.
+        workspace_name_or_id: Optional name or ID of the workspace.
+
+    Returns:
+        The created model version.
+    """
+    if workspace_name_or_id:
+        workspace = zen_store().get_workspace(workspace_name_or_id)
+        model_version.workspace = workspace.id
+
+    if model_id:
+        model_version.model = model_id
+
+    return verify_permissions_and_create_entity(
+        request_model=model_version,
+        create_method=zen_store().create_model_version,
+    )
+
+
+@model_router.get(
+    "/{model_name_or_id}" + MODEL_VERSIONS,
+    responses={401: error_response, 404: error_response, 422: error_response},
+)
 @router.get(
     "",
-    response_model=Page[ModelVersionResponse],
     responses={401: error_response, 404: error_response, 422: error_response},
 )
 @handle_exceptions
@@ -76,6 +133,7 @@ def list_model_versions(
     model_version_filter_model: ModelVersionFilter = Depends(
         make_dependable(ModelVersionFilter)
     ),
+    model_name_or_id: Optional[Union[str, UUID]] = None,
     hydrate: bool = False,
     auth_context: AuthContext = Security(authorize),
 ) -> Page[ModelVersionResponse]:
@@ -84,19 +142,38 @@ def list_model_versions(
     Args:
         model_version_filter_model: Filter model used for pagination, sorting,
             filtering.
+        model_name_or_id: Optional name or ID of the model.
         hydrate: Flag deciding whether to hydrate the output model(s)
             by including metadata fields in the response.
         auth_context: The authentication context.
 
     Returns:
         The model versions according to query filters.
+
+    Raises:
+        ValueError: If the model is missing from the filter.
     """
-    allowed_model_ids = get_allowed_resource_ids(
-        resource_type=ResourceType.MODEL
+    if model_name_or_id:
+        model_version_filter_model.model = model_name_or_id
+
+    if not model_version_filter_model.model:
+        raise ValueError("Model missing from the filter")
+
+    # A workspace scoped request must always be scoped to a specific
+    # workspace. This is required for the RBAC check to work.
+    set_filter_workspace_scope(model_version_filter_model)
+    assert isinstance(model_version_filter_model.workspace, UUID)
+
+    model = zen_store().get_model_by_name_or_id(
+        model_version_filter_model.model,
+        workspace=model_version_filter_model.workspace,
     )
+
+    # Check read permissions on the model
+    verify_permission_for_model(model, action=Action.READ)
+
     model_version_filter_model.configure_rbac(
-        authenticated_user_id=auth_context.user.id,
-        model_id=allowed_model_ids,
+        authenticated_user_id=auth_context.user.id
     )
 
     model_versions = zen_store().list_model_versions(
@@ -108,7 +185,6 @@ def list_model_versions(
 
 @router.get(
     "/{model_version_id}",
-    response_model=ModelVersionResponse,
     responses={401: error_response, 404: error_response, 422: error_response},
 )
 @handle_exceptions
@@ -137,7 +213,6 @@ def get_model_version(
 
 @router.put(
     "/{model_version_id}",
-    response_model=ModelVersionResponse,
     responses={401: error_response, 404: error_response, 422: error_response},
 )
 @handle_exceptions
@@ -184,9 +259,11 @@ def delete_model_version(
     Args:
         model_version_id: The name or ID of the model version to delete.
     """
-    model_version = zen_store().get_model_version(model_version_id)
-    verify_permission_for_model(model_version, action=Action.DELETE)
-    zen_store().delete_model_version(model_version_id)
+    verify_permissions_and_delete_entity(
+        id=model_version_id,
+        get_method=zen_store().get_model_version,
+        delete_method=zen_store().delete_model_version,
+    )
 
 
 ##########################
@@ -220,17 +297,18 @@ def create_model_version_artifact_link(
     model_version = zen_store().get_model_version(
         model_version_artifact_link.model_version
     )
-    verify_permission_for_model(model_version, action=Action.UPDATE)
 
-    mv = zen_store().create_model_version_artifact_link(
-        model_version_artifact_link
+    return verify_permissions_and_create_entity(
+        request_model=model_version_artifact_link,
+        create_method=zen_store().create_model_version_artifact_link,
+        # Check for UPDATE permissions on the model version instead of the
+        # model version artifact link
+        surrogate_models=[model_version],
     )
-    return mv
 
 
 @model_version_artifacts_router.get(
     "",
-    response_model=Page[ModelVersionArtifactResponse],
     responses={401: error_response, 404: error_response, 422: error_response},
 )
 @handle_exceptions
@@ -342,17 +420,18 @@ def create_model_version_pipeline_run_link(
     model_version = zen_store().get_model_version(
         model_version_pipeline_run_link.model_version, hydrate=False
     )
-    verify_permission_for_model(model_version, action=Action.UPDATE)
 
-    mv = zen_store().create_model_version_pipeline_run_link(
-        model_version_pipeline_run_link
+    return verify_permissions_and_create_entity(
+        request_model=model_version_pipeline_run_link,
+        create_method=zen_store().create_model_version_pipeline_run_link,
+        # Check for UPDATE permissions on the model version instead of the
+        # model version pipeline run link
+        surrogate_models=[model_version],
     )
-    return mv
 
 
 @model_version_pipeline_runs_router.get(
     "",
-    response_model=Page[ModelVersionPipelineRunResponse],
     responses={401: error_response, 404: error_response, 422: error_response},
 )
 @handle_exceptions