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

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",

zenml/zen_stores/schemas/action_schemas.py

@@ -0,0 +1,192 @@
+#  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.
+"""SQL Model Implementations for Actions."""
+
+import base64
+import json
+from datetime import datetime
+from typing import TYPE_CHECKING, Any, List, Optional
+from uuid import UUID
+
+from pydantic.json import pydantic_encoder
+from sqlalchemy import TEXT, Column
+from sqlmodel import Field, Relationship
+
+from zenml.models import (
+    ActionRequest,
+    ActionResponse,
+    ActionResponseBody,
+    ActionResponseMetadata,
+    ActionResponseResources,
+    ActionUpdate,
+)
+from zenml.zen_stores.schemas.base_schemas import NamedSchema
+from zenml.zen_stores.schemas.schema_utils import build_foreign_key_field
+from zenml.zen_stores.schemas.user_schemas import UserSchema
+from zenml.zen_stores.schemas.workspace_schemas import WorkspaceSchema
+
+if TYPE_CHECKING:
+    from zenml.zen_stores.schemas import TriggerSchema
+
+
+class ActionSchema(NamedSchema, table=True):
+    """SQL Model for actions."""
+
+    __tablename__ = "action"
+
+    workspace_id: UUID = build_foreign_key_field(
+        source=__tablename__,
+        target=WorkspaceSchema.__tablename__,
+        source_column="workspace_id",
+        target_column="id",
+        ondelete="CASCADE",
+        nullable=False,
+    )
+    workspace: "WorkspaceSchema" = Relationship(back_populates="actions")
+
+    user_id: Optional[UUID] = build_foreign_key_field(
+        source=__tablename__,
+        target=UserSchema.__tablename__,
+        source_column="user_id",
+        target_column="id",
+        ondelete="SET NULL",
+        nullable=True,
+    )
+    user: Optional["UserSchema"] = Relationship(
+        back_populates="actions",
+        sa_relationship_kwargs={"foreign_keys": "[ActionSchema.user_id]"},
+    )
+
+    triggers: List["TriggerSchema"] = Relationship(back_populates="action")
+
+    service_account_id: UUID = build_foreign_key_field(
+        source=__tablename__,
+        target=UserSchema.__tablename__,
+        source_column="service_account_id",
+        target_column="id",
+        ondelete="CASCADE",
+        nullable=False,
+    )
+    service_account: UserSchema = Relationship(
+        back_populates="auth_actions",
+        sa_relationship_kwargs={
+            "foreign_keys": "[ActionSchema.service_account_id]"
+        },
+    )
+    auth_window: int
+
+    flavor: str = Field(nullable=False)
+    plugin_subtype: str = Field(nullable=False)
+    description: Optional[str] = Field(sa_column=Column(TEXT, nullable=True))
+
+    configuration: bytes
+
+    @classmethod
+    def from_request(cls, request: "ActionRequest") -> "ActionSchema":
+        """Convert a `ActionRequest` to a `ActionSchema`.
+
+        Args:
+            request: The request model to convert.
+
+        Returns:
+            The converted schema.
+        """
+        return cls(
+            name=request.name,
+            workspace_id=request.workspace,
+            user_id=request.user,
+            configuration=base64.b64encode(
+                json.dumps(
+                    request.configuration, default=pydantic_encoder
+                ).encode("utf-8"),
+            ),
+            flavor=request.flavor,
+            plugin_subtype=request.plugin_subtype,
+            description=request.description,
+            service_account_id=request.service_account_id,
+            auth_window=request.auth_window,
+        )
+
+    def update(self, action_update: "ActionUpdate") -> "ActionSchema":
+        """Updates a action schema with a action update model.
+
+        Args:
+            action_update: `ActionUpdate` to update the action with.
+
+        Returns:
+            The updated ActionSchema.
+        """
+        for field, value in action_update.dict(
+            exclude_unset=True,
+            exclude_none=True,
+        ).items():
+            if field == "configuration":
+                self.configuration = base64.b64encode(
+                    json.dumps(
+                        action_update.configuration, default=pydantic_encoder
+                    ).encode("utf-8")
+                )
+            else:
+                setattr(self, field, value)
+
+        self.updated = datetime.utcnow()
+        return self
+
+    def to_model(
+        self,
+        include_metadata: bool = False,
+        include_resources: bool = False,
+        **kwargs: Any,
+    ) -> "ActionResponse":
+        """Converts the action schema to a model.
+
+        Args:
+            include_metadata: Flag deciding whether to include the output model(s)
+                metadata fields in the response.
+            include_resources: Flag deciding whether to include the output model(s)
+                metadata fields in the response.
+            **kwargs: Keyword arguments to allow schema specific logic
+
+        Returns:
+            The converted model.
+        """
+        body = ActionResponseBody(
+            user=self.user.to_model() if self.user else None,
+            created=self.created,
+            updated=self.updated,
+            flavor=self.flavor,
+            plugin_subtype=self.plugin_subtype,
+        )
+        metadata = None
+        if include_metadata:
+            metadata = ActionResponseMetadata(
+                workspace=self.workspace.to_model(),
+                configuration=json.loads(
+                    base64.b64decode(self.configuration).decode()
+                ),
+                description=self.description,
+                auth_window=self.auth_window,
+            )
+        resources = None
+        if include_resources:
+            resources = ActionResponseResources(
+                service_account=self.service_account.to_model(),
+            )
+        return ActionResponse(
+            id=self.id,
+            name=self.name,
+            body=body,
+            metadata=metadata,
+            resources=resources,
+        )

