zenml-nightly 0.70.0.dev20241201__py3-none-any.whl → 0.71.0.dev20241223__py3-none-any.whl

zenml/models/v2/base/scoped.py

@@ -23,6 +23,7 @@ from typing import (
     Optional,
     Type,
     TypeVar,
+    Union,
 )
 from uuid import UUID
 
@@ -151,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.
@@ -170,6 +187,73 @@ 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
+            )
+
+            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,
@@ -240,21 +324,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.
@@ -264,6 +364,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,
@@ -291,6 +420,44 @@ 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,
+            )
+
+            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."""
@@ -304,6 +471,11 @@ class WorkspaceScopedTaggableFilter(WorkspaceScopedFilter):
         "tag",
     ]
 
+    CUSTOM_SORTING_OPTIONS: ClassVar[List[str]] = [
+        *WorkspaceScopedFilter.CUSTOM_SORTING_OPTIONS,
+        "tag",
+    ]
+
     def apply_filter(
         self,
         query: AnyQuery,
@@ -330,15 +502,20 @@ class WorkspaceScopedTaggableFilter(WorkspaceScopedFilter):
 
         return query
 
-    def get_custom_filters(self) -> List["ColumnElement[bool]"]:
+    def get_custom_filters(
+        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(
                 self.generate_custom_query_conditions_for_column(
@@ -347,3 +524,79 @@ class WorkspaceScopedTaggableFilter(WorkspaceScopedFilter):
             )
 
         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 == "tag":
+            from sqlmodel import and_, asc, desc, func
+
+            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,
+            }
+
+            query = (
+                query.outerjoin(
+                    TagResourceSchema,
+                    and_(
+                        table.id == TagResourceSchema.resource_id,
+                        TagResourceSchema.resource_type
+                        == resource_type_mapping[table],
+                    ),
+                )
+                .outerjoin(TagSchema, TagResourceSchema.tag_id == TagSchema.id)
+                .group_by(table.id)
+            )
+
+            if operand == SorterOps.ASCENDING:
+                query = query.order_by(
+                    asc(
+                        func.group_concat(TagSchema.name, ",").label(
+                            "tags_list"
+                        )
+                    )
+                )
+            else:
+                query = query.order_by(
+                    desc(
+                        func.group_concat(TagSchema.name, ",").label(
+                            "tags_list"
+                        )
+                    )
+                )
+
+            return query
+
+        return super().apply_sorting(query=query, table=table)

zenml/models/v2/core/artifact_version.py

@@ -20,6 +20,8 @@ from typing import (
     Dict,
     List,
     Optional,
+    Type,
+    TypeVar,
     Union,
 )
 from uuid import UUID
@@ -58,6 +60,10 @@ if TYPE_CHECKING:
     )
     from zenml.models.v2.core.pipeline_run import PipelineRunResponse
     from zenml.models.v2.core.step_run import StepRunResponse
+    from zenml.zen_stores.schemas.base_schemas import BaseSchema
+
+    AnySchema = TypeVar("AnySchema", bound=BaseSchema)
+
 
 logger = get_logger(__name__)
 
@@ -471,7 +477,6 @@ class ArtifactVersionFilter(WorkspaceScopedTaggableFilter):
         "name",
         "only_unused",
         "has_custom_name",
-        "user",
         "model",
         "pipeline_run",
         "model_version_id",
@@ -516,19 +521,10 @@ class ArtifactVersionFilter(WorkspaceScopedTaggableFilter):
         description="Artifact store for this artifact",
         union_mode="left_to_right",
     )
-    workspace_id: Optional[Union[UUID, str]] = Field(
-        default=None,
-        description="Workspace for this artifact",
-        union_mode="left_to_right",
-    )
-    user_id: Optional[Union[UUID, str]] = Field(
-        default=None,
-        description="User that produced this artifact",
-        union_mode="left_to_right",
-    )
     model_version_id: Optional[Union[UUID, str]] = Field(
         default=None,
-        description="ID of the model version that is associated with this artifact version.",
+        description="ID of the model version that is associated with this "
+        "artifact version.",
         union_mode="left_to_right",
     )
     only_unused: Optional[bool] = Field(
@@ -559,13 +555,18 @@ class ArtifactVersionFilter(WorkspaceScopedTaggableFilter):
 
     model_config = ConfigDict(protected_namespaces=())
 
-    def get_custom_filters(self) -> List[Union["ColumnElement[bool]"]]:
+    def get_custom_filters(
+        self, table: Type["AnySchema"]
+    ) -> List[Union["ColumnElement[bool]"]]:
         """Get custom filters.
 
