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

zenml/models/v2/base/page.py

@@ -15,8 +15,7 @@
 
 from typing import Generator, Generic, List, TypeVar
 
-from pydantic import SecretStr
-from pydantic.generics import GenericModel
+from pydantic import BaseModel
 from pydantic.types import NonNegativeInt, PositiveInt
 
 from zenml.models.v2.base.base import BaseResponse
@@ -25,7 +24,7 @@ from zenml.models.v2.base.filter import BaseFilter
 B = TypeVar("B", bound=BaseResponse)  # type: ignore[type-arg]
 
 
-class Page(GenericModel, Generic[B]):
+class Page(BaseModel, Generic[B]):
     """Return Model for List Models to accommodate pagination."""
 
     index: PositiveInt
@@ -91,12 +90,3 @@ class Page(GenericModel, Generic[B]):
             Whether the item is in the page.
         """
         return item in self.items
-
-    class Config:
-        """Pydantic configuration class."""
-
-        # This is needed to allow the REST API server to unpack SecretStr
-        # values correctly before sending them to the client.
-        json_encoders = {
-            SecretStr: lambda v: v.get_secret_value() if v else None
-        }

zenml/models/v2/base/scoped.py

@@ -28,18 +28,18 @@ from typing import (
 from uuid import UUID
 
 from pydantic import Field
-from sqlmodel import col
 
 from zenml.models.v2.base.base import (
+    BaseDatedResponseBody,
+    BaseIdentifiedResponse,
     BaseRequest,
-    BaseResponse,
-    BaseResponseBody,
     BaseResponseMetadata,
+    BaseResponseResources,
 )
 from zenml.models.v2.base.filter import AnyQuery, BaseFilter
 
 if TYPE_CHECKING:
-    from sqlalchemy.sql.elements import BinaryExpression, BooleanClauseList
+    from sqlalchemy.sql.elements import ColumnElement
 
     from zenml.models.v2.core.user import UserResponse
     from zenml.models.v2.core.workspace import WorkspaceResponse
@@ -94,11 +94,11 @@ class WorkspaceScopedRequest(UserScopedRequest):
 
 
 # User-scoped models
-class UserScopedResponseBody(BaseResponseBody):
+class UserScopedResponseBody(BaseDatedResponseBody):
     """Base user-owned body."""
 
     user: Optional["UserResponse"] = Field(
-        title="The user who created this resource."
+        title="The user who created this resource.", default=None
     )
 
 
@@ -106,12 +106,18 @@ class UserScopedResponseMetadata(BaseResponseMetadata):
     """Base user-owned metadata."""
 
 
+class UserScopedResponseResources(BaseResponseResources):
+    """Base class for all resource models associated with the user."""
+
+
 UserBody = TypeVar("UserBody", bound=UserScopedResponseBody)
 UserMetadata = TypeVar("UserMetadata", bound=UserScopedResponseMetadata)
+UserResources = TypeVar("UserResources", bound=UserScopedResponseResources)
 
 
 class UserScopedResponse(
-    BaseResponse[UserBody, UserMetadata], Generic[UserBody, UserMetadata]
+    BaseIdentifiedResponse[UserBody, UserMetadata, UserResources],
+    Generic[UserBody, UserMetadata, UserResources],
 ):
     """Base user-owned model.
 
@@ -146,16 +152,32 @@ class UserScopedFilter(BaseFilter):
 
     FILTER_EXCLUDE_FIELDS: ClassVar[List[str]] = [
         *BaseFilter.FILTER_EXCLUDE_FIELDS,
+        "user",
         "scope_user",
     ]
     CLI_EXCLUDE_FIELDS: ClassVar[List[str]] = [
         *BaseFilter.CLI_EXCLUDE_FIELDS,
+        "user_id",
         "scope_user",
     ]
+    CUSTOM_SORTING_OPTIONS: ClassVar[List[str]] = [
+        *BaseFilter.CUSTOM_SORTING_OPTIONS,
+        "user",
+    ]
+
     scope_user: Optional[UUID] = Field(
         default=None,
         description="The user to scope this query to.",
     )
