zenml-nightly 0.75.0.dev20250311__py3-none-any.whl → 0.75.0.dev20250313__py3-none-any.whl

zenml/models/v2/base/scoped.py

@@ -27,8 +27,9 @@ from typing import (
 )
 from uuid import UUID
 
-from pydantic import Field
+from pydantic import Field, model_validator
 
+from zenml.logger import get_logger
 from zenml.models.v2.base.base import (
     BaseDatedResponseBody,
     BaseIdentifiedResponse,
@@ -47,6 +48,9 @@ if TYPE_CHECKING:
 
     AnySchema = TypeVar("AnySchema", bound=BaseSchema)
 
+logger = get_logger(__name__)
+
+
 # ---------------------- Request Models ----------------------
 
 
@@ -56,7 +60,14 @@ class UserScopedRequest(BaseRequest):
     Used as a base class for all domain models that are "owned" by a user.
     """
 
-    user: UUID = Field(title="The id of the user that created this resource.")
+    user: Optional[UUID] = Field(
+        default=None,
+        title="The id of the user that created this resource. Set "
+        "automatically by the server.",
+        # This field is set automatically by the server, so the client doesn't
+        # need to set it and it will not be serialized.
+        exclude=True,
+    )
 
     def get_analytics_metadata(self) -> Dict[str, Any]:
         """Fetches the analytics metadata for user scoped models.
@@ -157,7 +168,6 @@ class UserScopedFilter(BaseFilter):
     ]
     CLI_EXCLUDE_FIELDS: ClassVar[List[str]] = [
         *BaseFilter.CLI_EXCLUDE_FIELDS,
-        "user_id",
         "scope_user",
     ]
     CUSTOM_SORTING_OPTIONS: ClassVar[List[str]] = [
@@ -169,14 +179,10 @@ class UserScopedFilter(BaseFilter):
         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.",
+        union_mode="left_to_right",
     )
 
     def set_scope_user(self, user_id: UUID) -> None:
@@ -316,6 +322,18 @@ class WorkspaceScopedResponse(
     Used as a base class for all domain models that are workspace-scoped.
     """
 
+    # Analytics
+    def get_analytics_metadata(self) -> Dict[str, Any]:
+        """Fetches the analytics metadata for workspace scoped models.
+
+        Returns:
+            The analytics metadata.
+        """
+        metadata = super().get_analytics_metadata()
+        if self.workspace is not None:
+            metadata["workspace_id"] = self.workspace.id
+        return metadata
+
     # Body and metadata properties
     @property
     def workspace(self) -> "WorkspaceResponse":
@@ -327,75 +345,25 @@ class WorkspaceScopedResponse(
         return self.get_metadata().workspace
 
 
+# ---------------------- Filter Models ----------------------
+
+
 class WorkspaceScopedFilter(UserScopedFilter):
     """Model to enable advanced scoping with workspace."""
 
     FILTER_EXCLUDE_FIELDS: ClassVar[List[str]] = [
         *UserScopedFilter.FILTER_EXCLUDE_FIELDS,
         "workspace",
-        "scope_workspace",
-    ]
-    CLI_EXCLUDE_FIELDS: ClassVar[List[str]] = [
-        *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.",
+        description="Name/ID of the workspace which the search is scoped to. "
+        "This field must always be set and is always applied in addition to "
+        "the other filters, regardless of the value of the "
+        "logical_operator field.",
+        union_mode="left_to_right",
     )
 
-    def set_scope_workspace(self, workspace_id: UUID) -> None:
-        """Set the workspace to scope this response.
-
-        Args:
-            workspace_id: The workspace to scope this response to.
-        """
-        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,
@@ -409,59 +377,35 @@ class WorkspaceScopedFilter(UserScopedFilter):
 
         Returns:
             The query with filter applied.
-        """
-        from sqlmodel import or_
-
-        query = super().apply_filter(query=query, table=table)
-
-        if self.scope_workspace:
-            scope_filter = or_(
-                getattr(table, "workspace_id") == self.scope_workspace,
-                getattr(table, "workspace_id").is_(None),
-            )
-            query = query.where(scope_filter)
-
-        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.
+        Raises:
+            ValueError: If the workspace scope is missing from the filter.
         """
-        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 = super().apply_filter(query=query, table=table)
 
-            query = query.join(
-                WorkspaceSchema,
-                getattr(table, "workspace_id") == WorkspaceSchema.id,
+        # The workspace scope must always be set and must be a UUID. If the
+        # client sets this to a string, the server will try to resolve it to a
+        # workspace ID.
+        #
+        # If not set by the client, the server will fall back to using the
+        # user's default workspace or even the server's default workspace, if
+        # they are configured. If this also fails to yield a workspace, this
+        # method will raise a ValueError.
+        #
+        # See: SqlZenStore._set_filter_workspace_id
+
+        if not self.workspace:
+            raise ValueError("Workspace scope missing from the filter.")
+
+        if not isinstance(self.workspace, UUID):
+            raise ValueError(
+                f"Workspace scope must be a UUID, got {type(self.workspace)}."
             )
 
-            query = query.add_columns(WorkspaceSchema.name)
+        scope_filter = getattr(table, "workspace_id") == self.workspace
+        query = query.where(scope_filter)
 
-            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)
+        return query
 
 
 class TaggableFilter(BaseFilter):
@@ -470,16 +414,44 @@ class TaggableFilter(BaseFilter):
     tag: Optional[str] = Field(
         description="Tag to apply to the filter query.", default=None
     )
+    tags: Optional[List[str]] = Field(
+        description="Tags to apply to the filter query.", default=None
+    )
 
+    CLI_EXCLUDE_FIELDS = [
+        *BaseFilter.CLI_EXCLUDE_FIELDS,
+    ]
     FILTER_EXCLUDE_FIELDS: ClassVar[List[str]] = [
         *BaseFilter.FILTER_EXCLUDE_FIELDS,
         "tag",
+        "tags",
     ]
     CUSTOM_SORTING_OPTIONS: ClassVar[List[str]] = [
         *BaseFilter.CUSTOM_SORTING_OPTIONS,
         "tags",
     ]
 
+    @model_validator(mode="after")
+    def add_tag_to_tags(self) -> "TaggableFilter":
+        """Deprecated the tag attribute in favor of the tags attribute.
+
+        Returns:
+            self
+        """
+        if self.tag is not None:
+            logger.warning(
+                "The `tag` attribute is deprecated in favor of the `tags` attribute. "
+                "Please update your code to use the `tags` attribute instead."
+            )
+            if self.tags is not None:
+                self.tags.append(self.tag)
+            else:
+                self.tags = [self.tag]
+
+            self.tag = None
+
+        return self
+
     def apply_filter(
         self,
         query: AnyQuery,
@@ -497,7 +469,8 @@ class TaggableFilter(BaseFilter):
         from zenml.zen_stores.schemas import TagResourceSchema, TagSchema
 
         query = super().apply_filter(query=query, table=table)
-        if self.tag:
+
+        if self.tags is not None:
             query = query.join(
                 TagResourceSchema,
                 TagResourceSchema.resource_id == getattr(table, "id"),
@@ -516,15 +489,25 @@ class TaggableFilter(BaseFilter):
         Returns:
             A list of custom filters.
         """
-        from zenml.zen_stores.schemas import TagSchema
-
         custom_filters = super().get_custom_filters(table)
-        if self.tag:
-            custom_filters.append(
-                self.generate_custom_query_conditions_for_column(
-                    value=self.tag, table=TagSchema, column="name"
+
+        if self.tags is not None:
+            from sqlmodel import exists, select
+
+            from zenml.zen_stores.schemas import TagResourceSchema, TagSchema
+
+            for tag in self.tags:
+                condition = self.generate_custom_query_conditions_for_column(
+                    value=tag, table=TagSchema, column="name"
                 )
-            )
+                exists_subquery = exists(
+                    select(TagResourceSchema)
+                    .join(TagSchema, TagSchema.id == TagResourceSchema.tag_id)  # type: ignore[arg-type]
+                    .where(
+                        TagResourceSchema.resource_id == table.id, condition
+                    )
+                )
+                custom_filters.append(exists_subquery)
 
         return custom_filters
 

zenml/models/v2/core/api_key.py

@@ -97,7 +97,7 @@ class APIKeyRequest(BaseRequest):
     )
 
 
-class APIKeyRotateRequest(BaseModel):
+class APIKeyRotateRequest(BaseRequest):
     """Request model for API key rotation."""
 
     retain_period_minutes: int = Field(

zenml/models/v2/core/artifact.py

@@ -25,17 +25,19 @@ from typing import (
 )
 from uuid import UUID
 
-from pydantic import BaseModel, Field
+from pydantic import Field
 
 from zenml.constants import SORT_BY_LATEST_VERSION_KEY, STR_FIELD_MAX_LENGTH
-from zenml.models.v2.base.base import (
-    BaseDatedResponseBody,
-    BaseIdentifiedResponse,
-    BaseRequest,
-    BaseResponseMetadata,
-    BaseResponseResources,
+from zenml.models.v2.base.base import BaseUpdate
+from zenml.models.v2.base.scoped import (
+    TaggableFilter,
+    WorkspaceScopedFilter,
+    WorkspaceScopedRequest,
+    WorkspaceScopedResponse,
+    WorkspaceScopedResponseBody,
+    WorkspaceScopedResponseMetadata,
+    WorkspaceScopedResponseResources,
 )
-from zenml.models.v2.base.scoped import TaggableFilter
 from zenml.models.v2.core.tag import TagResponse
 
 if TYPE_CHECKING:
@@ -49,7 +51,7 @@ AnyQuery = TypeVar("AnyQuery", bound=Any)
 # ------------------ Request Model ------------------
 
 
-class ArtifactRequest(BaseRequest):
+class ArtifactRequest(WorkspaceScopedRequest):
     """Artifact request model."""
 
     name: str = Field(
@@ -70,7 +72,7 @@ class ArtifactRequest(BaseRequest):
 # ------------------ Update Model ------------------
 
 
-class ArtifactUpdate(BaseModel):
+class ArtifactUpdate(BaseUpdate):
     """Artifact update model."""
 
     name: Optional[str] = None
@@ -82,7 +84,7 @@ class ArtifactUpdate(BaseModel):
 # ------------------ Response Model ------------------
 
 
-class ArtifactResponseBody(BaseDatedResponseBody):
+class ArtifactResponseBody(WorkspaceScopedResponseBody):
     """Response body for artifacts."""
 
     tags: List[TagResponse] = Field(
@@ -92,7 +94,7 @@ class ArtifactResponseBody(BaseDatedResponseBody):
     latest_version_id: Optional[UUID] = None
 
 
-class ArtifactResponseMetadata(BaseResponseMetadata):
+class ArtifactResponseMetadata(WorkspaceScopedResponseMetadata):
     """Response metadata for artifacts."""
 
     has_custom_name: bool = Field(
@@ -101,12 +103,12 @@ class ArtifactResponseMetadata(BaseResponseMetadata):
     )
 
 
-class ArtifactResponseResources(BaseResponseResources):
+class ArtifactResponseResources(WorkspaceScopedResponseResources):
     """Class for all resource models associated with the Artifact Entity."""
 
 
 class ArtifactResponse(
-    BaseIdentifiedResponse[
+    WorkspaceScopedResponse[
         ArtifactResponseBody,
         ArtifactResponseMetadata,
         ArtifactResponseResources,
@@ -176,24 +178,35 @@ class ArtifactResponse(
         """
         from zenml.client import Client
 
-        responses = Client().list_artifact_versions(name=self.name)
+        responses = Client().list_artifact_versions(artifact=self.name)
         return {str(response.version): response for response in responses}
 
 
 # ------------------ Filter Model ------------------
 
 
-class ArtifactFilter(TaggableFilter):
+class ArtifactFilter(WorkspaceScopedFilter, TaggableFilter):
     """Model to enable advanced filtering of artifacts."""
 
-    name: Optional[str] = None
-    has_custom_name: Optional[bool] = None
+    FILTER_EXCLUDE_FIELDS: ClassVar[List[str]] = [
+        *WorkspaceScopedFilter.FILTER_EXCLUDE_FIELDS,
+        *TaggableFilter.FILTER_EXCLUDE_FIELDS,
+    ]
 
     CUSTOM_SORTING_OPTIONS: ClassVar[List[str]] = [
+        *WorkspaceScopedFilter.CUSTOM_SORTING_OPTIONS,
         *TaggableFilter.CUSTOM_SORTING_OPTIONS,
         SORT_BY_LATEST_VERSION_KEY,
     ]
 
+    CLI_EXCLUDE_FIELDS: ClassVar[List[str]] = [
+        *WorkspaceScopedFilter.CLI_EXCLUDE_FIELDS,
+        *TaggableFilter.CLI_EXCLUDE_FIELDS,
+    ]
+
+    name: Optional[str] = None
+    has_custom_name: Optional[bool] = None
+
     def apply_sorting(
         self,
         query: AnyQuery,

zenml/models/v2/core/artifact_version.py

@@ -27,7 +27,6 @@ from typing import (
 from uuid import UUID
 
 from pydantic import (
-    BaseModel,
     ConfigDict,
     Field,
     field_validator,
@@ -36,10 +35,11 @@ from pydantic import (
 
 from zenml.config.source import Source, SourceWithValidator
 from zenml.constants import STR_FIELD_MAX_LENGTH, TEXT_FIELD_MAX_LENGTH
-from zenml.enums import ArtifactSaveType, ArtifactType, GenericFilterOps
+from zenml.enums import ArtifactSaveType, ArtifactType
 from zenml.logger import get_logger
 from zenml.metadata.metadata_types import MetadataType
-from zenml.models.v2.base.filter import FilterGenerator, StrFilter
+from zenml.models.v2.base.base import BaseUpdate
+from zenml.models.v2.base.filter import FilterGenerator
 from zenml.models.v2.base.scoped import (
     TaggableFilter,
     WorkspaceScopedFilter,
@@ -165,7 +165,7 @@ class ArtifactVersionRequest(WorkspaceScopedRequest):
 # ------------------ Update Model ------------------
 
 
-class ArtifactVersionUpdate(BaseModel):
+class ArtifactVersionUpdate(BaseUpdate):
     """Artifact version update model."""
 
     name: Optional[str] = None
@@ -476,7 +476,8 @@ class ArtifactVersionFilter(WorkspaceScopedFilter, TaggableFilter):
     FILTER_EXCLUDE_FIELDS: ClassVar[List[str]] = [
         *WorkspaceScopedFilter.FILTER_EXCLUDE_FIELDS,
         *TaggableFilter.FILTER_EXCLUDE_FIELDS,
-        "name",
+        "artifact_id",
+        "artifact",
         "only_unused",
         "has_custom_name",
         "model",
@@ -484,23 +485,28 @@ class ArtifactVersionFilter(WorkspaceScopedFilter, TaggableFilter):
         "model_version_id",
         "run_metadata",
     ]
-    CUSTOM_SORTING_OPTIONS = [
+    CUSTOM_SORTING_OPTIONS: ClassVar[List[str]] = [
         *WorkspaceScopedFilter.CUSTOM_SORTING_OPTIONS,
         *TaggableFilter.CUSTOM_SORTING_OPTIONS,
     ]
-    CLI_EXCLUDE_FIELDS = [
+    CLI_EXCLUDE_FIELDS: ClassVar[List[str]] = [
         *WorkspaceScopedFilter.CLI_EXCLUDE_FIELDS,
         *TaggableFilter.CLI_EXCLUDE_FIELDS,
+        "artifact_id",
     ]
 
-    artifact_id: Optional[Union[UUID, str]] = Field(
+    artifact: Optional[Union[UUID, str]] = Field(
         default=None,
-        description="ID of the artifact to which this version belongs.",
+        description="The name or ID of the artifact which the search is scoped "
+        "to. This field must always be set and is always applied in addition "
+        "to the other filters, regardless of the value of the "
+        "logical_operator field.",
         union_mode="left_to_right",
     )
-    name: Optional[str] = Field(
+    artifact_id: Optional[Union[UUID, str]] = Field(
         default=None,
-        description="Name of the artifact to which this version belongs.",
+        description="[Deprecated] Use 'artifact' instead. ID of the artifact to which this version belongs.",
+        union_mode="left_to_right",
     )
     version: Optional[str] = Field(
         default=None,
@@ -545,10 +551,6 @@ class ArtifactVersionFilter(WorkspaceScopedFilter, TaggableFilter):
         default=None,
         description="Filter only artifacts with/without custom names.",
     )
-    user: Optional[Union[UUID, str]] = Field(
-        default=None,
-        description="Name/ID of the user that created the artifact version.",
-    )
     model: Optional[Union[UUID, str]] = Field(
         default=None,
         description="Name/ID of the model that is associated with this "
@@ -595,18 +597,15 @@ class ArtifactVersionFilter(WorkspaceScopedFilter, TaggableFilter):
             StepRunSchema,
         )
 
-        if self.name:
-            value, filter_operator = self._resolve_operator(self.name)
-            filter_ = StrFilter(
-                operation=GenericFilterOps(filter_operator),
-                column="name",
-                value=value,
-            )
-            artifact_name_filter = and_(
+        if self.artifact:
+            value, operator = self._resolve_operator(self.artifact)
+            artifact_filter = and_(
                 ArtifactVersionSchema.artifact_id == ArtifactSchema.id,
-                filter_.generate_query_conditions(ArtifactSchema),
+                self.generate_name_or_id_query_conditions(
+                    value=self.artifact, table=ArtifactSchema
+                ),
             )
-            custom_filters.append(artifact_name_filter)
+            custom_filters.append(artifact_filter)
 
         if self.only_unused:
             unused_filter = and_(
@@ -702,6 +701,24 @@ class ArtifactVersionFilter(WorkspaceScopedFilter, TaggableFilter):
 
         return custom_filters
 
+    @model_validator(mode="after")
+    def _migrate_artifact_id(self) -> "ArtifactVersionFilter":
+        """Migrate value from the deprecated artifact_id attribute.
+
+        Returns:
+            The filter with migrated value.
+        """
+        # Handle deprecated artifact_id field
+        if self.artifact_id is not None:
+            logger.warning(
+                "The 'ArtifactVersionFilter.artifact_id' field is deprecated "
+                "and will be removed in a future version. Please use "
+                "'ArtifactVersionFilter.artifact' instead."
+            )
+            self.artifact = self.artifact or self.artifact_id
+
+        return self
+
 
 # -------------------- Lazy Loader --------------------