codekingpro/portable-devtools

1from __future__ import annotations2 3import typing as t4 5import sqlalchemy.exc as sa_exc6import sqlalchemy.orm as sa_orm7from flask import abort8 9from .pagination import Pagination10from .pagination import QueryPagination11 12 13class Query(sa_orm.Query):  # type: ignore[type-arg]14    """SQLAlchemy :class:`~sqlalchemy.orm.query.Query` subclass with some extra methods15    useful for querying in a web application.16 17    This is the default query class for :attr:`.Model.query`.18 19    .. versionchanged:: 3.020        Renamed to ``Query`` from ``BaseQuery``.21    """22 23    def get_or_404(self, ident: t.Any, description: str | None = None) -> t.Any:24        """Like :meth:`~sqlalchemy.orm.Query.get` but aborts with a ``404 Not Found``25        error instead of returning ``None``.26 27        :param ident: The primary key to query.28        :param description: A custom message to show on the error page.29        """30        rv = self.get(ident)31 32        if rv is None:33            abort(404, description=description)34 35        return rv36 37    def first_or_404(self, description: str | None = None) -> t.Any:38        """Like :meth:`~sqlalchemy.orm.Query.first` but aborts with a ``404 Not Found``39        error instead of returning ``None``.40 41        :param description: A custom message to show on the error page.42        """43        rv = self.first()44 45        if rv is None:46            abort(404, description=description)47 48        return rv49 50    def one_or_404(self, description: str | None = None) -> t.Any:51        """Like :meth:`~sqlalchemy.orm.Query.one` but aborts with a ``404 Not Found``52        error instead of raising ``NoResultFound`` or ``MultipleResultsFound``.53 54        :param description: A custom message to show on the error page.55 56        .. versionadded:: 3.057        """58        try:59            return self.one()60        except (sa_exc.NoResultFound, sa_exc.MultipleResultsFound):61            abort(404, description=description)62 63    def paginate(64        self,65        *,66        page: int | None = None,67        per_page: int | None = None,68        max_per_page: int | None = None,69        error_out: bool = True,70        count: bool = True,71    ) -> Pagination:72        """Apply an offset and limit to the query based on the current page and number73        of items per page, returning a :class:`.Pagination` object.74 75        :param page: The current page, used to calculate the offset. Defaults to the76            ``page`` query arg during a request, or 1 otherwise.77        :param per_page: The maximum number of items on a page, used to calculate the78            offset and limit. Defaults to the ``per_page`` query arg during a request,79            or 20 otherwise.80        :param max_per_page: The maximum allowed value for ``per_page``, to limit a81            user-provided value. Use ``None`` for no limit. Defaults to 100.82        :param error_out: Abort with a ``404 Not Found`` error if no items are returned83            and ``page`` is not 1, or if ``page`` or ``per_page`` is less than 1, or if84            either are not ints.85        :param count: Calculate the total number of values by issuing an extra count86            query. For very complex queries this may be inaccurate or slow, so it can be87            disabled and set manually if necessary.88 89        .. versionchanged:: 3.090            All parameters are keyword-only.91 92        .. versionchanged:: 3.093            The ``count`` query is more efficient.94 95        .. versionchanged:: 3.096            ``max_per_page`` defaults to 100.97        """98        return QueryPagination(99            query=self,100            page=page,101            per_page=per_page,102            max_per_page=max_per_page,103            error_out=error_out,104            count=count,105        )106