async-orm-kit 1.0.0__py3-none-any.whl

async_orm_kit/query/expressions/functions.py

@@ -0,0 +1,202 @@
+"""Typed, parameterized SQL function expressions with a portable core API.
+
+User-defined raw function names are intentionally not accepted: each wrapper
+chooses a reviewed SQLAlchemy function/operator and resolves operands through
+the same field resolver as ``F``, ``Case`` and aggregate expressions.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from sqlalchemy import cast, extract, func
+from sqlalchemy.sql.elements import ClauseElement
+from sqlalchemy.types import TypeEngine
+
+from async_orm_kit.query.lookups.resolver import ResolverContext
+from async_orm_kit.query.mixins.combinable import CombinableMixin
+from async_orm_kit.query.mixins.orderable import OrderableMixin
+
+
+def _operand(value: Any, field_resolver: Any, context: ResolverContext):
+    """Resolve ORM operands, preserving SQLAlchemy SQL and bound literals."""
+    from sqlalchemy import literal
+
+    from async_orm_kit.types import SAFE_SCALAR_TYPES
+
+    if isinstance(value, CombinableMixin) and hasattr(value, "resolve"):
+        return value.resolve(field_resolver=field_resolver, context=context)
+    if isinstance(value, ClauseElement):
+        return value
+    if isinstance(value, SAFE_SCALAR_TYPES):
+        return literal(value)
+    raise TypeError(f"Unsupported SQL function operand: {type(value).__name__}")
+
+
+class SQLFunction(CombinableMixin, OrderableMixin):
+    """Base class for a fixed, reviewed SQL function with immutable operands."""
+
+    _name: str
+    _arguments: tuple[Any, ...]
+    __slots__ = ("_name", "_arguments")
+    _allowed = frozenset(
+        {"coalesce", "nullif", "lower", "upper", "length", "trim", "round", "abs"}
+    )
+
+    def __init__(self, name: str, *arguments: Any) -> None:
+        if name not in self._allowed:
+            raise ValueError(f"Unsupported SQL function: {name!r}")
+        object.__setattr__(self, "_name", name)
+        object.__setattr__(self, "_arguments", tuple(arguments))
+
+    def __setattr__(self, key, value):
+        raise AttributeError("SQL function expressions are immutable")
+
+    def resolve(self, field_resolver, context: ResolverContext = ResolverContext()):
+        operands = [_operand(arg, field_resolver, context) for arg in self._arguments]
+        return getattr(func, self._name)(*operands)
+
+
+class Function(CombinableMixin, OrderableMixin):
+    """Extensible SQL function with a trusted SQLAlchemy expression factory.
+
+    A factory receives resolved, parameterized SQLAlchemy operands and must
+    return a SQLAlchemy ColumnElement. Calling application-owned functions is
+    intentional; the package never interpolates untrusted function names or SQL.
+
+    Example::
+        Function(lambda col: func.reverse(col), F("name"))
+    """
+
+    __slots__ = ("factory", "arguments")
+
+    def __init__(self, factory: Any, *arguments: Any):
+        if not callable(factory):
+            raise TypeError("Function factory must be callable")
+        self.factory = factory
+        self.arguments = tuple(arguments)
+
+    def resolve(self, field_resolver, context: ResolverContext = ResolverContext()):
+        from sqlalchemy.sql.elements import ColumnElement
+
+        values = [_operand(arg, field_resolver, context) for arg in self.arguments]
+        result = self.factory(*values)
+        if not isinstance(result, ColumnElement):
+            raise TypeError("Function factory must return a SQLAlchemy ColumnElement")
+        return result
+
+
+class Coalesce(SQLFunction):
+    """Return first non-NULL expression."""
+
+    def __init__(self, *expressions: Any):
+        if len(expressions) < 2:
+            raise ValueError("Coalesce requires at least two expressions")
+        super().__init__("coalesce", *expressions)
+
+
+class NullIf(SQLFunction):
+    """Return NULL when expressions compare equal."""
+
+    def __init__(self, first: Any, second: Any):
+        super().__init__("nullif", first, second)
+
+
+class Lower(SQLFunction):
+    def __init__(self, expression: Any):
+        super().__init__("lower", expression)
+
+
+class Upper(SQLFunction):
+    def __init__(self, expression: Any):
+        super().__init__("upper", expression)
+
+
+class Length(SQLFunction):
+    def __init__(self, expression: Any):
+        super().__init__("length", expression)
+
+
+class Trim(SQLFunction):
+    def __init__(self, expression: Any):
+        super().__init__("trim", expression)
+
+
+class Round(SQLFunction):
+    def __init__(self, expression: Any, precision: Any = 0):
+        super().__init__("round", expression, precision)
+
+
+class Abs(SQLFunction):
+    def __init__(self, expression: Any):
+        super().__init__("abs", expression)
+
+
+class Concat(CombinableMixin, OrderableMixin):
+    """Concatenate strings using the dialect's string concatenation operator."""
+
+    __slots__ = ("arguments",)
+
+    def __init__(self, *expressions: Any):
+        if len(expressions) < 2:
+            raise ValueError("Concat requires at least two expressions")
+        self.arguments = tuple(expressions)
+
+    def resolve(self, field_resolver, context: ResolverContext = ResolverContext()):
+        values = [_operand(arg, field_resolver, context) for arg in self.arguments]
+        out = values[0]
+        for value in values[1:]:
+            out = out.concat(value)
+        return out
+
+
+class Cast(CombinableMixin, OrderableMixin):
+    """Cast an expression to a validated SQLAlchemy SQL type."""
+
+    __slots__ = ("expression", "type_")
+
+    def __init__(self, expression: Any, type_: TypeEngine | type[TypeEngine]):
+        if isinstance(type_, type) and issubclass(type_, TypeEngine):
+            type_ = type_()
+        if not isinstance(type_, TypeEngine):
+            raise TypeError("Cast type_ must be a SQLAlchemy TypeEngine")
+        self.expression = expression
+        self.type_ = type_
+
+    def resolve(self, field_resolver, context: ResolverContext = ResolverContext()):
+        return cast(_operand(self.expression, field_resolver, context), self.type_)
+
+
+class Extract(CombinableMixin, OrderableMixin):
+    """Extract a supported date/time part with SQLAlchemy dialect compilation."""
+
+    __slots__ = ("part", "expression")
+    _parts = frozenset(
+        {"year", "month", "day", "hour", "minute", "second", "dow", "week", "quarter"}
+    )
+
+    def __init__(self, part: str, expression: Any):
+        if part not in self._parts:
+            raise ValueError(f"Unsupported Extract part: {part!r}")
+        self.part = part
+        self.expression = expression
+
+    def resolve(self, field_resolver, context: ResolverContext = ResolverContext()):
+        return extract(self.part, _operand(self.expression, field_resolver, context))
+
+
+__all__ = [
+    "Abs",
+    "Cast",
+    "Coalesce",
+    "Concat",
+    "Extract",
+    "Length",
+    "Lower",
+    "Function",
+    "NullIf",
+    "Round",
+    "SQLFunction",
+    "Trim",
+    "Upper",
+]

