from __future__ import annotations import typing as t import sqlalchemy as sa import sqlalchemy.exc import sqlalchemy.orm from flask import abort from .pagination import Pagination from .pagination import QueryPagination class Query(sa.orm.Query): # type: ignore[type-arg] """SQLAlchemy :class:`~sqlalchemy.orm.query.Query` subclass with some extra methods useful for querying in a web application. This is the default query class for :attr:`.Model.query`. .. versionchanged:: 3.0 Renamed to ``Query`` from ``BaseQuery``. """ def get_or_404(self, ident: t.Any, description: str | None = None) -> t.Any: """Like :meth:`~sqlalchemy.orm.Query.get` but aborts with a ``404 Not Found`` error instead of returning ``None``. :param ident: The primary key to query. :param description: A custom message to show on the error page. """ rv = self.get(ident) if rv is None: abort(404, description=description) return rv def first_or_404(self, description: str | None = None) -> t.Any: """Like :meth:`~sqlalchemy.orm.Query.first` but aborts with a ``404 Not Found`` error instead of returning ``None``. :param description: A custom message to show on the error page. """ rv = self.first() if rv is None: abort(404, description=description) return rv def one_or_404(self, description: str | None = None) -> t.Any: """Like :meth:`~sqlalchemy.orm.Query.one` but aborts with a ``404 Not Found`` error instead of raising ``NoResultFound`` or ``MultipleResultsFound``. :param description: A custom message to show on the error page. .. versionadded:: 3.0 """ try: return self.one() except (sa.exc.NoResultFound, sa.exc.MultipleResultsFound): abort(404, description=description) def paginate( self, *, page: int | None = None, per_page: int | None = None, max_per_page: int | None = None, error_out: bool = True, count: bool = True, ) -> Pagination: """Apply an offset and limit to the query based on the current page and number of items per page, returning a :class:`.Pagination` object. :param page: The current page, used to calculate the offset. Defaults to the ``page`` query arg during a request, or 1 otherwise. :param per_page: The maximum number of items on a page, used to calculate the offset and limit. Defaults to the ``per_page`` query arg during a request, or 20 otherwise. :param max_per_page: The maximum allowed value for ``per_page``, to limit a user-provided value. Use ``None`` for no limit. Defaults to 100. :param error_out: Abort with a ``404 Not Found`` error if no items are returned and ``page`` is not 1, or if ``page`` or ``per_page`` is less than 1, or if either are not ints. :param count: Calculate the total number of values by issuing an extra count query. For very complex queries this may be inaccurate or slow, so it can be disabled and set manually if necessary. .. versionchanged:: 3.0 All parameters are keyword-only. .. versionchanged:: 3.0 The ``count`` query is more efficient. .. versionchanged:: 3.0 ``max_per_page`` defaults to 100. """ return QueryPagination( query=self, page=page, per_page=per_page, max_per_page=max_per_page, error_out=error_out, count=count, )