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

zenml/models/v2/base/filter.py

@@ -54,7 +54,7 @@ from zenml.utils.typing_utils import get_args
 if TYPE_CHECKING:
     from sqlalchemy.sql.elements import ColumnElement
 
-    from zenml.zen_stores.schemas import BaseSchema
+    from zenml.zen_stores.schemas import BaseSchema, NamedSchema
 
     AnySchema = TypeVar("AnySchema", bound=BaseSchema)
 
@@ -142,7 +142,10 @@ class Filter(BaseModel, ABC):
 class BoolFilter(Filter):
     """Filter for all Boolean fields."""
 
-    ALLOWED_OPS: ClassVar[List[str]] = [GenericFilterOps.EQUALS]
+    ALLOWED_OPS: ClassVar[List[str]] = [
+        GenericFilterOps.EQUALS,
+        GenericFilterOps.NOT_EQUALS,
+    ]
 
     def generate_query_conditions_from_column(self, column: Any) -> Any:
         """Generate query conditions for a boolean column.
@@ -153,6 +156,9 @@ class BoolFilter(Filter):
         Returns:
             A list of query conditions.
         """
+        if self.operation == GenericFilterOps.NOT_EQUALS:
+            return column != self.value
+
         return column == self.value
 
 
@@ -161,6 +167,7 @@ class StrFilter(Filter):
 
     ALLOWED_OPS: ClassVar[List[str]] = [
         GenericFilterOps.EQUALS,
+        GenericFilterOps.NOT_EQUALS,
         GenericFilterOps.STARTSWITH,
         GenericFilterOps.CONTAINS,
         GenericFilterOps.ENDSWITH,
@@ -181,12 +188,31 @@ class StrFilter(Filter):
             return column.startswith(f"{self.value}")
         if self.operation == GenericFilterOps.ENDSWITH:
             return column.endswith(f"{self.value}")
+        if self.operation == GenericFilterOps.NOT_EQUALS:
+            return column != self.value
+
         return column == self.value
 
 
 class UUIDFilter(StrFilter):
     """Filter for all uuid fields which are mostly treated like strings."""
 
+    @field_validator("value", mode="before")
+    @classmethod
+    def _remove_hyphens_from_value(cls, value: Any) -> Any:
+        """Remove hyphens from the value to enable string comparisons.
+
+        Args:
+            value: The filter value.
+
+        Returns:
+            The filter value with removed hyphens.
+        """
+        if isinstance(value, str):
+            return value.replace("-", "")
+
+        return value
+
     def generate_query_conditions_from_column(self, column: Any) -> Any:
         """Generate query conditions for a UUID column.
 
@@ -203,6 +229,9 @@ class UUIDFilter(StrFilter):
         if self.operation == GenericFilterOps.EQUALS:
             return column == self.value
 
+        if self.operation == GenericFilterOps.NOT_EQUALS:
+            return column != self.value
+
         # For all other operations, cast and handle the column as string
         return super().generate_query_conditions_from_column(
             column=cast_if(column, sqlalchemy.String)
@@ -216,6 +245,7 @@ class NumericFilter(Filter):
 
     ALLOWED_OPS: ClassVar[List[str]] = [
         GenericFilterOps.EQUALS,
+        GenericFilterOps.NOT_EQUALS,
         GenericFilterOps.GT,
         GenericFilterOps.GTE,
         GenericFilterOps.LT,
@@ -223,14 +253,59 @@ class NumericFilter(Filter):
     ]
 
     def generate_query_conditions_from_column(self, column: Any) -> Any:
-        """Generate query conditions for a UUID column.
+        """Generate query conditions for a numeric column.
 
         Args:
-            column: The UUID column of an SQLModel table on which to filter.
+            column: The numeric column of an SQLModel table on which to filter.
+
+        Returns:
+            A list of query conditions.
+        """
+        if self.operation == GenericFilterOps.GTE:
+            return column >= self.value
+        if self.operation == GenericFilterOps.GT:
+            return column > self.value
+        if self.operation == GenericFilterOps.LTE:
+            return column <= self.value
+        if self.operation == GenericFilterOps.LT:
+            return column < self.value
+        if self.operation == GenericFilterOps.NOT_EQUALS:
+            return column != self.value
+        return column == self.value
+
+
+class DatetimeFilter(Filter):
+    """Filter for all datetime fields."""
+
+    value: Union[datetime, Tuple[datetime, datetime]] = Field(
+        union_mode="left_to_right"
+    )
+
+    ALLOWED_OPS: ClassVar[List[str]] = [
+        GenericFilterOps.EQUALS,
+        GenericFilterOps.NOT_EQUALS,
+        GenericFilterOps.GT,
+        GenericFilterOps.GTE,
+        GenericFilterOps.LT,
+        GenericFilterOps.LTE,
+        GenericFilterOps.IN,
+    ]
+
+    def generate_query_conditions_from_column(self, column: Any) -> Any:
+        """Generate query conditions for a datetime column.
+
+        Args:
+            column: The datetime column of an SQLModel table on which to filter.
 
         Returns:
             A list of query conditions.
         """
+        if self.operation == GenericFilterOps.IN:
+            assert isinstance(self.value, tuple)
+            lower_bound, upper_bound = self.value
+            return column.between(lower_bound, upper_bound)
+
+        assert isinstance(self.value, datetime)
         if self.operation == GenericFilterOps.GTE:
             return column >= self.value
         if self.operation == GenericFilterOps.GT:
@@ -239,6 +314,8 @@ class NumericFilter(Filter):
             return column <= self.value
         if self.operation == GenericFilterOps.LT:
             return column < self.value
+        if self.operation == GenericFilterOps.NOT_EQUALS:
+            return column != self.value
         return column == self.value
 
 
@@ -490,7 +567,7 @@ class BaseFilter(BaseModel):
             value, operator = cls._resolve_operator(value)
 
             # Define the filter
-            filter = cls._define_filter(
+            filter = FilterGenerator(cls).define_filter(
                 column=key, value=value, operator=operator
             )
             list_of_filters.append(filter)
@@ -523,9 +600,185 @@ class BaseFilter(BaseModel):
                 operator = GenericFilterOps(split_value[0])
         return value, operator
 
-    @classmethod
-    def _define_filter(
-        cls, column: str, value: Any, operator: GenericFilterOps
+    def generate_name_or_id_query_conditions(
+        self,
+        value: Union[UUID, str],
+        table: Type["NamedSchema"],
+    ) -> "ColumnElement[bool]":
+        """Generate filter conditions for name or id of a table.
+
+        Args:
+            value: The filter value.
+            table: The table to filter.
+
+        Returns:
+            The query conditions.
+        """
+        from sqlmodel import or_
+
+        value, operator = BaseFilter._resolve_operator(value)
+        value = str(value)
+
+        conditions = []
+
+        try:
+            filter_ = FilterGenerator(table).define_filter(
+                column="id", value=value, operator=operator
+            )
+            conditions.append(filter_.generate_query_conditions(table=table))
+        except ValueError:
+            # UUID filter with equal operators and no full UUID fail with
+            # a ValueError. In this case, we already know that the filter
+            # will not produce any result and can simply ignore it.
+            pass
+
+        filter_ = FilterGenerator(table).define_filter(
+            column="name", value=value, operator=operator
+        )
+        conditions.append(filter_.generate_query_conditions(table=table))
+
+        return or_(*conditions)
+
+    def generate_custom_query_conditions_for_column(
+        self,
+        value: Any,
+        table: Type[SQLModel],
+        column: str,
+    ) -> "ColumnElement[bool]":
+        """Generate custom filter conditions for a column of a table.
+
+        Args:
+            value: The filter value.
+            table: The table which contains the column.
+            column: The column name.
+
+        Returns:
+            The query conditions.
+        """
+        value, operator = BaseFilter._resolve_operator(value)
+        filter_ = FilterGenerator(table).define_filter(
+            column=column, value=value, operator=operator
+        )
+        return filter_.generate_query_conditions(table=table)
+
+    @property
+    def offset(self) -> int:
+        """Returns the offset needed for the query on the data persistence layer.
+
+        Returns:
+            The offset for the query.
+        """
+        return self.size * (self.page - 1)
+
+    def generate_filter(
+        self, table: Type[SQLModel]
+    ) -> Union["ColumnElement[bool]"]:
+        """Generate the filter for the query.
+
+        Args:
+            table: The Table that is being queried from.
+
+        Returns:
+            The filter expression for the query.
+
+        Raises:
+            RuntimeError: If a valid logical operator is not supplied.
+        """
+        from sqlmodel import and_, or_
+
+        filters = []
+        for column_filter in self.list_of_filters:
+            filters.append(
+                column_filter.generate_query_conditions(table=table)
+            )
+        for custom_filter in self.get_custom_filters():
+            filters.append(custom_filter)
+        if self.logical_operator == LogicalOperators.OR:
+            return or_(False, *filters)
+        elif self.logical_operator == LogicalOperators.AND:
+            return and_(True, *filters)
+        else:
+            raise RuntimeError("No valid logical operator was supplied.")
+
+    def get_custom_filters(self) -> List["ColumnElement[bool]"]:
+        """Get custom filters.
+
+        This can be overridden by subclasses to define custom filters that are
+        not based on the columns of the underlying table.
+
+        Returns:
+            A list of custom filters.
+        """
+        return []
+
+    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.
+        """
+        rbac_filter = self.generate_rbac_filter(table=table)
+
+        if rbac_filter is not None:
+            query = query.where(rbac_filter)
+
+        filters = self.generate_filter(table=table)
+
+        if filters is not None:
+            query = query.where(filters)
+
+        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.
+        """
+        column, operand = self.sorting_params
+
+        if operand == SorterOps.DESCENDING:
+            sort_clause = desc(getattr(table, column))  # type: ignore[var-annotated]
+        else:
+            sort_clause = asc(getattr(table, column))
+
+        # We always add the `id` column as a tiebreaker to ensure a stable,
+        # repeatable order of items, otherwise subsequent pages might contain
+        # the same items.
+        query = query.order_by(sort_clause, asc(table.id))  # type: ignore[arg-type]
+
+        return query
+
+
+class FilterGenerator:
+    """Helper class to define filters for a class."""
+
+    def __init__(self, model_class: Type[BaseModel]) -> None:
+        """Initialize the object.
+
+        Args:
+            model_class: The model class for which to define filters.
+        """
+        self._model_class = model_class
+
+    def define_filter(
+        self, column: str, value: Any, operator: GenericFilterOps
     ) -> Filter:
         """Define a filter for a given column.
 
@@ -538,23 +791,23 @@ class BaseFilter(BaseModel):
             A Filter object.
         """
         # Create datetime filters
-        if cls.is_datetime_field(column):
-            return cls._define_datetime_filter(
+        if self.is_datetime_field(column):
+            return self._define_datetime_filter(
                 column=column,
                 value=value,
                 operator=operator,
             )
 
         # Create UUID filters
-        if cls.is_uuid_field(column):
-            return cls._define_uuid_filter(
+        if self.is_uuid_field(column):
+            return self._define_uuid_filter(
                 column=column,
                 value=value,
                 operator=operator,
             )
 
         # Create int filters
-        if cls.is_int_field(column):
+        if self.is_int_field(column):
             return NumericFilter(
                 operation=GenericFilterOps(operator),
                 column=column,
@@ -562,15 +815,15 @@ class BaseFilter(BaseModel):
             )
 
         # Create bool filters
-        if cls.is_bool_field(column):
-            return cls._define_bool_filter(
+        if self.is_bool_field(column):
+            return self._define_bool_filter(
                 column=column,
                 value=value,
                 operator=operator,
             )
 
         # Create str filters
-        if cls.is_str_field(column):
+        if self.is_str_field(column):
             return StrFilter(
                 operation=GenericFilterOps(operator),
                 column=column,
@@ -579,8 +832,8 @@ class BaseFilter(BaseModel):
 
         # Handle unsupported datatypes
         logger.warning(
-            f"The Datatype {cls.model_fields[column].annotation} might not be "
-            "supported for filtering. Defaulting to a string filter."
+            f"The Datatype {self._model_class.model_fields[column].annotation} might "
+            "not be supported for filtering. Defaulting to a string filter."
         )
         return StrFilter(
             operation=GenericFilterOps(operator),
@@ -588,8 +841,7 @@ class BaseFilter(BaseModel):
             value=str(value),
         )
 
-    @classmethod
-    def check_field_annotation(cls, k: str, type_: Any) -> bool:
+    def check_field_annotation(self, k: str, type_: Any) -> bool:
         """Checks whether a model field has a certain annotation.
 
         Args:
@@ -604,7 +856,7 @@ class BaseFilter(BaseModel):
             otherwise.
         """
         try:
-            annotation = cls.model_fields[k].annotation
+            annotation = self._model_class.model_fields[k].annotation
 
             if annotation is not None:
                 return (
@@ -613,14 +865,13 @@ class BaseFilter(BaseModel):
                 )
             else:
                 raise ValueError(
-                    f"The field '{k}' inside the model {cls.__name__} "
+                    f"The field '{k}' inside the model {self._model_class.__name__} "
                     "does not have an annotation."
                 )
         except TypeError:
             return False
 
-    @classmethod
-    def is_datetime_field(cls, k: str) -> bool:
+    def is_datetime_field(self, k: str) -> bool:
         """Checks if it's a datetime field.
 
         Args:
@@ -629,10 +880,9 @@ class BaseFilter(BaseModel):
         Returns:
             True if the field is a datetime field, False otherwise.
         """
-        return cls.check_field_annotation(k=k, type_=datetime)
+        return self.check_field_annotation(k=k, type_=datetime)
 
-    @classmethod
-    def is_uuid_field(cls, k: str) -> bool:
+    def is_uuid_field(self, k: str) -> bool:
         """Checks if it's a UUID field.
 
         Args:
@@ -641,10 +891,9 @@ class BaseFilter(BaseModel):
         Returns:
             True if the field is a UUID field, False otherwise.
         """
-        return cls.check_field_annotation(k=k, type_=UUID)
+        return self.check_field_annotation(k=k, type_=UUID)
 
-    @classmethod
-    def is_int_field(cls, k: str) -> bool:
+    def is_int_field(self, k: str) -> bool:
         """Checks if it's an int field.
 
         Args:
@@ -653,10 +902,9 @@ class BaseFilter(BaseModel):
         Returns:
             True if the field is an int field, False otherwise.
         """
-        return cls.check_field_annotation(k=k, type_=int)
+        return self.check_field_annotation(k=k, type_=int)
 
-    @classmethod
-    def is_bool_field(cls, k: str) -> bool:
+    def is_bool_field(self, k: str) -> bool:
         """Checks if it's a bool field.
 
         Args:
@@ -665,10 +913,9 @@ class BaseFilter(BaseModel):
         Returns:
             True if the field is a bool field, False otherwise.
         """
-        return cls.check_field_annotation(k=k, type_=bool)
+        return self.check_field_annotation(k=k, type_=bool)
 
-    @classmethod
-    def is_str_field(cls, k: str) -> bool:
+    def is_str_field(self, k: str) -> bool:
         """Checks if it's a string field.
 
         Args:
@@ -677,10 +924,9 @@ class BaseFilter(BaseModel):
         Returns:
             True if the field is a string field, False otherwise.
         """
-        return cls.check_field_annotation(k=k, type_=str)
+        return self.check_field_annotation(k=k, type_=str)
 
-    @classmethod
-    def is_sort_by_field(cls, k: str) -> bool:
+    def is_sort_by_field(self, k: str) -> bool:
         """Checks if it's a sort by field.
 
         Args:
@@ -689,12 +935,12 @@ class BaseFilter(BaseModel):
         Returns:
             True if the field is a sort by field, False otherwise.
         """
-        return cls.check_field_annotation(k=k, type_=str) and k == "sort_by"
+        return self.check_field_annotation(k=k, type_=str) and k == "sort_by"
 
     @staticmethod
     def _define_datetime_filter(
         column: str, value: Any, operator: GenericFilterOps
-    ) -> NumericFilter:
+    ) -> DatetimeFilter:
         """Define a datetime filter for a given column.
 
         Args:
@@ -709,10 +955,17 @@ class BaseFilter(BaseModel):
             ValueError: If the value is not a valid datetime.
         """
         try:
+            filter_value: Union[datetime, Tuple[datetime, datetime]]
             if isinstance(value, datetime):
-                datetime_value = value
+                filter_value = value
+            elif "," in value:
+                lower_bound, upper_bound = value.split(",", 1)
+                filter_value = (
+                    datetime.strptime(lower_bound, FILTERING_DATETIME_FORMAT),
+                    datetime.strptime(upper_bound, FILTERING_DATETIME_FORMAT),
+                )
             else:
-                datetime_value = datetime.strptime(
+                filter_value = datetime.strptime(
                     value, FILTERING_DATETIME_FORMAT
                 )
         except ValueError as e:
@@ -720,10 +973,27 @@ class BaseFilter(BaseModel):
                 "The datetime filter only works with values in the following "
                 f"format: {FILTERING_DATETIME_FORMAT}"
             ) from e
-        datetime_filter = NumericFilter(
+
+        if operator == GenericFilterOps.IN and not isinstance(
+            filter_value, tuple
+        ):
+            raise ValueError(
+                "Two comma separated datetime values are required for the `in` "
+                "operator."
+            )
+
+        if operator != GenericFilterOps.IN and not isinstance(
+            filter_value, datetime
+        ):
+            raise ValueError(
+                "Only a single datetime value is allowed for operator "
+                f"{operator}."
+            )
+
+        datetime_filter = DatetimeFilter(
             operation=GenericFilterOps(operator),
             column=column,
-            value=datetime_value,
+            value=filter_value,
         )
         return datetime_filter
 
@@ -789,107 +1059,3 @@ class BaseFilter(BaseModel):
             column=column,
             value=bool(value),
         )
-
-    @property
-    def offset(self) -> int:
-        """Returns the offset needed for the query on the data persistence layer.
-
-        Returns:
-            The offset for the query.
-        """
-        return self.size * (self.page - 1)
-
-    def generate_filter(
-        self, table: Type[SQLModel]
-    ) -> Union["ColumnElement[bool]"]:
-        """Generate the filter for the query.
-
-        Args:
-            table: The Table that is being queried from.
-
-        Returns:
-            The filter expression for the query.
-
-        Raises:
-            RuntimeError: If a valid logical operator is not supplied.
-        """
-        from sqlmodel import and_, or_
-
-        filters = []
-        for column_filter in self.list_of_filters:
-            filters.append(
-                column_filter.generate_query_conditions(table=table)
-            )
-        for custom_filter in self.get_custom_filters():
-            filters.append(custom_filter)
-        if self.logical_operator == LogicalOperators.OR:
-            return or_(False, *filters)
-        elif self.logical_operator == LogicalOperators.AND:
-            return and_(True, *filters)
-        else:
-            raise RuntimeError("No valid logical operator was supplied.")
-
-    def get_custom_filters(self) -> List["ColumnElement[bool]"]:
-        """Get custom filters.
-
-        This can be overridden by subclasses to define custom filters that are
-        not based on the columns of the underlying table.
-
-        Returns:
-            A list of custom filters.
-        """
-        return []
-
-    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.
-        """
-        rbac_filter = self.generate_rbac_filter(table=table)
-
-        if rbac_filter is not None:
-            query = query.where(rbac_filter)
-
-        filters = self.generate_filter(table=table)
-
-        if filters is not None:
-            query = query.where(filters)
-
-        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.
-        """
-        column, operand = self.sorting_params
-
-        if operand == SorterOps.DESCENDING:
-            sort_clause = desc(getattr(table, column))  # type: ignore[var-annotated]
-        else:
-            sort_clause = asc(getattr(table, column))
-
-        # We always add the `id` column as a tiebreaker to ensure a stable,
-        # repeatable order of items, otherwise subsequent pages might contain
-        # the same items.
-        query = query.order_by(sort_clause, asc(table.id))  # type: ignore[arg-type]
-
-        return query

zenml/models/v2/base/scoped.py

@@ -27,7 +27,6 @@ from typing import (
 from uuid import UUID
 
 from pydantic import Field
-from sqlmodel import col
 
 from zenml.models.v2.base.base import (
     BaseDatedResponseBody,
@@ -341,6 +340,10 @@ class WorkspaceScopedTaggableFilter(WorkspaceScopedFilter):
 
         custom_filters = super().get_custom_filters()
         if self.tag:
-            custom_filters.append(col(TagSchema.name) == self.tag)
+            custom_filters.append(
+                self.generate_custom_query_conditions_for_column(
+                    value=self.tag, table=TagSchema, column="name"
+                )
+            )
 
         return custom_filters