async_orm_kit/query/expressions/subquery.py

@@ -0,0 +1,206 @@
+"""Safe correlated subquery expressions for Django-inspired QuerySets.
+
+OuterRef markers are never rendered as user-supplied SQL. They are replaced by
+validated, mapped outer columns while a Subquery or Exists expression resolves.
+Outside a containing subquery, unresolved markers are rejected by QueryCompiler.
+"""
+
+from __future__ import annotations
+
+from sqlalchemy import column, inspect as sa_inspect
+from sqlalchemy.sql import visitors
+
+from async_orm_kit.exceptions import QueryStateError
+from async_orm_kit.query.lookups.resolver import ResolverContext
+from async_orm_kit.query.mixins.combinable import CombinableMixin
+from async_orm_kit.query.mixins.orderable import OrderableMixin
+
+_OUTER_REF = "_async_orm_kit_outer_ref"
+
+
+class OuterRef(CombinableMixin, OrderableMixin):
+    """Reference a mapped field on the *containing* queryset's model.
+
+    Example::
+
+        posts = Post.objects.filter(session, author_id=OuterRef('id'))
+        authors = Author.objects.annotate(session, newest=Subquery(
+            posts.order_by('-id').values_list('title', flat=True)[:1]
+        ))
+    """
+
+    name: str
+    depth: int
+    __slots__ = ("name", "depth")
+
+    def __init__(self, name: str | OuterRef):
+        if isinstance(name, OuterRef):
+            # Django-style OuterRef(OuterRef("id")) skips the immediate
+            # containing queryset and resolves against its parent instead.
+            self.name = name.name
+            self.depth = name.depth + 1
+        elif isinstance(name, str) and name and all(name.split("__")):
+            self.name = name
+            self.depth = 1
+        else:
+            raise ValueError(
+                "OuterRef requires a non-empty mapped field path or another OuterRef"
+            )
+
+    def resolve(self, field_resolver, context: ResolverContext = ResolverContext()):
+        # Store the depth in SQLAlchemy's annotation metadata, never SQL text.
+        # Each containing Subquery/Exists consumes one correlation boundary.
+        return _marker(self.name, self.depth)
+
+
+def _marker(name: str, depth: int):
+    return column("__async_orm_outer_ref__")._annotate({_OUTER_REF: (name, depth)})
+
+
+class _QueryExpression(CombinableMixin, OrderableMixin):
+    __slots__ = ("queryset",)
+
+    def __init__(self, queryset):
+        from async_orm_kit.queryset.queryset import QuerySet
+
+        if not isinstance(queryset, QuerySet):
+            raise TypeError("Subquery/Exists requires a QuerySet")
+        self.queryset = queryset
+
+    def _compile(self, outer_resolver, context: ResolverContext, *, scalar: bool):
+        state = self.queryset._state.clone()
+        state.select_related.clear()
+        state.prefetch_related.clear()
+        state.for_update = None
+        state.order_by = state.order_by if scalar else []
+        if scalar:
+            if state.values_select is not None:
+                if len(state.values_select) != 1:
+                    raise QueryStateError("Subquery requires a single selected column")
+                expression = state.values_select[0][1]
+            elif state.values_list_select is not None:
+                if len(state.values_list_select) != 1:
+                    raise QueryStateError("Subquery requires a single selected column")
+                expression = state.values_list_select[0][1]
+            else:
+                if state.annotations or state.has_aggregates:
+                    raise QueryStateError(
+                        "Annotated Subquery requires explicit scalar projection"
+                    )
+                primary_key = sa_inspect(self.queryset.model).primary_key
+                if len(primary_key) != 1:
+                    raise QueryStateError(
+                        "Composite-key Subquery requires explicit scalar projection"
+                    )
+                expression = getattr(self.queryset.model, primary_key[0].key)
+        else:
+            expression = None
+            state.order_by.clear()
+
+        # Referring to the same mapped class on both sides of a correlation
+        # produces SQL such as "WHERE authors.id = authors.id". SQLAlchemy
+        # cannot distinguish the inner/outer identity unless the inner model
+        # is explicitly aliased. Reject the ambiguous form before rendering it.
+        if self.queryset.model is outer_resolver.model:
+            clauses = (
+                list(state.filters)
+                + list(state.having_conditions)
+                + list(state.annotations.values())
+                + list(state.order_by)
+            )
+            clauses.extend(expr for _, expr in (state.values_select or ()))
+            clauses.extend(expr for _, expr in (state.values_list_select or ()))
+            for clause in clauses:
+                for element in visitors.iterate(clause):
+                    reference = getattr(element, "_annotations", {}).get(_OUTER_REF)
+                    if reference is None:
+                        continue
+                    depth = reference[1] if isinstance(reference, tuple) else 1
+                    if depth == 1:
+                        raise QueryStateError(
+                            "Self-correlated queries require a distinct inner "
+                            "SQLAlchemy aliased(Model) entity"
+                        )
+
+        def replace(element):
+            reference = getattr(element, "_annotations", {}).get(_OUTER_REF)
+            if reference is None:
+                return None
+            name, depth = reference if isinstance(reference, tuple) else (reference, 1)
+            if depth > 1:
+                return _marker(name, depth - 1)
+            return outer_resolver.resolve_to_column(name, context=context).column
+
+        def resolve_in(expression):
+            # Replacing references on the entire compiled SELECT forces
+            # SQLAlchemy to clone ORM JOIN aliases. A selected aggregate could
+            # then reference one alias while FROM uses another, causing a
+            # Cartesian product or an ambiguous column error. Resolve markers
+            # in the query state BEFORE the compiler constructs its FROM/JOIN
+            # graph, and leave unrelated column expressions intact.
+            if not any(
+                getattr(element, "_annotations", {}).get(_OUTER_REF) is not None
+                for element in visitors.iterate(expression)
+            ):
+                return expression
+            return visitors.replacement_traverse(expression, {}, replace)
+
+        state.filters = [resolve_in(expr) for expr in state.filters]
+        state.having_conditions = [resolve_in(expr) for expr in state.having_conditions]
+        state.annotations = {
+            key: resolve_in(expr) for key, expr in state.annotations.items()
+        }
+        state.values_select = (
+            [(key, resolve_in(expr)) for key, expr in state.values_select]
+            if state.values_select is not None
+            else None
+        )
+        state.values_list_select = (
+            [(key, resolve_in(expr)) for key, expr in state.values_list_select]
+            if state.values_list_select is not None
+            else None
+        )
+        state.order_by = [resolve_in(expr) for expr in state.order_by]
+        state.group_by_fields = [resolve_in(expr) for expr in state.group_by_fields]
+        state.distinct_on = [
+            resolve_in(expr) if not isinstance(expr, str) else expr
+            for expr in state.distinct_on
+        ]
+
+        statement = self.queryset._compiler.compile(state, allow_outer_refs=True)
+        if scalar:
+            if state.values_select is not None:
+                expression = state.values_select[0][1]
+            elif state.values_list_select is not None:
+                expression = state.values_list_select[0][1]
+            statement = statement.with_only_columns(
+                expression, maintain_column_froms=True
+            )
+        return statement
+
+
+class Subquery(_QueryExpression):
+    """Use a single-column QuerySet as a scalar SQL subquery."""
+
+    def resolve(self, field_resolver, context: ResolverContext = ResolverContext()):
+        return self._compile(field_resolver, context, scalar=True).scalar_subquery()
+
+
+class Exists(_QueryExpression):
+    """Use a QuerySet as an automatically correlated EXISTS predicate."""
+
+    __slots__ = ("negated",)
+
+    def __init__(self, queryset, *, negated: bool = False):
+        super().__init__(queryset)
+        self.negated = negated
+
+    def __invert__(self):
+        return Exists(self.queryset, negated=not self.negated)
+
+    def resolve(self, field_resolver, context: ResolverContext = ResolverContext()):
+        expression = self._compile(field_resolver, context, scalar=False).exists()
+        return ~expression if self.negated else expression
+
+
+__all__ = ["Exists", "OuterRef", "Subquery"]