+        Args:
+            table: The query table.
+
         Returns:
             A list of custom filters.
         """
-        custom_filters = super().get_custom_filters()
+        custom_filters = super().get_custom_filters(table)
 
         from sqlmodel import and_, or_, select
 
@@ -581,7 +582,6 @@ class ArtifactVersionFilter(WorkspaceScopedTaggableFilter):
             StepRunInputArtifactSchema,
             StepRunOutputArtifactSchema,
             StepRunSchema,
-            UserSchema,
         )
 
         if self.name:
@@ -629,17 +629,6 @@ class ArtifactVersionFilter(WorkspaceScopedTaggableFilter):
             )
             custom_filters.append(custom_name_filter)
 
-        if self.user:
-            user_filter = and_(
-                ArtifactVersionSchema.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)
-
         if self.model:
             model_filter = and_(
                 ArtifactVersionSchema.id

zenml/models/v2/core/code_repository.py

@@ -13,8 +13,7 @@
 #  permissions and limitations under the License.
 """Models representing code repositories."""
 
-from typing import Any, Dict, Optional, Union
-from uuid import UUID
+from typing import Any, Dict, Optional
 
 from pydantic import Field
 
@@ -189,13 +188,3 @@ class CodeRepositoryFilter(WorkspaceScopedFilter):
         description="Name of the code repository.",
         default=None,
     )
-    workspace_id: Optional[Union[UUID, str]] = Field(
-        description="Workspace of the code repository.",
-        default=None,
-        union_mode="left_to_right",
-    )
-    user_id: Optional[Union[UUID, str]] = Field(
-        description="User that created the code repository.",
-        default=None,
-        union_mode="left_to_right",
-    )

zenml/models/v2/core/component.py

@@ -21,6 +21,7 @@ from typing import (
     List,
     Optional,
     Type,
+    TypeVar,
     Union,
 )
 from uuid import UUID
@@ -42,9 +43,11 @@ from zenml.utils import secret_utils
 
 if TYPE_CHECKING:
     from sqlalchemy.sql.elements import ColumnElement
-    from sqlmodel import SQLModel
 
     from zenml.models import FlavorResponse, ServiceConnectorResponse
+    from zenml.zen_stores.schemas.base_schemas import BaseSchema
+
+    AnySchema = TypeVar("AnySchema", bound=BaseSchema)
 
 # ------------------ Base Model ------------------
 
@@ -356,7 +359,6 @@ class ComponentFilter(WorkspaceScopedFilter):
         *WorkspaceScopedFilter.FILTER_EXCLUDE_FIELDS,
         "scope_type",
         "stack_id",
-        "user",
     ]
     CLI_EXCLUDE_FIELDS: ClassVar[List[str]] = [
         *WorkspaceScopedFilter.CLI_EXCLUDE_FIELDS,
@@ -366,7 +368,6 @@ class ComponentFilter(WorkspaceScopedFilter):
         default=None,
         description="The type to scope this query to.",
     )
-
     name: Optional[str] = Field(
         default=None,
         description="Name of the stack component",
@@ -379,16 +380,6 @@ class ComponentFilter(WorkspaceScopedFilter):
         default=None,
         description="Type of the stack component",
     )
-    workspace_id: Optional[Union[UUID, str]] = Field(
-        default=None,
-        description="Workspace of the stack component",
-        union_mode="left_to_right",
-    )
-    user_id: Optional[Union[UUID, str]] = Field(
-        default=None,
-        description="User of the stack component",
-        union_mode="left_to_right",
-    )
     connector_id: Optional[Union[UUID, str]] = Field(
         default=None,
         description="Connector linked to the stack component",
@@ -399,10 +390,6 @@ class ComponentFilter(WorkspaceScopedFilter):
         description="Stack of the stack component",
         union_mode="left_to_right",
     )
-    user: Optional[Union[UUID, str]] = Field(
-        default=None,
-        description="Name/ID of the user that created the component.",
-    )
 
     def set_scope_type(self, component_type: str) -> None:
         """Set the type of component on which to perform the filtering to scope the response.
