sqlspec 0.10.1__py3-none-any.whl → 0.11.1__py3-none-any.whl

sqlspec/filters.py

@@ -1,16 +1,20 @@
 """Collection filter datastructures."""
 
-from abc import ABC
+from abc import ABC, abstractmethod
 from collections import abc
 from dataclasses import dataclass
 from datetime import datetime
-from typing import Generic, Literal, Optional, Protocol, Union
+from typing import Any, Generic, Literal, Optional, Protocol, Union, cast, runtime_checkable
 
-from typing_extensions import TypeVar
+from sqlglot import exp
+from typing_extensions import TypeAlias, TypeVar
+
+from sqlspec.statement import SQLStatement
 
 __all__ = (
     "BeforeAfter",
     "CollectionFilter",
+    "FilterTypes",
     "InAnyFilter",
     "LimitOffset",
     "NotInCollectionFilter",
@@ -20,18 +24,27 @@ __all__ = (
     "PaginationFilter",
     "SearchFilter",
     "StatementFilter",
+    "apply_filter",
 )
 
 T = TypeVar("T")
-StatementT = TypeVar("StatementT", bound="str")
 
 
+@runtime_checkable
 class StatementFilter(Protocol):
     """Protocol for filters that can be appended to a statement."""
 
-    def append_to_statement(self, statement: StatementT) -> StatementT:
-        """Append the filter to the statement."""
-        return statement
+    @abstractmethod
+    def append_to_statement(self, statement: SQLStatement) -> SQLStatement:
+        """Append the filter to the statement.
+
+        Args:
+            statement: The SQL statement to modify.
+
+        Returns:
+            The modified statement.
+        """
+        raise NotImplementedError
 
 
 @dataclass
@@ -45,6 +58,27 @@ class BeforeAfter(StatementFilter):
     after: Optional[datetime] = None
     """Filter results where field later than this."""
 
+    def append_to_statement(self, statement: SQLStatement) -> SQLStatement:
+        conditions = []
+        params: dict[str, Any] = {}
+        col_expr = exp.column(self.field_name)
+
+        if self.before:
+            param_name = statement.generate_param_name(f"{self.field_name}_before")
+            conditions.append(exp.LT(this=col_expr, expression=exp.Placeholder(this=param_name)))
+            params[param_name] = self.before
+        if self.after:
+            param_name = statement.generate_param_name(f"{self.field_name}_after")
+            conditions.append(exp.GT(this=col_expr, expression=exp.Placeholder(this=param_name)))  # type: ignore[arg-type]
+            params[param_name] = self.after
+
+        if conditions:
+            final_condition = conditions[0]
+            for cond in conditions[1:]:
+                final_condition = exp.And(this=final_condition, expression=cond)  # type: ignore[assignment]
+            statement.add_condition(final_condition, params)
+        return statement
+
 
 @dataclass
 class OnBeforeAfter(StatementFilter):
@@ -57,13 +91,38 @@ class OnBeforeAfter(StatementFilter):
     on_or_after: Optional[datetime] = None
     """Filter results where field on or later than this."""
 
+    def append_to_statement(self, statement: SQLStatement) -> SQLStatement:
+        conditions = []
+        params: dict[str, Any] = {}
+        col_expr = exp.column(self.field_name)
+
+        if self.on_or_before:
+            param_name = statement.generate_param_name(f"{self.field_name}_on_or_before")
+            conditions.append(exp.LTE(this=col_expr, expression=exp.Placeholder(this=param_name)))
+            params[param_name] = self.on_or_before
+        if self.on_or_after:
+            param_name = statement.generate_param_name(f"{self.field_name}_on_or_after")
+            conditions.append(exp.GTE(this=col_expr, expression=exp.Placeholder(this=param_name)))  # type: ignore[arg-type]
+            params[param_name] = self.on_or_after
+
+        if conditions:
+            final_condition = conditions[0]
+            for cond in conditions[1:]:
+                final_condition = exp.And(this=final_condition, expression=cond)  # type: ignore[assignment]
+            statement.add_condition(final_condition, params)
+        return statement
 
-class InAnyFilter(StatementFilter, ABC):
+
+class InAnyFilter(StatementFilter, ABC, Generic[T]):
     """Subclass for methods that have a `prefer_any` attribute."""
 
+    @abstractmethod
+    def append_to_statement(self, statement: SQLStatement) -> SQLStatement:
+        raise NotImplementedError
+
 
 @dataclass
-class CollectionFilter(InAnyFilter, Generic[T]):
+class CollectionFilter(InAnyFilter[T]):
     """Data required to construct a ``WHERE ... IN (...)`` clause."""
 
     field_name: str
@@ -73,9 +132,30 @@ class CollectionFilter(InAnyFilter, Generic[T]):
 
     An empty list will return an empty result set, however, if ``None``, the filter is not applied to the query, and all rows are returned. """
 
+    def append_to_statement(self, statement: SQLStatement) -> SQLStatement:
+        if self.values is None:
+            return statement
+
+        if not self.values:  # Empty collection
+            # Add a condition that is always false
+            statement.add_condition(exp.false())
+            return statement
+
+        placeholder_expressions: list[exp.Placeholder] = []
+        current_params: dict[str, Any] = {}
+
+        for i, value_item in enumerate(self.values):
+            param_key = statement.generate_param_name(f"{self.field_name}_in_{i}")
+            placeholder_expressions.append(exp.Placeholder(this=param_key))
+            current_params[param_key] = value_item
+
+        in_condition = exp.In(this=exp.column(self.field_name), expressions=placeholder_expressions)
+        statement.add_condition(in_condition, current_params)
+        return statement
+
 
 @dataclass
-class NotInCollectionFilter(InAnyFilter, Generic[T]):
+class NotInCollectionFilter(InAnyFilter[T]):
     """Data required to construct a ``WHERE ... NOT IN (...)`` clause."""
 
     field_name: str
@@ -85,10 +165,31 @@ class NotInCollectionFilter(InAnyFilter, Generic[T]):
 
     An empty list or ``None`` will return all rows."""
 
+    def append_to_statement(self, statement: SQLStatement) -> SQLStatement:
+        if self.values is None or not self.values:  # Empty list or None, no filter applied
+            return statement
+
+        placeholder_expressions: list[exp.Placeholder] = []
+        current_params: dict[str, Any] = {}
+
+        for i, value_item in enumerate(self.values):
+            param_key = statement.generate_param_name(f"{self.field_name}_notin_{i}")
+            placeholder_expressions.append(exp.Placeholder(this=param_key))
+            current_params[param_key] = value_item
+
+        in_expr = exp.In(this=exp.column(self.field_name), expressions=placeholder_expressions)
+        not_in_condition = exp.Not(this=in_expr)
+        statement.add_condition(not_in_condition, current_params)
+        return statement
+
 
 class PaginationFilter(StatementFilter, ABC):
     """Subclass for methods that function as a pagination type."""
 
+    @abstractmethod
+    def append_to_statement(self, statement: SQLStatement) -> SQLStatement:
+        raise NotImplementedError
+
 
 @dataclass
 class LimitOffset(PaginationFilter):
@@ -99,6 +200,16 @@ class LimitOffset(PaginationFilter):
     offset: int
     """Value for ``OFFSET`` clause of query."""
 
+    def append_to_statement(self, statement: SQLStatement) -> SQLStatement:
+        # Generate parameter names for limit and offset
+        limit_param_name = statement.generate_param_name("limit_val")
+        offset_param_name = statement.generate_param_name("offset_val")
+
+        statement.add_limit(self.limit, param_name=limit_param_name)
+        statement.add_offset(self.offset, param_name=offset_param_name)
+
+        return statement
+
 
 @dataclass
 class OrderBy(StatementFilter):
@@ -109,6 +220,16 @@ class OrderBy(StatementFilter):
     sort_order: Literal["asc", "desc"] = "asc"
     """Sort ascending or descending"""
 
+    def append_to_statement(self, statement: SQLStatement) -> SQLStatement:
+        # Basic validation for sort_order, though Literal helps at type checking time
+        normalized_sort_order = self.sort_order.lower()
+        if normalized_sort_order not in {"asc", "desc"}:
+            normalized_sort_order = "asc"
+
+        statement.add_order_by(self.field_name, direction=cast("Literal['asc', 'desc']", normalized_sort_order))
+
+        return statement
+
 
 @dataclass
 class SearchFilter(StatementFilter):
@@ -121,7 +242,90 @@ class SearchFilter(StatementFilter):
     ignore_case: Optional[bool] = False
     """Should the search be case insensitive."""
 
+    def append_to_statement(self, statement: SQLStatement) -> SQLStatement:
+        if not self.value:
+            return statement
+
+        search_val_param_name = statement.generate_param_name("search_val")
+
+        # The pattern %value% needs to be handled carefully.
+        params = {search_val_param_name: f"%{self.value}%"}
+        pattern_expr = exp.Placeholder(this=search_val_param_name)
+
+        like_op = exp.ILike if self.ignore_case else exp.Like
+
+        if isinstance(self.field_name, str):
+            condition = like_op(this=exp.column(self.field_name), expression=pattern_expr)
+            statement.add_condition(condition, params)
+        elif isinstance(self.field_name, set) and self.field_name:
+            field_conditions = [like_op(this=exp.column(field), expression=pattern_expr) for field in self.field_name]
+            if not field_conditions:
+                return statement
+
+            final_condition = field_conditions[0]
+            for cond in field_conditions[1:]:
+                final_condition = exp.Or(this=final_condition, expression=cond)  # type: ignore[assignment]
+            statement.add_condition(final_condition, params)
+
+        return statement
+
 
 @dataclass
-class NotInSearchFilter(SearchFilter):
+class NotInSearchFilter(SearchFilter):  # Inherits field_name, value, ignore_case
     """Data required to construct a ``WHERE field_name NOT LIKE '%' || :value || '%'`` clause."""
+
+    def append_to_statement(self, statement: SQLStatement) -> SQLStatement:
+        if not self.value:
+            return statement
+
+        search_val_param_name = statement.generate_param_name("not_search_val")
+
+        params = {search_val_param_name: f"%{self.value}%"}
+        pattern_expr = exp.Placeholder(this=search_val_param_name)
+
+        like_op = exp.ILike if self.ignore_case else exp.Like
+
+        if isinstance(self.field_name, str):
+            condition = exp.Not(this=like_op(this=exp.column(self.field_name), expression=pattern_expr))
+            statement.add_condition(condition, params)
+        elif isinstance(self.field_name, set) and self.field_name:
+            field_conditions = [
+                exp.Not(this=like_op(this=exp.column(field), expression=pattern_expr)) for field in self.field_name
+            ]
+            if not field_conditions:
+                return statement
+
+            # Combine with AND: (field1 NOT LIKE pattern) AND (field2 NOT LIKE pattern) ...
+            final_condition = field_conditions[0]
+            for cond in field_conditions[1:]:
+                final_condition = exp.And(this=final_condition, expression=cond)  # type: ignore[assignment]
+            statement.add_condition(final_condition, params)
+
+        return statement
+
+
+# Function to be imported in SQLStatement module
+def apply_filter(statement: SQLStatement, filter_obj: StatementFilter) -> SQLStatement:
+    """Apply a statement filter to a SQL statement.
+
+    Args:
+        statement: The SQL statement to modify.
+        filter_obj: The filter to apply.
+
+    Returns:
+        The modified statement.
+    """
+    return filter_obj.append_to_statement(statement)
+
+
+FilterTypes: TypeAlias = Union[
+    BeforeAfter,
+    OnBeforeAfter,
+    CollectionFilter[Any],
+    LimitOffset,
+    OrderBy,
+    SearchFilter,
+    NotInCollectionFilter[Any],
+    NotInSearchFilter,
+]
+"""Aggregate type alias of the types supported for collection filtering."""

sqlspec/mixins.py

@@ -1,19 +1,40 @@
+import datetime
 from abc import abstractmethod
+from collections.abc import Sequence
+from enum import Enum
+from functools import partial
+from pathlib import Path, PurePath
 from typing import (
     TYPE_CHECKING,
     Any,
+    Callable,
     ClassVar,
     Generic,
     Optional,
+    Union,
+    cast,
+    overload,
 )
+from uuid import UUID
 
 from sqlglot import parse_one
 from sqlglot.dialects.dialect import DialectType
 
-from sqlspec.exceptions import SQLConversionError, SQLParsingError
-from sqlspec.typing import ConnectionT, StatementParameterType
+from sqlspec.exceptions import SQLConversionError, SQLParsingError, SQLSpecError
+from sqlspec.typing import (
+    ConnectionT,
+    ModelDTOT,
+    ModelT,
+    StatementParameterType,
+    convert,
+    get_type_adapter,
+    is_dataclass,
+    is_msgspec_struct,
+    is_pydantic_model,
+)
 
 if TYPE_CHECKING:
+    from sqlspec.filters import StatementFilter
     from sqlspec.typing import ArrowTable
 
 __all__ = (
@@ -31,12 +52,11 @@ class SyncArrowBulkOperationsMixin(Generic[ConnectionT]):
     __supports_arrow__: "ClassVar[bool]" = True
 
     @abstractmethod
-    def select_arrow(  # pyright: ignore[reportUnknownParameterType]
+    def select_arrow(
         self,
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
-        /,
-        *,
+        *filters: "StatementFilter",
         connection: "Optional[ConnectionT]" = None,
         **kwargs: Any,
     ) -> "ArrowTable":  # pyright: ignore[reportUnknownReturnType]
@@ -45,6 +65,7 @@ class SyncArrowBulkOperationsMixin(Generic[ConnectionT]):
         Args:
             sql: The SQL query string.
             parameters: Parameters for the query.
+            filters: Optional filters to apply to the query.
             connection: Optional connection override.
             **kwargs: Additional keyword arguments to merge with parameters if parameters is a dict.
 
@@ -60,12 +81,11 @@ class AsyncArrowBulkOperationsMixin(Generic[ConnectionT]):
     __supports_arrow__: "ClassVar[bool]" = True
 
     @abstractmethod
-    async def select_arrow(  # pyright: ignore[reportUnknownParameterType]
+    async def select_arrow(
         self,
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
-        /,
-        *,
+        *filters: "StatementFilter",
         connection: "Optional[ConnectionT]" = None,
         **kwargs: Any,
     ) -> "ArrowTable":  # pyright: ignore[reportUnknownReturnType]
@@ -74,6 +94,7 @@ class AsyncArrowBulkOperationsMixin(Generic[ConnectionT]):
         Args:
             sql: The SQL query string.
             parameters: Parameters for the query.
+            filters: Optional filters to apply to the query.
             connection: Optional connection override.
             **kwargs: Additional keyword arguments to merge with parameters if parameters is a dict.
 
@@ -91,8 +112,7 @@ class SyncParquetExportMixin(Generic[ConnectionT]):
         self,
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
-        /,
-        *,
+        *filters: "StatementFilter",
         connection: "Optional[ConnectionT]" = None,
         **kwargs: Any,
     ) -> None:
@@ -108,8 +128,7 @@ class AsyncParquetExportMixin(Generic[ConnectionT]):
         self,
         sql: str,
         parameters: "Optional[StatementParameterType]" = None,
-        /,
-        *,
+        *filters: "StatementFilter",
         connection: "Optional[ConnectionT]" = None,
         **kwargs: Any,
     ) -> None:
@@ -154,3 +173,133 @@ class SQLTranslatorMixin(Generic[ConnectionT]):
         except Exception as e:
             error_msg = f"Failed to convert SQL to {to_dialect}: {e!s}"
             raise SQLConversionError(error_msg) from e
+
+
+_DEFAULT_TYPE_DECODERS = [  # pyright: ignore[reportUnknownVariableType]
+    (lambda x: x is UUID, lambda t, v: t(v.hex)),  # pyright: ignore[reportUnknownLambdaType,reportUnknownMemberType]
+    (lambda x: x is datetime.datetime, lambda t, v: t(v.isoformat())),  # pyright: ignore[reportUnknownLambdaType,reportUnknownMemberType]
+    (lambda x: x is datetime.date, lambda t, v: t(v.isoformat())),  # pyright: ignore[reportUnknownLambdaType,reportUnknownMemberType]
+    (lambda x: x is datetime.time, lambda t, v: t(v.isoformat())),  # pyright: ignore[reportUnknownLambdaType,reportUnknownMemberType]
+    (lambda x: x is Enum, lambda t, v: t(v.value)),  # pyright: ignore[reportUnknownLambdaType,reportUnknownMemberType]
+]
+
+
+def _default_msgspec_deserializer(
+    target_type: Any,
+    value: Any,
+    type_decoders: "Union[Sequence[tuple[Callable[[Any], bool], Callable[[Any, Any], Any]]], None]" = None,
+) -> Any:  # pragma: no cover
+    """Transform values non-natively supported by ``msgspec``
+
+    Args:
+        target_type: Encountered type
+        value: Value to coerce
+        type_decoders: Optional sequence of type decoders
+
+    Raises:
+        TypeError: If the value cannot be coerced to the target type
+
+    Returns:
+        A ``msgspec``-supported type
+    """
+
+    if isinstance(value, target_type):
+        return value
+
+    if type_decoders:
+        for predicate, decoder in type_decoders:
+            if predicate(target_type):
+                return decoder(target_type, value)
+
+    if issubclass(target_type, (Path, PurePath, UUID)):
+        return target_type(value)
+
+    try:
+        return target_type(value)
+    except Exception as e:
+        msg = f"Unsupported type: {type(value)!r}"
+        raise TypeError(msg) from e
+
+
+class ResultConverter:
+    """Simple mixin to help convert to dictionary or list of dictionaries to specified schema type.
+
+    Single objects are transformed to the supplied schema type, and lists of objects are transformed into a list of the supplied schema type.
+
+    Args:
+        data: A database model instance or row mapping.
+              Type: :class:`~sqlspec.typing.ModelDictT`
+
+    Returns:
+        The converted schema object.
+    """
+
+    @overload
+    @staticmethod
+    def to_schema(data: "ModelT", *, schema_type: None = None) -> "ModelT": ...
+    @overload
+    @staticmethod
+    def to_schema(data: "dict[str, Any]", *, schema_type: "type[ModelDTOT]") -> "ModelDTOT": ...
+    @overload
+    @staticmethod
+    def to_schema(data: "Sequence[ModelT]", *, schema_type: None = None) -> "Sequence[ModelT]": ...
+    @overload
+    @staticmethod
+    def to_schema(data: "Sequence[dict[str, Any]]", *, schema_type: "type[ModelDTOT]") -> "Sequence[ModelDTOT]": ...
+
+    @staticmethod
+    def to_schema(
+        data: "Union[ModelT, Sequence[ModelT], dict[str, Any], Sequence[dict[str, Any]], ModelDTOT, Sequence[ModelDTOT]]",
+        *,
+        schema_type: "Optional[type[ModelDTOT]]" = None,
+    ) -> "Union[ModelT, Sequence[ModelT], ModelDTOT, Sequence[ModelDTOT]]":
+        if schema_type is None:
+            if not isinstance(data, Sequence):
+                return cast("ModelT", data)
+            return cast("Sequence[ModelT]", data)
+        if is_dataclass(schema_type):
+            if not isinstance(data, Sequence):
+                # data is assumed to be dict[str, Any] as per the method's overloads
+                return cast("ModelDTOT", schema_type(**data))  # type: ignore[operator]
+            # data is assumed to be Sequence[dict[str, Any]]
+            return cast("Sequence[ModelDTOT]", [schema_type(**item) for item in data])  # type: ignore[operator]
+        if is_msgspec_struct(schema_type):
+            if not isinstance(data, Sequence):
+                return cast(
+                    "ModelDTOT",
+                    convert(
+                        obj=data,
+                        type=schema_type,
+                        from_attributes=True,
+                        dec_hook=partial(
+                            _default_msgspec_deserializer,
+                            type_decoders=_DEFAULT_TYPE_DECODERS,
+                        ),
+                    ),
+                )
+            return cast(
+                "Sequence[ModelDTOT]",
+                convert(
+                    obj=data,
+                    type=list[schema_type],  # type: ignore[valid-type]
+                    from_attributes=True,
+                    dec_hook=partial(
+                        _default_msgspec_deserializer,
+                        type_decoders=_DEFAULT_TYPE_DECODERS,
+                    ),
+                ),
+            )
+
+        if schema_type is not None and is_pydantic_model(schema_type):
+            if not isinstance(data, Sequence):
+                return cast(
+                    "ModelDTOT",
+                    get_type_adapter(schema_type).validate_python(data, from_attributes=True),  # pyright: ignore
+                )
+            return cast(
+                "Sequence[ModelDTOT]",
+                get_type_adapter(list[schema_type]).validate_python(data, from_attributes=True),  # type: ignore[valid-type] # pyright: ignore[reportUnknownArgumentType]
+            )
+
+        msg = "`schema_type` should be a valid Dataclass, Pydantic model or Msgspec struct"
+        raise SQLSpecError(msg)