zenml-nightly 0.66.0.dev20240923__py3-none-any.whl → 0.66.0.dev20240927__py3-none-any.whl

zenml/models/v2/core/artifact_version.py

@@ -24,7 +24,7 @@ from typing import (
 )
 from uuid import UUID
 
-from pydantic import BaseModel, Field, field_validator
+from pydantic import BaseModel, ConfigDict, Field, field_validator
 
 from zenml.config.source import Source, SourceWithValidator
 from zenml.constants import STR_FIELD_MAX_LENGTH, TEXT_FIELD_MAX_LENGTH
@@ -430,14 +430,14 @@ class ArtifactVersionResponse(
 class ArtifactVersionFilter(WorkspaceScopedTaggableFilter):
     """Model to enable advanced filtering of artifact versions."""
 
-    # `name` and `only_unused` refer to properties related to other entities
-    #  rather than a field in the db, hence they need to be handled
-    #  explicitly
     FILTER_EXCLUDE_FIELDS: ClassVar[List[str]] = [
         *WorkspaceScopedTaggableFilter.FILTER_EXCLUDE_FIELDS,
         "name",
         "only_unused",
         "has_custom_name",
+        "user",
+        "model",
+        "pipeline_run",
     ]
     artifact_id: Optional[Union[UUID, str]] = Field(
         default=None,
@@ -495,6 +495,22 @@ class ArtifactVersionFilter(WorkspaceScopedTaggableFilter):
         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 "
+        "artifact version.",
+    )
+    pipeline_run: Optional[Union[UUID, str]] = Field(
+        default=None,
+        description="Name/ID of a pipeline run that is associated with this "
+        "artifact version.",
+    )
+
+    model_config = ConfigDict(protected_namespaces=())
 
     def get_custom_filters(self) -> List[Union["ColumnElement[bool]"]]:
         """Get custom filters.
@@ -504,15 +520,18 @@ class ArtifactVersionFilter(WorkspaceScopedTaggableFilter):
         """
         custom_filters = super().get_custom_filters()
 
-        from sqlmodel import and_, select
+        from sqlmodel import and_, or_, select
 
-        from zenml.zen_stores.schemas.artifact_schemas import (
+        from zenml.zen_stores.schemas import (
             ArtifactSchema,
             ArtifactVersionSchema,
-        )
-        from zenml.zen_stores.schemas.step_run_schemas import (
+            ModelSchema,
+            ModelVersionArtifactSchema,
+            PipelineRunSchema,
             StepRunInputArtifactSchema,
             StepRunOutputArtifactSchema,
+            StepRunSchema,
+            UserSchema,
         )
 
         if self.name:
@@ -546,6 +565,48 @@ 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
+                ),
+            )
+            custom_filters.append(user_filter)
+
+        if self.model:
+            model_filter = and_(
+                ArtifactVersionSchema.id
+                == ModelVersionArtifactSchema.artifact_version_id,
+                ModelVersionArtifactSchema.model_id == ModelSchema.id,
+                self.generate_name_or_id_query_conditions(
+                    value=self.model, table=ModelSchema
+                ),
+            )
+            custom_filters.append(model_filter)
+
+        if self.pipeline_run:
+            pipeline_run_filter = and_(
+                or_(
+                    and_(
+                        ArtifactVersionSchema.id
+                        == StepRunOutputArtifactSchema.artifact_id,
+                        StepRunOutputArtifactSchema.step_id
+                        == StepRunSchema.id,
+                    ),
+                    and_(
+                        ArtifactVersionSchema.id
+                        == StepRunInputArtifactSchema.artifact_id,
+                        StepRunInputArtifactSchema.step_id == StepRunSchema.id,
+                    ),
+                ),
+                StepRunSchema.pipeline_run_id == PipelineRunSchema.id,
+                self.generate_name_or_id_query_conditions(
+                    value=self.pipeline_run, table=PipelineRunSchema
+                ),
+            )
+            custom_filters.append(pipeline_run_filter)
+
         return custom_filters
 
 

zenml/models/v2/core/model.py

@@ -30,10 +30,11 @@ 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
 
-
 # ------------------ Request Model ------------------
 
 
@@ -316,6 +317,16 @@ 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",
@@ -330,9 +341,35 @@ class ModelFilter(WorkspaceScopedTaggableFilter):
         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.",
+    )
 
-    CLI_EXCLUDE_FIELDS: ClassVar[List[str]] = [
-        *WorkspaceScopedTaggableFilter.CLI_EXCLUDE_FIELDS,
-        "workspace_id",
-        "user_id",
-    ]
+    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
+                ),
+            )
+            custom_filters.append(user_filter)
+
+        return custom_filters

zenml/models/v2/core/model_version.py

@@ -13,7 +13,16 @@
 #  permissions and limitations under the License.
 """Models representing model versions."""
 
-from typing import TYPE_CHECKING, Dict, List, Optional, Type, TypeVar, Union
+from typing import (
+    TYPE_CHECKING,
+    ClassVar,
+    Dict,
+    List,
+    Optional,
+    Type,
+    TypeVar,
+    Union,
+)
 from uuid import UUID
 
 from pydantic import BaseModel, ConfigDict, Field, PrivateAttr, field_validator
