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

zenml/zen_server/routers/stacks_endpoints.py

@@ -13,12 +13,19 @@
 #  permissions and limitations under the License.
 """Endpoint definitions for stacks."""
 
+from typing import Optional, Union
 from uuid import UUID
 
 from fastapi import APIRouter, Depends, Security
 
 from zenml.constants import API, STACKS, VERSION_1
-from zenml.models import Page, StackFilter, StackResponse, StackUpdate
+from zenml.models import (
+    Page,
+    StackFilter,
+    StackRequest,
+    StackResponse,
+    StackUpdate,
+)
 from zenml.zen_server.auth import AuthContext, authorize
 from zenml.zen_server.exceptions import error_response
 from zenml.zen_server.rbac.endpoint_utils import (
@@ -28,7 +35,14 @@ from zenml.zen_server.rbac.endpoint_utils import (
     verify_permissions_and_update_entity,
 )
 from zenml.zen_server.rbac.models import Action, ResourceType
-from zenml.zen_server.rbac.utils import batch_verify_permissions_for_models
+from zenml.zen_server.rbac.utils import (
+    batch_verify_permissions_for_models,
+    verify_permission,
+    verify_permission_for_model,
+)
+from zenml.zen_server.routers.workspaces_endpoints import (
+    router as workspace_router,
+)
 from zenml.zen_server.utils import (
     handle_exceptions,
     make_dependable,
@@ -42,27 +56,106 @@ 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}" + STACKS,
+    responses={401: error_response, 409: error_response, 422: error_response},
+    deprecated=True,
+    tags=["stacks"],
+)
+@handle_exceptions
+def create_stack(
+    stack: StackRequest,
+    workspace_name_or_id: Optional[Union[str, UUID]] = None,
+    auth_context: AuthContext = Security(authorize),
+) -> StackResponse:
+    """Creates a stack, optionally in a specific workspace.
+
+    Args:
+        stack: Stack to register.
+        workspace_name_or_id: Optional name or ID of the workspace.
+        auth_context: Authentication context.
+
+    Returns:
+        The created stack.
+    """
+    # Check the service connector creation
+    is_connector_create_needed = False
+    for connector_id_or_info in stack.service_connectors:
+        if isinstance(connector_id_or_info, UUID):
+            service_connector = zen_store().get_service_connector(
+                connector_id_or_info, hydrate=False
+            )
+            verify_permission_for_model(
+                model=service_connector, action=Action.READ
+            )
+        else:
+            is_connector_create_needed = True
+
+    # Check the component creation
+    if is_connector_create_needed:
+        verify_permission(
+            resource_type=ResourceType.SERVICE_CONNECTOR, action=Action.CREATE
+        )
+    is_component_create_needed = False
+    for components in stack.components.values():
+        for component_id_or_info in components:
+            if isinstance(component_id_or_info, UUID):
+                component = zen_store().get_stack_component(
+                    component_id_or_info, hydrate=False
+                )
+                verify_permission_for_model(
+                    model=component, action=Action.READ
+                )
+            else:
+                is_component_create_needed = True
+    if is_component_create_needed:
+        verify_permission(
+            resource_type=ResourceType.STACK_COMPONENT,
+            action=Action.CREATE,
+        )
+
+    # Check the stack creation
+    verify_permission_for_model(model=stack, action=Action.CREATE)
+
+    return zen_store().create_stack(stack)
+
+
 @router.get(
     "",
-    response_model=Page[StackResponse],
     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}" + STACKS,
+    responses={401: error_response, 404: error_response, 422: error_response},
+    deprecated=True,
+    tags=["stacks"],
+)
 @handle_exceptions
 def list_stacks(
+    workspace_name_or_id: Optional[Union[str, UUID]] = None,
     stack_filter_model: StackFilter = Depends(make_dependable(StackFilter)),
     hydrate: bool = False,
     _: AuthContext = Security(authorize),
 ) -> Page[StackResponse]:
-    """Returns all stacks.
+    """Returns all stacks, optionally filtered by workspace.
 
     Args:
