zenml-nightly 0.58.2.dev20240623__py3-none-any.whl → 0.61.0.dev20240712__py3-none-any.whl

zenml/zen_server/routers/stack_deployment_endpoints.py

@@ -0,0 +1,158 @@
+#  Copyright (c) ZenML GmbH 2024. 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 stack deployments."""
+
+import datetime
+from typing import Optional
+
+from fastapi import APIRouter, Request, Security
+
+from zenml.constants import (
+    API,
+    CONFIG,
+    INFO,
+    STACK,
+    STACK_DEPLOYMENT,
+    STACK_DEPLOYMENT_API_TOKEN_EXPIRATION,
+    VERSION_1,
+)
+from zenml.enums import StackDeploymentProvider
+from zenml.models import (
+    DeployedStack,
+    StackDeploymentConfig,
+    StackDeploymentInfo,
+)
+from zenml.stack_deployments.utils import get_stack_deployment_class
+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.utils import verify_permission
+from zenml.zen_server.utils import (
+    handle_exceptions,
+)
+
+router = APIRouter(
+    prefix=API + VERSION_1 + STACK_DEPLOYMENT,
+    tags=["stacks"],
+    responses={401: error_response, 403: error_response},
+)
+
+
+@router.get(
+    INFO,
+)
+@handle_exceptions
+def get_stack_deployment_info(
+    provider: StackDeploymentProvider,
+    _: AuthContext = Security(authorize),
+) -> StackDeploymentInfo:
+    """Get information about a stack deployment provider.
+
+    Args:
+        provider: The stack deployment provider.
+
+    Returns:
+        Information about the stack deployment provider.
+    """
+    stack_deployment_class = get_stack_deployment_class(provider)
+    return stack_deployment_class.get_deployment_info()
+
+
+@router.get(
+    CONFIG,
+)
+@handle_exceptions
+def get_stack_deployment_config(
+    request: Request,
+    provider: StackDeploymentProvider,
+    stack_name: str,
+    location: Optional[str] = None,
+    auth_context: AuthContext = Security(authorize),
+) -> StackDeploymentConfig:
+    """Return the URL to deploy the ZenML stack to the specified cloud provider.
+
+    Args:
+        request: The FastAPI request object.
+        provider: The stack deployment provider.
+        stack_name: The name of the stack.
+        location: The location where the stack should be deployed.
+        auth_context: The authentication context.
+
+    Returns:
+        The cloud provider console URL where the stack will be deployed and
+        the configuration for the stack deployment.
+    """
+    verify_permission(
+        resource_type=ResourceType.SERVICE_CONNECTOR, action=Action.CREATE
+    )
+    verify_permission(
+        resource_type=ResourceType.STACK_COMPONENT,
+        action=Action.CREATE,
+    )
+    verify_permission(resource_type=ResourceType.STACK, action=Action.CREATE)
+
+    stack_deployment_class = get_stack_deployment_class(provider)
+    # Get the base server URL used to call this FastAPI endpoint
+    url = request.url.replace(path="").replace(query="")
+    # Use HTTPS for the URL
+    url = url.replace(scheme="https")
+
+    token = auth_context.access_token
+    assert token is not None
+
+    # A new API token is generated for the stack deployment
+    expires = datetime.datetime.utcnow() + datetime.timedelta(
+        minutes=STACK_DEPLOYMENT_API_TOKEN_EXPIRATION
+    )
+    api_token = token.encode(expires=expires)
+
+    return stack_deployment_class(
+        stack_name=stack_name,
+        location=location,
+        zenml_server_url=str(url),
+        zenml_server_api_token=api_token,
+    ).get_deployment_config()
+
+
+@router.get(
+    STACK,
+)
+@handle_exceptions
+def get_deployed_stack(
+    provider: StackDeploymentProvider,
+    stack_name: str,
+    location: Optional[str] = None,
+    date_start: Optional[datetime.datetime] = None,
+    _: AuthContext = Security(authorize),
+) -> Optional[DeployedStack]:
+    """Return a matching ZenML stack that was deployed and registered.
+
+    Args:
+        provider: The stack deployment provider.
+        stack_name: The name of the stack.
+        location: The location where the stack should be deployed.
+        date_start: The date when the deployment started.
+
+    Returns:
+        The ZenML stack that was deployed and registered or None if the stack
+        was not found.
+    """
+    stack_deployment_class = get_stack_deployment_class(provider)
+    return stack_deployment_class(
+        stack_name=stack_name,
+        location=location,
+        # These fields are not needed for this operation
+        zenml_server_url="",
+        zenml_server_api_token="",
+    ).get_stack(date_start=date_start)

