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

zenml/models/v2/core/trigger.py

@@ -13,15 +13,14 @@
 #  permissions and limitations under the License.
 """Collection of all models concerning triggers."""
 
-import copy
-from typing import TYPE_CHECKING, Any, Dict, Optional, Union
+from typing import TYPE_CHECKING, Any, ClassVar, Dict, List, Optional, Union
 from uuid import UUID
 
-from pydantic import BaseModel, Field
+from pydantic import Field, model_validator
 
+from zenml.config.schedule import Schedule
 from zenml.constants import STR_FIELD_MAX_LENGTH
-from zenml.enums import PluginSubType
-from zenml.models.v2.base.base import BaseZenModel
+from zenml.models.v2.base.base import BaseUpdate
 from zenml.models.v2.base.page import Page
 from zenml.models.v2.base.scoped import (
     WorkspaceScopedFilter,
@@ -32,138 +31,129 @@ from zenml.models.v2.base.scoped import (
     WorkspaceScopedResponseResources,
 )
 from zenml.models.v2.core.trigger_execution import TriggerExecutionResponse
-from zenml.models.v2.core.user import UserResponse
 
 if TYPE_CHECKING:
+    from sqlalchemy.sql.elements import ColumnElement
+
+    from zenml.models.v2.core.action import (
+        ActionResponse,
+    )
     from zenml.models.v2.core.event_source import EventSourceResponse
 
-# ------------------ Base Model ------------------
+
+# ------------------ Request Model ------------------
 
 
-class TriggerBase(BaseModel):
-    """Base model for triggers."""
+class TriggerRequest(WorkspaceScopedRequest):
+    """Model for creating a new trigger."""
 
     name: str = Field(
-        title="The name of the Trigger.", max_length=STR_FIELD_MAX_LENGTH
+        title="The name of the trigger.", max_length=STR_FIELD_MAX_LENGTH
     )
     description: str = Field(
         default="",
         title="The description of the trigger",
         max_length=STR_FIELD_MAX_LENGTH,
     )
-    event_source_id: UUID = Field(
-        title="The event source that activates this trigger.",
-    )
-    event_filter: Dict[str, Any] = Field(
-        title="Filter applied to events that activate this trigger.",
-    )
-
-    action: Dict[str, Any] = Field(
-        title="The configuration for the action that is executed by this "
-        "trigger.",
-    )
-    action_flavor: str = Field(
-        title="The flavor of the action that is executed by this trigger.",
-        max_length=STR_FIELD_MAX_LENGTH,
+    action_id: UUID = Field(
+        title="The action that is executed by this trigger.",
     )
-    action_subtype: PluginSubType = Field(
-        title="The subtype of the action that is executed by this trigger.",
+    schedule: Optional[Schedule] = Field(
+        default=None,
+        title="The schedule for the trigger. Either a schedule or an event "
+        "source is required.",
     )
-    service_account_id: UUID = Field(
-        title="The service account that is used to execute the action.",
+    event_source_id: Optional[UUID] = Field(
+        default=None,
+        title="The event source that activates this trigger. Either a schedule "
+        "or an event source is required.",
     )
-    auth_window: Optional[int] = Field(
+    event_filter: Optional[Dict[str, Any]] = Field(
         default=None,
-        title="The time window in minutes for which the service account is "
-        "authorized to execute the action. Set this to 0 to authorize the "
-        "service account indefinitely (not recommended). If not set, a "
-        "default value defined for each individual action type is used.",
+        title="Filter applied to events that activate this trigger. Only "
+        "set if the trigger is activated by an event source.",
     )
 
+    @model_validator(mode="after")
+    def _validate_schedule_or_event_source(self) -> "TriggerRequest":
+        """Validate that either a schedule or an event source is provided.
 
-# ------------------ Request Model ------------------
-class TriggerRequest(TriggerBase, WorkspaceScopedRequest):
-    """Model for creating a new Trigger."""
+        Returns:
+            The validated request.
+
+        Raises:
+            ValueError: If neither a schedule nor an event source is provided,
+                or if both are provided.
+        """
+        if not self.schedule and not self.event_source_id:
+            raise ValueError(
+                "Either a schedule or an event source is required."
+            )
+
+        if self.schedule and self.event_source_id:
+            raise ValueError("Only a schedule or an event source is allowed.")
+
+        return self
 
 
 # ------------------ Update Model ------------------
 
 
-class TriggerUpdate(BaseZenModel):
+class TriggerUpdate(BaseUpdate):
     """Update model for triggers."""
 
     name: Optional[str] = Field(
         default=None,
-        title="The new name for the Trigger.",
+        title="The new name for the trigger.",
         max_length=STR_FIELD_MAX_LENGTH,
     )
     description: Optional[str] = Field(
         default=None,
-        title="The new description for the trigger",
+        title="The new description for the trigger.",
         max_length=STR_FIELD_MAX_LENGTH,
     )
     event_filter: Optional[Dict[str, Any]] = Field(
         default=None,
-        title="New filter applied to events that activate this trigger.",
+        title="New filter applied to events that activate this trigger. Only "
+        "valid if the trigger is already configured to be activated by an "
+        "event source.",
     )
-    action: Optional[Dict[str, Any]] = Field(
+    schedule: Optional[Schedule] = Field(
         default=None,
-        title="The new configuration for the action that is executed by this "
-        "trigger.",
+        title="The updated schedule for the trigger. Only valid if the trigger "
+        "is already configured to be activated by a schedule.",
     )
     is_active: Optional[bool] = Field(
         default=None,
         title="The new status of the trigger.",
     )
-    service_account_id: Optional[UUID] = Field(
-        default=None,
-        title="The service account that is used to execute the action.",
-    )
-    auth_window: Optional[int] = Field(
-        default=None,
-        title="The time window in minutes for which the service account is "
-        "authorized to execute the action. Set this to 0 to authorize the "
-        "service account indefinitely (not recommended). If not set, a "
-        "default value defined for each individual action type is used.",
-    )
-
-    @classmethod
-    def from_response(cls, response: "TriggerResponse") -> "TriggerUpdate":
-        """Create an update model from a response model.
-
-        Args:
-            response: The response model to create the update model from.
-
-        Returns:
-            The update model.
-        """
-        return TriggerUpdate(
-            name=response.name,
-            description=response.description,
-            action=copy.deepcopy(response.action),
-            event_filter=copy.deepcopy(response.event_filter),
-            is_active=response.is_active,
-            service_account_id=response.get_resources().service_account.id,
-        )
 
 
 # ------------------ Response Model ------------------
 
 
 class TriggerResponseBody(WorkspaceScopedResponseBody):
-    """ResponseBody for triggers."""
+    """Response body for triggers."""
 
-    event_source_flavor: str = Field(
-        title="The flavor of the event source that activates this trigger.",
-        max_length=STR_FIELD_MAX_LENGTH,
-    )
     action_flavor: str = Field(
         title="The flavor of the action that is executed by this trigger.",
         max_length=STR_FIELD_MAX_LENGTH,
     )
-    action_subtype: PluginSubType = Field(
+    action_subtype: str = Field(
         title="The subtype of the action that is executed by this trigger.",
     )
+    event_source_flavor: Optional[str] = Field(
+        default=None,
+        title="The flavor of the event source that activates this trigger. Not "
+        "set if the trigger is activated by a schedule.",
+        max_length=STR_FIELD_MAX_LENGTH,
+    )
+    event_source_subtype: Optional[str] = Field(
+        default=None,
+        title="The subtype of the event source that activates this trigger. "
+        "Not set if the trigger is activated by a schedule.",
+        max_length=STR_FIELD_MAX_LENGTH,
+    )
     is_active: bool = Field(
         title="Whether the trigger is active.",
     )
@@ -172,33 +162,33 @@ class TriggerResponseBody(WorkspaceScopedResponseBody):
 class TriggerResponseMetadata(WorkspaceScopedResponseMetadata):
     """Response metadata for triggers."""
 
-    event_filter: Dict[str, Any] = Field(
-        title="The event that activates this trigger.",
-    )
-    action: Dict[str, Any] = Field(
-        title="The action that is executed by this trigger.",
-    )
     description: str = Field(
         default="",
-        title="The description of the trigger",
+        title="The description of the trigger.",
         max_length=STR_FIELD_MAX_LENGTH,
     )
-    auth_window: int = Field(
-        title="The time window in minutes for which the service account is "
-        "authorized to execute the action. Set this to 0 to authorize the "
-        "service account indefinitely (not recommended). If not set, a "
-        "default value defined for each individual action type is used.",
+    event_filter: Optional[Dict[str, Any]] = Field(
+        default=None,
+        title="The event that activates this trigger. Not set if the trigger "
+        "is activated by a schedule.",
+    )
+    schedule: Optional[Schedule] = Field(
+        default=None,
+        title="The schedule that activates this trigger. Not set if the "
+        "trigger is activated by an event source.",
     )
 
 
 class TriggerResponseResources(WorkspaceScopedResponseResources):
     """Class for all resource models associated with the trigger entity."""
 
-    event_source: "EventSourceResponse" = Field(
-        title="The event source that activates this trigger.",
+    action: "ActionResponse" = Field(
+        title="The action that is executed by this trigger.",
     )
-    service_account: UserResponse = Field(
-        title="The service account that is used to execute the action.",
+    event_source: Optional["EventSourceResponse"] = Field(
+        default=None,
+        title="The event source that activates this trigger. Not set if the "
+        "trigger is activated by a schedule.",
     )
     executions: Page[TriggerExecutionResponse] = Field(
         title="The executions of this trigger.",
@@ -213,7 +203,7 @@ class TriggerResponse(
     """Response model for models."""
 
     name: str = Field(
-        title="The name of the model",
+        title="The name of the trigger",
         max_length=STR_FIELD_MAX_LENGTH,
     )
 
@@ -227,15 +217,6 @@ class TriggerResponse(
 
         return Client().zen_store.get_trigger(self.id)
 
-    @property
-    def event_source_flavor(self) -> str:
-        """The `event_source_flavor` property.
-
-        Returns:
-            the value of the property.
-        """
-        return self.get_body().event_source_flavor
-
     @property
     def action_flavor(self) -> str:
         """The `action_flavor` property.
@@ -246,7 +227,7 @@ class TriggerResponse(
         return self.get_body().action_flavor
 
     @property
-    def action_subtype(self) -> PluginSubType:
+    def action_subtype(self) -> str:
         """The `action_subtype` property.
 
         Returns:
@@ -255,48 +236,40 @@ class TriggerResponse(
         return self.get_body().action_subtype
 
     @property
-    def is_active(self) -> bool:
-        """The `is_active` property.
+    def event_source_flavor(self) -> Optional[str]:
+        """The `event_source_flavor` property.
 
         Returns:
             the value of the property.
         """
-        return self.get_body().is_active
+        return self.get_body().event_source_flavor
 
     @property
-    def event_filter(self) -> Dict[str, Any]:
-        """The `event_filter` property.
+    def event_source_subtype(self) -> Optional[str]:
+        """The `event_source_subtype` property.
 
         Returns:
             the value of the property.
         """
-        return self.get_metadata().event_filter
+        return self.get_body().event_source_subtype
 
     @property
-    def action(self) -> Dict[str, Any]:
-        """The `action` property.
+    def is_active(self) -> bool:
+        """The `is_active` property.
 
         Returns:
             the value of the property.
         """
-        return self.get_metadata().action
-
-    def set_action(self, action: Dict[str, Any]) -> None:
-        """Set the `action` property.
-
-        Args:
-            action: The value to set.
-        """
-        self.get_metadata().action = action
+        return self.get_body().is_active
 
     @property
-    def event_source(self) -> "EventSourceResponse":
-        """The `event_source` property.
+    def event_filter(self) -> Optional[Dict[str, Any]]:
+        """The `event_filter` property.
 
         Returns:
             the value of the property.
         """
-        return self.get_resources().event_source
+        return self.get_metadata().event_filter
 
     @property
     def description(self) -> str:
@@ -308,37 +281,59 @@ class TriggerResponse(
         return self.get_metadata().description
 
     @property
-    def service_account(self) -> UserResponse:
-        """The `service_account` property.
+    def action(self) -> "ActionResponse":
+        """The `action` property.
 
         Returns:
             the value of the property.
         """
-        return self.get_resources().service_account
+        return self.get_resources().action
 
     @property
-    def auth_window(self) -> int:
-        """The `auth_window` property.
+    def event_source(self) -> Optional["EventSourceResponse"]:
+        """The `event_source` property.
 
         Returns:
             the value of the property.
         """
-        return self.get_metadata().auth_window
+        return self.get_resources().event_source
+
+    @property
+    def executions(self) -> Page[TriggerExecutionResponse]:
+        """The `event_source` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_resources().executions
 
 
 # ------------------ Filter Model ------------------
 
 
 class TriggerFilter(WorkspaceScopedFilter):
-    """Model to enable advanced filtering of all TriggerModels."""
+    """Model to enable advanced filtering of all triggers."""
+
+    FILTER_EXCLUDE_FIELDS: ClassVar[List[str]] = [
+        *WorkspaceScopedFilter.FILTER_EXCLUDE_FIELDS,
+        "action_flavor",
+        "action_subtype",
+        "event_source_flavor",
+        "event_source_subtype",
+    ]
 
     name: Optional[str] = Field(
         default=None,
-        description="Name of the trigger",
+        description="Name of the trigger.",
     )
     event_source_id: Optional[Union[UUID, str]] = Field(
         default=None,
-        description="By the event source this trigger is attached to.",
+        description="The event source this trigger is attached to.",
+        union_mode="left_to_right",
+    )
+    action_id: Optional[Union[UUID, str]] = Field(
+        default=None,
+        description="The action this trigger is attached to.",
         union_mode="left_to_right",
     )
     is_active: Optional[bool] = Field(
@@ -353,13 +348,59 @@ class TriggerFilter(WorkspaceScopedFilter):
         default=None,
         title="The subtype of the action that is executed by this trigger.",
     )
-    # TODO: Ignore these in normal filter and handle in sqlzenstore
-    resource_id: Optional[Union[UUID, str]] = Field(
+    event_source_flavor: Optional[str] = Field(
         default=None,
-        description="By the resource this trigger references.",
-        union_mode="left_to_right",
+        title="The flavor of the event source that activates this trigger.",
     )
-    resource_type: Optional[str] = Field(
+    event_source_subtype: Optional[str] = Field(
         default=None,
-        description="By the resource type this trigger references.",
+        title="The subtype of the event source that activates this trigger.",
     )
+
+    def get_custom_filters(
+        self,
+    ) -> List["ColumnElement[bool]"]:
+        """Get custom filters.
+
+        Returns:
+            A list of custom filters.
+        """
+        from sqlmodel import and_
+
+        from zenml.zen_stores.schemas import (
+            ActionSchema,
+            EventSourceSchema,
+            TriggerSchema,
+        )
+
+        custom_filters = super().get_custom_filters()
+
+        if self.event_source_flavor:
+            event_source_flavor_filter = and_(
+                EventSourceSchema.id == TriggerSchema.event_source_id,
+                EventSourceSchema.flavor == self.event_source_flavor,
+            )
+            custom_filters.append(event_source_flavor_filter)
+
+        if self.event_source_subtype:
+            event_source_subtype_filter = and_(
+                EventSourceSchema.id == TriggerSchema.event_source_id,
+                EventSourceSchema.plugin_subtype == self.event_source_subtype,
+            )
+            custom_filters.append(event_source_subtype_filter)
+
+        if self.action_flavor:
+            action_flavor_filter = and_(
+                ActionSchema.id == TriggerSchema.action_id,
+                ActionSchema.flavor == self.action_flavor,
+            )
+            custom_filters.append(action_flavor_filter)
+
+        if self.action_subtype:
+            action_subtype_filter = and_(
+                ActionSchema.id == TriggerSchema.action_id,
+                ActionSchema.plugin_subtype == self.action_subtype,
+            )
+            custom_filters.append(action_subtype_filter)
+
+        return custom_filters

zenml/models/v2/misc/full_stack.py

@@ -0,0 +1,97 @@
+#  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.
+"""Models representing full stack requests."""
+
+from typing import Any, Dict, List, Optional, Union
+from uuid import UUID
+
+from pydantic import BaseModel, Field, model_validator
+
+from zenml.constants import STR_FIELD_MAX_LENGTH
+from zenml.enums import StackComponentType
+from zenml.models.v2.base.base import BaseRequest
+
+
+class ServiceConnectorInfo(BaseModel):
+    """Information about the service connector when creating a full stack."""
+
+    type: str
+    auth_method: str
+    configuration: Dict[str, Any] = {}
+
+
+class ComponentInfo(BaseModel):
+    """Information about each stack components when creating a full stack."""
+
+    flavor: str
+    service_connector_index: Optional[int] = Field(
+        default=None,
+        title="The id of the service connector from the list "
+        "`service_connectors`.",
+        description="The id of the service connector from the list "
+        "`service_connectors` from `FullStackRequest`.",
+    )
+    service_connector_resource_id: Optional[str] = None
+    configuration: Dict[str, Any] = {}
+
+
+class FullStackRequest(BaseRequest):
+    """Request model for a full-stack."""
+
+    user: Optional[UUID] = None
+    workspace: Optional[UUID] = None
+
+    name: str = Field(
+        title="The name of the stack.", max_length=STR_FIELD_MAX_LENGTH
+    )
+    description: Optional[str] = Field(
+        default="",
+        title="The description of the stack",
+        max_length=STR_FIELD_MAX_LENGTH,
+    )
+    labels: Optional[Dict[str, Any]] = Field(
+        default=None,
+        title="The stack labels.",
+    )
+    service_connectors: List[Union[UUID, ServiceConnectorInfo]] = Field(
+        default=[],
+        title="The service connectors dictionary for the full stack "
+        "registration.",
+        description="The UUID of an already existing service connector or "
+        "request information to create a service connector from "
+        "scratch.",
+    )
+    components: Dict[StackComponentType, Union[UUID, ComponentInfo]] = Field(
+        title="The mapping for the components of the full stack registration.",
+        description="The mapping from component types to either UUIDs of "
+        "existing components or request information for brand new "
+        "components.",
+    )
+
+    @model_validator(mode="after")
+    def _validate_indexes_in_components(self) -> "FullStackRequest":
+        for component in self.components.values():
+            if isinstance(component, ComponentInfo):
+                if component.service_connector_index is not None:
+                    if (
+                        component.service_connector_index < 0
+                        or component.service_connector_index
+                        >= len(self.service_connectors)
+                    ):
+                        raise ValueError(
+                            f"Service connector index {component.service_connector_index} "
+                            "is out of range. Please provide a valid index referring to "
+                            "the position in the list of service connectors."
+                        )
+        return self

zenml/models/v2/misc/stack_deployment.py

@@ -0,0 +1,86 @@
+#  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.
+"""Models related to cloud stack deployments."""
+
+from typing import Dict, List, Optional
+
+from pydantic import BaseModel, Field
+
+from zenml.enums import StackDeploymentProvider
+from zenml.models.v2.core.service_connector import ServiceConnectorResponse
+from zenml.models.v2.core.stack import StackResponse
+
+
+class StackDeploymentInfo(BaseModel):
+    """Information about a stack deployment."""
+
+    provider: StackDeploymentProvider = Field(
+        title="The provider of the stack deployment."
+    )
+    description: str = Field(
+        title="The description of the stack deployment.",
+        description="The description of the stack deployment.",
+    )
+    instructions: str = Field(
+        title="The instructions for deploying the stack.",
+        description="The instructions for deploying the stack.",
+    )
+    post_deploy_instructions: str = Field(
+        title="The instructions for post-deployment.",
+        description="The instructions for post-deployment.",
+    )
+    integrations: List[str] = Field(
+        title="ZenML integrations required for the stack.",
+        description="The list of ZenML integrations that need to be installed "
+        "for the stack to be usable.",
+    )
+    permissions: Dict[str, List[str]] = Field(
+        title="The permissions granted to ZenML to access the cloud resources.",
+        description="The permissions granted to ZenML to access the cloud "
+        "resources, as a dictionary grouping permissions by resource.",
+    )
+    locations: Dict[str, str] = Field(
+        title="The locations where the stack can be deployed.",
+        description="The locations where the stack can be deployed, as a "
+        "dictionary mapping location names to descriptions.",
+    )
+
+
+class StackDeploymentConfig(BaseModel):
+    """Configuration about a stack deployment."""
+
+    deployment_url: str = Field(
+        title="The cloud provider console URL where the stack will be deployed.",
+    )
+    deployment_url_text: str = Field(
+        title="A textual description for the cloud provider console URL.",
+    )
+    configuration: Optional[str] = Field(
+        title="Configuration for the stack deployment that the user must "
+        "manually configure into the cloud provider console.",
+    )
+
+
+class DeployedStack(BaseModel):
+    """Information about a deployed stack."""
+
+    stack: StackResponse = Field(
+        title="The stack that was deployed.",
+        description="The stack that was deployed.",
+    )
+    service_connector: Optional[ServiceConnectorResponse] = Field(
+        default=None,
+        title="The service connector for the deployed stack.",
+        description="The service connector for the deployed stack.",
+    )