async_orm_kit/query/expressions/value.py

@@ -0,0 +1,101 @@
+from __future__ import annotations
+
+from typing import Any
+
+from sqlalchemy import literal
+
+from async_orm_kit.query.lookups.resolver import FieldResolver, ResolverContext
+from async_orm_kit.query.mixins.combinable import CombinableMixin
+from async_orm_kit.query.mixins.orderable import OrderableMixin
+from async_orm_kit.types import SAFE_SCALAR_TYPES
+
+
+class Value(CombinableMixin, OrderableMixin):
+    """Wraps a literal scalar value for use in database expressions.
+
+    Value objects convert Python scalars (str, int, float, bool, None, datetime, etc.)
+    to SQLAlchemy literal() expressions, enabling their use in annotations, Case/When
+    expressions, and F expressions. They ensure type-safe SQL generation by wrapping
+    literals that would otherwise be ambiguous or require special handling.
+
+    Common use cases include providing default values in Case expressions, creating
+    constant annotations, and mixing literals with F objects in arithmetic operations.
+    Value objects participate in the same expression system as F objects, supporting
+    arithmetic operations and comparisons through CombinableMixin.
+
+    Args:
+        value: Python scalar value to wrap. Supports str, int, float, bool, None,
+               datetime, date, time, Decimal, and other SQLAlchemy-compatible types.
+
+    Note:
+        Value objects are immutable and hashable, making them safe to reuse across
+        multiple queries. They resolve to SQLAlchemy literal() expressions at query
+        compilation time, ensuring proper SQL escaping and type handling.
+
+        Unlike F objects which reference database columns, Value objects represent
+        constant values that are embedded directly in the SQL query. This distinction
+        is important for understanding query behavior and optimization.
+    """
+
+    _value: Any
+    __slots__ = ("_value",)
+
+    @property
+    def value(self) -> Any:
+        return self._value
+
+    def __setattr__(self, name: str, value: Any) -> None:
+        raise AttributeError("Value expressions are immutable.")
+
+    def __init__(self, value: Any):
+        """Initialize Value with a Python scalar.
+
+        Args:
+            value: Scalar value to wrap. Common types: str, int, float, bool, None,
+                   datetime, date, Decimal.
+        Example:
+            >>> Value("Active")
+            >>> Value(42)
+            >>> Value(3.14)
+            >>> Value(None)
+        """
+
+        if not isinstance(value, SAFE_SCALAR_TYPES):
+            raise TypeError(
+                f"Value() only accepts supported scalar values, got {type(value).__name__}."
+            )
+        object.__setattr__(self, "_value", value)
+
+    def resolve(self, field_resolver: FieldResolver, context: ResolverContext):
+        """Resolve to SQLAlchemy literal expression.
+
+        Converts the wrapped Python value to a SQLAlchemy literal() expression,
+        which handles proper SQL escaping and type conversion based on the database
+        dialect. The literal is embedded directly in the generated SQL.
+
+        Args:
+            field_resolver: FieldResolver instance (unused, for interface consistency).
+            context: ResolverContext (unused, for interface consistency).
+
+        Returns:
+            SQLAlchemy literal expression representing the wrapped value.
+
+        Example:
+            >>> # Internal usage during query compilation
+            >>> v = Value("Active")
+            >>> literal_expr = v.resolve(field_resolver, context)
+            >>> # Returns: literal("Active") -> SQL: 'Active'
+        """
+        return literal(self.value)
+
+    def __repr__(self):
+        """Return string representation showing wrapped value."""
+        return f"Value({self.value!r})"
+
+    def __eq__(self, other):
+        """Equality based on wrapped value for comparison and deduplication."""
+        return isinstance(other, Value) and self.value == other.value
+
+    def __hash__(self):
+        """Return hash based on wrapped value for use in sets and dicts."""
+        return hash(("Value", self.value))