zenml/zen_stores/schemas/stack_schemas.py

@@ -13,6 +13,8 @@
 #  permissions and limitations under the License.
 """SQL Model Implementations for Stacks."""
 
+import base64
+import json
 from datetime import datetime
 from typing import TYPE_CHECKING, Any, List, Optional
 from uuid import UUID
@@ -75,6 +77,7 @@ class StackSchema(NamedSchema, table=True):
 
     __tablename__ = "stack"
     stack_spec_path: Optional[str]
+    labels: Optional[bytes]
 
     workspace_id: UUID = build_foreign_key_field(
         source=__tablename__,
@@ -124,6 +127,10 @@ class StackSchema(NamedSchema, table=True):
         ).items():
             if field == "components":
                 self.components = components
+            elif field == "labels":
+                self.labels = base64.b64encode(
+                    json.dumps(stack_update.labels).encode("utf-8")
+                )
             else:
                 setattr(self, field, value)
 
@@ -158,6 +165,9 @@ class StackSchema(NamedSchema, table=True):
                 workspace=self.workspace.to_model(),
                 components={c.type: [c.to_model()] for c in self.components},
                 stack_spec_path=self.stack_spec_path,
+                labels=json.loads(base64.b64decode(self.labels).decode())
+                if self.labels
+                else None,
             )
 
         return StackResponse(

zenml/zen_stores/schemas/step_run_schemas.py

@@ -192,6 +192,7 @@ class StepRunSchema(NamedSchema, table=True):
             The created StepRunResponse.
 
         Raises:
+            ValueError: In case the step run configuration can not be loaded.
             RuntimeError: If the step run schema does not have a deployment_id
                 or a step_configuration.
         """
@@ -212,19 +213,34 @@ class StepRunSchema(NamedSchema, table=True):
             for artifact in self.output_artifacts
         }
 
+        full_step_config = None
         if self.deployment is not None:
-            full_step_config = Step.model_validate(
-                json.loads(self.deployment.step_configurations)[self.name]
-            )
-        elif self.step_configuration is not None:
-            full_step_config = Step.model_validate_json(
-                self.step_configuration
-            )
-        else:
-            raise RuntimeError(
-                "Step run model creation has failed. Each step run entry "
-                "should either have a deployment_id or step_configuration."
+            step_configuration = json.loads(
+                self.deployment.step_configurations
             )
+            if self.name in step_configuration:
+                full_step_config = Step.model_validate(
+                    step_configuration[self.name]
+                )
+            elif not self.step_configuration:
+                raise ValueError(
+                    f"Unable to load the configuration for step `{self.name}` from the"
+                    f"database. To solve this please delete the pipeline run that this"
+                    f"step run belongs to. Pipeline Run ID: `{self.pipeline_run_id}`."
+                )
+
+        # the step configuration moved into the deployment - the following case is to ensure
+        # backwards compatibility
+        if full_step_config is None:
+            if self.step_configuration:
+                full_step_config = Step.model_validate_json(
+                    self.step_configuration
+                )
+            else:
+                raise RuntimeError(
+                    "Step run model creation has failed. Each step run entry "
+                    "should either have a deployment_id or step_configuration."
+                )
 
         body = StepRunResponseBody(
             user=self.user.to_model() if self.user else None,

zenml/zen_stores/schemas/trigger_schemas.py

@@ -22,13 +22,14 @@ from uuid import UUID
 from sqlalchemy import TEXT, Column
 from sqlmodel import Field, Relationship
 
-from zenml import TriggerExecutionResponseResources
+from zenml.config.schedule import Schedule
 from zenml.models import (
     Page,
     TriggerExecutionRequest,
     TriggerExecutionResponse,
     TriggerExecutionResponseBody,
     TriggerExecutionResponseMetadata,
+    TriggerExecutionResponseResources,
     TriggerRequest,
     TriggerResponse,
     TriggerResponseBody,
@@ -37,6 +38,7 @@ from zenml.models import (
     TriggerUpdate,
 )
 from zenml.utils.json_utils import pydantic_encoder
+from zenml.zen_stores.schemas.action_schemas import ActionSchema
 from zenml.zen_stores.schemas.base_schemas import BaseSchema, NamedSchema
 from zenml.zen_stores.schemas.event_source_schemas import EventSourceSchema
 from zenml.zen_stores.schemas.schema_utils import build_foreign_key_field
@@ -78,36 +80,33 @@ class TriggerSchema(NamedSchema, table=True):
         target=EventSourceSchema.__tablename__,
         source_column="event_source_id",
         target_column="id",
-        ondelete="CASCADE",  # TODO: this should be set null and the trigger should be deactivated
-        nullable=False,
+        # This won't happen because the SQL zen store prevents an event source
+        # from being deleted if it has associated triggers
+        ondelete="SET NULL",
+        nullable=True,
     )
-    event_source: "EventSourceSchema" = Relationship(back_populates="triggers")
-
-    executions: List["TriggerExecutionSchema"] = Relationship(
-        back_populates="trigger"
+    event_source: Optional["EventSourceSchema"] = Relationship(
+        back_populates="triggers"
     )
 
-    service_account_id: UUID = build_foreign_key_field(
+    action_id: UUID = build_foreign_key_field(
         source=__tablename__,
-        target=UserSchema.__tablename__,
-        source_column="service_account_id",
+        target=ActionSchema.__tablename__,
+        source_column="action_id",
         target_column="id",
-        ondelete="CASCADE",  # TODO: this should be set null and the trigger should be deactivated
+        # This won't happen because the SQL zen store prevents an action
+        # from being deleted if it has associated triggers
+        ondelete="CASCADE",
         nullable=False,
     )
-    service_account: UserSchema = Relationship(
-        back_populates="auth_triggers",
-        sa_relationship_kwargs={
-            "foreign_keys": "[TriggerSchema.service_account_id]"
-        },
+    action: "ActionSchema" = Relationship(back_populates="triggers")
+
+    executions: List["TriggerExecutionSchema"] = Relationship(
+        back_populates="trigger"
     )
-    auth_window: int
 
     event_filter: bytes
-
-    action: bytes
-    action_flavor: str  # <- "builtin"
-    action_subtype: str  # <- "PipelineRun"
+    schedule: Optional[bytes] = Field(nullable=True)
 
     description: str = Field(sa_column=Column(TEXT, nullable=True))
     is_active: bool = Field(nullable=False)
@@ -131,12 +130,6 @@ class TriggerSchema(NamedSchema, table=True):
                         trigger_update.event_filter, default=pydantic_encoder
                     ).encode("utf-8")
                 )
-            elif field == "action":
-                self.action = base64.b64encode(
-                    json.dumps(
-                        trigger_update.action, default=pydantic_encoder
-                    ).encode("utf-8")
-                )
             else:
                 setattr(self, field, value)
 
@@ -152,31 +145,21 @@ class TriggerSchema(NamedSchema, table=True):
 
         Returns:
             The converted schema.
-
-        Raises:
-            ValueError: If `auth_window` is not set.
         """
-        if request.auth_window is None:
-            raise ValueError("auth_window must be set")
         return cls(
             name=request.name,
             workspace_id=request.workspace,
             user_id=request.user,
-            action=base64.b64encode(
-                json.dumps(request.action, default=pydantic_encoder).encode(
-                    "utf-8"
-                ),
-            ),
-            action_flavor=request.action_flavor,
-            action_subtype=request.action_subtype,
+            action_id=request.action_id,
             event_source_id=request.event_source_id,
-            service_account_id=request.service_account_id,
-            auth_window=request.auth_window,
             event_filter=base64.b64encode(
                 json.dumps(
                     request.event_filter, default=pydantic_encoder
                 ).encode("utf-8")
             ),
+            schedule=base64.b64encode(request.schedule.json().encode("utf-8"))
+            if request.schedule
+            else None,
             description=request.description,
             is_active=True,  # Makes no sense for it to be created inactive
         )
@@ -205,9 +188,14 @@ class TriggerSchema(NamedSchema, table=True):
             user=self.user.to_model() if self.user else None,
             created=self.created,
             updated=self.updated,
-            action_flavor=self.action_flavor,
-            action_subtype=self.action_subtype,
-            event_source_flavor=self.event_source.flavor,
+            action_flavor=self.action.flavor,
+            action_subtype=self.action.plugin_subtype,
+            event_source_flavor=self.event_source.flavor
+            if self.event_source
+            else None,
+            event_source_subtype=self.event_source.plugin_subtype
+            if self.event_source
+            else None,
             is_active=self.is_active,
         )
         metadata = None
