Queries¶

class Query(Generic[T]):
    """Build and execute fluent ORM queries.

    Attributes:
        model_cls: Model class used to hydrate results.
        where_clause: Accumulated filter nodes for the query.
        order_by_clause: Sort definitions sent to the Rust core.
    """

    def __init__(
        self, model_cls: Type[T], using: str | None = None, session: Any | None = None
    ):
        """Initialize a query for a model class.

        Args:
            model_cls: Model class that defines the target table.

        Examples:
            >>> query = Query(User)
            >>> query.model_cls is User
            True
        """
        self.model_cls = model_cls
        self._using = using
        self._session = session
        self.where_clause: list["QueryNode"] = []
        self.order_by_clause: list[OrderByEntry] = []
        self._limit: int | None = None
        self._offset: int | None = None
        self._m2m_context: M2mContext | None = None
        # Relation paths that must render a join, insertion-ordered (full path
        # tuple -> registered join_type). Populated by where()/order_by()
        # traversal ("inner") and by the explicit join()/left_join() chainers
        # ("inner"/"left"). Serialized into the QueryIR ``joins`` section by
        # all()/count() (#270, #272).
        self._joins: dict[tuple[str, ...], str] = {}
        # Edges (path prefixes, length ≥ 1) whose join type was fixed by an
        # explicit chainer, insertion-ordered (edge tuple -> "inner"|"left").
        # ``.left_join`` marks every edge of its path "left" (whole-path rule,
        # ADR-0006); ``.join`` marks them "inner". The single source of truth
        # for LEFT on the wire — implicit where()/order_by() traversal never
        # touches this, so explicit always beats implicit (#272).
        self._explicit_edges: dict[tuple[str, ...], str] = {}
        # Relation paths to populate (#286, ADR-0008), insertion-ordered and
        # deduped by path identity (dict-as-ordered-set). CRITICAL: include
        # paths never enter ``_joins``/the wire ``joins`` section — they ride the
        # ``instances`` materialization plan, so the ``joins`` wire section
        # keeps its stage-1 semantics untouched and include contributes no
        # join-type opinion on any edge another clause references.
        self._includes: dict[tuple[str, ...], None] = {}

    async def _transaction_or_using(self) -> "RouteHandle":
        from .. import _ensure_rust_registration_synced_for_operation
        from ..state import resolve_operation_scope

        await _ensure_rust_registration_synced_for_operation()
        return resolve_operation_scope(using=self._using, session=self._session)

    def _clone(self) -> Self:
        """Return a copy of this query with no shared mutable state.

        ``copy.copy`` preserves the concrete class (``Relation`` stays
        ``Relation``); the mutable containers are then replaced so chained
        queries never alias the originals (FF-F F-1).
        """
        new = copy.copy(self)
        new.where_clause = list(self.where_clause)
        new.order_by_clause = list(self.order_by_clause)
        # M2mContext is frozen, so clones can share it safely.
        new._m2m_context = self._m2m_context
        new._joins = dict(self._joins)
        new._explicit_edges = dict(self._explicit_edges)
        new._includes = dict(self._includes)
        return new

    def _m2m(
        self, join_table: str, source_col: str, target_col: str, source_id: Any
    ) -> Self:
        """Store many-to-many linkage context for relationship operations.

        A new ``Query`` with the m2m context set; ``self`` is unchanged.
        """
        new = self._clone()
        new._m2m_context = M2mContext(
            join_table=join_table,
            source_col=source_col,
            target_col=target_col,
            source_id=source_id,
        )
        return new

    def where(self, predicate: "Predicate[T]") -> Self:
        """Add a filter condition to the query.

        ``predicate`` is a lambda of shape ``Callable[[QueryProxy[T]], QueryNode]``.
        The lambda receives a fresh :class:`QueryProxy` whose attributes
        return :class:`FieldProxy` instances, so
        ``lambda user: user.archived == False`` builds a comparison. Column
        names are validated at build time against the model's declared
        fields (plus shadow ``{fk}_id`` columns): a misspelled column raises
        ``AttributeError`` naming the closest valid match, before any query
        is sent to the database.

        Attribute access on a declared forward-FK field traverses the relation
        (``lambda t: t.account.ledger_id == lid``): each hop resolves against
        the related model, and every distinct traversed path renders one INNER
        join (ADR-0006) when the query runs. Every hop is validated at build
        time with the same did-you-mean naming the hop's model.

        Args:
            predicate: A callable that takes a :class:`QueryProxy` and
                returns a :class:`QueryNode`.

        Returns:
            A new ``Query`` with the clause added; ``self`` is unchanged.

        Raises:
            TypeError: If ``predicate`` is not callable, or if it does not
                return a ``QueryNode``.

        Examples:
            >>> q1 = User.where(lambda user: user.archived == False)  # noqa: E712
            >>> q2 = User.where(lambda user: user.id == 1)
            >>> isinstance(q1, Query) and isinstance(q2, Query)
            True
        """
        new = self._clone()
        node = _resolve_where_node(predicate, self.model_cls)
        new.where_clause.append(node)
        _register_join_paths(node, new._joins)
        return new

    @overload
    def select(self) -> Self: ...

    @overload
    def select(self, selector: "RowSelector[T]") -> "ProjectedQuery[T]": ...

    @overload
    def select(self, *columns: str) -> "ProjectedQuery[T]": ...

    def select(
        self, *selectors: "RowSelector[T] | str"
    ) -> "Self | ProjectedQuery[T]":
        """Project the query to a column subset, or pass through unchanged.

        Bare ``select()`` keeps meaning "full query": it returns an equivalent
        query of complete model instances. With a lambda selector — the
        documented style — the query becomes a projection: its results are
        :class:`Row` records in the list-like :class:`Rows` container, never
        model instances (the complete-instance invariant, ADR-0007). The
        selector returns one field (``select(lambda t: t.amount)``), a tuple
        (``select(lambda t: (t.id, t.amount))``), or a dict whose keys name
        the output fields (``select(lambda t: {"account_name":
        t.account.name})`` — output aliases). Fields may traverse declared
        forward-FK relations at any depth (``t.account.name`` — traversed
        projection); traversal narrows exactly like a ``where()`` predicate
        on the same path (INNER, one join per path, ADR-0006) and
        ``left_join()`` is the opt-out — kept rows decode the traversed
        fields as ``None``. Unaliased traversed fields take the bare leaf
        column name; two fields resolving to the same output name raise at
        build time. Column-name strings (``select("id", "amount")``) follow
        ``order_by``'s string contract exactly: root columns only, never
        traversal, never mixed with a lambda in one call. All forms validate
        at build time with did-you-mean.

        A projected query composes like any other query: ``where()``
        (relation traversal included), ``order_by()`` (by unselected columns
        too), ``limit()``/``offset()``, ``first()``, ``count()``, and
        ``exists()`` all work; ``update()``/``delete()`` and a second
        ``select()`` raise at build time.

        Args:
            *selectors: Nothing (full query), one lambda selector, or
                column-name strings.

        Returns:
            A new query; ``self`` is unchanged. Projected when a selector is
            given.

        Raises:
            TypeError: If the selector is not callable/strings, selects
                non-fields (a bare relation, a comparison), nests shapes
                (a dict inside a tuple, a tuple inside a dict), uses a
                non-string dict key, or mixes strings with a lambda in one
                call.
            ValueError: If the selection is empty, two fields resolve to the
                same output name, a string is a dotted path (strings never
                traverse), or the query already carries an ``include()``
                (one materialization plan per query — record results are
                flat, permanently).

        Examples:
            >>> rows = await Transaction.select(lambda t: (t.id, t.amount)).all()  # doctest: +SKIP
            >>> rows[0].amount  # doctest: +SKIP
            Decimal('12.50')
            >>> named = await Transaction.select(  # doctest: +SKIP
            ...     lambda t: {"id": t.id, "account_name": t.account.name}
            ... ).all()
        """
        if not selectors:
            return self._clone()
        if self._includes:
            raise ValueError(
                "select() with a projection cannot be combined with "
                "include(): a query carries exactly one materialization plan "
                "— projected records or populated instances, never both, and "
                "record results are flat (ADR-0009). Reach across the "
                "relation with traversed projection instead (e.g. "
                'select(lambda t: {"account_name": t.account.name})), or '
                "populate instances with include() on an unprojected query."
            )
        if any(isinstance(selector, str) for selector in selectors):
            if not all(isinstance(selector, str) for selector in selectors):
                raise TypeError(
                    "select() cannot mix column-name strings and a lambda "
                    "selector in one call; use one form "
                    '(select("id", "amount") or '
                    "select(lambda t: (t.id, t.amount)))."
                )
            names: tuple[str, ...] = selectors  # type: ignore[assignment]  # ty: ignore[invalid-assignment]
            columns = _resolve_projection_strings(names, self.model_cls)
            fields = tuple(
                _ProjectedField(name=column, column=column, path=())
                for column in columns
            )
            return ProjectedQuery(self, fields)
        if len(selectors) > 1:
            raise TypeError(
                "select() takes a single lambda selector naming the columns "
                "(e.g. `select(lambda t: (t.id, t.amount))`), got "
                f"{len(selectors)} arguments"
            )
        selector = selectors[0]
        if not callable(selector):
            raise TypeError(
                "select() expected a selector callable or column-name strings "
                f"(e.g. `lambda t: (t.id, t.amount)`), got {type(selector).__name__}"
            )
        # Strings were dispatched above, so the lone selector is the lambda
        # form; the tuple element type is too wide for the checker to see it.
        fields = _resolve_projection_selector(
            cast("RowSelector[T]", selector), self.model_cls
        )
        return ProjectedQuery(self, fields)

    def order_by(
        self,
        field: "str | Callable[[QueryProxy[T]], FieldProxy[Any]]",
        direction: str = "asc",
    ) -> Self:
        """Add an ordering clause and return a new query.

        Accepts a lambda naming the column (``order_by(lambda u: u.created_at,
        "desc")`` — the documented style, matching ``where`` predicates) or a
        column-name string (``order_by("created_at", "desc")``). Both forms are
        validated against the model's queryable columns at build time.

        A lambda selector may traverse a declared forward-FK relation
        (``order_by(lambda t: t.account.name)``) exactly like ``where()``: each
        hop resolves against the related model, and the traversed path renders
        one INNER join (ADR-0006), shared with any ``where()`` traversal of the
        same path in the same query — the same path referenced in both yields
        exactly one join. String selectors do not traverse: ``"account.name"``
        is looked up as a literal (unqualified) column name on the queried
        model.

        Args:
            field: Column selector — lambda receiving a :class:`QueryProxy`,
                or a column-name string.
            direction: ``"asc"`` (default) or ``"desc"``.

        Returns:
            A new ``Query`` with the ordering added; ``self`` is unchanged.

        Raises:
            AttributeError: If the column is not a queryable column.
            TypeError: If a lambda selector returns something other than a
                single column reference (a bare relation, e.g.
                ``lambda t: t.account``, is meaningless as a sort key).
            ValueError: If ``direction`` is not ``"asc"`` or ``"desc"``.

        Examples:
            >>> newest = await Post.select().order_by(lambda p: p.created_at, "desc").all()
        """
        if direction.lower() not in ("asc", "desc"):
            raise ValueError("direction must be 'asc' or 'desc'")

        path: tuple[str, ...] = ()
        if isinstance(field, str):
            col_name = validate_query_column(self.model_cls, field)
        elif callable(field):
            selected = field(QueryProxy(self.model_cls))
            # Ordering by a bare relation (no column selected) is meaningless —
            # reject it loudly rather than emitting an order_by the SELECT
            # walker cannot render. Ordering by a related COLUMN (a FieldProxy
            # with a non-empty path) is valid traversal (#271).
            if isinstance(selected, RelationProxy):
                relation = selected._path[-1]
                raise TypeError(
                    f"order_by() selector returned the bare relation {relation!r}, "
                    "not a column; order by a column on it instead "
                    f"(e.g. t.{relation}.<column>)."
                )
            if not isinstance(selected, FieldProxy):
                raise TypeError(
                    "order_by() selector must return a FieldProxy "
                    f"(e.g. `lambda u: u.created_at`), got {type(selected).__name__}"
                )
            col_name = selected.column
            path = selected.path
        else:
            raise TypeError(
                "order_by() expected a column-name string or a lambda selector, "
                f"got {type(field).__name__}"
            )

        new = self._clone()
        new.order_by_clause.append(
            OrderByEntry(column=col_name, direction=direction.lower(), path=path)
        )
        if path:
            new._joins.setdefault(path, "inner")
        return new

    def join(self, selector: "Callable[[QueryProxy[T]], Any]") -> Self:
        """Force an INNER join on a relation path and return a new query.

        ``selector`` is a lambda naming a RELATION path
        (``lambda t: t.account``, ``lambda t: t.account.owner``) — the same
        traversal syntax as ``where()``, but resolving to the relation itself,
        not a column on it. A bare ``.join(lambda t: t.account)`` with no
        predicate is a meaningful **existence filter** on a nullable relation:
        it narrows the result to rows where the relation exists (ADR-0006).

        Every edge of the path is marked explicit-INNER; combining it with an
        explicit ``.left_join`` on the same edge is a build-time error (see
        :meth:`left_join`).

        Args:
            selector: A lambda receiving a :class:`QueryProxy` and returning a
                relation path (a ``RelationProxy``).

        Returns:
            A new ``Query`` with the join registered; ``self`` is unchanged.

        Raises:
            TypeError: If ``selector`` is not callable or does not resolve to a
                relation path (a column selector is rejected — join names a
                relation, not a column).
            ValueError: If an edge of the path is already marked explicit-LEFT.

        Examples:
            >>> with_account = QJTransaction.select().join(lambda t: t.account)
        """
        return self._add_explicit_join(selector, "inner")

    def left_join(self, selector: "Callable[[QueryProxy[T]], Any]") -> Self:
        """Mark a relation path LEFT (whole-path) and return a new query.

        ``selector`` names a RELATION path exactly like :meth:`join`. Every edge
        of the path is marked LEFT (the **whole-path rule**, ADR-0006), so a
        left-marked 2-hop path retains rows missing the relation at either hop.
        Relation-less rows are retained (NULL retention), observable both in
        ordered results and in traversal predicates on related columns (e.g.
        ``where(lambda t: t.account.name == None)``).

        Conflict rules: an explicit LEFT beats implicit ``where()``/``order_by()``
        traversal on a shared edge (the path renders LEFT); an explicit ``.join``
        plus an explicit ``.left_join`` on the same edge is a build-time error.

        Args:
            selector: A lambda receiving a :class:`QueryProxy` and returning a
                relation path (a ``RelationProxy``).

        Returns:
            A new ``Query`` with the LEFT join registered; ``self`` is unchanged.

        Raises:
            TypeError: If ``selector`` is not callable or does not resolve to a
                relation path.
            ValueError: If an edge of the path is already marked explicit-INNER.

        Examples:
            >>> keep_orphans = QJNote.select().left_join(lambda n: n.account)
        """
        return self._add_explicit_join(selector, "left")

    def include(self, selector: "Callable[[QueryProxy[T]], Any]") -> Self:
        """Populate a relation path on every result and return a new query.

        ``selector`` is a lambda naming a forward-FK RELATION path
        (``lambda t: t.account``, ``lambda t: t.account.owner``) — the same
        traversal syntax as :meth:`join`. Results are the same ``list[T]``
        the query always returned, in one SQL statement, with each included
        relation **populated**: access is a plain attribute
        (``txn.account.label`` — no await, no query) returning the complete
        related instance, exactly as the field's declared annotation claims
        (ADR-0008). Unpopulated relations keep the awaitable contract
        unchanged; a nullable FK with no target populates as ``None`` with
        the root row retained.

        Include decides attached data only — joins decide membership,
        projection decides shape. Adding ``.include(...)`` to any query
        returns exactly the rows that query returned without it: include
        contributes no join-type opinion on any edge a predicate or explicit
        chainer references, and renders LEFT only on edges nothing else
        touches. ``count()``/``exists()`` are unaffected.

        Cumulative, order-free, and idempotent: chained includes accumulate,
        shared prefixes dedup by path identity, and including a path
        populates every hop along it.

        Args:
            selector: A lambda receiving a :class:`QueryProxy` and returning
                a relation path (a ``RelationProxy``).

        Returns:
            A new ``Query`` with the population registered; ``self`` is
            unchanged.

        Raises:
            TypeError: If ``selector`` is a string or not callable, resolves
                to a column rather than a relation, or names a BackRef/M2M
                relation (reverse population is a separate future mechanism).

        Examples:
            >>> txns = await Transaction.select().include(lambda t: t.account).all()  # doctest: +SKIP
            >>> txns[0].account.label  # doctest: +SKIP
            'checking'
        """
        path = _resolve_include_selector(selector, self.model_cls)
        new = self._clone()
        new._includes.setdefault(path, None)
        return new

    def _add_explicit_join(
        self, selector: "Callable[[QueryProxy[T]], Any]", join_type: str
    ) -> Self:
        """Shared body of :meth:`join`/:meth:`left_join` (#272).

        Resolves the selector to a relation path, checks every edge for a
        contradictory explicit mark (INNER vs LEFT), then records the marks and
        registers the full path so it renders. Re-marking an edge the same
        direction is idempotent.
        """
        path = _resolve_join_selector(selector, self.model_cls)
        edges = path_edges(path)
        new = self._clone()
        for edge in edges:
            existing = new._explicit_edges.get(edge)
            if existing is not None and existing != join_type:
                relation_path = ".".join(edge)
                raise ValueError(
                    f"conflicting explicit join types on relation edge "
                    f"{relation_path!r}: already marked {existing!r} by a prior "
                    f"join()/left_join(), cannot re-mark {join_type!r}. Use one "
                    "join type per edge."
                )
        for edge in edges:
            new._explicit_edges[edge] = join_type
        new._joins.setdefault(path, join_type)
        return new

    def limit(self, value: int) -> Self:
        """Limit the number of records returned

        Args:
            value: The maximum number of records to return.

        Returns:
            A new ``Query`` with the clause added; ``self`` is unchanged.

        Examples:
            >>> query = User.select().limit(10)
            >>> query._limit
            10
        """
        new = self._clone()
        new._limit = value
        return new

    def offset(self, value: int) -> Self:
        """Skip a specific number of records

        Args:
            value: The number of records to skip.

        Returns:
            A new ``Query`` with the clause added; ``self`` is unchanged.

        Examples:
            >>> query = User.select().offset(20)
            >>> query._offset
            20
        """
        new = self._clone()
        new._offset = value
        return new

    async def all(self) -> list[T]:
        """Return all model instances that match the current query

        Returns:
            A list of model instances.

        Examples:
            >>> users = await User.where(lambda user: user.active == True).all()  # noqa: E712
            >>> isinstance(users, list)
            True
        """
        compiled = compile_query(self, "fetch")
        route = await self._transaction_or_using()
        return await fetch_filtered(
            self.model_cls,
            compiled.wire_json,
            route,
            hop_classes=compiled.hop_classes,
        )

    async def count(self) -> int:
        """Return the number of records that match the current query

        Returns:
            The count of matching records.

        Examples:
            >>> total = await User.where(lambda user: user.active == True).count()  # noqa: E712
            >>> isinstance(total, int)
            True
        """
        compiled = compile_query(self, "count")
        route = await self._transaction_or_using()
        return await count_filtered(
            model_identity(self.model_cls),
            compiled.wire_json,
            route,
        )

    async def update(self, **fields) -> int:
        """Update all records matching the current query

        Args:
            **fields: Field names and values to update.

        Returns:
            The number of records updated.

        Raises:
            ValueError: If ``limit()`` or ``offset()`` was set on this query, or
                if the query traverses a relation (a ``where()`` predicate on a
                related column, or an explicit ``join()``/``left_join()``) —
                multi-table mutation has no portable SQL. Filter by a column on
                the target model, or resolve the related primary keys first and
                update by primary-key set.

        Examples:
            >>> updated = await User.where(lambda user: user.id == 1).update(name="Taylor")
            >>> isinstance(updated, int)
            True
        """
        compiled = compile_query(self, "update")
        route = await self._transaction_or_using()
        return await update_filtered(
            model_identity(self.model_cls),
            compiled.wire_json,
            update_bind_payload(fields),
            route,
        )

    async def first(self) -> T | None:
        """Return the first matching record, or None

        Returns:
            A model instance or None.

        Examples:
            >>> user = await User.select().order_by("id").first()
            >>> user is None or isinstance(user, User)
            True
        """
        results = await self.limit(1).all()
        return results[0] if results else None

    async def delete(self) -> int:
        """Delete all records matching the current query

        Returns:
            The number of records deleted.

        Raises:
            ValueError: If ``limit()`` or ``offset()`` was set on this query, or
                if the query traverses a relation (a ``where()`` predicate on a
                related column, or an explicit ``join()``/``left_join()``) —
                multi-table mutation has no portable SQL. Filter by a column on
                the target model, or resolve the related primary keys first and
                delete by primary-key set.

        Examples:
            >>> deleted = await User.where(lambda user: user.disabled == True).delete()  # noqa: E712
            >>> isinstance(deleted, int)
            True
        """
        compiled = compile_query(self, "delete")
        route = await self._transaction_or_using()
        return await delete_filtered(
            model_identity(self.model_cls),
            compiled.wire_json,
            route,
        )

    async def exists(self) -> bool:
        """Return whether at least one record matches the current query

        Returns:
            True if records exist, otherwise False.

        Examples:
            >>> found = await User.where(lambda user: user.email == "[email protected]").exists()
            >>> isinstance(found, bool)
            True
        """
        return await self.count() > 0

    async def add(self, *instances: Any) -> None:
        """Add links to a many-to-many relationship

        Args:
            *instances: Target model instances that provide an ``id`` attribute.

        Raises:
            RuntimeError: If the query is not bound to a many-to-many context.

        Examples:
            >>> user = await User.create(email="[email protected]")
            >>> admin = await Group.create(name="admin")
            >>> staff = await Group.create(name="staff")
            >>> await user.groups.add(admin, staff)
        """
        if not self._m2m_context:
            raise RuntimeError(
                "'.add()' can only be used on Many-to-Many relationships"
            )

        ids = []
        for inst in instances:
            # Assume 'id' for now
            ids.append(getattr(inst, "id"))

        route = await self._transaction_or_using()
        await add_m2m_links(
            self._m2m_context.join_table,
            self._m2m_context.source_col,
            self._m2m_context.target_col,
            self._m2m_context.source_id,
            ids,
            route,
        )

    async def remove(self, *instances: Any) -> None:
        """Remove links from a many-to-many relationship

        Args:
            *instances: Target model instances that provide an ``id`` attribute.

        Raises:
            RuntimeError: If the query is not bound to a many-to-many context.

        Examples:
            >>> user = await User.create(email="[email protected]")
            >>> admin = await Group.create(name="admin")
            >>> await user.groups.remove(admin)
        """
        if not self._m2m_context:
            raise RuntimeError(
                "'.remove()' can only be used on Many-to-Many relationships"
            )

        ids = []
        for inst in instances:
            ids.append(getattr(inst, "id"))

        route = await self._transaction_or_using()
        await remove_m2m_links(
            self._m2m_context.join_table,
            self._m2m_context.source_col,
            self._m2m_context.target_col,
            self._m2m_context.source_id,
            ids,
            route,
        )

    async def clear(self) -> None:
        """Clear all links in a many-to-many relationship

        Raises:
            RuntimeError: If the query is not bound to a many-to-many context.

        Examples:
            >>> user = await User.create(email="[email protected]")
            >>> await user.groups.clear()
        """
        if not self._m2m_context:
            raise RuntimeError(
                "'.clear()' can only be used on Many-to-Many relationships"
            )

        route = await self._transaction_or_using()
        await clear_m2m_links(
            self._m2m_context.join_table,
            self._m2m_context.source_col,
            self._m2m_context.source_id,
            route,
        )

    def __repr__(self):
        """Return a developer-friendly representation of the query"""
        return f"<Query model={self.model_cls.__name__} where={self.where_clause}>"