zenml-nightly 0.58.2.dev20240623__py3-none-any.whl → 0.58.2.dev20240626__py3-none-any.whl

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/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,
@@ -262,6 +263,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)

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

@@ -47,6 +47,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,
@@ -108,6 +109,10 @@ 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,
@@ -480,6 +485,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 +1533,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 +1637,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 +1645,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  --------------------
 

zenml/zen_stores/schemas/__init__.py

@@ -13,6 +13,7 @@
 #  permissions and limitations under the License.
 """SQL Model Implementations."""
 
+from zenml.zen_stores.schemas.action_schemas import ActionSchema
 from zenml.zen_stores.schemas.api_key_schemas import APIKeySchema
 from zenml.zen_stores.schemas.artifact_schemas import (
     ArtifactSchema,
@@ -71,6 +72,7 @@ from zenml.zen_stores.schemas.model_schemas import (
 from zenml.zen_stores.schemas.server_settings_schemas import ServerSettingsSchema
 
 __all__ = [
+    "ActionSchema",
     "APIKeySchema",
     "ArtifactSchema",
     "ArtifactVersionSchema",