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

zenml/zen_server/routers/run_templates_endpoints.py

@@ -13,7 +13,7 @@
 #  permissions and limitations under the License.
 """Endpoint definitions for run templates."""
 
-from typing import Optional
+from typing import Optional, Union
 from uuid import UUID
 
 from fastapi import APIRouter, BackgroundTasks, Depends, Security
@@ -26,12 +26,14 @@ from zenml.models import (
     Page,
     PipelineRunResponse,
     RunTemplateFilter,
+    RunTemplateRequest,
     RunTemplateResponse,
     RunTemplateUpdate,
 )
 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,
@@ -39,6 +41,9 @@ from zenml.zen_server.rbac.endpoint_utils import (
 )
 from zenml.zen_server.rbac.models import Action, ResourceType
 from zenml.zen_server.rbac.utils import verify_permission
+from zenml.zen_server.routers.workspaces_endpoints import (
+    router as workspace_router,
+)
 from zenml.zen_server.utils import (
     handle_exceptions,
     make_dependable,
@@ -53,15 +58,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}" + RUN_TEMPLATES,
+    responses={401: error_response, 409: error_response, 422: error_response},
+    deprecated=True,
+    tags=["run_templates"],
+)
+@handle_exceptions
+def create_run_template(
+    run_template: RunTemplateRequest,
+    workspace_name_or_id: Optional[Union[str, UUID]] = None,
+    _: AuthContext = Security(authorize),
+) -> RunTemplateResponse:
+    """Create a run template.
+
+    Args:
+        run_template: Run template to create.
+        workspace_name_or_id: Optional name or ID of the workspace.
+
+    Returns:
+        The created run template.
+    """
+    if workspace_name_or_id:
+        workspace = zen_store().get_workspace(workspace_name_or_id)
+        run_template.workspace = workspace.id
+
+    return verify_permissions_and_create_entity(
+        request_model=run_template,
+        create_method=zen_store().create_run_template,
+    )
+
+
 @router.get(
     "",
     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}" + RUN_TEMPLATES,
+    responses={401: error_response, 404: error_response, 422: error_response},
+    deprecated=True,
+    tags=["run_templates"],
+)
 @handle_exceptions
 def list_run_templates(
     filter_model: RunTemplateFilter = Depends(
         make_dependable(RunTemplateFilter)
     ),
+    workspace_name_or_id: Optional[Union[str, UUID]] = None,
     hydrate: bool = False,
     _: AuthContext = Security(authorize),
 ) -> Page[RunTemplateResponse]:
@@ -70,12 +121,16 @@ def list_run_templates(
     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 run templates.
     """
+    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.RUN_TEMPLATE,
@@ -189,19 +244,27 @@ if server_config().workload_manager_enabled:
         """
         from zenml.zen_server.template_execution.utils import run_template
 
-        with track_handler(event=AnalyticsEvent.EXECUTED_RUN_TEMPLATE):
+        with track_handler(
+            event=AnalyticsEvent.EXECUTED_RUN_TEMPLATE,
+        ) as analytics_handler:
             template = verify_permissions_and_get_entity(
                 id=template_id,
                 get_method=zen_store().get_run_template,
                 hydrate=True,
             )
+            analytics_handler.metadata = {
+                "workspace_id": template.workspace.id,
+            }
 
             verify_permission(
                 resource_type=ResourceType.PIPELINE_DEPLOYMENT,
                 action=Action.CREATE,
+                workspace_id=template.workspace.id,
             )
             verify_permission(
-                resource_type=ResourceType.PIPELINE_RUN, action=Action.CREATE
+                resource_type=ResourceType.PIPELINE_RUN,
+                action=Action.CREATE,
+                workspace_id=template.workspace.id,
             )
 
             return run_template(

zenml/zen_server/routers/runs_endpoints.py

@@ -13,7 +13,7 @@
 #  permissions and limitations under the License.
 """Endpoint definitions for pipeline runs."""
 
-from typing import Any, Dict
+from typing import Any, Dict, Optional, Tuple, Union
 from uuid import UUID
 
 from fastapi import APIRouter, Depends, Security
@@ -32,6 +32,7 @@ from zenml.logger import get_logger
 from zenml.models import (
     Page,
     PipelineRunFilter,
+    PipelineRunRequest,
     PipelineRunResponse,
     PipelineRunUpdate,
     StepRunFilter,
@@ -42,11 +43,17 @@ from zenml.zen_server.exceptions import error_response
 from zenml.zen_server.rbac.endpoint_utils import (
     verify_permissions_and_delete_entity,
     verify_permissions_and_get_entity,
+    verify_permissions_and_get_or_create_entity,
     verify_permissions_and_list_entities,
     verify_permissions_and_update_entity,
 )
 from zenml.zen_server.rbac.models import Action, ResourceType
-from zenml.zen_server.rbac.utils import verify_permission_for_model
+from zenml.zen_server.rbac.utils import (
+    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,
@@ -63,16 +70,62 @@ router = APIRouter(
 logger = get_logger(__name__)
 
 
+@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}" + RUNS,
+    responses={401: error_response, 409: error_response, 422: error_response},
+    deprecated=True,
+    tags=["runs"],
+)
+@handle_exceptions
+def get_or_create_pipeline_run(
+    pipeline_run: PipelineRunRequest,
+    workspace_name_or_id: Optional[Union[str, UUID]] = None,
+    _: AuthContext = Security(authorize),
+) -> Tuple[PipelineRunResponse, bool]:
+    """Get or create a pipeline run.
+
+    Args:
+        pipeline_run: Pipeline run to create.
+        workspace_name_or_id: Optional name or ID of the workspace.
+
+    Returns:
+        The pipeline run and a boolean indicating whether the run was created
+        or not.
+    """
+    if workspace_name_or_id:
+        workspace = zen_store().get_workspace(workspace_name_or_id)
+        pipeline_run.workspace = workspace.id
+
+    return verify_permissions_and_get_or_create_entity(
+        request_model=pipeline_run,
+        get_or_create_method=zen_store().get_or_create_run,
+    )
+
+
 @router.get(
     "",
-    response_model=Page[PipelineRunResponse],
     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}" + RUNS,
+    responses={401: error_response, 404: error_response, 422: error_response},
+    deprecated=True,
+    tags=["runs"],
+)
 @handle_exceptions
 def list_runs(
     runs_filter_model: PipelineRunFilter = Depends(
         make_dependable(PipelineRunFilter)
     ),
+    workspace_name_or_id: Optional[Union[str, UUID]] = None,
     hydrate: bool = False,
     _: AuthContext = Security(authorize),
 ) -> Page[PipelineRunResponse]:
@@ -80,12 +133,16 @@ def list_runs(
 
     Args:
         runs_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:
         The pipeline runs according to query filters.
     """
+    if workspace_name_or_id:
+        runs_filter_model.workspace = workspace_name_or_id
+
     return verify_permissions_and_list_entities(
         filter_model=runs_filter_model,
         resource_type=ResourceType.PIPELINE_RUN,
@@ -96,7 +153,6 @@ def list_runs(
 
 @router.get(
     "/{run_id}",
-    response_model=PipelineRunResponse,
     responses={401: error_response, 404: error_response, 422: error_response},
 )
 @handle_exceptions
@@ -157,7 +213,6 @@ def get_run(
 
 @router.put(
     "/{run_id}",
-    response_model=PipelineRunResponse,
     responses={401: error_response, 404: error_response, 422: error_response},
 )
 @handle_exceptions
@@ -206,7 +261,6 @@ def delete_run(
 
 @router.get(
     "/{run_id}" + STEPS,
-    response_model=Page[StepRunResponse],
     responses={401: error_response, 404: error_response, 422: error_response},
 )
 @handle_exceptions
@@ -236,7 +290,6 @@ def get_run_steps(
 
 @router.get(
     "/{run_id}" + PIPELINE_CONFIGURATION,
-    response_model=Dict[str, Any],
     responses={401: error_response, 404: error_response, 422: error_response},
 )
 @handle_exceptions
@@ -260,7 +313,6 @@ def get_pipeline_configuration(
 
 @router.get(
     "/{run_id}" + STATUS,
-    response_model=ExecutionStatus,
     responses={401: error_response, 404: error_response, 422: error_response},
 )
 @handle_exceptions

zenml/zen_server/routers/schedule_endpoints.py

@@ -13,6 +13,7 @@
 #  permissions and limitations under the License.
 """Endpoint definitions for pipeline run schedules."""
 
+from typing import Optional, Union
 from uuid import UUID
 
 from fastapi import APIRouter, Depends, Security
@@ -21,11 +22,18 @@ from zenml.constants import API, SCHEDULES, VERSION_1
 from zenml.models import (
     Page,
     ScheduleFilter,
+    ScheduleRequest,
     ScheduleResponse,
     ScheduleUpdate,
 )
 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.routers.workspaces_endpoints import (
+    router as workspace_router,
+)
 from zenml.zen_server.utils import (
     handle_exceptions,
     make_dependable,
@@ -39,16 +47,64 @@ 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}" + SCHEDULES,
+    responses={401: error_response, 409: error_response, 422: error_response},
+    deprecated=True,
+    tags=["schedules"],
+)
+@handle_exceptions
+def create_schedule(
+    schedule: ScheduleRequest,
+    workspace_name_or_id: Optional[Union[str, UUID]] = None,
+    auth_context: AuthContext = Security(authorize),
+) -> ScheduleResponse:
+    """Creates a schedule.
+
+    Args:
+        schedule: Schedule to create.
+        workspace_name_or_id: Optional name or ID of the workspace.
+        auth_context: Authentication context.
+
+    Returns:
+        The created schedule.
+    """
+    if workspace_name_or_id:
+        workspace = zen_store().get_workspace(workspace_name_or_id)
+        schedule.workspace = workspace.id
+
+    # NOTE: no RBAC is enforced currently for schedules, but we're
+    # keeping the RBAC checks here for consistency
+    return verify_permissions_and_create_entity(
+        request_model=schedule,
+        create_method=zen_store().create_schedule,
+    )
+
+
 @router.get(
     "",
-    response_model=Page[ScheduleResponse],
     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}" + SCHEDULES,
+    responses={401: error_response, 404: error_response, 422: error_response},
+    deprecated=True,
+    tags=["schedules"],
+)
 @handle_exceptions
 def list_schedules(
     schedule_filter_model: ScheduleFilter = Depends(
         make_dependable(ScheduleFilter)
     ),
+    workspace_name_or_id: Optional[Union[str, UUID]] = None,
     hydrate: bool = False,
     _: AuthContext = Security(authorize),
 ) -> Page[ScheduleResponse]:
@@ -57,20 +113,24 @@ def list_schedules(
     Args:
         schedule_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:
         List of schedule objects.
     """
+    if workspace_name_or_id:
+        schedule_filter_model.workspace = workspace_name_or_id
+
     return zen_store().list_schedules(
-        schedule_filter_model=schedule_filter_model, hydrate=hydrate
+        schedule_filter_model=schedule_filter_model,
+        hydrate=hydrate,
     )
 
 
 @router.get(
     "/{schedule_id}",
-    response_model=ScheduleResponse,
     responses={401: error_response, 404: error_response, 422: error_response},
 )
 @handle_exceptions
@@ -89,12 +149,14 @@ def get_schedule(
     Returns:
         A specific schedule object.
     """
-    return zen_store().get_schedule(schedule_id=schedule_id, hydrate=hydrate)
+    return zen_store().get_schedule(
+        schedule_id=schedule_id,
+        hydrate=hydrate,
+    )
 
 
 @router.put(
     "/{schedule_id}",
-    response_model=ScheduleResponse,
     responses={401: error_response, 404: error_response, 422: error_response},
 )
 @handle_exceptions
@@ -113,7 +175,8 @@ def update_schedule(
         The updated schedule object.
     """
     return zen_store().update_schedule(
-        schedule_id=schedule_id, schedule_update=schedule_update
+        schedule_id=schedule_id,
+        schedule_update=schedule_update,
     )
 
 

zenml/zen_server/routers/secrets_endpoints.py

@@ -13,7 +13,7 @@
 #  permissions and limitations under the License.
 """Endpoint definitions for pipeline run secrets."""
 
-from typing import Optional
+from typing import Optional, Union
 from uuid import UUID
 
 from fastapi import APIRouter, Depends, Security
@@ -29,12 +29,14 @@ from zenml.constants import (
 from zenml.models import (
     Page,
     SecretFilter,
+    SecretRequest,
     SecretResponse,
     SecretUpdate,
 )
 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,
@@ -47,6 +49,9 @@ from zenml.zen_server.rbac.utils import (
     is_owned_by_authenticated_user,
     verify_permission,
 )
+from zenml.zen_server.routers.workspaces_endpoints import (
+    router as workspace_router,
+)
 from zenml.zen_server.utils import (
     handle_exceptions,
     make_dependable,
@@ -66,9 +71,41 @@ op_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}" + SECRETS,
+    responses={401: error_response, 409: error_response, 422: error_response},
+    deprecated=True,
+    tags=["secrets"],
+)
+@handle_exceptions
+def create_secret(
+    secret: SecretRequest,
+    workspace_name_or_id: Optional[Union[str, UUID]] = None,
+    _: AuthContext = Security(authorize),
+) -> SecretResponse:
+    """Creates a secret.
+
+    Args:
+        secret: Secret to create.
+        workspace_name_or_id: Optional name or ID of the workspace.
+
+    Returns:
+        The created secret.
+    """
+    return verify_permissions_and_create_entity(
+        request_model=secret,
+        create_method=zen_store().create_secret,
+    )
+
+
 @router.get(
     "",
-    response_model=Page[SecretResponse],
     responses={401: error_response, 404: error_response, 422: error_response},
 )
 @handle_exceptions
@@ -116,7 +153,6 @@ def list_secrets(
 
 @router.get(
     "/{secret_id}",
-    response_model=SecretResponse,
     responses={401: error_response, 404: error_response, 422: error_response},
 )
 @handle_exceptions
@@ -140,6 +176,7 @@ def get_secret(
         get_method=zen_store().get_secret,
         hydrate=hydrate,
     )
+
     if not has_permissions_for_model(secret, action=Action.READ_SECRET_VALUE):
         secret.remove_secrets()
 
@@ -148,7 +185,6 @@ def get_secret(
 
 @router.put(
     "/{secret_id}",
-    response_model=SecretResponse,
     responses={401: error_response, 404: error_response, 422: error_response},
 )
 @handle_exceptions

zenml/zen_server/routers/server_endpoints.py

@@ -25,20 +25,27 @@ from zenml.constants import (
     LOAD_INFO,
     ONBOARDING_STATE,
     SERVER_SETTINGS,
+    STATISTICS,
     VERSION_1,
 )
 from zenml.enums import AuthScheme
 from zenml.exceptions import IllegalOperationError
 from zenml.models import (
+    ComponentFilter,
     ServerActivationRequest,
     ServerLoadInfo,
     ServerModel,
     ServerSettingsResponse,
     ServerSettingsUpdate,
+    ServerStatistics,
+    StackFilter,
     UserResponse,
+    WorkspaceFilter,
 )
 from zenml.zen_server.auth import AuthContext, authorize
 from zenml.zen_server.exceptions import error_response
+from zenml.zen_server.rbac.models import ResourceType
+from zenml.zen_server.rbac.utils import get_allowed_resource_ids
 from zenml.zen_server.utils import handle_exceptions, server_config, zen_store
 
 router = APIRouter(
@@ -60,7 +67,6 @@ def version() -> str:
 
 @router.get(
     INFO,
-    response_model=ServerModel,
     responses={401: error_response, 404: error_response, 422: error_response},
 )
 @handle_exceptions
@@ -234,3 +240,49 @@ if server_config().auth_scheme != AuthScheme.EXTERNAL:
             The default admin user that was created during activation, if any.
         """
         return zen_store().activate_server(activate_request)
+
+
+@router.get(
+    STATISTICS,
+    responses={401: error_response, 404: error_response, 422: error_response},
+)
+@handle_exceptions
+def get_server_statistics(
+    auth_context: AuthContext = Security(authorize),
+) -> ServerStatistics:
+    """Gets server statistics.
+
+    Args:
+        auth_context: Authentication context.
+
+    Returns:
+        Statistics of the server.
+    """
+    user_id = auth_context.user.id
+    component_filter = ComponentFilter()
+    component_filter.configure_rbac(
+        authenticated_user_id=user_id,
+        id=get_allowed_resource_ids(
+            resource_type=ResourceType.STACK_COMPONENT
+        ),
+    )
+
+    workspace_filter = WorkspaceFilter()
+    workspace_filter.configure_rbac(
+        authenticated_user_id=user_id,
+        id=get_allowed_resource_ids(resource_type=ResourceType.WORKSPACE),
+    )
+
+    stack_filter = StackFilter()
+    stack_filter.configure_rbac(
+        authenticated_user_id=user_id,
+        id=get_allowed_resource_ids(resource_type=ResourceType.STACK),
+    )
+
+    return ServerStatistics(
+        stacks=zen_store().count_stacks(filter_model=stack_filter),
+        components=zen_store().count_stack_components(
+            filter_model=component_filter
+        ),
+        workspaces=zen_store().count_workspaces(filter_model=workspace_filter),
+    )