+    user_id: Optional[Union[UUID, str]] = Field(
+        default=None,
+        description="UUID of the user that created the entity.",
+        union_mode="left_to_right",
+    )
+    user: Optional[Union[UUID, str]] = Field(
+        default=None,
+        description="Name/ID of the user that created the entity.",
+    )
 
     def set_scope_user(self, user_id: UUID) -> None:
         """Set the user that is performing the filtering to scope the response.
@@ -165,6 +187,75 @@ class UserScopedFilter(BaseFilter):
         """
         self.scope_user = user_id
 
+    def get_custom_filters(
+        self, table: Type["AnySchema"]
+    ) -> List["ColumnElement[bool]"]:
+        """Get custom filters.
+
+        Args:
+            table: The query table.
+
+        Returns:
+            A list of custom filters.
+        """
+        custom_filters = super().get_custom_filters(table)
+
+        from sqlmodel import and_
+
+        from zenml.zen_stores.schemas import UserSchema
+
+        if self.user:
+            user_filter = and_(
+                getattr(table, "user_id") == UserSchema.id,
+                self.generate_name_or_id_query_conditions(
+                    value=self.user,
+                    table=UserSchema,
+                    additional_columns=["full_name"],
+                ),
+            )
+            custom_filters.append(user_filter)
+
+        return custom_filters
+
+    def apply_sorting(
+        self,
+        query: AnyQuery,
+        table: Type["AnySchema"],
+    ) -> AnyQuery:
+        """Apply sorting to the query.
+
+        Args:
+            query: The query to which to apply the sorting.
+            table: The query table.
+
+        Returns:
+            The query with sorting applied.
+        """
+        from sqlmodel import asc, desc
+
+        from zenml.enums import SorterOps
+        from zenml.zen_stores.schemas import UserSchema
+
+        sort_by, operand = self.sorting_params
+
+        if sort_by == "user":
+            column = UserSchema.name
+
+            query = query.join(
+                UserSchema, getattr(table, "user_id") == UserSchema.id
+            )
+
+            query = query.add_columns(UserSchema.name)
+
+            if operand == SorterOps.ASCENDING:
+                query = query.order_by(asc(column))
+            else:
+                query = query.order_by(desc(column))
+
+            return query
+
+        return super().apply_sorting(query=query, table=table)
+
     def apply_filter(
         self,
         query: AnyQuery,
@@ -202,15 +293,22 @@ class WorkspaceScopedResponseMetadata(UserScopedResponseMetadata):
     )
 
 
+class WorkspaceScopedResponseResources(UserScopedResponseResources):
+    """Base workspace-scoped resources."""
+
+
 WorkspaceBody = TypeVar("WorkspaceBody", bound=WorkspaceScopedResponseBody)
 WorkspaceMetadata = TypeVar(
     "WorkspaceMetadata", bound=WorkspaceScopedResponseMetadata
 )
+WorkspaceResources = TypeVar(
+    "WorkspaceResources", bound=WorkspaceScopedResponseResources
+)
 
 
 class WorkspaceScopedResponse(
-    UserScopedResponse[WorkspaceBody, WorkspaceMetadata],
-    Generic[WorkspaceBody, WorkspaceMetadata],
+    UserScopedResponse[WorkspaceBody, WorkspaceMetadata, WorkspaceResources],
+    Generic[WorkspaceBody, WorkspaceMetadata, WorkspaceResources],
 ):
     """Base workspace-scoped domain model.
 
@@ -228,21 +326,37 @@ class WorkspaceScopedResponse(
         return self.get_metadata().workspace
 
 
-class WorkspaceScopedFilter(BaseFilter):
+class WorkspaceScopedFilter(UserScopedFilter):
     """Model to enable advanced scoping with workspace."""
 
     FILTER_EXCLUDE_FIELDS: ClassVar[List[str]] = [
-        *BaseFilter.FILTER_EXCLUDE_FIELDS,
+        *UserScopedFilter.FILTER_EXCLUDE_FIELDS,
+        "workspace",
         "scope_workspace",
     ]
     CLI_EXCLUDE_FIELDS: ClassVar[List[str]] = [
-        *BaseFilter.CLI_EXCLUDE_FIELDS,
+        *UserScopedFilter.CLI_EXCLUDE_FIELDS,
+        "workspace_id",
+        "workspace",
         "scope_workspace",
     ]
+    CUSTOM_SORTING_OPTIONS: ClassVar[List[str]] = [
+        *UserScopedFilter.CUSTOM_SORTING_OPTIONS,
+        "workspace",
+    ]
     scope_workspace: Optional[UUID] = Field(
         default=None,
         description="The workspace to scope this query to.",
     )
+    workspace_id: Optional[Union[UUID, str]] = Field(
+        default=None,
+        description="UUID of the workspace that this entity belongs to.",
+        union_mode="left_to_right",
+    )
+    workspace: Optional[Union[UUID, str]] = Field(
+        default=None,
+        description="Name/ID of the workspace that this entity belongs to.",
+    )
 
     def set_scope_workspace(self, workspace_id: UUID) -> None:
         """Set the workspace to scope this response.