+        workspace_name_or_id: Optional name or ID of the workspace to filter by.
         stack_filter_model: Filter model used for pagination, sorting,
             filtering.
         hydrate: Flag deciding whether to hydrate the output model(s)
             by including metadata fields in the response.
 
     Returns:
-        All stacks.
+        All stacks matching the filter criteria.
     """
     return verify_permissions_and_list_entities(
         filter_model=stack_filter_model,
@@ -74,7 +167,6 @@ def list_stacks(
 
 @router.get(
     "/{stack_id}",
-    response_model=StackResponse,
     responses={401: error_response, 404: error_response, 422: error_response},
 )
 @handle_exceptions
@@ -100,7 +192,6 @@ def get_stack(
 
 @router.put(
     "/{stack_id}",
-    response_model=StackResponse,
     responses={401: error_response, 404: error_response, 422: error_response},
 )
 @handle_exceptions

zenml/zen_server/routers/steps_endpoints.py

@@ -37,6 +37,9 @@ 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,
+)
 from zenml.zen_server.rbac.models import Action, ResourceType
 from zenml.zen_server.rbac.utils import (
     dehydrate_page,
@@ -47,6 +50,7 @@ from zenml.zen_server.rbac.utils import (
 from zenml.zen_server.utils import (
     handle_exceptions,
     make_dependable,
+    set_filter_workspace_scope,
     zen_store,
 )
 
@@ -59,7 +63,6 @@ router = APIRouter(
 
 @router.get(
     "",
-    response_model=Page[StepRunResponse],
     responses={401: error_response, 404: error_response, 422: error_response},
 )
 @handle_exceptions
@@ -82,8 +85,14 @@ def list_run_steps(
     Returns:
         The run steps according to query filters.
     """
