zenml-nightly 0.82.1.dev20250521__py3-none-any.whl → 0.82.1.dev20250524__py3-none-any.whl

zenml/zen_server/routers/run_templates_endpoints.py

@@ -55,7 +55,7 @@ from zenml.zen_server.rbac.models import Action, ResourceType
 from zenml.zen_server.rbac.utils import verify_permission
 from zenml.zen_server.routers.projects_endpoints import workspace_router
 from zenml.zen_server.utils import (
-    handle_exceptions,
+    async_fastapi_endpoint_wrapper,
     make_dependable,
     server_config,
     zen_store,
@@ -80,7 +80,7 @@ router = APIRouter(
     deprecated=True,
     tags=["run_templates"],
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def create_run_template(
     run_template: RunTemplateRequest,
     project_name_or_id: Optional[Union[str, UUID]] = None,
@@ -117,7 +117,7 @@ def create_run_template(
     deprecated=True,
     tags=["run_templates"],
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def list_run_templates(
     filter_model: RunTemplateFilter = Depends(
         make_dependable(RunTemplateFilter)
@@ -153,7 +153,7 @@ def list_run_templates(
     "/{template_id}",
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def get_run_template(
     template_id: UUID,
     hydrate: bool = True,
@@ -180,7 +180,7 @@ def get_run_template(
     "/{template_id}",
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def update_run_template(
     template_id: UUID,
     update: RunTemplateUpdate,
@@ -207,7 +207,7 @@ def update_run_template(
     "/{template_id}",
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def delete_run_template(
     template_id: UUID,
     _: AuthContext = Security(authorize),
@@ -235,7 +235,7 @@ if server_config().workload_manager_enabled:
             429: error_response,
         },
     )
-    @handle_exceptions
+    @async_fastapi_endpoint_wrapper
     def create_template_run(
         template_id: UUID,
         config: Optional[PipelineRunConfiguration] = None,

zenml/zen_server/routers/runs_endpoints.py

@@ -54,7 +54,7 @@ from zenml.zen_server.rbac.utils import (
 )
 from zenml.zen_server.routers.projects_endpoints import workspace_router
 from zenml.zen_server.utils import (
-    handle_exceptions,
+    async_fastapi_endpoint_wrapper,
     make_dependable,
     server_config,
     workload_manager,
@@ -83,7 +83,7 @@ logger = get_logger(__name__)
     deprecated=True,
     tags=["runs"],
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def get_or_create_pipeline_run(
     pipeline_run: PipelineRunRequest,
     project_name_or_id: Optional[Union[str, UUID]] = None,
@@ -121,7 +121,7 @@ def get_or_create_pipeline_run(
     deprecated=True,
     tags=["runs"],
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def list_runs(
     runs_filter_model: PipelineRunFilter = Depends(
         make_dependable(PipelineRunFilter)
@@ -156,7 +156,7 @@ def list_runs(
     "/{run_id}",
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def get_run(
     run_id: UUID,
     hydrate: bool = True,
@@ -222,7 +222,7 @@ def get_run(
     "/{run_id}",
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def update_run(
     run_id: UUID,
     run_model: PipelineRunUpdate,
@@ -249,7 +249,7 @@ def update_run(
     "/{run_id}",
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def delete_run(
     run_id: UUID,
     _: AuthContext = Security(authorize),
@@ -270,7 +270,7 @@ def delete_run(
     "/{run_id}" + STEPS,
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def get_run_steps(
     run_id: UUID,
     step_run_filter_model: StepRunFilter = Depends(
@@ -299,7 +299,7 @@ def get_run_steps(
     "/{run_id}" + PIPELINE_CONFIGURATION,
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def get_pipeline_configuration(
     run_id: UUID,
     _: AuthContext = Security(authorize),
@@ -322,7 +322,7 @@ def get_pipeline_configuration(
     "/{run_id}" + STATUS,
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def get_run_status(
     run_id: UUID,
     _: AuthContext = Security(authorize),
@@ -345,7 +345,7 @@ def get_run_status(
     "/{run_id}" + REFRESH,
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def refresh_run_status(
     run_id: UUID,
     _: AuthContext = Security(authorize),
@@ -394,7 +394,7 @@ def refresh_run_status(
         422: error_response,
     },
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def run_logs(
     run_id: UUID,
     offset: int = 0,

zenml/zen_server/routers/schedule_endpoints.py

@@ -33,7 +33,7 @@ from zenml.zen_server.rbac.endpoint_utils import (
 )
 from zenml.zen_server.routers.projects_endpoints import workspace_router
 from zenml.zen_server.utils import (
-    handle_exceptions,
+    async_fastapi_endpoint_wrapper,
     make_dependable,
     zen_store,
 )
@@ -57,7 +57,7 @@ router = APIRouter(
     deprecated=True,
     tags=["schedules"],
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def create_schedule(
     schedule: ScheduleRequest,
     project_name_or_id: Optional[Union[str, UUID]] = None,
@@ -97,7 +97,7 @@ def create_schedule(
     deprecated=True,
     tags=["schedules"],
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def list_schedules(
     schedule_filter_model: ScheduleFilter = Depends(
         make_dependable(ScheduleFilter)
@@ -131,7 +131,7 @@ def list_schedules(
     "/{schedule_id}",
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def get_schedule(
     schedule_id: UUID,
     hydrate: bool = True,
@@ -157,7 +157,7 @@ def get_schedule(
     "/{schedule_id}",
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def update_schedule(
     schedule_id: UUID,
     schedule_update: ScheduleUpdate,
@@ -182,7 +182,7 @@ def update_schedule(
     "/{schedule_id}",
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def delete_schedule(
     schedule_id: UUID,
     _: AuthContext = Security(authorize),

zenml/zen_server/routers/secrets_endpoints.py

@@ -51,7 +51,7 @@ from zenml.zen_server.rbac.utils import (
 )
 from zenml.zen_server.routers.projects_endpoints import workspace_router
 from zenml.zen_server.utils import (
-    handle_exceptions,
+    async_fastapi_endpoint_wrapper,
     make_dependable,
     zen_store,
 )
@@ -81,7 +81,7 @@ op_router = APIRouter(
     deprecated=True,
     tags=["secrets"],
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def create_secret(
     secret: SecretRequest,
     workspace_name_or_id: Optional[Union[str, UUID]] = None,
@@ -106,7 +106,7 @@ def create_secret(
     "",
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def list_secrets(
     secret_filter_model: SecretFilter = Depends(make_dependable(SecretFilter)),
     hydrate: bool = False,
@@ -153,7 +153,7 @@ def list_secrets(
     "/{secret_id}",
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def get_secret(
     secret_id: UUID,
     hydrate: bool = True,
@@ -185,7 +185,7 @@ def get_secret(
     "/{secret_id}",
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def update_secret(
     secret_id: UUID,
     secret_update: SecretUpdate,
@@ -225,7 +225,7 @@ def update_secret(
     "/{secret_id}",
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def delete_secret(
     secret_id: UUID,
     _: AuthContext = Security(authorize),
@@ -246,7 +246,7 @@ def delete_secret(
     SECRETS_BACKUP,
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def backup_secrets(
     ignore_errors: bool = True,
     delete_secrets: bool = False,
@@ -276,7 +276,7 @@ def backup_secrets(
     SECRETS_RESTORE,
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def restore_secrets(
     ignore_errors: bool = False,
     delete_secrets: bool = False,

zenml/zen_server/routers/server_endpoints.py

@@ -46,7 +46,11 @@ 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
+from zenml.zen_server.utils import (
+    async_fastapi_endpoint_wrapper,
+    server_config,
+    zen_store,
+)
 
 router = APIRouter(
     prefix=API + VERSION_1,
@@ -56,7 +60,7 @@ router = APIRouter(
 
 
 @router.get("/version")
-def version() -> str:
+async def version() -> str:
     """Get version of the server.
 
     Returns:
@@ -69,7 +73,7 @@ def version() -> str:
     INFO,
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def server_info() -> ServerModel:
     """Get information about the server.
 
@@ -83,7 +87,7 @@ def server_info() -> ServerModel:
     LOAD_INFO,
     response_model=ServerLoadInfo,
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def server_load_info(_: AuthContext = Security(authorize)) -> ServerLoadInfo:
     """Get information about the server load.
 
@@ -132,7 +136,7 @@ def server_load_info(_: AuthContext = Security(authorize)) -> ServerLoadInfo:
         422: error_response,
     },
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def get_onboarding_state(
     _: AuthContext = Security(authorize),
 ) -> List[str]:
@@ -159,7 +163,7 @@ if server_config().external_server_id is None:
             422: error_response,
         },
     )
-    @handle_exceptions
+    @async_fastapi_endpoint_wrapper
     def get_settings(
         _: AuthContext = Security(authorize),
         hydrate: bool = True,
@@ -182,7 +186,7 @@ if server_config().external_server_id is None:
             422: error_response,
         },
     )
-    @handle_exceptions
+    @async_fastapi_endpoint_wrapper
     def update_server_settings(
         settings_update: ServerSettingsUpdate,
         auth_context: AuthContext = Security(authorize),
@@ -227,7 +231,7 @@ if server_config().auth_scheme != AuthScheme.EXTERNAL:
             422: error_response,
         },
     )
-    @handle_exceptions
+    @async_fastapi_endpoint_wrapper
     def activate_server(
         activate_request: ServerActivationRequest,
     ) -> Optional[UserResponse]:
@@ -246,7 +250,7 @@ if server_config().auth_scheme != AuthScheme.EXTERNAL:
     STATISTICS,
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def get_server_statistics(
     auth_context: AuthContext = Security(authorize),
 ) -> ServerStatistics:

zenml/zen_server/routers/service_accounts_endpoints.py

@@ -49,7 +49,7 @@ 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_for_model
 from zenml.zen_server.utils import (
-    handle_exceptions,
+    async_fastapi_endpoint_wrapper,
     make_dependable,
     zen_store,
 )
@@ -73,7 +73,7 @@ router = APIRouter(
         422: error_response,
     },
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def create_service_account(
     service_account: ServiceAccountRequest,
     _: AuthContext = Security(authorize),
@@ -96,7 +96,7 @@ def create_service_account(
     "/{service_account_name_or_id}",
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def get_service_account(
     service_account_name_or_id: Union[str, UUID],
     _: AuthContext = Security(authorize),
@@ -123,7 +123,7 @@ def get_service_account(
     "",
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def list_service_accounts(
     filter_model: ServiceAccountFilter = Depends(
         make_dependable(ServiceAccountFilter)
@@ -158,7 +158,7 @@ def list_service_accounts(
         422: error_response,
     },
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def update_service_account(
     service_account_name_or_id: Union[str, UUID],
     service_account_update: ServiceAccountUpdate,
@@ -185,7 +185,7 @@ def update_service_account(
     "/{service_account_name_or_id}",
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def delete_service_account(
     service_account_name_or_id: Union[str, UUID],
     _: AuthContext = Security(authorize),
@@ -211,7 +211,7 @@ def delete_service_account(
     "/{service_account_id}" + API_KEYS,
     responses={401: error_response, 409: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def create_api_key(
     service_account_id: UUID,
     api_key: APIKeyRequest,
@@ -249,7 +249,7 @@ def create_api_key(
     "/{service_account_id}" + API_KEYS + "/{api_key_name_or_id}",
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def get_api_key(
     service_account_id: UUID,
     api_key_name_or_id: Union[str, UUID],
@@ -282,7 +282,7 @@ def get_api_key(
     "/{service_account_id}" + API_KEYS,
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def list_api_keys(
     service_account_id: UUID,
     filter_model: APIKeyFilter = Depends(make_dependable(APIKeyFilter)),
@@ -316,7 +316,7 @@ def list_api_keys(
     "/{service_account_id}" + API_KEYS + "/{api_key_name_or_id}",
     responses={401: error_response, 409: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def update_api_key(
     service_account_id: UUID,
     api_key_name_or_id: Union[str, UUID],
@@ -350,7 +350,7 @@ def update_api_key(
     + API_KEY_ROTATE,
     responses={401: error_response, 409: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def rotate_api_key(
     service_account_id: UUID,
     api_key_name_or_id: Union[str, UUID],
@@ -381,7 +381,7 @@ def rotate_api_key(
     "/{service_account_id}" + API_KEYS + "/{api_key_name_or_id}",
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def delete_api_key(
     service_account_id: UUID,
     api_key_name_or_id: Union[str, UUID],

zenml/zen_server/routers/service_connectors_endpoints.py

@@ -61,7 +61,7 @@ from zenml.zen_server.rbac.utils import (
 )
 from zenml.zen_server.routers.projects_endpoints import workspace_router
 from zenml.zen_server.utils import (
-    handle_exceptions,
+    async_fastapi_endpoint_wrapper,
     make_dependable,
     zen_store,
 )
@@ -91,7 +91,7 @@ types_router = APIRouter(
     deprecated=True,
     tags=["service_connectors"],
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def create_service_connector(
     connector: ServiceConnectorRequest,
     project_name_or_id: Optional[Union[str, UUID]] = None,
@@ -124,7 +124,7 @@ def create_service_connector(
     deprecated=True,
     tags=["service_connectors"],
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def list_service_connectors(
     connector_filter_model: ServiceConnectorFilter = Depends(
         make_dependable(ServiceConnectorFilter)
@@ -198,7 +198,7 @@ def list_service_connectors(
     deprecated=True,
     tags=["service_connectors"],
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def list_service_connector_resources(
     filter_model: ServiceConnectorFilter = Depends(
         make_dependable(ServiceConnectorFilter)
@@ -234,7 +234,7 @@ def list_service_connector_resources(
     "/{connector_id}",
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def get_service_connector(
     connector_id: UUID,
     expand_secrets: bool = True,
@@ -276,7 +276,7 @@ def get_service_connector(
     "/{connector_id}",
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def update_service_connector(
     connector_id: UUID,
     connector_update: ServiceConnectorUpdate,
@@ -303,7 +303,7 @@ def update_service_connector(
     "/{connector_id}",
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def delete_service_connector(
     connector_id: UUID,
     _: AuthContext = Security(authorize),
@@ -324,7 +324,7 @@ def delete_service_connector(
     SERVICE_CONNECTOR_VERIFY,
     responses={401: error_response, 409: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def validate_and_verify_service_connector_config(
     connector: ServiceConnectorRequest,
     list_resources: bool = True,
@@ -357,7 +357,7 @@ def validate_and_verify_service_connector_config(
     "/{connector_id}" + SERVICE_CONNECTOR_VERIFY,
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def validate_and_verify_service_connector(
     connector_id: UUID,
     resource_type: Optional[str] = None,
@@ -398,7 +398,7 @@ def validate_and_verify_service_connector(
     "/{connector_id}" + SERVICE_CONNECTOR_CLIENT,
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def get_service_connector_client(
     connector_id: UUID,
     resource_type: Optional[str] = None,
@@ -435,7 +435,7 @@ def get_service_connector_client(
     "",
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def list_service_connector_types(
     connector_type: Optional[str] = None,
     resource_type: Optional[str] = None,
@@ -465,7 +465,7 @@ def list_service_connector_types(
     "/{connector_type}",
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def get_service_connector_type(
     connector_type: str,
     _: AuthContext = Security(authorize),
@@ -485,7 +485,7 @@ def get_service_connector_type(
     SERVICE_CONNECTOR_FULL_STACK,
     responses={401: error_response, 409: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def get_resources_based_on_service_connector_info(
     connector_info: Optional[ServiceConnectorInfo] = None,
     connector_uuid: Optional[UUID] = None,

zenml/zen_server/routers/service_endpoints.py

@@ -38,7 +38,7 @@ from zenml.zen_server.rbac.endpoint_utils import (
 from zenml.zen_server.rbac.models import ResourceType
 from zenml.zen_server.routers.projects_endpoints import workspace_router
 from zenml.zen_server.utils import (
-    handle_exceptions,
+    async_fastapi_endpoint_wrapper,
     make_dependable,
     zen_store,
 )
@@ -62,7 +62,7 @@ router = APIRouter(
     deprecated=True,
     tags=["services"],
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def create_service(
     service: ServiceRequest,
     project_name_or_id: Optional[Union[str, UUID]] = None,
@@ -91,7 +91,7 @@ def create_service(
     "",
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def list_services(
     filter_model: ServiceFilter = Depends(make_dependable(ServiceFilter)),
     hydrate: bool = False,
@@ -120,7 +120,7 @@ def list_services(
     "/{service_id}",
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def get_service(
     service_id: UUID,
     hydrate: bool = True,
@@ -147,7 +147,7 @@ def get_service(
     "/{service_id}",
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def update_service(
     service_id: UUID,
     update: ServiceUpdate,
@@ -174,7 +174,7 @@ def update_service(
     "/{service_id}",
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def delete_service(
     service_id: UUID,
     _: AuthContext = Security(authorize),