@@ -217,9 +205,12 @@ class TriggerSchema(NamedSchema, table=True):
                 event_filter=json.loads(
                     base64.b64decode(self.event_filter).decode()
                 ),
-                action=json.loads(base64.b64decode(self.action).decode()),
+                schedule=Schedule.parse_raw(
+                    base64.b64decode(self.schedule).decode()
+                )
+                if self.schedule
+                else None,
                 description=self.description,
-                auth_window=self.auth_window,
             )
         resources = None
         if include_resources:
@@ -228,13 +219,15 @@ class TriggerSchema(NamedSchema, table=True):
                 get_page_from_list(
                     items_list=self.executions,
                     response_model=TriggerExecutionResponse,
-                    include_resources=include_resources,
-                    include_metadata=include_metadata,
+                    include_resources=False,
+                    include_metadata=False,
                 ),
             )
             resources = TriggerResponseResources(
-                event_source=self.event_source.to_model(),
-                service_account=self.service_account.to_model(),
+                action=self.action.to_model(),
+                event_source=self.event_source.to_model()
+                if self.event_source
+                else None,
                 executions=executions,
             )
         return TriggerResponse(

zenml/zen_stores/schemas/user_schemas.py

@@ -37,6 +37,7 @@ from zenml.zen_stores.schemas.base_schemas import NamedSchema
 
 if TYPE_CHECKING:
     from zenml.zen_stores.schemas import (
+        ActionSchema,
         APIKeySchema,
         ArtifactVersionSchema,
         CodeRepositorySchema,
@@ -87,6 +88,13 @@ class UserSchema(NamedSchema, table=True):
         back_populates="user",
     )
     flavors: List["FlavorSchema"] = Relationship(back_populates="user")
+    actions: List["ActionSchema"] = Relationship(
+        back_populates="user",
+        sa_relationship_kwargs={
+            "cascade": "delete",
+            "primaryjoin": "UserSchema.id==ActionSchema.user_id",
+        },
+    )
     event_sources: List["EventSourceSchema"] = Relationship(
         back_populates="user"
     )
@@ -114,11 +122,11 @@ class UserSchema(NamedSchema, table=True):
             "primaryjoin": "UserSchema.id==TriggerSchema.user_id",
         },
     )