zenml/zen_server/routers/triggers_endpoints.py

@@ -18,7 +18,6 @@ from uuid import UUID
 from fastapi import APIRouter, Depends, Security
 
 from zenml import TriggerRequest
-from zenml.actions.base_action import BaseActionHandler
 from zenml.constants import API, TRIGGER_EXECUTIONS, TRIGGERS, VERSION_1
 from zenml.enums import PluginType
 from zenml.event_sources.base_event_source import BaseEventSourceHandler
@@ -81,54 +80,11 @@ def list_triggers(
     Returns:
         All triggers.
     """
-
-    def list_triggers_fn(
-        filter_model: TriggerFilter,
-    ) -> Page[TriggerResponse]:
-        """List triggers through their associated plugins.
-
-        Args:
-            filter_model: Filter model used for pagination, sorting,
-                filtering.
-
-        Returns:
-            All triggers.
-
-        Raises:
-            ValueError: If the plugin for a trigger action is not a valid action
-                plugin.
-        """
-        triggers = zen_store().list_triggers(
-            trigger_filter_model=filter_model, hydrate=hydrate
-        )
-
-        # Process the triggers through their associated plugins
-        for idx, trigger in enumerate(triggers.items):
-            action_handler = plugin_flavor_registry().get_plugin(
-                name=trigger.action_flavor,
-                _type=PluginType.ACTION,
-                subtype=trigger.action_subtype,
-            )
-
-            # Validate that the flavor and plugin_type correspond to an action
-            # handler implementation
-            if not isinstance(action_handler, BaseActionHandler):
-                raise ValueError(
-                    f"Action handler plugin {trigger.action_subtype} "
-                    f"for flavor {trigger.action_flavor} is not a valid action "
-                    "handler plugin."
-                )
-
-            triggers.items[idx] = action_handler.get_trigger(
-                trigger, hydrate=hydrate
-            )
-
-        return triggers
-
     return verify_permissions_and_list_entities(
         filter_model=trigger_filter_model,
         resource_type=ResourceType.TRIGGER,
-        list_method=list_triggers_fn,
+        list_method=zen_store().list_triggers,
+        hydrate=hydrate,
     )
 
 
@@ -152,31 +108,9 @@ def get_trigger(
 
     Returns:
         The requested trigger.
-
-    Raises:
-        ValueError: If the action flavor/subtype combination is not actually a webhook event source
     """
     trigger = zen_store().get_trigger(trigger_id=trigger_id, hydrate=hydrate)
-
     verify_permission_for_model(trigger, action=Action.READ)
-
-    action_handler = plugin_flavor_registry().get_plugin(
-        name=trigger.action_flavor,
-        _type=PluginType.ACTION,
-        subtype=trigger.action_subtype,
-    )
-
-    # Validate that the flavor and plugin_type correspond to an action
-    # handler implementation
-    if not isinstance(action_handler, BaseActionHandler):
-        raise ValueError(
-            f"Action handler plugin {trigger.action_subtype} "
-            f"for flavor {trigger.action_flavor} is not a valid action "
-            "handler plugin."
-        )
-
-    trigger = action_handler.get_trigger(trigger, hydrate=hydrate)
-
     return dehydrate_response_model(trigger)
 
 
@@ -201,55 +135,35 @@ def create_trigger(
     Raises:
         ValueError: If the action flavor/subtype combination is not actually a webhook event source
     """
-    if trigger.service_account_id:
-        service_account = zen_store().get_service_account(
-            service_account_name_or_id=trigger.service_account_id
+    if trigger.event_source_id and trigger.event_filter:
+        event_source = zen_store().get_event_source(
+            event_source_id=trigger.event_source_id
         )
-        verify_permission_for_model(service_account, action=Action.READ)
-
-    event_source = zen_store().get_event_source(
-        event_source_id=trigger.event_source_id
-    )
 
-    event_source_handler = plugin_flavor_registry().get_plugin(
-        name=event_source.flavor,
-        _type=PluginType.EVENT_SOURCE,
-        subtype=event_source.plugin_subtype,
-    )
-
-    # Validate that the flavor and plugin_type correspond to an event source
-    # implementation
-    if not isinstance(event_source_handler, BaseEventSourceHandler):
-        raise ValueError(
-            f"Event source plugin {event_source.plugin_subtype} "
-            f"for flavor {event_source.flavor} is not a valid event source "
-            "handler implementation."
+        event_source_handler = plugin_flavor_registry().get_plugin(
+            name=event_source.flavor,
+            _type=PluginType.EVENT_SOURCE,
+            subtype=event_source.plugin_subtype,
         )
 
-    # Validate the trigger event filter
-    event_source_handler.validate_event_filter_configuration(
-        trigger.event_filter
-    )
-
-    action_handler = plugin_flavor_registry().get_plugin(
-        name=trigger.action_flavor,
-        _type=PluginType.ACTION,
-        subtype=trigger.action_subtype,
-    )
+        # Validate that the flavor and plugin_type correspond to an event source
+        # implementation
+        if not isinstance(event_source_handler, BaseEventSourceHandler):
+            raise ValueError(
+                f"Event source plugin {event_source.plugin_subtype} "
+                f"for flavor {event_source.flavor} is not a valid event source "
+                "handler implementation."
+            )
 
-    # Validate that the flavor and plugin_type correspond to an action
-    # handler implementation
-    if not isinstance(action_handler, BaseActionHandler):
-        raise ValueError(
-            f"Action handler plugin {trigger.action_subtype} "
-            f"for flavor {trigger.action_flavor} is not a valid action "
-            "handler plugin."
+        # Validate the trigger event filter
+        event_source_handler.validate_event_filter_configuration(
+            trigger.event_filter
         )
 
     return verify_permissions_and_create_entity(
         request_model=trigger,
         resource_type=ResourceType.TRIGGER,
-        create_method=action_handler.create_trigger,
+        create_method=zen_store().create_trigger,
     )
 
 
@@ -278,13 +192,12 @@ def update_trigger(
     """
     trigger = zen_store().get_trigger(trigger_id=trigger_id)
 
-    if trigger_update.service_account_id:
-        service_account = zen_store().get_service_account(
-            service_account_name_or_id=trigger_update.service_account_id
-        )
-        verify_permission_for_model(service_account, action=Action.READ)
-
     if trigger_update.event_filter:
+        if not trigger.event_source:
+            raise ValueError(
+                "Trying to set event filter for trigger without event source."
+            )
+
         event_source = zen_store().get_event_source(
             event_source_id=trigger.event_source.id
         )
@@ -306,29 +219,13 @@ def update_trigger(
 
         # Validate the trigger event filter
         event_source_handler.validate_event_filter_configuration(
-            trigger.event_filter
+            trigger_update.event_filter
         )
 
     verify_permission_for_model(trigger, action=Action.UPDATE)
 
-    action_handler = plugin_flavor_registry().get_plugin(
-        name=trigger.action_flavor,
-        _type=PluginType.ACTION,
-        subtype=trigger.action_subtype,
-    )
-
-    # Validate that the flavor and plugin_type correspond to an action
-    # handler implementation
-    if not isinstance(action_handler, BaseActionHandler):
-        raise ValueError(
-            f"Action handler plugin {trigger.action_subtype} "
-            f"for flavor {trigger.action_flavor} is not a valid action "
-            "handler plugin."
-        )
-
-    updated_trigger = action_handler.update_trigger(
-        trigger=trigger,
-        trigger_update=trigger_update,
+    updated_trigger = zen_store().update_trigger(
+        trigger_id=trigger_id, trigger_update=trigger_update
     )
 
     return dehydrate_response_model(updated_trigger)
@@ -341,41 +238,16 @@ def update_trigger(
 @handle_exceptions
 def delete_trigger(
     trigger_id: UUID,
-    force: bool = False,
     _: AuthContext = Security(authorize),
 ) -> None:
     """Deletes a trigger.
 
     Args:
         trigger_id: Name of the trigger.
-        force: Flag deciding whether to force delete the trigger.
-
-    Raises:
-        ValueError: If the action flavor/subtype combination is not actually a webhook event source
     """
     trigger = zen_store().get_trigger(trigger_id=trigger_id)
-
     verify_permission_for_model(trigger, action=Action.DELETE)
-
-    action_handler = plugin_flavor_registry().get_plugin(
-        name=trigger.action_flavor,
-        _type=PluginType.ACTION,
-        subtype=trigger.action_subtype,
-    )
-
-    # Validate that the flavor and plugin_type correspond to an action
-    # handler implementation
-    if not isinstance(action_handler, BaseActionHandler):
-        raise ValueError(
-            f"Action handler plugin {trigger.action_subtype} "
-            f"for flavor {trigger.action_flavor} is not a valid action "
-            "handler plugin."
-        )
-
-    action_handler.delete_trigger(
-        trigger=trigger,
-        force=force,
-    )
+    zen_store().delete_trigger(trigger_id=trigger_id)
 
 
 executions_router = APIRouter(

zenml/zen_server/routers/workspaces_endpoints.py

@@ -22,6 +22,7 @@ from zenml.constants import (
     API,
     ARTIFACTS,
     CODE_REPOSITORIES,
+    FULL_STACK,
     GET_OR_CREATE,
     MODEL_VERSIONS,
     MODELS,
@@ -51,6 +52,7 @@ from zenml.models import (
     ComponentFilter,
     ComponentRequest,
     ComponentResponse,
+    FullStackRequest,
     ModelRequest,
     ModelResponse,
     ModelVersionArtifactRequest,
@@ -349,6 +351,68 @@ def create_stack(
     )
 
 
+@router.post(
+    WORKSPACES + "/{workspace_name_or_id}" + FULL_STACK,
+    response_model=StackResponse,
+    responses={401: error_response, 409: error_response, 422: error_response},
+)
+@handle_exceptions
+def create_full_stack(
+    workspace_name_or_id: Union[str, UUID],
+    full_stack: FullStackRequest,
+    auth_context: AuthContext = Security(authorize),
+) -> StackResponse:
+    """Creates a stack for a particular workspace.
+
+    Args:
+        workspace_name_or_id: Name or ID of the workspace.
+        full_stack: Stack to register.
+        auth_context: Authentication context.
+
+    Returns:
+        The created stack.
+    """
+    workspace = zen_store().get_workspace(workspace_name_or_id)
+
+    is_connector_create_needed = False
+    for connector_id_or_info in full_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
+    if is_connector_create_needed:
+        verify_permission(
+            resource_type=ResourceType.SERVICE_CONNECTOR, action=Action.CREATE
+        )
+
+    is_component_create_needed = False
+    for component_id_or_info in full_stack.components.values():
+        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,
+        )
+
+    verify_permission(resource_type=ResourceType.STACK, action=Action.CREATE)
+
+    full_stack.user = auth_context.user.id
+    full_stack.workspace = workspace.id
+
+    return zen_store().create_full_stack(full_stack)
+
+
 @router.get(
     WORKSPACES + "/{workspace_name_or_id}" + STACK_COMPONENTS,
     response_model=Page[ComponentResponse],

zenml/zen_server/zen_server_api.py

@@ -40,6 +40,7 @@ from zenml.constants import API, HEALTH
 from zenml.enums import AuthScheme, SourceContextTypes
 from zenml.zen_server.exceptions import error_detail
 from zenml.zen_server.routers import (
+    actions_endpoints,
     artifact_endpoint,
     artifact_version_endpoints,
     auth_endpoints,
@@ -62,6 +63,7 @@ from zenml.zen_server.routers import (
     service_connectors_endpoints,
     service_endpoints,
     stack_components_endpoints,
+    stack_deployment_endpoints,
     stacks_endpoints,
     steps_endpoints,
     tags_endpoints,
@@ -262,6 +264,7 @@ async def dashboard(request: Request) -> Any:
     return templates.TemplateResponse("index.html", {"request": request})
 
 
+app.include_router(actions_endpoints.router)
 app.include_router(artifact_endpoint.artifact_router)
 app.include_router(artifact_version_endpoints.artifact_version_router)
 app.include_router(auth_endpoints.router)
@@ -288,6 +291,7 @@ app.include_router(service_accounts_endpoints.router)
 app.include_router(service_connectors_endpoints.router)
 app.include_router(service_connectors_endpoints.types_router)
 app.include_router(service_endpoints.router)
+app.include_router(stack_deployment_endpoints.router)
 app.include_router(stacks_endpoints.router)
 app.include_router(stack_components_endpoints.router)
 app.include_router(stack_components_endpoints.types_router)

zenml/zen_stores/migrations/utils.py

@@ -364,7 +364,7 @@ class MigrationUtils(BaseModel):
                         # Convert column values to the correct type
                         for column in table.columns:
                             # Blob columns are stored as binary strings
-                            if column.type.python_type == bytes and isinstance(
+                            if column.type.python_type is bytes and isinstance(
                                 row[column.name], str
                             ):
                                 # Convert the string to bytes

zenml/zen_stores/migrations/versions/0.60.0_release.py

@@ -0,0 +1,23 @@
+"""Release [0.60.0].
+
+Revision ID: 0.60.0
+Revises: 25155145c545
+Create Date: 2024-06-25 14:26:48.867005
+
+"""
+
+# revision identifiers, used by Alembic.
+revision = "0.60.0"
+down_revision = "25155145c545"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    """Upgrade database schema and/or data, creating a new revision."""
+    pass
+
+
+def downgrade() -> None:
+    """Downgrade database schema and/or data back to the previous revision."""
+    pass

zenml/zen_stores/migrations/versions/0.61.0_release.py

@@ -0,0 +1,23 @@
+"""Release [0.61.0].
+
+Revision ID: 0.61.0
+Revises: 0d707865f404
+Create Date: 2024-07-08 15:52:54.765307
+
+"""
+
+# revision identifiers, used by Alembic.
+revision = "0.61.0"
+down_revision = "0d707865f404"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    """Upgrade database schema and/or data, creating a new revision."""
+    pass
+
+
+def downgrade() -> None:
+    """Downgrade database schema and/or data back to the previous revision."""
+    pass

zenml/zen_stores/migrations/versions/0d707865f404_adding_labels_to_stacks.py

@@ -0,0 +1,30 @@
+"""adding labels to stacks [0d707865f404].
+
+Revision ID: 0d707865f404
+Revises: 0.60.0
+Create Date: 2024-07-04 16:10:07.709184
+
+"""
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "0d707865f404"
+down_revision = "0.60.0"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    """Upgrade database schema and/or data, creating a new revision."""
+    with op.batch_alter_table("stack", schema=None) as batch_op:
+        batch_op.add_column(
+            sa.Column("labels", sa.LargeBinary(), nullable=True)
+        )
+
+
+def downgrade() -> None:
+    """Downgrade database schema and/or data back to the previous revision."""
+    with op.batch_alter_table("stack", schema=None) as batch_op:
+        batch_op.drop_column("labels")