@@ -34,6 +43,8 @@ from zenml.models.v2.core.service import ServiceResponse
 from zenml.models.v2.core.tag import TagResponse
 
 if TYPE_CHECKING:
+    from sqlalchemy.sql.elements import ColumnElement
+
     from zenml.model.model import Model
     from zenml.models.v2.core.artifact_version import ArtifactVersionResponse
     from zenml.models.v2.core.model import ModelResponse
@@ -582,6 +593,11 @@ class ModelVersionResponse(
 class ModelVersionFilter(WorkspaceScopedTaggableFilter):
     """Filter model for model versions."""
 
+    FILTER_EXCLUDE_FIELDS: ClassVar[List[str]] = [
+        *WorkspaceScopedTaggableFilter.FILTER_EXCLUDE_FIELDS,
+        "user",
+    ]
+
     name: Optional[str] = Field(
         default=None,
         description="The name of the Model Version",
@@ -605,6 +621,10 @@ class ModelVersionFilter(WorkspaceScopedTaggableFilter):
         default=None,
         union_mode="left_to_right",
     )
+    user: Optional[Union[UUID, str]] = Field(
+        default=None,
+        description="Name/ID of the user that created the model version.",
+    )
 
     _model_id: UUID = PrivateAttr(None)
 
@@ -623,6 +643,34 @@ class ModelVersionFilter(WorkspaceScopedTaggableFilter):
 
         self._model_id = model_id
 
+    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 (
+            ModelVersionSchema,
+            UserSchema,
+        )
+
+        if self.user:
+            user_filter = and_(
+                ModelVersionSchema.user_id == UserSchema.id,
+                self.generate_name_or_id_query_conditions(
+                    value=self.user, table=UserSchema
+                ),
+            )
+            custom_filters.append(user_filter)
+
+        return custom_filters
+
     def apply_filter(
         self,
         query: AnyQuery,

zenml/models/v2/core/model_version_artifact.py

@@ -166,6 +166,7 @@ class ModelVersionArtifactFilter(WorkspaceScopedFilter):
         "only_model_artifacts",
         "only_deployment_artifacts",
         "has_custom_name",
+        "user",
     ]
     CLI_EXCLUDE_FIELDS = [
         *WorkspaceScopedFilter.CLI_EXCLUDE_FIELDS,
@@ -214,6 +215,10 @@ class ModelVersionArtifactFilter(WorkspaceScopedFilter):
     only_model_artifacts: Optional[bool] = False
     only_deployment_artifacts: Optional[bool] = False
     has_custom_name: Optional[bool] = None
+    user: Optional[Union[UUID, str]] = Field(
+        default=None,
+        description="Name/ID of the user that created the artifact.",
+    )
 
     # TODO: In Pydantic v2, the `model_` is a protected namespaces for all
     #  fields defined under base models. If not handled, this raises a warning.
@@ -233,12 +238,11 @@ class ModelVersionArtifactFilter(WorkspaceScopedFilter):
 
         from sqlmodel import and_
 
-        from zenml.zen_stores.schemas.artifact_schemas import (
+        from zenml.zen_stores.schemas import (
             ArtifactSchema,
             ArtifactVersionSchema,
-        )
-        from zenml.zen_stores.schemas.model_schemas import (
             ModelVersionArtifactSchema,
+            UserSchema,
         )
 
         if self.artifact_name:
@@ -284,4 +288,15 @@ class ModelVersionArtifactFilter(WorkspaceScopedFilter):
             )
             custom_filters.append(custom_name_filter)
 
+        if self.user:
+            user_filter = and_(
+                ModelVersionArtifactSchema.artifact_version_id
+                == ArtifactVersionSchema.id,
+                ArtifactVersionSchema.user_id == UserSchema.id,
+                self.generate_name_or_id_query_conditions(
+                    value=self.user, table=UserSchema
+                ),
+            )
+            custom_filters.append(user_filter)
+
         return custom_filters

zenml/models/v2/core/model_version_pipeline_run.py

@@ -123,10 +123,10 @@ class ModelVersionPipelineRunResponse(
 class ModelVersionPipelineRunFilter(WorkspaceScopedFilter):
     """Model version pipeline run links filter model."""
 
-    # Pipeline run name is not a DB field and needs to be handled separately
     FILTER_EXCLUDE_FIELDS = [
         *WorkspaceScopedFilter.FILTER_EXCLUDE_FIELDS,
         "pipeline_run_name",
+        "user",
     ]
     CLI_EXCLUDE_FIELDS = [
         *WorkspaceScopedFilter.CLI_EXCLUDE_FIELDS,
@@ -167,6 +167,10 @@ class ModelVersionPipelineRunFilter(WorkspaceScopedFilter):
         default=None,
         description="Name of the pipeline run",
     )
+    user: Optional[Union[UUID, str]] = Field(
+        default=None,
+        description="Name/ID of the user that created the pipeline run.",
+    )
 
     # TODO: In Pydantic v2, the `model_` is a protected namespaces for all
     #  fields defined under base models. If not handled, this raises a warning.
@@ -186,11 +190,10 @@ class ModelVersionPipelineRunFilter(WorkspaceScopedFilter):
 
         from sqlmodel import and_
 
-        from zenml.zen_stores.schemas.model_schemas import (
+        from zenml.zen_stores.schemas import (
             ModelVersionPipelineRunSchema,
-        )
-        from zenml.zen_stores.schemas.pipeline_run_schemas import (
             PipelineRunSchema,
+            UserSchema,
         )
 
         if self.pipeline_run_name:
@@ -209,4 +212,15 @@ class ModelVersionPipelineRunFilter(WorkspaceScopedFilter):
             )
             custom_filters.append(pipeline_run_name_filter)
 
+        if self.user:
+            user_filter = and_(
+                ModelVersionPipelineRunSchema.pipeline_run_id
+                == PipelineRunSchema.id,
+                PipelineRunSchema.user_id == UserSchema.id,
+                self.generate_name_or_id_query_conditions(
+                    value=self.user, table=UserSchema
+                ),
+            )
+            custom_filters.append(user_filter)
+
         return custom_filters

zenml/models/v2/core/pipeline.py

@@ -13,7 +13,16 @@
 #  permissions and limitations under the License.
 """Models representing pipelines."""
 
-from typing import TYPE_CHECKING, Any, List, Optional, Type, TypeVar, Union
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    ClassVar,
+    List,
+    Optional,
+    Type,
+    TypeVar,
+    Union,
+)
 from uuid import UUID
 
 from pydantic import Field
@@ -36,6 +45,8 @@ from zenml.models.v2.base.scoped import (
 from zenml.models.v2.core.tag import TagResponse
 
 if TYPE_CHECKING:
+    from sqlalchemy.sql.elements import ColumnElement
+
     from zenml.models.v2.core.pipeline_run import PipelineRunResponse
     from zenml.zen_stores.schemas import BaseSchema
 
@@ -248,11 +259,21 @@ class PipelineFilter(WorkspaceScopedTaggableFilter):
     """Pipeline filter model."""
 
     CUSTOM_SORTING_OPTIONS = [SORT_PIPELINES_BY_LATEST_RUN_KEY]
+    FILTER_EXCLUDE_FIELDS: ClassVar[List[str]] = [
+        *WorkspaceScopedTaggableFilter.FILTER_EXCLUDE_FIELDS,
+        "user",
+        "latest_run_status",
+    ]
 
     name: Optional[str] = Field(
         default=None,
         description="Name of the Pipeline",
     )
+    latest_run_status: Optional[str] = Field(
+        default=None,
+        description="Filter by the status of the latest run of a pipeline. "
+        "This will always be applied as an `AND` filter for now.",
+    )
     workspace_id: Optional[Union[UUID, str]] = Field(
         default=None,
         description="Workspace of the Pipeline",
@@ -263,6 +284,92 @@ class PipelineFilter(WorkspaceScopedTaggableFilter):
         description="User of the Pipeline",
         union_mode="left_to_right",
     )
+    user: Optional[Union[UUID, str]] = Field(
+        default=None,
+        description="Name/ID of the user that created the pipeline.",
+    )
+
+    def apply_filter(
+        self, query: AnyQuery, table: Type["AnySchema"]
+    ) -> AnyQuery:
+        """Applies the filter to a query.
+
+        Args:
+            query: The query to which to apply the filter.
+            table: The query table.
+
+        Returns:
+            The query with filter applied.
+        """
+        query = super().apply_filter(query, table)
+
+        from sqlmodel import and_, col, func, select
+
+        from zenml.zen_stores.schemas import PipelineRunSchema, PipelineSchema
+
+        if self.latest_run_status:
+            latest_pipeline_run_subquery = (
+                select(
+                    PipelineRunSchema.pipeline_id,
+                    func.max(PipelineRunSchema.created).label("created"),
+                )
+                .where(col(PipelineRunSchema.pipeline_id).is_not(None))
+                .group_by(col(PipelineRunSchema.pipeline_id))
+                .subquery()
+            )
+
+            query = (
+                query.join(
+                    PipelineRunSchema,
+                    PipelineSchema.id == PipelineRunSchema.pipeline_id,
+                )
+                .join(
+                    latest_pipeline_run_subquery,
+                    and_(
+                        PipelineRunSchema.pipeline_id
+                        == latest_pipeline_run_subquery.c.pipeline_id,
+                        PipelineRunSchema.created
+                        == latest_pipeline_run_subquery.c.created,
+                    ),
+                )
+                .where(
+                    self.generate_custom_query_conditions_for_column(
+                        value=self.latest_run_status,
+                        table=PipelineRunSchema,
+                        column="status",
+                    )
+                )
+            )
+
+        return query
+
+    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 (
+            PipelineSchema,
+            UserSchema,
+        )
+
+        if self.user:
+            user_filter = and_(
+                PipelineSchema.user_id == UserSchema.id,
+                self.generate_name_or_id_query_conditions(
+                    value=self.user, table=UserSchema
+                ),
+            )
+            custom_filters.append(user_filter)
+
+        return custom_filters
 
     def apply_sorting(
         self,