async_orm_kit/query/expressions/window.py

@@ -0,0 +1,95 @@
+"""First-class SQL window expressions, without accidental GROUP BY injection."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from sqlalchemy import func
+
+from async_orm_kit.query.expressions.functions import _operand
+from async_orm_kit.query.lookups.resolver import ResolverContext
+from async_orm_kit.query.mixins.combinable import CombinableMixin
+from async_orm_kit.query.mixins.orderable import OrderableMixin, OrderedExpression
+
+
+class _WindowFunction(CombinableMixin, OrderableMixin):
+    __slots__ = ("name", "arguments")
+    _allowed = frozenset({"row_number", "rank", "dense_rank", "lag", "lead"})
+
+    def __init__(self, name: str, *arguments: Any):
+        if name not in self._allowed:
+            raise ValueError(f"Unsupported window function: {name!r}")
+        self.name = name
+        self.arguments = tuple(arguments)
+
+    def resolve(self, field_resolver, context: ResolverContext = ResolverContext()):
+        return getattr(func, self.name)(
+            *[_operand(arg, field_resolver, context) for arg in self.arguments]
+        )
+
+
+class RowNumber(_WindowFunction):
+    def __init__(self):
+        super().__init__("row_number")
+
+
+class Rank(_WindowFunction):
+    def __init__(self):
+        super().__init__("rank")
+
+
+class DenseRank(_WindowFunction):
+    def __init__(self):
+        super().__init__("dense_rank")
+
+
+class Lag(_WindowFunction):
+    def __init__(self, expression: Any, offset: int = 1, default: Any = None):
+        if isinstance(offset, bool) or not isinstance(offset, int) or offset < 1:
+            raise ValueError("Lag offset must be a positive integer")
+        super().__init__("lag", expression, offset, default)
+
+
+class Lead(_WindowFunction):
+    def __init__(self, expression: Any, offset: int = 1, default: Any = None):
+        if isinstance(offset, bool) or not isinstance(offset, int) or offset < 1:
+            raise ValueError("Lead offset must be a positive integer")
+        super().__init__("lead", expression, offset, default)
+
+
+class Window(CombinableMixin, OrderableMixin):
+    """Apply OVER/PARTITION BY/ORDER BY to a SQL expression.
+
+    ``order_by`` requires expressions with explicit ordering (`.asc()`/`.desc()`)
+    or raw SQLAlchemy order expressions. A string field name can be specified
+    using ``F('field').asc()`` rather than trusted raw SQL.
+    """
+
+    __slots__ = ("expression", "partition_by", "order_by")
+
+    def __init__(self, expression: Any, *, partition_by=(), order_by=()):
+        from sqlalchemy.sql.elements import ClauseElement
+
+        if not (isinstance(expression, (CombinableMixin, ClauseElement))):
+            raise TypeError("Window requires an ORM or SQLAlchemy SQL expression")
+        self.expression = expression
+        self.partition_by = tuple(partition_by)
+        self.order_by = tuple(order_by)
+        if isinstance(expression, (RowNumber, Rank, DenseRank)) and not self.order_by:
+            raise ValueError("Ranking window expressions require order_by")
+
+    def resolve(self, field_resolver, context: ResolverContext = ResolverContext()):
+        partition = [_operand(x, field_resolver, context) for x in self.partition_by]
+        ordering = [
+            x.resolve(field_resolver, context)
+            if isinstance(x, OrderedExpression)
+            else _operand(x, field_resolver, context)
+            for x in self.order_by
+        ]
+        return _operand(self.expression, field_resolver, context).over(
+            partition_by=partition or None,
+            order_by=ordering or None,
+        )
+
+
+__all__ = ["DenseRank", "Lag", "Lead", "Rank", "RowNumber", "Window"]