-    auth_triggers: List["TriggerSchema"] = Relationship(
+    auth_actions: List["ActionSchema"] = Relationship(
         back_populates="service_account",
         sa_relationship_kwargs={
             "cascade": "delete",
-            "primaryjoin": "UserSchema.id==TriggerSchema.service_account_id",
+            "primaryjoin": "UserSchema.id==ActionSchema.service_account_id",
         },
     )
     deployments: List["PipelineDeploymentSchema"] = Relationship(

zenml/zen_stores/schemas/workspace_schemas.py

@@ -29,6 +29,7 @@ from zenml.zen_stores.schemas.base_schemas import NamedSchema
 
 if TYPE_CHECKING:
     from zenml.zen_stores.schemas import (
+        ActionSchema,
         ArtifactVersionSchema,
         CodeRepositorySchema,
         EventSourceSchema,
@@ -104,6 +105,10 @@ class WorkspaceSchema(NamedSchema, table=True):
         back_populates="workspace",
         sa_relationship_kwargs={"cascade": "delete"},
     )
+    actions: List["ActionSchema"] = Relationship(
+        back_populates="workspace",
+        sa_relationship_kwargs={"cascade": "delete"},
+    )
     triggers: List["TriggerSchema"] = Relationship(
         back_populates="workspace",
         sa_relationship_kwargs={"cascade": "delete"},