@@ -252,6 +366,35 @@ class WorkspaceScopedFilter(BaseFilter):
         """
         self.scope_workspace = workspace_id
 
+    def get_custom_filters(
+        self, table: Type["AnySchema"]
+    ) -> List["ColumnElement[bool]"]:
+        """Get custom filters.
+
+        Args:
+            table: The query table.
+
+        Returns:
+            A list of custom filters.
+        """
+        custom_filters = super().get_custom_filters(table)
+
+        from sqlmodel import and_
+
+        from zenml.zen_stores.schemas import WorkspaceSchema
+
+        if self.workspace:
+            workspace_filter = and_(
+                getattr(table, "workspace_id") == WorkspaceSchema.id,
+                self.generate_name_or_id_query_conditions(
+                    value=self.workspace,
+                    table=WorkspaceSchema,
+                ),
+            )
+            custom_filters.append(workspace_filter)
+
+        return custom_filters
+
     def apply_filter(
         self,
         query: AnyQuery,
@@ -279,6 +422,46 @@ class WorkspaceScopedFilter(BaseFilter):
 
         return query
 
+    def apply_sorting(
+        self,
+        query: AnyQuery,
+        table: Type["AnySchema"],
+    ) -> AnyQuery:
+        """Apply sorting to the query.
+
+        Args:
+            query: The query to which to apply the sorting.
+            table: The query table.
+
+        Returns:
+            The query with sorting applied.
+        """
+        from sqlmodel import asc, desc
+
+        from zenml.enums import SorterOps
+        from zenml.zen_stores.schemas import WorkspaceSchema
+
+        sort_by, operand = self.sorting_params
+
+        if sort_by == "workspace":
+            column = WorkspaceSchema.name
+
+            query = query.join(
+                WorkspaceSchema,
+                getattr(table, "workspace_id") == WorkspaceSchema.id,
+            )
+
+            query = query.add_columns(WorkspaceSchema.name)
+
+            if operand == SorterOps.ASCENDING:
+                query = query.order_by(asc(column))
+            else:
+                query = query.order_by(desc(column))
+
+            return query
+
+        return super().apply_sorting(query=query, table=table)
+
 
 class WorkspaceScopedTaggableFilter(WorkspaceScopedFilter):
     """Model to enable advanced scoping with workspace and tagging."""
@@ -291,6 +474,10 @@ class WorkspaceScopedTaggableFilter(WorkspaceScopedFilter):
         *WorkspaceScopedFilter.FILTER_EXCLUDE_FIELDS,
         "tag",
     ]
+    CUSTOM_SORTING_OPTIONS: ClassVar[List[str]] = [
+        *WorkspaceScopedFilter.CUSTOM_SORTING_OPTIONS,
+        "tags",
+    ]
 
     def apply_filter(
         self,
@@ -306,30 +493,113 @@ class WorkspaceScopedTaggableFilter(WorkspaceScopedFilter):
         Returns:
             The query with filter applied.
         """