@@ -413,7 +400,7 @@ class ComponentFilter(WorkspaceScopedFilter):
         self.scope_type = component_type
 
     def generate_filter(
-        self, table: Type["SQLModel"]
+        self, table: Type["AnySchema"]
     ) -> Union["ColumnElement[bool]"]:
         """Generate the filter for the query.
 
@@ -449,31 +436,3 @@ class ComponentFilter(WorkspaceScopedFilter):
             base_filter = operator(base_filter, stack_filter)
 
         return base_filter
-
-    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 (
-            StackComponentSchema,
-            UserSchema,
-        )
-
-        custom_filters = super().get_custom_filters()
-
-        if self.user:
-            user_filter = and_(
-                StackComponentSchema.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

zenml/models/v2/core/flavor.py

@@ -13,7 +13,7 @@
 #  permissions and limitations under the License.
 """Models representing flavors."""
 
-from typing import TYPE_CHECKING, Any, ClassVar, Dict, List, Optional, Union
+from typing import TYPE_CHECKING, Any, ClassVar, Dict, List, Optional
 from uuid import UUID
 
 from pydantic import Field
@@ -428,13 +428,3 @@ class FlavorFilter(WorkspaceScopedFilter):
         default=None,
         description="Integration associated with the flavor",
     )
-    workspace_id: Optional[Union[UUID, str]] = Field(
-        default=None,
-        description="Workspace of the stack",
-        union_mode="left_to_right",
-    )
-    user_id: Optional[Union[UUID, str]] = Field(
-        default=None,
-        description="User of the stack",
-        union_mode="left_to_right",
-    )

zenml/models/v2/core/model.py

@@ -13,7 +13,7 @@
 #  permissions and limitations under the License.
 """Models representing models."""
 
-from typing import TYPE_CHECKING, ClassVar, List, Optional, Union
+from typing import TYPE_CHECKING, List, Optional
 from uuid import UUID
 
 from pydantic import BaseModel, Field
@@ -30,8 +30,6 @@ from zenml.models.v2.base.scoped import (
 from zenml.utils.pagination_utils import depaginate
 
 if TYPE_CHECKING:
-    from sqlalchemy.sql.elements import ColumnElement
-
     from zenml.model.model import Model
     from zenml.models.v2.core.tag import TagResponse
 
@@ -318,61 +316,7 @@ class ModelResponse(
 class ModelFilter(WorkspaceScopedTaggableFilter):
     """Model to enable advanced filtering of all Workspaces."""
 
-    CLI_EXCLUDE_FIELDS: ClassVar[List[str]] = [
-        *WorkspaceScopedTaggableFilter.CLI_EXCLUDE_FIELDS,
-        "workspace_id",
-        "user_id",
-    ]
-    FILTER_EXCLUDE_FIELDS: ClassVar[List[str]] = [
-        *WorkspaceScopedTaggableFilter.FILTER_EXCLUDE_FIELDS,
-        "user",
-    ]
-
     name: Optional[str] = Field(
         default=None,
         description="Name of the Model",
     )
-    workspace_id: Optional[Union[UUID, str]] = Field(
-        default=None,
-        description="Workspace of the Model",
-        union_mode="left_to_right",
-    )
-    user_id: Optional[Union[UUID, str]] = Field(
-        default=None,
-        description="User of the Model",
-        union_mode="left_to_right",
-    )
-    user: Optional[Union[UUID, str]] = Field(
-        default=None,
-        description="Name/ID of the user that created the model.",
-    )
-
-    def get_custom_filters(
-        self,
-    ) -> List["ColumnElement[bool]"]:
-        """Get custom filters.
-
-        Returns:
-            A list of custom filters.
-        """
-        custom_filters = super().get_custom_filters()
-
-        from sqlmodel import and_
-
-        from zenml.zen_stores.schemas import (
-            ModelSchema,
-            UserSchema,
-        )
-
-        if self.user:
-            user_filter = and_(
-                ModelSchema.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