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

zenml/zen_stores/migrations/versions/25155145c545_separate_actions_and_triggers.py

@@ -0,0 +1,228 @@
+"""Separate actions and triggers [25155145c545].
+
+Revision ID: 25155145c545
+Revises: 0.58.2
+Create Date: 2024-05-16 11:29:53.341275
+
+"""
+
+from datetime import datetime
+from uuid import uuid4
+
+import sqlalchemy as sa
+import sqlmodel
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "25155145c545"
+down_revision = "0.58.2"
+branch_labels = None
+depends_on = None
+
+
+def migrate_actions() -> None:
+    """Migrate actions from the trigger table."""
+    conn = op.get_bind()
+    meta = sa.MetaData()
+    meta.reflect(only=("trigger", "action"), bind=op.get_bind())
+    trigger_table = sa.Table("trigger", meta)
+    action_table = sa.Table("action", meta)
+
+    triggers = conn.execute(
+        sa.select(
+            trigger_table.c.id,
+            trigger_table.c.name,
+            trigger_table.c.user_id,
+            trigger_table.c.workspace_id,
+            trigger_table.c.service_account_id,
+            trigger_table.c.auth_window,
+            trigger_table.c.action,
+            trigger_table.c.action_subtype,
+            trigger_table.c.action_flavor,
+        )
+    ).fetchall()
+
+    now = datetime.utcnow()
+
+    actions_to_insert = []
+    trigger_updates = {}
+    for trigger in triggers:
+        action_id = str(uuid4()).replace("-", "")
+
+        actions_to_insert.append(
+            {
+                "id": action_id,
+                "workspace_id": trigger.workspace_id,
+                "user_id": trigger.user_id,
+                "created": now,
+                "updated": now,
+                "name": f"{trigger.name}_action",
+                "description": f"Automatically migrated action for trigger {trigger.name}",
+                "service_account_id": trigger.service_account_id,
+                "auth_window": trigger.auth_window,
+                "configuration": trigger.action,
+                "flavor": trigger.action_flavor,
+                "plugin_subtype": trigger.action_subtype,
+            }
+        )
+
+        trigger_updates[trigger.id] = action_id
+
+    op.bulk_insert(action_table, actions_to_insert)
+
+    for trigger_id, action_id in trigger_updates.items():
+        query = (
+            sa.update(trigger_table)
+            .where(trigger_table.c.id.is_(trigger_id))
+            .values(action_id=action_id)
+        )
+        conn.execute(query)
+
+
+def upgrade() -> None:
+    """Upgrade database schema and/or data, creating a new revision."""
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table(
+        "action",
+        sa.Column(
+            "workspace_id", sqlmodel.sql.sqltypes.GUID(), nullable=False
+        ),
+        sa.Column("user_id", sqlmodel.sql.sqltypes.GUID(), nullable=True),
+        sa.Column(
+            "service_account_id", sqlmodel.sql.sqltypes.GUID(), nullable=False
+        ),
+        sa.Column("description", sa.TEXT(), nullable=True),
+        sa.Column("id", sqlmodel.sql.sqltypes.GUID(), nullable=False),
+        sa.Column("created", sa.DateTime(), nullable=False),
+        sa.Column("updated", sa.DateTime(), nullable=False),
+        sa.Column("name", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
+        sa.Column("auth_window", sa.Integer(), nullable=False),
+        sa.Column(
+            "flavor", sqlmodel.sql.sqltypes.AutoString(), nullable=False
+        ),
+        sa.Column(
+            "plugin_subtype",
+            sqlmodel.sql.sqltypes.AutoString(),
+            nullable=False,
+        ),
+        sa.Column("configuration", sa.LargeBinary(), nullable=False),
+        sa.ForeignKeyConstraint(
+            ["service_account_id"],
+            ["user.id"],
+            name="fk_action_service_account_id_user",
+            ondelete="CASCADE",
+        ),
+        sa.ForeignKeyConstraint(
+            ["user_id"],
+            ["user.id"],
+            name="fk_action_user_id_user",
+            ondelete="SET NULL",
+        ),
+        sa.ForeignKeyConstraint(
+            ["workspace_id"],
+            ["workspace.id"],
+            name="fk_action_workspace_id_workspace",
+            ondelete="CASCADE",
+        ),
+        sa.PrimaryKeyConstraint("id"),
+    )
+
+    # Add the action_id column as nullable until we migrate the actions into
+    # a separate table
+    with op.batch_alter_table("trigger", schema=None) as batch_op:
+        batch_op.add_column(
+            sa.Column("action_id", sqlmodel.sql.sqltypes.GUID(), nullable=True)
+        )
+
+    migrate_actions()
+
+    with op.batch_alter_table("trigger", schema=None) as batch_op:
+        # Make the action_id column non-nullable now that we've inserted the
+        # actions and the value is set for each row
+        batch_op.alter_column(
+            "action_id",
+            existing_type=sqlmodel.sql.sqltypes.GUID(),
+            existing_nullable=True,
+            nullable=False,
+        )
+
+        batch_op.add_column(
+            sa.Column("schedule", sa.LargeBinary(), nullable=True)
+        )
+        batch_op.alter_column(
+            "event_source_id", existing_type=sa.CHAR(length=32), nullable=True
+        )
+        batch_op.drop_constraint(
+            "fk_trigger_event_source_id_event_source", type_="foreignkey"
+        )
+        batch_op.drop_constraint(
+            "fk_trigger_service_account_id_user", type_="foreignkey"
+        )
+        batch_op.create_foreign_key(
+            "fk_trigger_event_source_id_event_source",
+            "event_source",
+            ["event_source_id"],
+            ["id"],
+            ondelete="SET NULL",
+        )
+        batch_op.create_foreign_key(
+            "fk_trigger_action_id_action",
+            "action",
+            ["action_id"],
+            ["id"],
+            ondelete="CASCADE",
+        )
+        batch_op.drop_column("action_subtype")
+        batch_op.drop_column("action_flavor")
+        batch_op.drop_column("action")
+        batch_op.drop_column("auth_window")
+        batch_op.drop_column("service_account_id")
+
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    """Downgrade database schema and/or data back to the previous revision."""
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table("trigger", schema=None) as batch_op:
+        batch_op.add_column(
+            sa.Column("service_account_id", sa.CHAR(length=32), nullable=False)
+        )
+        batch_op.add_column(
+            sa.Column("auth_window", sa.INTEGER(), nullable=False)
+        )
+        batch_op.add_column(sa.Column("action", sa.BLOB(), nullable=False))
+        batch_op.add_column(
+            sa.Column("action_flavor", sa.VARCHAR(), nullable=False)
+        )
+        batch_op.add_column(
+            sa.Column("action_subtype", sa.VARCHAR(), nullable=False)
+        )
+        batch_op.drop_constraint(
+            "fk_trigger_action_id_action", type_="foreignkey"
+        )
+        batch_op.drop_constraint(
+            "fk_trigger_event_source_id_event_source", type_="foreignkey"
+        )
+        batch_op.create_foreign_key(
+            "fk_trigger_service_account_id_user",
+            "user",
+            ["service_account_id"],
+            ["id"],
+            ondelete="CASCADE",
+        )
+        batch_op.create_foreign_key(
+            "fk_trigger_event_source_id_event_source",
+            "event_source",
+            ["event_source_id"],
+            ["id"],
+            ondelete="CASCADE",
+        )
+        batch_op.alter_column(
+            "event_source_id", existing_type=sa.CHAR(length=32), nullable=False
+        )
+        batch_op.drop_column("schedule")
+        batch_op.drop_column("action_id")
+
+    op.drop_table("action")
+    # ### end Alembic commands ###

zenml/zen_stores/rest_zen_store.py

@@ -15,6 +15,7 @@
 
 import os
 import re
+from datetime import datetime
 from pathlib import Path
 from typing import (
     Any,
@@ -47,6 +48,7 @@ from zenml.config.global_config import GlobalConfiguration
 from zenml.config.pipeline_run_configuration import PipelineRunConfiguration
 from zenml.config.store_config import StoreConfiguration
 from zenml.constants import (
+    ACTIONS,
     API,
     API_KEY_ROTATE,
     API_KEYS,
@@ -56,6 +58,7 @@ from zenml.constants import (
     ARTIFACTS,
     CODE_REFERENCES,
     CODE_REPOSITORIES,
+    CONFIG,
     CURRENT_USER,
     DEACTIVATE,
     DEFAULT_HTTP_TIMEOUT,
@@ -64,6 +67,7 @@ from zenml.constants import (
     ENV_ZENML_DISABLE_CLIENT_SERVER_MISMATCH_WARNING,
     EVENT_SOURCES,
     FLAVORS,
+    FULL_STACK,
     GET_OR_CREATE,
     INFO,
     LOGIN,
@@ -88,9 +92,12 @@ from zenml.constants import (
     SERVICE_CONNECTOR_RESOURCES,
     SERVICE_CONNECTOR_TYPES,
     SERVICE_CONNECTOR_VERIFY,
+    SERVICE_CONNECTOR_VERIFY_REQUEST_TIMEOUT,
     SERVICE_CONNECTORS,
     SERVICES,
+    STACK,
     STACK_COMPONENTS,
+    STACK_DEPLOYMENT,
     STACKS,
     STEPS,
     TAGS,
@@ -102,12 +109,17 @@ from zenml.constants import (
 )
 from zenml.enums import (
     OAuthGrantTypes,
+    StackDeploymentProvider,
     StoreType,
 )
 from zenml.exceptions import AuthorizationException, MethodNotAllowedError
 from zenml.io import fileio
 from zenml.logger import get_logger
 from zenml.models import (
+    ActionFilter,
+    ActionRequest,
+    ActionResponse,
+    ActionUpdate,
     APIKeyFilter,
     APIKeyRequest,
     APIKeyResponse,
@@ -134,6 +146,7 @@ from zenml.models import (
     ComponentRequest,
     ComponentResponse,
     ComponentUpdate,
+    DeployedStack,
     EventSourceFilter,
     EventSourceRequest,
     EventSourceResponse,
@@ -142,6 +155,7 @@ from zenml.models import (
     FlavorRequest,
     FlavorResponse,
     FlavorUpdate,
+    FullStackRequest,
     LogsResponse,
     ModelFilter,
     ModelRequest,
@@ -203,6 +217,8 @@ from zenml.models import (
     ServiceRequest,
     ServiceResponse,
     ServiceUpdate,
+    StackDeploymentConfig,
+    StackDeploymentInfo,
     StackFilter,
     StackRequest,
     StackResponse,
@@ -480,6 +496,100 @@ class RestZenStore(BaseZenStore):
         response_body = self.put(SERVER_SETTINGS, body=settings_update)
         return ServerSettingsResponse.model_validate(response_body)
 
+    # -------------------- Actions  --------------------
+
+    def create_action(self, action: ActionRequest) -> ActionResponse:
+        """Create an action.
+
+        Args:
+            action: The action to create.
+
+        Returns:
+            The created action.
+        """
+        return self._create_resource(
+            resource=action,
+            route=ACTIONS,
+            response_model=ActionResponse,
+        )
+
+    def get_action(
+        self,
+        action_id: UUID,
+        hydrate: bool = True,
+    ) -> ActionResponse:
+        """Get an action by ID.
+
+        Args:
+            action_id: The ID of the action to get.
+            hydrate: Flag deciding whether to hydrate the output model(s)
+                by including metadata fields in the response.
+
+        Returns:
+            The action.
+        """
+        return self._get_resource(
+            resource_id=action_id,
+            route=ACTIONS,
+            response_model=ActionResponse,
+            params={"hydrate": hydrate},
+        )
+
+    def list_actions(
+        self,
+        action_filter_model: ActionFilter,
+        hydrate: bool = False,
+    ) -> Page[ActionResponse]:
+        """List all actions matching the given filter criteria.
+
+        Args:
+            action_filter_model: All filter parameters including pagination
+                params.
+            hydrate: Flag deciding whether to hydrate the output model(s)
+                by including metadata fields in the response.
+
+        Returns:
+            A list of all actions matching the filter criteria.
+        """
+        return self._list_paginated_resources(
+            route=ACTIONS,
+            response_model=ActionResponse,
+            filter_model=action_filter_model,
+            params={"hydrate": hydrate},
+        )
+
+    def update_action(
+        self,
+        action_id: UUID,
+        action_update: ActionUpdate,
+    ) -> ActionResponse:
+        """Update an existing action.
+
+        Args:
+            action_id: The ID of the action to update.
+            action_update: The update to be applied to the action.
+
+        Returns:
+            The updated action.
+        """
+        return self._update_resource(
+            resource_id=action_id,
+            resource_update=action_update,
+            route=ACTIONS,
+            response_model=ActionResponse,
+        )
+
+    def delete_action(self, action_id: UUID) -> None:
+        """Delete an action.
+
+        Args:
+            action_id: The ID of the action to delete.
+        """
+        self._delete_resource(
+            resource_id=action_id,
+            route=ACTIONS,
+        )
+
     # ----------------------------- API Keys -----------------------------
 
     def create_api_key(
@@ -1434,14 +1544,14 @@ class RestZenStore(BaseZenStore):
         run_configuration = run_configuration or PipelineRunConfiguration()
         try:
             response_body = self.post(
-                f"{PIPELINE_BUILDS}/{build_id}/run", body=run_configuration
+                f"{PIPELINE_BUILDS}/{build_id}/runs", body=run_configuration
             )
         except MethodNotAllowedError as e:
             raise RuntimeError(
                 "Running a build is not supported for this server."
             ) from e
 
-        return PipelineRunResponse.parse_obj(response_body)
+        return PipelineRunResponse.model_validate(response_body)
 
     # -------------------------- Pipeline Deployments --------------------------
 
@@ -1538,7 +1648,7 @@ class RestZenStore(BaseZenStore):
 
         try:
             response_body = self.post(
-                f"{PIPELINE_DEPLOYMENTS}/{deployment_id}/run",
+                f"{PIPELINE_DEPLOYMENTS}/{deployment_id}/runs",
                 body=run_configuration,
             )
         except MethodNotAllowedError as e:
@@ -1546,7 +1656,7 @@ class RestZenStore(BaseZenStore):
                 "Running a deployment is not supported for this server."
             ) from e
 
-        return PipelineRunResponse.parse_obj(response_body)
+        return PipelineRunResponse.model_validate(response_body)
 
     # -------------------- Event Sources  --------------------
 
@@ -2368,6 +2478,10 @@ class RestZenStore(BaseZenStore):
             f"{SERVICE_CONNECTORS}{SERVICE_CONNECTOR_VERIFY}",
             body=service_connector,
             params={"list_resources": list_resources},
+            timeout=max(
+                self.config.http_timeout,
+                SERVICE_CONNECTOR_VERIFY_REQUEST_TIMEOUT,
+            ),
         )
 
         resources = ServiceConnectorResourcesModel.model_validate(
@@ -2405,6 +2519,10 @@ class RestZenStore(BaseZenStore):
         response_body = self.put(
             f"{SERVICE_CONNECTORS}/{str(service_connector_id)}{SERVICE_CONNECTOR_VERIFY}",
             params=params,
+            timeout=max(
+                self.config.http_timeout,
+                SERVICE_CONNECTOR_VERIFY_REQUEST_TIMEOUT,
+            ),
         )
 
         resources = ServiceConnectorResourcesModel.model_validate(
@@ -2645,6 +2763,23 @@ class RestZenStore(BaseZenStore):
             response_model=StackResponse,
         )
 
+    def create_full_stack(self, full_stack: FullStackRequest) -> StackResponse:
+        """Register a full-stack.
+
+        Args:
+            full_stack: The full stack configuration.
+
+        Returns:
+            The registered stack.
+        """
+        assert full_stack.workspace is not None
+
+        return self._create_resource(
+            resource=full_stack,
+            response_model=StackResponse,
+            route=f"{WORKSPACES}/{str(full_stack.workspace)}{FULL_STACK}",
+        )
+
     def get_stack(self, stack_id: UUID, hydrate: bool = True) -> StackResponse:
         """Get a stack by its unique ID.
 
@@ -2714,6 +2849,88 @@ class RestZenStore(BaseZenStore):
             route=STACKS,
         )
 
+    # ---------------- Stack deployments-----------------
+
+    def get_stack_deployment_info(
+        self,
+        provider: StackDeploymentProvider,
+    ) -> StackDeploymentInfo:
+        """Get information about a stack deployment provider.
+
+        Args:
+            provider: The stack deployment provider.
+
+        Returns:
+            Information about the stack deployment provider.
+        """
+        body = self.get(
+            f"{STACK_DEPLOYMENT}{INFO}",
+            params={"provider": provider.value},
+        )
+        return StackDeploymentInfo.model_validate(body)
+
+    def get_stack_deployment_config(
+        self,
+        provider: StackDeploymentProvider,
+        stack_name: str,
+        location: Optional[str] = None,
+    ) -> StackDeploymentConfig:
+        """Return the cloud provider console URL and configuration needed to deploy the ZenML stack.
+
+        Args:
+            provider: The stack deployment provider.
+            stack_name: The name of the stack.
+            location: The location where the stack should be deployed.
+
+        Returns:
+            The cloud provider console URL and configuration needed to deploy
+            the ZenML stack to the specified cloud provider.
+        """
+        params = {
+            "provider": provider.value,
+            "stack_name": stack_name,
+        }
+        if location:
+            params["location"] = location
+        body = self.get(f"{STACK_DEPLOYMENT}{CONFIG}", params=params)
+        return StackDeploymentConfig.model_validate(body)
+
+    def get_stack_deployment_stack(
+        self,
+        provider: StackDeploymentProvider,
+        stack_name: str,
+        location: Optional[str] = None,
+        date_start: Optional[datetime] = None,
+    ) -> 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.
+        """
+        params = {
+            "provider": provider.value,
+            "stack_name": stack_name,
+        }
+        if location:
+            params["location"] = location
+        if date_start:
+            params["date_start"] = str(date_start)
+        body = self.get(
+            f"{STACK_DEPLOYMENT}{STACK}",
+            params=params,
+        )
+        if body:
+            return DeployedStack.model_validate(body)
+
+        return None
+
     # ----------------------------- Step runs -----------------------------
 
     def create_run_step(self, step_run: StepRunRequest) -> StepRunResponse:
@@ -3865,6 +4082,7 @@ class RestZenStore(BaseZenStore):
         method: str,
         url: str,
         params: Optional[Dict[str, Any]] = None,
+        timeout: Optional[int] = None,
         **kwargs: Any,
     ) -> Json:
         """Make a request to the REST API.
@@ -3873,6 +4091,7 @@ class RestZenStore(BaseZenStore):
             method: The HTTP method to use.
             url: The URL to request.
             params: The query parameters to pass to the endpoint.
+            timeout: The request timeout in seconds.
             kwargs: Additional keyword arguments to pass to the request.
 
         Returns:
@@ -3895,7 +4114,7 @@ class RestZenStore(BaseZenStore):
                     url,
                     params=params,
                     verify=self.config.verify_ssl,
-                    timeout=self.config.http_timeout,
+                    timeout=timeout or self.config.http_timeout,
                     **kwargs,
                 )
             )
@@ -3924,13 +4143,18 @@ class RestZenStore(BaseZenStore):
             raise
 
     def get(
-        self, path: str, params: Optional[Dict[str, Any]] = None, **kwargs: Any
+        self,
+        path: str,
+        params: Optional[Dict[str, Any]] = None,
+        timeout: Optional[int] = None,
+        **kwargs: Any,
     ) -> Json:
         """Make a GET request to the given endpoint path.
 
         Args:
             path: The path to the endpoint.
             params: The query parameters to pass to the endpoint.
+            timeout: The request timeout in seconds.
             kwargs: Additional keyword arguments to pass to the request.
 
         Returns:
@@ -3938,17 +4162,26 @@ class RestZenStore(BaseZenStore):
         """
         logger.debug(f"Sending GET request to {path}...")
         return self._request(
-            "GET", self.url + API + VERSION_1 + path, params=params, **kwargs
+            "GET",
+            self.url + API + VERSION_1 + path,
+            params=params,
+            timeout=timeout,
+            **kwargs,
         )
 
     def delete(
-        self, path: str, params: Optional[Dict[str, Any]] = None, **kwargs: Any
+        self,
+        path: str,
+        params: Optional[Dict[str, Any]] = None,
+        timeout: Optional[int] = None,
+        **kwargs: Any,
     ) -> Json:
         """Make a DELETE request to the given endpoint path.
 
         Args:
             path: The path to the endpoint.
             params: The query parameters to pass to the endpoint.
+            timeout: The request timeout in seconds.
             kwargs: Additional keyword arguments to pass to the request.
 
         Returns:
@@ -3959,6 +4192,7 @@ class RestZenStore(BaseZenStore):
             "DELETE",
             self.url + API + VERSION_1 + path,
             params=params,
+            timeout=timeout,
             **kwargs,
         )
 
@@ -3967,6 +4201,7 @@ class RestZenStore(BaseZenStore):
         path: str,
         body: BaseModel,
         params: Optional[Dict[str, Any]] = None,
+        timeout: Optional[int] = None,
         **kwargs: Any,
     ) -> Json:
         """Make a POST request to the given endpoint path.
@@ -3975,6 +4210,7 @@ class RestZenStore(BaseZenStore):
             path: The path to the endpoint.
             body: The body to send.
             params: The query parameters to pass to the endpoint.
+            timeout: The request timeout in seconds.
             kwargs: Additional keyword arguments to pass to the request.
 
         Returns:
@@ -3986,6 +4222,7 @@ class RestZenStore(BaseZenStore):
             self.url + API + VERSION_1 + path,
             data=body.model_dump_json(),
             params=params,
+            timeout=timeout,
             **kwargs,
         )
 
@@ -3994,6 +4231,7 @@ class RestZenStore(BaseZenStore):
         path: str,
         body: Optional[BaseModel] = None,
         params: Optional[Dict[str, Any]] = None,
+        timeout: Optional[int] = None,
         **kwargs: Any,
     ) -> Json:
         """Make a PUT request to the given endpoint path.
@@ -4002,6 +4240,7 @@ class RestZenStore(BaseZenStore):
             path: The path to the endpoint.
             body: The body to send.
             params: The query parameters to pass to the endpoint.
+            timeout: The request timeout in seconds.
             kwargs: Additional keyword arguments to pass to the request.
 
         Returns:
@@ -4014,6 +4253,7 @@ class RestZenStore(BaseZenStore):
             self.url + API + VERSION_1 + path,
             data=data,
             params=params,
+            timeout=timeout,
             **kwargs,
         )