-        from zenml.zen_stores.schemas import TagResourceSchema
+        from zenml.zen_stores.schemas import TagResourceSchema, TagSchema
 
         query = super().apply_filter(query=query, table=table)
         if self.tag:
-            query = (
-                query.join(getattr(table, "tags"))
-                .join(TagResourceSchema.tag)
-                .distinct()
-            )
+            query = query.join(
+                TagResourceSchema,
+                TagResourceSchema.resource_id == getattr(table, "id"),
+            ).join(TagSchema, TagSchema.id == TagResourceSchema.tag_id)
 
         return query
 
     def get_custom_filters(
-        self,
-    ) -> List[Union["BinaryExpression[Any]", "BooleanClauseList[Any]"]]:
+        self, table: Type["AnySchema"]
+    ) -> List["ColumnElement[bool]"]:
         """Get custom tag filters.
 
+        Args:
+            table: The query table.
+
         Returns:
             A list of custom filters.
         """
         from zenml.zen_stores.schemas import TagSchema
 
-        custom_filters = super().get_custom_filters()
+        custom_filters = super().get_custom_filters(table)
         if self.tag:
-            custom_filters.append(col(TagSchema.name) == self.tag)  # type: ignore[arg-type]
+            custom_filters.append(
+                self.generate_custom_query_conditions_for_column(
+                    value=self.tag, table=TagSchema, column="name"
+                )
+            )
 
         return custom_filters
+
+    def apply_sorting(
+        self,
+        query: AnyQuery,
+        table: Type["AnySchema"],
+    ) -> AnyQuery:
+        """Apply sorting to the query.
+
+        Args:
+            query: The query to which to apply the sorting.
+            table: The query table.
+
+        Returns:
+            The query with sorting applied.
+        """
+        sort_by, operand = self.sorting_params
+
+        if sort_by == "tags":
+            from sqlmodel import asc, desc, func, select
+
+            from zenml.enums import SorterOps, TaggableResourceTypes
+            from zenml.zen_stores.schemas import (
+                ArtifactSchema,
+                ArtifactVersionSchema,
+                ModelSchema,
+                ModelVersionSchema,
+                PipelineRunSchema,
+                PipelineSchema,
+                RunTemplateSchema,
+                TagResourceSchema,
+                TagSchema,
+            )
+
+            resource_type_mapping = {
+                ArtifactSchema: TaggableResourceTypes.ARTIFACT,
+                ArtifactVersionSchema: TaggableResourceTypes.ARTIFACT_VERSION,
+                ModelSchema: TaggableResourceTypes.MODEL,
+                ModelVersionSchema: TaggableResourceTypes.MODEL_VERSION,
+                PipelineSchema: TaggableResourceTypes.PIPELINE,
+                PipelineRunSchema: TaggableResourceTypes.PIPELINE_RUN,
+                RunTemplateSchema: TaggableResourceTypes.RUN_TEMPLATE,
+            }
+
+            sorted_tags = (
+                select(TagResourceSchema.resource_id, TagSchema.name)
+                .join(TagSchema, TagResourceSchema.tag_id == TagSchema.id)  # type: ignore[arg-type]
+                .filter(
+                    TagResourceSchema.resource_type  # type: ignore[arg-type]
+                    == resource_type_mapping[table]
+                )
+                .order_by(
+                    asc(TagResourceSchema.resource_id), asc(TagSchema.name)
+                )
+            ).alias("sorted_tags")
+
+            tags_subquery = (
+                select(
+                    sorted_tags.c.resource_id,
+                    func.group_concat(sorted_tags.c.name, ", ").label(
+                        "tags_list"
+                    ),
+                ).group_by(sorted_tags.c.resource_id)
+            ).alias("tags_subquery")
+
+            query = query.add_columns(tags_subquery.c.tags_list).outerjoin(
+                tags_subquery, table.id == tags_subquery.c.resource_id
+            )
+
+            # Apply ordering based on the tags list
+            if operand == SorterOps.ASCENDING:
+                query = query.order_by(asc("tags_list"))
+            else:
+                query = query.order_by(desc("tags_list"))
+
+            return query
+
+        return super().apply_sorting(query=query, table=table)