+    # 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(step_run_filter_model)
+    assert isinstance(step_run_filter_model.workspace, UUID)
+
     allowed_pipeline_run_ids = get_allowed_resource_ids(
-        resource_type=ResourceType.PIPELINE_RUN
+        resource_type=ResourceType.PIPELINE_RUN,
+        workspace_id=step_run_filter_model.workspace,
     )
     step_run_filter_model.configure_rbac(
         authenticated_user_id=auth_context.user.id,
@@ -98,7 +107,6 @@ def list_run_steps(
 
 @router.post(
     "",
-    response_model=StepRunResponse,
     responses={401: error_response, 409: error_response, 422: error_response},
 )
 @handle_exceptions
@@ -110,20 +118,22 @@ def create_run_step(
 
     Args:
         step: The run step to create.
+        _: Authentication context.
 
     Returns:
         The created run step.
     """
     pipeline_run = zen_store().get_run(step.pipeline_run_id)
-    verify_permission_for_model(pipeline_run, action=Action.UPDATE)
 
-    step_response = zen_store().create_run_step(step_run=step)
-    return dehydrate_response_model(step_response)
+    return verify_permissions_and_create_entity(
+        request_model=step,
+        create_method=zen_store().create_run_step,
+        surrogate_models=[pipeline_run],
+    )
 
 
 @router.get(
     "/{step_id}",
-    response_model=StepRunResponse,
     responses={401: error_response, 404: error_response, 422: error_response},
 )
 @handle_exceptions
@@ -157,7 +167,6 @@ def get_step(
 
 @router.put(
     "/{step_id}",
-    response_model=StepRunResponse,
     responses={401: error_response, 404: error_response, 422: error_response},
 )
 @handle_exceptions
@@ -187,7 +196,6 @@ def update_step(
 
 @router.get(
     "/{step_id}" + STEP_CONFIGURATION,
-    response_model=Dict[str, Any],
     responses={401: error_response, 404: error_response, 422: error_response},
 )
 @handle_exceptions
@@ -212,7 +220,6 @@ def get_step_configuration(
 
 @router.get(
     "/{step_id}" + STATUS,
-    response_model=ExecutionStatus,
     responses={401: error_response, 404: error_response, 422: error_response},
 )
 @handle_exceptions
@@ -237,7 +244,6 @@ def get_step_status(
 
 @router.get(
     "/{step_id}" + LOGS,
-    response_model=str,
     responses={401: error_response, 404: error_response, 422: error_response},
 )
 @handle_exceptions

zenml/zen_server/routers/tag_resource_endpoints.py

@@ -0,0 +1,115 @@
+#  Copyright (c) ZenML GmbH 2023. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at:
+#
+#       https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+#  or implied. See the License for the specific language governing
+#  permissions and limitations under the License.
+"""Endpoint definitions for the link between tags and resources."""
+
+from typing import List
+
+from fastapi import APIRouter, Security
+
+from zenml.constants import (
+    API,
+    BATCH,
+    TAG_RESOURCES,
+    VERSION_1,
+)
+from zenml.models import TagResourceRequest, TagResourceResponse
+from zenml.zen_server.auth import AuthContext, authorize
+from zenml.zen_server.exceptions import error_response
+from zenml.zen_server.utils import (
+    handle_exceptions,
+    zen_store,
+)
+
+router = APIRouter(
+    prefix=API + VERSION_1 + TAG_RESOURCES,
+    tags=["tag_resources"],
+    responses={401: error_response, 403: error_response},
+)
+
+
+@router.post(
+    "",
+    responses={401: error_response, 404: error_response, 422: error_response},
+)
+@handle_exceptions
+def create_tag_resource(
+    tag_resource: TagResourceRequest,
+    _: AuthContext = Security(authorize),
+) -> TagResourceResponse:
+    """Attach different tags to different resources.
+
+    Args:
+        tag_resource: A tag resource request.
+
+    Returns:
+        A tag resource response.
+    """
+    return zen_store().create_tag_resource(tag_resource=tag_resource)
+
+
+@router.post(
+    BATCH,
+    responses={401: error_response, 409: error_response, 422: error_response},
+)
+@handle_exceptions
+def batch_create_tag_resource(
+    tag_resources: List[TagResourceRequest],
+    _: AuthContext = Security(authorize),
+) -> List[TagResourceResponse]:
+    """Attach different tags to different resources.
+
+    Args:
+        tag_resources: A list of tag resource requests.
+
+    Returns:
+        A list of tag resource responses.
+    """
+    return [
+        zen_store().create_tag_resource(tag_resource=tag_resource)
+        for tag_resource in tag_resources
+    ]
+
+
+@router.delete(
+    "",
+    responses={401: error_response, 404: error_response, 422: error_response},
+)
+@handle_exceptions
+def delete_tag_resource(
+    tag_resource: TagResourceRequest,
+    _: AuthContext = Security(authorize),
+) -> None:
+    """Detach a tag from a resource.
+
+    Args:
+        tag_resource: The tag resource relationship to delete.
+    """
+    zen_store().delete_tag_resource(tag_resource=tag_resource)
+
+
+@router.delete(
+    BATCH,
+    responses={401: error_response, 404: error_response, 422: error_response},
+)
+@handle_exceptions
+def batch_delete_tag_resource(
+    tag_resources: List[TagResourceRequest],
+    _: AuthContext = Security(authorize),
+) -> None:
+    """Detach different tags from different resources.
+
+    Args:
+        tag_resources: A list of tag resource requests.
+    """
+    zen_store().batch_delete_tag_resource(tag_resources=tag_resources)

zenml/zen_server/routers/tags_endpoints.py

@@ -33,13 +33,10 @@ 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.utils import (
     handle_exceptions,
     make_dependable,
@@ -59,7 +56,6 @@ router = APIRouter(
 
 @router.post(
     "",
-    response_model=TagResponse,
     responses={401: error_response, 409: error_response, 422: error_response},
 )
 @handle_exceptions
@@ -75,16 +71,11 @@ def create_tag(
     Returns:
         The created tag.
     """
-    return verify_permissions_and_create_entity(
-        request_model=tag,
-        resource_type=ResourceType.TAG,
-        create_method=zen_store().create_tag,
-    )
+    return zen_store().create_tag(tag=tag)
 
 
 @router.get(
     "",
-    response_model=Page[TagResponse],
     responses={401: error_response, 404: error_response, 422: error_response},
 )
 @handle_exceptions
@@ -104,17 +95,14 @@ def list_tags(
     Returns:
         The tags according to query filters.
     """
-    return verify_permissions_and_list_entities(
-        filter_model=tag_filter_model,
-        resource_type=ResourceType.TAG,
-        list_method=zen_store().list_tags,
+    return zen_store().list_tags(
+        tag_filter_model=tag_filter_model,
         hydrate=hydrate,
     )
 
 
 @router.get(
     "/{tag_name_or_id}",
-    response_model=TagResponse,
     responses={401: error_response, 404: error_response, 422: error_response},
 )
 @handle_exceptions
@@ -134,13 +122,14 @@ def get_tag(
         The tag with the given name or ID.
     """
     return verify_permissions_and_get_entity(
-        id=tag_name_or_id, get_method=zen_store().get_tag, hydrate=hydrate
+        id=tag_name_or_id,
+        get_method=zen_store().get_tag,
+        hydrate=hydrate,
     )
 
 
 @router.put(
     "/{tag_id}",
-    response_model=TagResponse,
     responses={401: error_response, 404: error_response, 422: error_response},
 )
 @handle_exceptions

zenml/zen_server/routers/triggers_endpoints.py

@@ -58,7 +58,6 @@ router = APIRouter(
 
 @router.get(
     "",
-    response_model=Page[TriggerResponse],
     responses={401: error_response, 404: error_response, 422: error_response},
 )
 @handle_exceptions
@@ -90,7 +89,6 @@ def list_triggers(
 
 @router.get(
     "/{trigger_id}",
-    response_model=TriggerResponse,
     responses={401: error_response, 404: error_response, 422: error_response},
 )
 @handle_exceptions
@@ -116,7 +114,6 @@ def get_trigger(
 
 @router.post(
     "",
-    response_model=TriggerResponse,
     responses={401: error_response, 409: error_response, 422: error_response},
 )
 @handle_exceptions
@@ -162,14 +159,12 @@ def create_trigger(
 
     return verify_permissions_and_create_entity(
         request_model=trigger,
-        resource_type=ResourceType.TRIGGER,
         create_method=zen_store().create_trigger,
     )
 
 
 @router.put(
     "/{trigger_id}",
-    response_model=TriggerResponse,
     responses={401: error_response, 404: error_response, 422: error_response},
 )
 @handle_exceptions
@@ -245,9 +240,11 @@ def delete_trigger(
     Args:
         trigger_id: Name of the trigger.
     """
-    trigger = zen_store().get_trigger(trigger_id=trigger_id)
-    verify_permission_for_model(trigger, action=Action.DELETE)
-    zen_store().delete_trigger(trigger_id=trigger_id)
+    verify_permissions_and_delete_entity(
+        id=trigger_id,
+        get_method=zen_store().get_trigger,
+        delete_method=zen_store().delete_trigger,
+    )
 
 
 executions_router = APIRouter(

zenml/zen_server/routers/users_endpoints.py

@@ -38,6 +38,7 @@ from zenml.models import (
     UserRequest,
     UserResponse,
     UserUpdate,
+    WorkspaceScopedResponse,
 )
 from zenml.zen_server.auth import (
     AuthContext,
@@ -46,6 +47,9 @@ from zenml.zen_server.auth import (
 )
 from zenml.zen_server.exceptions import error_response
 from zenml.zen_server.rate_limit import RequestLimiter
+from zenml.zen_server.rbac.endpoint_utils import (
+    verify_permissions_and_get_entity,
+)
 from zenml.zen_server.rbac.models import Action, Resource, ResourceType
 from zenml.zen_server.rbac.utils import (
     dehydrate_page,
@@ -87,7 +91,6 @@ current_user_router = APIRouter(
 
 @router.get(
     "",
-    response_model=Page[UserResponse],
     responses={401: error_response, 404: error_response, 422: error_response},
 )
 @handle_exceptions
@@ -133,7 +136,6 @@ if server_config().auth_scheme != AuthScheme.EXTERNAL:
 
     @router.post(
         "",
-        response_model=UserResponse,
         responses={
             401: error_response,
             409: error_response,
@@ -174,7 +176,6 @@ if server_config().auth_scheme != AuthScheme.EXTERNAL:
 
         # new_user = verify_permissions_and_create_entity(
         #     request_model=user,
-        #     resource_type=ResourceType.USER,
         #     create_method=zen_store().create_user,
         # )
         new_user = zen_store().create_user(user)
@@ -188,7 +189,6 @@ if server_config().auth_scheme != AuthScheme.EXTERNAL:
 
 @router.get(
     "/{user_name_or_id}",
-    response_model=UserResponse,
     responses={401: error_response, 404: error_response, 422: error_response},
 )
 @handle_exceptions
@@ -233,7 +233,6 @@ if server_config().auth_scheme != AuthScheme.EXTERNAL:
 
     @router.put(
         "/{user_name_or_id}",
-        response_model=UserResponse,
         responses={
             401: error_response,
             404: error_response,
@@ -403,7 +402,6 @@ if server_config().auth_scheme != AuthScheme.EXTERNAL:
 
     @activation_router.put(
         "/{user_name_or_id}" + ACTIVATE,
-        response_model=UserResponse,
         responses={
             401: error_response,
             404: error_response,
@@ -464,7 +462,6 @@ if server_config().auth_scheme != AuthScheme.EXTERNAL:
 
     @router.put(
         "/{user_name_or_id}" + DEACTIVATE,
-        response_model=UserResponse,
         responses={
             401: error_response,
             404: error_response,
@@ -555,7 +552,6 @@ if server_config().auth_scheme != AuthScheme.EXTERNAL:
 
     @router.put(
         "/{user_name_or_id}" + EMAIL_ANALYTICS,
-        response_model=UserResponse,
         responses={
             401: error_response,
             404: error_response,
@@ -608,7 +604,6 @@ if server_config().auth_scheme != AuthScheme.EXTERNAL:
 
 @current_user_router.get(
     "/current-user",
-    response_model=UserResponse,
     responses={401: error_response, 404: error_response, 422: error_response},
 )
 @handle_exceptions
@@ -632,7 +627,6 @@ if server_config().auth_scheme != AuthScheme.EXTERNAL:
 
     @current_user_router.put(
         "/current-user",
-        response_model=UserResponse,
         responses={
             401: error_response,
             404: error_response,
@@ -752,8 +746,6 @@ if server_config().rbac_enabled:
             )
 
         resource_type = ResourceType(resource_type)
-        resource = Resource(type=resource_type, id=resource_id)
-
         schema_class = get_schema_for_resource_type(resource_type)
         model = zen_store().get_entity_by_id(
             entity_id=resource_id, schema_class=schema_class
@@ -765,6 +757,14 @@ if server_config().rbac_enabled:
                 "not exist."
             )
 
+        workspace_id = None
+        if isinstance(model, WorkspaceScopedResponse):
+            workspace_id = model.workspace.id
+
+        resource = Resource(
+            type=resource_type, id=resource_id, workspace_id=workspace_id
+        )
+
         verify_permission_for_model(model=model, action=Action.SHARE)
         for action in actions:
             # Make sure users aren't able to share permissions they don't have
@@ -776,3 +776,38 @@ if server_config().rbac_enabled:
             resource=resource,
             actions=[Action(action) for action in actions],
         )
+
+
+@current_user_router.put(
+    "/default-workspace",
+    responses={
+        401: error_response,
+        404: error_response,
+        422: error_response,
+    },
+)
+@handle_exceptions
+def update_user_default_workspace(
+    workspace_name_or_id: Union[str, UUID],
+    auth_context: AuthContext = Security(authorize),
+) -> UserResponse:
+    """Updates the default workspace of the current user.
+
+    Args:
+        workspace_name_or_id: Name or ID of the workspace.
+        auth_context: Authentication context.
+
+    Returns:
+        The updated user.
+    """
+    workspace = verify_permissions_and_get_entity(
+        id=workspace_name_or_id,
+        get_method=zen_store().get_workspace,
+    )
+
+    user = zen_store().update_user(
+        user_id=auth_context.user.id,
+        user_update=UserUpdate(default_workspace_id=workspace.id),
+    )
+
+    return dehydrate_response_model(user)