zenml-nightly 0.55.0.dev20240124__py3-none-any.whl → 0.72.0.dev20250116__py3-none-any.whl

zenml/zen_stores/schemas/trigger_schemas.py

@@ -0,0 +1,316 @@
+#  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 Triggers."""
+
+import base64
+import json
+from datetime import datetime
+from typing import Any, List, Optional, cast
+from uuid import UUID
+
+from sqlalchemy import TEXT, Column
+from sqlmodel import Field, Relationship
+
+from zenml.config.schedule import Schedule
+from zenml.models import (
+    Page,
+    TriggerExecutionRequest,
+    TriggerExecutionResponse,
+    TriggerExecutionResponseBody,
+    TriggerExecutionResponseMetadata,
+    TriggerExecutionResponseResources,
+    TriggerRequest,
+    TriggerResponse,
+    TriggerResponseBody,
+    TriggerResponseMetadata,
+    TriggerResponseResources,
+    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
+from zenml.zen_stores.schemas.user_schemas import UserSchema
+from zenml.zen_stores.schemas.utils import get_page_from_list
+from zenml.zen_stores.schemas.workspace_schemas import WorkspaceSchema
+
+
+class TriggerSchema(NamedSchema, table=True):
+    """SQL Model for triggers."""
+
+    __tablename__ = "trigger"
+
+    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="triggers")
+
+    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="triggers",
+        sa_relationship_kwargs={"foreign_keys": "[TriggerSchema.user_id]"},
+    )
+
+    event_source_id: Optional[UUID] = build_foreign_key_field(
+        source=__tablename__,
+        target=EventSourceSchema.__tablename__,  # type: ignore[has-type]
+        source_column="event_source_id",
+        target_column="id",
+        # 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: Optional["EventSourceSchema"] = Relationship(
+        back_populates="triggers"
+    )
+
+    action_id: UUID = build_foreign_key_field(
+        source=__tablename__,
+        target=ActionSchema.__tablename__,
+        source_column="action_id",
+        target_column="id",
+        # This won't happen because the SQL zen store prevents an action
+        # from being deleted if it has associated triggers
+        ondelete="CASCADE",
+        nullable=False,
+    )
+    action: "ActionSchema" = Relationship(back_populates="triggers")
+
+    executions: List["TriggerExecutionSchema"] = Relationship(
+        back_populates="trigger"
+    )
+
+    event_filter: bytes
+    schedule: Optional[bytes] = Field(nullable=True)
+
+    description: str = Field(sa_column=Column(TEXT, nullable=True))
+    is_active: bool = Field(nullable=False)
+
+    def update(self, trigger_update: "TriggerUpdate") -> "TriggerSchema":
+        """Updates a trigger schema with a trigger update model.
+
+        Args:
+            trigger_update: `TriggerUpdate` to update the trigger with.
+
+        Returns:
+            The updated TriggerSchema.
+        """
+        for field, value in trigger_update.model_dump(
+            exclude_unset=True,
+            exclude_none=True,
+        ).items():
+            if field == "event_filter":
+                self.event_filter = base64.b64encode(
+                    json.dumps(
+                        trigger_update.event_filter, default=pydantic_encoder
+                    ).encode("utf-8")
+                )
+            else:
+                setattr(self, field, value)
+
+        self.updated = datetime.utcnow()
+        return self
+
+    @classmethod
+    def from_request(cls, request: "TriggerRequest") -> "TriggerSchema":
+        """Convert a `TriggerRequest` to a `TriggerSchema`.
+
+        Args:
+            request: The request model to convert.
+
+        Returns:
+            The converted schema.
+        """
+        return cls(
+            name=request.name,
+            workspace_id=request.workspace,
+            user_id=request.user,
+            action_id=request.action_id,
+            event_source_id=request.event_source_id,
+            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
+        )
+
+    def to_model(
+        self,
+        include_metadata: bool = False,
+        include_resources: bool = False,
+        **kwargs: Any,
+    ) -> "TriggerResponse":
+        """Converts the 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.
+        """
+        from zenml.models import TriggerExecutionResponse
+
+        body = TriggerResponseBody(
+            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.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
+        if include_metadata:
+            metadata = TriggerResponseMetadata(
+                workspace=self.workspace.to_model(),
+                event_filter=json.loads(
+                    base64.b64decode(self.event_filter).decode()
+                ),
+                schedule=Schedule.parse_raw(
+                    base64.b64decode(self.schedule).decode()
+                )
+                if self.schedule
+                else None,
+                description=self.description,
+            )
+        resources = None
+        if include_resources:
+            executions = cast(
+                Page[TriggerExecutionResponse],
+                get_page_from_list(
+                    items_list=self.executions,
+                    response_model=TriggerExecutionResponse,
+                    include_resources=False,
+                    include_metadata=False,
+                ),
+            )
+            resources = TriggerResponseResources(
+                action=self.action.to_model(),
+                event_source=self.event_source.to_model()
+                if self.event_source
+                else None,
+                executions=executions,
+            )
+        return TriggerResponse(
+            id=self.id,
+            name=self.name,
+            body=body,
+            metadata=metadata,
+            resources=resources,
+        )
+
+
+class TriggerExecutionSchema(BaseSchema, table=True):
+    """SQL Model for trigger executions."""
+
+    __tablename__ = "trigger_execution"
+
+    trigger_id: UUID = build_foreign_key_field(
+        source=__tablename__,
+        target=TriggerSchema.__tablename__,
+        source_column="trigger_id",
+        target_column="id",
+        ondelete="CASCADE",
+        nullable=False,
+    )
+    trigger: TriggerSchema = Relationship(back_populates="executions")
+
+    event_metadata: Optional[bytes] = None
+
+    @classmethod
+    def from_request(
+        cls, request: "TriggerExecutionRequest"
+    ) -> "TriggerExecutionSchema":
+        """Convert a `TriggerExecutionRequest` to a `TriggerExecutionSchema`.
+
+        Args:
+            request: The request model to convert.
+
+        Returns:
+            The converted schema.
+        """
+        return cls(
+            trigger_id=request.trigger,
+            event_metadata=base64.b64encode(
+                json.dumps(request.event_metadata).encode("utf-8")
+            ),
+        )
+
+    def to_model(
+        self,
+        include_metadata: bool = False,
+        include_resources: bool = False,
+        **kwargs: Any,
+    ) -> "TriggerExecutionResponse":
+        """Converts the schema to a model.
+
+        Args:
+            include_metadata: Whether the metadata will be filled.
+            include_resources: Whether the resources will be filled.
+            **kwargs: Keyword arguments to allow schema specific logic
+
+
+        Returns:
+            The converted model.
+        """
+        body = TriggerExecutionResponseBody(
+            created=self.created,
+            updated=self.updated,
+        )
+        metadata = None
+        if include_metadata:
+            metadata = TriggerExecutionResponseMetadata(
+                event_metadata=json.loads(
+                    base64.b64decode(self.event_metadata).decode()
+                )
+                if self.event_metadata
+                else {},
+            )
+        resources = None
+        if include_resources:
+            resources = TriggerExecutionResponseResources(
+                trigger=self.trigger.to_model(),
+            )
+
+        return TriggerExecutionResponse(
+            id=self.id, body=body, metadata=metadata, resources=resources
+        )

zenml/zen_stores/schemas/user_schemas.py

@@ -13,11 +13,12 @@
 #  permissions and limitations under the License.
 """SQLModel implementation of user tables."""
 
+import json
 from datetime import datetime
-from typing import TYPE_CHECKING, List, Optional
+from typing import TYPE_CHECKING, Any, List, Optional
 from uuid import UUID
 
-from sqlalchemy import TEXT, Column
+from sqlalchemy import TEXT, Column, UniqueConstraint
 from sqlmodel import Field, Relationship
 
 from zenml.models import (
@@ -36,13 +37,13 @@ from zenml.zen_stores.schemas.base_schemas import NamedSchema
 
 if TYPE_CHECKING:
     from zenml.zen_stores.schemas import (
+        ActionSchema,
         APIKeySchema,
         ArtifactVersionSchema,
         CodeRepositorySchema,
+        EventSourceSchema,
         FlavorSchema,
         ModelSchema,
-        ModelVersionArtifactSchema,
-        ModelVersionPipelineRunSchema,
         ModelVersionSchema,
         OAuthDeviceSchema,
         PipelineBuildSchema,
@@ -53,9 +54,11 @@ if TYPE_CHECKING:
         ScheduleSchema,
         SecretSchema,
         ServiceConnectorSchema,
+        ServiceSchema,
         StackComponentSchema,
         StackSchema,
         StepRunSchema,
+        TriggerSchema,
     )
 
 
@@ -63,6 +66,7 @@ class UserSchema(NamedSchema, table=True):
     """SQL Model for users."""
 
     __tablename__ = "user"
+    __table_args__ = (UniqueConstraint("name", "is_service_account"),)
 
     is_service_account: bool = Field(default=False)
     full_name: str
@@ -71,15 +75,26 @@ class UserSchema(NamedSchema, table=True):
     active: bool
     password: Optional[str] = Field(nullable=True)
     activation_token: Optional[str] = Field(nullable=True)
-    hub_token: Optional[str] = Field(nullable=True)
     email_opted_in: Optional[bool] = Field(nullable=True)
     external_user_id: Optional[UUID] = Field(nullable=True)
+    is_admin: bool = Field(default=False)
+    user_metadata: Optional[str] = Field(nullable=True)
 
     stacks: List["StackSchema"] = Relationship(back_populates="user")
     components: List["StackComponentSchema"] = Relationship(
         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"
+    )
     pipelines: List["PipelineSchema"] = Relationship(back_populates="user")
     schedules: List["ScheduleSchema"] = Relationship(
         back_populates="user",
@@ -97,12 +112,27 @@ class UserSchema(NamedSchema, table=True):
         back_populates="user",
         sa_relationship_kwargs={"cascade": "delete"},
     )
+    triggers: List["TriggerSchema"] = Relationship(
+        back_populates="user",
+        sa_relationship_kwargs={
+            "cascade": "delete",
+            "primaryjoin": "UserSchema.id==TriggerSchema.user_id",
+        },
+    )
+    auth_actions: List["ActionSchema"] = Relationship(
+        back_populates="service_account",
+        sa_relationship_kwargs={
+            "cascade": "delete",
+            "primaryjoin": "UserSchema.id==ActionSchema.service_account_id",
+        },
+    )
     deployments: List["PipelineDeploymentSchema"] = Relationship(
         back_populates="user",
     )
     code_repositories: List["CodeRepositorySchema"] = Relationship(
         back_populates="user",
     )
+    services: List["ServiceSchema"] = Relationship(back_populates="user")
     service_connectors: List["ServiceConnectorSchema"] = Relationship(
         back_populates="user",
     )
@@ -112,12 +142,6 @@ class UserSchema(NamedSchema, table=True):
     model_versions: List["ModelVersionSchema"] = Relationship(
         back_populates="user",
     )
-    model_versions_artifacts_links: List[
-        "ModelVersionArtifactSchema"
-    ] = Relationship(back_populates="user")
-    model_versions_pipeline_runs_links: List[
-        "ModelVersionPipelineRunSchema"
-    ] = Relationship(back_populates="user")
     auth_devices: List["OAuthDeviceSchema"] = Relationship(
         back_populates="user",
         sa_relationship_kwargs={"cascade": "delete"},
@@ -147,6 +171,10 @@ class UserSchema(NamedSchema, table=True):
             email_opted_in=model.email_opted_in,
             email=model.email,
             is_service_account=False,
+            is_admin=model.is_admin,
+            user_metadata=json.dumps(model.user_metadata)
+            if model.user_metadata
+            else None,
         )
 
     @classmethod
@@ -169,6 +197,7 @@ class UserSchema(NamedSchema, table=True):
             is_service_account=True,
             email_opted_in=False,
             full_name="",
+            is_admin=False,
         )
 
     def update_user(self, user_update: UserUpdate) -> "UserSchema":
@@ -180,13 +209,19 @@ class UserSchema(NamedSchema, table=True):
         Returns:
             The updated `UserSchema`.
         """
-        for field, value in user_update.dict(exclude_unset=True).items():
+        for field, value in user_update.model_dump(exclude_unset=True).items():
+            if field == "old_password":
+                continue
+
             if field == "password":
                 setattr(self, field, user_update.create_hashed_password())
             elif field == "activation_token":
                 setattr(
                     self, field, user_update.create_hashed_activation_token()
                 )
+            elif field == "user_metadata":
+                if value is not None:
+                    self.user_metadata = json.dumps(value)
             else:
                 setattr(self, field, value)
 
@@ -205,7 +240,7 @@ class UserSchema(NamedSchema, table=True):
         Returns:
             The updated `UserSchema`.
         """
-        for field, value in service_account_update.dict(
+        for field, value in service_account_update.model_dump(
             exclude_none=True
         ).items():
             setattr(self, field, value)
@@ -214,13 +249,18 @@ class UserSchema(NamedSchema, table=True):
         return self
 
     def to_model(
-        self, hydrate: bool = False, include_private: bool = False
+        self,
+        include_metadata: bool = False,
+        include_resources: bool = False,
+        include_private: bool = False,
+        **kwargs: Any,
     ) -> UserResponse:
         """Convert a `UserSchema` to a `UserResponse`.
 
         Args:
-            hydrate: bool to decide whether to return a hydrated version of the
-                model.
+            include_metadata: Whether the metadata will be filled.
+            include_resources: Whether the resources will be filled.
+            **kwargs: Keyword arguments to allow schema specific logic
             include_private: Whether to include the user private information
                              this is to limit the amount of data one can get
                              about other users
@@ -229,11 +269,13 @@ class UserSchema(NamedSchema, table=True):
             The converted `UserResponse`.
         """
         metadata = None
-        if hydrate:
+        if include_metadata:
             metadata = UserResponseMetadata(
                 email=self.email if include_private else None,
-                hub_token=self.hub_token if include_private else None,
                 external_user_id=self.external_user_id,
+                user_metadata=json.loads(self.user_metadata)
+                if self.user_metadata
+                else {},
             )
 
         return UserResponse(
@@ -246,24 +288,25 @@ class UserSchema(NamedSchema, table=True):
                 is_service_account=self.is_service_account,
                 created=self.created,
                 updated=self.updated,
+                is_admin=self.is_admin,
             ),
             metadata=metadata,
         )
 
     def to_service_account_model(
-        self, hydrate: bool = False
+        self, include_metadata: bool = False, include_resources: bool = False
     ) -> ServiceAccountResponse:
         """Convert a `UserSchema` to a `ServiceAccountResponse`.
 
         Args:
-            hydrate: bool to decide whether to return a hydrated version of the
-                model.
+             include_metadata: Whether the metadata will be filled.
+             include_resources: Whether the resources will be filled.
 
         Returns:
-            The converted `ServiceAccountResponse`.
+             The converted `ServiceAccountResponse`.
         """
         metadata = None
-        if hydrate:
+        if include_metadata:
             metadata = ServiceAccountResponseMetadata(
                 description=self.description or "",
             )

zenml/zen_stores/schemas/utils.py

@@ -0,0 +1,112 @@
+#  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.
+"""Utils for schemas."""
+
+import json
+import math
+from typing import Dict, List, Type, TypeVar
+
+from sqlmodel import Relationship
+
+from zenml.metadata.metadata_types import MetadataType
+from zenml.models import BaseResponse, Page, RunMetadataEntry
+from zenml.zen_stores.schemas.base_schemas import BaseSchema
+
+S = TypeVar("S", bound=BaseSchema)
+
+
+def get_page_from_list(
+    items_list: List[S],
+    response_model: Type[BaseResponse],  # type: ignore[type-arg]
+    size: int = 5,
+    page: int = 1,
+    include_resources: bool = False,
+    include_metadata: bool = False,
+) -> Page[BaseResponse]:  # type: ignore[type-arg]
+    """Converts list of schemas into page of response models.
+
+    Args:
+        items_list: List of schemas
+        response_model: Response model
+        size: Page size
+        page: Page number
+        include_metadata: Whether metadata should be included in response models
+        include_resources: Whether resources should be included in response models
+
+    Returns:
+        A page of list items.
+    """
+    total = len(items_list)
+    if total == 0:
+        total_pages = 1
+    else:
+        total_pages = math.ceil(total / size)
+
+    start = (page - 1) * size
+    end = start + size
+
+    page_items = [
+        item.to_model(
+            include_metadata=include_metadata,
+            include_resources=include_resources,
+        )
+        for item in items_list[start:end]
+    ]
+    return Page[response_model](  # type: ignore[valid-type]
+        index=page,
+        max_size=size,
+        total_pages=total_pages,
+        total=total,
+        items=page_items,
+    )
+
+
+class RunMetadataInterface:
+    """The interface for entities with run metadata."""
+
+    run_metadata = Relationship()
+
+    def fetch_metadata_collection(self) -> Dict[str, List[RunMetadataEntry]]:
+        """Fetches all the metadata entries related to the entity.
+
+        Returns:
+            A dictionary, where the key is the key of the metadata entry
+                and the values represent the list of entries with this key.
+        """
+        metadata_collection: Dict[str, List[RunMetadataEntry]] = {}
+
+        for rm in self.run_metadata:
+            if rm.key not in metadata_collection:
+                metadata_collection[rm.key] = []
+            metadata_collection[rm.key].append(
+                RunMetadataEntry(
+                    value=json.loads(rm.value),
+                    created=rm.created,
+                )
+            )
+
+        return metadata_collection
+
+    def fetch_metadata(self) -> Dict[str, MetadataType]:
+        """Fetches the latest metadata entry related to the entity.
+
+        Returns:
+            A dictionary, where the key is the key of the metadata entry
+                and the values represent the latest entry with this key.
+        """
+        metadata_collection = self.fetch_metadata_collection()
+        return {
+            k: sorted(v, key=lambda x: x.created, reverse=True)[0].value
+            for k, v in metadata_collection.items()
+        }