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

zenml/zen_server/routers/stack_components_endpoints.py

@@ -40,7 +40,7 @@ from zenml.zen_server.rbac.models import Action, ResourceType
 from zenml.zen_server.rbac.utils import verify_permission_for_model
 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,
 )
@@ -70,7 +70,7 @@ types_router = APIRouter(
     deprecated=True,
     tags=["stack_components"],
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def create_stack_component(
     component: ComponentRequest,
     project_name_or_id: Optional[Union[str, UUID]] = None,
@@ -93,7 +93,7 @@ def create_stack_component(
 
     from zenml.stack.utils import validate_stack_component_config
 
-    validate_stack_component_config(
+    validated_config = validate_stack_component_config(
         configuration_dict=component.configuration,
         flavor=component.flavor,
         component_type=component.type,
@@ -102,6 +102,11 @@ def create_stack_component(
         validate_custom_flavors=False,
     )
 
+    if validated_config:
+        component.configuration = validated_config.model_dump(
+            mode="json", exclude_unset=True
+        )
+
     return verify_permissions_and_create_entity(
         request_model=component,
         create_method=zen_store().create_stack_component,
@@ -120,7 +125,7 @@ def create_stack_component(
     deprecated=True,
     tags=["stack_components"],
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def list_stack_components(
     component_filter_model: ComponentFilter = Depends(
         make_dependable(ComponentFilter)
@@ -153,7 +158,7 @@ def list_stack_components(
     "/{component_id}",
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def get_stack_component(
     component_id: UUID,
     hydrate: bool = True,
@@ -180,7 +185,7 @@ def get_stack_component(
     "/{component_id}",
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def update_stack_component(
     component_id: UUID,
     component_update: ComponentUpdate,
@@ -199,7 +204,7 @@ def update_stack_component(
         from zenml.stack.utils import validate_stack_component_config
 
         existing_component = zen_store().get_stack_component(component_id)
-        validate_stack_component_config(
+        validated_config = validate_stack_component_config(
             configuration_dict=component_update.configuration,
             flavor=existing_component.flavor_name,
             component_type=existing_component.type,
@@ -207,6 +212,10 @@ def update_stack_component(
             # We allow custom flavors to fail import on the server side.
             validate_custom_flavors=False,
         )
+        if validated_config:
+            component_update.configuration = validated_config.model_dump(
+                mode="json", exclude_unset=True
+            )
 
     if component_update.connector:
         service_connector = zen_store().get_service_connector(
@@ -226,7 +235,7 @@ def update_stack_component(
     "/{component_id}",
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def deregister_stack_component(
     component_id: UUID,
     _: AuthContext = Security(authorize),
@@ -247,7 +256,7 @@ def deregister_stack_component(
     "",
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def get_stack_component_types(
     _: AuthContext = Security(authorize),
 ) -> List[str]:

zenml/zen_server/routers/stack_deployment_endpoints.py

@@ -39,7 +39,7 @@ from zenml.zen_server.exceptions import error_response
 from zenml.zen_server.rbac.models import Action, ResourceType
 from zenml.zen_server.rbac.utils import verify_permission
 from zenml.zen_server.utils import (
-    handle_exceptions,
+    async_fastapi_endpoint_wrapper,
     server_config,
 )
 
@@ -53,7 +53,7 @@ router = APIRouter(
 @router.get(
     INFO,
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def get_stack_deployment_info(
     provider: StackDeploymentProvider,
     _: AuthContext = Security(authorize),
@@ -73,7 +73,7 @@ def get_stack_deployment_info(
 @router.get(
     CONFIG,
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def get_stack_deployment_config(
     request: Request,
     provider: StackDeploymentProvider,
@@ -139,7 +139,7 @@ def get_stack_deployment_config(
 @router.get(
     STACK,
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def get_deployed_stack(
     provider: StackDeploymentProvider,
     stack_name: str,

zenml/zen_server/routers/stacks_endpoints.py

@@ -42,7 +42,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,
 )
@@ -66,7 +66,7 @@ router = APIRouter(
     deprecated=True,
     tags=["stacks"],
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def create_stack(
     stack: StackRequest,
     project_name_or_id: Optional[Union[str, UUID]] = None,
@@ -136,7 +136,7 @@ def create_stack(
     deprecated=True,
     tags=["stacks"],
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def list_stacks(
     project_name_or_id: Optional[Union[str, UUID]] = None,
     stack_filter_model: StackFilter = Depends(make_dependable(StackFilter)),
@@ -167,7 +167,7 @@ def list_stacks(
     "/{stack_id}",
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def get_stack(
     stack_id: UUID,
     hydrate: bool = True,
@@ -192,7 +192,7 @@ def get_stack(
     "/{stack_id}",
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def update_stack(
     stack_id: UUID,
     stack_update: StackUpdate,
@@ -230,7 +230,7 @@ def update_stack(
     "/{stack_id}",
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def delete_stack(
     stack_id: UUID,
     _: AuthContext = Security(authorize),

zenml/zen_server/routers/steps_endpoints.py

@@ -48,7 +48,7 @@ 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,
     set_filter_project_scope,
     zen_store,
@@ -65,7 +65,7 @@ router = APIRouter(
     "",
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def list_run_steps(
     step_run_filter_model: StepRunFilter = Depends(
         make_dependable(StepRunFilter)
@@ -109,7 +109,7 @@ def list_run_steps(
     "",
     responses={401: error_response, 409: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def create_run_step(
     step: StepRunRequest,
     _: AuthContext = Security(authorize),
@@ -136,7 +136,7 @@ def create_run_step(
     "/{step_id}",
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def get_step(
     step_id: UUID,
     hydrate: bool = True,
@@ -169,7 +169,7 @@ def get_step(
     "/{step_id}",
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def update_step(
     step_id: UUID,
     step_model: StepRunUpdate,
@@ -198,7 +198,7 @@ def update_step(
     "/{step_id}" + STEP_CONFIGURATION,
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def get_step_configuration(
     step_id: UUID,
     _: AuthContext = Security(authorize),
@@ -222,7 +222,7 @@ def get_step_configuration(
     "/{step_id}" + STATUS,
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def get_step_status(
     step_id: UUID,
     _: AuthContext = Security(authorize),
@@ -246,7 +246,7 @@ def get_step_status(
     "/{step_id}" + LOGS,
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def get_step_logs(
     step_id: UUID,
     offset: int = 0,

zenml/zen_server/routers/tag_resource_endpoints.py

@@ -27,7 +27,7 @@ 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,
+    async_fastapi_endpoint_wrapper,
     zen_store,
 )
 
@@ -42,7 +42,7 @@ router = APIRouter(
     "",
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def create_tag_resource(
     tag_resource: TagResourceRequest,
     _: AuthContext = Security(authorize),
@@ -62,7 +62,7 @@ def create_tag_resource(
     BATCH,
     responses={401: error_response, 409: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def batch_create_tag_resource(
     tag_resources: List[TagResourceRequest],
     _: AuthContext = Security(authorize),
@@ -85,7 +85,7 @@ def batch_create_tag_resource(
     "",
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def delete_tag_resource(
     tag_resource: TagResourceRequest,
     _: AuthContext = Security(authorize),
@@ -102,7 +102,7 @@ def delete_tag_resource(
     BATCH,
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def batch_delete_tag_resource(
     tag_resources: List[TagResourceRequest],
     _: AuthContext = Security(authorize),

zenml/zen_server/routers/tags_endpoints.py

@@ -38,7 +38,7 @@ from zenml.zen_server.rbac.endpoint_utils import (
     verify_permissions_and_update_entity,
 )
 from zenml.zen_server.utils import (
-    handle_exceptions,
+    async_fastapi_endpoint_wrapper,
     make_dependable,
     zen_store,
 )
@@ -58,7 +58,7 @@ router = APIRouter(
     "",
     responses={401: error_response, 409: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def create_tag(
     tag: TagRequest,
     _: AuthContext = Security(authorize),
@@ -78,7 +78,7 @@ def create_tag(
     "",
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def list_tags(
     tag_filter_model: TagFilter = Depends(make_dependable(TagFilter)),
     hydrate: bool = False,
@@ -105,7 +105,7 @@ def list_tags(
     "/{tag_name_or_id}",
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def get_tag(
     tag_name_or_id: Union[str, UUID],
     hydrate: bool = True,
@@ -132,7 +132,7 @@ def get_tag(
     "/{tag_id}",
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def update_tag(
     tag_id: UUID,
     tag_update_model: TagUpdate,
@@ -159,7 +159,7 @@ def update_tag(
     "/{tag_name_or_id}",
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def delete_tag(
     tag_name_or_id: Union[str, UUID],
     _: AuthContext = Security(authorize),

zenml/zen_server/routers/triggers_endpoints.py

@@ -43,7 +43,7 @@ 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,
     plugin_flavor_registry,
     zen_store,
@@ -60,7 +60,7 @@ router = APIRouter(
     "",
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def list_triggers(
     trigger_filter_model: TriggerFilter = Depends(
         make_dependable(TriggerFilter)
@@ -91,7 +91,7 @@ def list_triggers(
     "/{trigger_id}",
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def get_trigger(
     trigger_id: UUID,
     hydrate: bool = True,
@@ -116,7 +116,7 @@ def get_trigger(
     "",
     responses={401: error_response, 409: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def create_trigger(
     trigger: TriggerRequest,
     _: AuthContext = Security(authorize),
@@ -167,7 +167,7 @@ def create_trigger(
     "/{trigger_id}",
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def update_trigger(
     trigger_id: UUID,
     trigger_update: TriggerUpdate,
@@ -230,7 +230,7 @@ def update_trigger(
     "/{trigger_id}",
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def delete_trigger(
     trigger_id: UUID,
     _: AuthContext = Security(authorize),
@@ -258,7 +258,7 @@ executions_router = APIRouter(
     "",
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def list_trigger_executions(
     trigger_execution_filter_model: TriggerExecutionFilter = Depends(
         make_dependable(TriggerExecutionFilter)
@@ -289,7 +289,7 @@ def list_trigger_executions(
     "/{trigger_execution_id}",
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def get_trigger_execution(
     trigger_execution_id: UUID,
     hydrate: bool = True,
@@ -316,7 +316,7 @@ def get_trigger_execution(
     "/{trigger_execution_id}",
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def delete_trigger_execution(
     trigger_execution_id: UUID,
     _: AuthContext = Security(authorize),

zenml/zen_server/routers/users_endpoints.py

@@ -56,7 +56,7 @@ 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,
     server_config,
     verify_admin_status_if_no_rbac,
@@ -95,7 +95,7 @@ current_user_router = APIRouter(
     "",
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def list_users(
     user_filter_model: UserFilter = Depends(make_dependable(UserFilter)),
     hydrate: bool = False,
@@ -144,7 +144,7 @@ if server_config().auth_scheme != AuthScheme.EXTERNAL:
             422: error_response,
         },
     )
-    @handle_exceptions
+    @async_fastapi_endpoint_wrapper
     def create_user(
         user: UserRequest,
         auth_context: AuthContext = Security(authorize),
@@ -193,7 +193,7 @@ if server_config().auth_scheme != AuthScheme.EXTERNAL:
     "/{user_name_or_id}",
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def get_user(
     user_name_or_id: Union[str, UUID],
     hydrate: bool = True,
@@ -237,7 +237,7 @@ if server_config().auth_scheme != AuthScheme.EXTERNAL:
             422: error_response,
         },
     )
-    @handle_exceptions
+    @async_fastapi_endpoint_wrapper
     def activate_user(
         user_name_or_id: Union[str, UUID],
         user_update: UserUpdate,
@@ -297,7 +297,7 @@ if server_config().auth_scheme != AuthScheme.EXTERNAL:
             422: error_response,
         },
     )
-    @handle_exceptions
+    @async_fastapi_endpoint_wrapper
     def deactivate_user(
         user_name_or_id: Union[str, UUID],
         auth_context: AuthContext = Security(authorize),
@@ -345,7 +345,7 @@ if server_config().auth_scheme != AuthScheme.EXTERNAL:
             422: error_response,
         },
     )
-    @handle_exceptions
+    @async_fastapi_endpoint_wrapper
     def delete_user(
         user_name_or_id: Union[str, UUID],
         auth_context: AuthContext = Security(authorize),
@@ -387,7 +387,7 @@ if server_config().auth_scheme != AuthScheme.EXTERNAL:
             422: error_response,
         },
     )
-    @handle_exceptions
+    @async_fastapi_endpoint_wrapper
     def email_opt_in_response(
         user_name_or_id: Union[str, UUID],
         user_response: UserUpdate,
@@ -439,7 +439,7 @@ if server_config().auth_scheme != AuthScheme.EXTERNAL:
         422: error_response,
     },
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def update_user(
     user_name_or_id: Union[str, UUID],
     user_update: UserUpdate,
@@ -599,7 +599,7 @@ def update_user(
     "/current-user",
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def get_current_user(
     auth_context: AuthContext = Security(authorize),
 ) -> UserResponse:
@@ -626,7 +626,7 @@ if server_config().auth_scheme != AuthScheme.EXTERNAL:
             422: error_response,
         },
     )
-    @handle_exceptions
+    @async_fastapi_endpoint_wrapper
     def update_myself(
         user: UserUpdate,
         request: Request,
@@ -705,7 +705,7 @@ if server_config().rbac_enabled:
             422: error_response,
         },
     )
-    @handle_exceptions
+    @async_fastapi_endpoint_wrapper
     def update_user_resource_membership(
         resource_type: str,
         resource_id: UUID,

zenml/zen_server/routers/webhook_endpoints.py

@@ -26,7 +26,7 @@ from zenml.event_sources.webhooks.base_webhook_event_source import (
 from zenml.exceptions import AuthorizationException, WebhookInactiveError
 from zenml.logger import get_logger
 from zenml.zen_server.utils import (
-    handle_exceptions,
+    async_fastapi_endpoint_wrapper,
     plugin_flavor_registry,
     zen_store,
 )
@@ -55,7 +55,7 @@ async def get_body(request: Request) -> bytes:
     "/{event_source_id}",
     response_model=Dict[str, str],
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def webhook(
     event_source_id: UUID,
     request: Request,