diff --git a/src/suou/sqlalchemy_async.py b/src/suou/sqlalchemy_async.py index 7a710bc..99f0777 100644 --- a/src/suou/sqlalchemy_async.py +++ b/src/suou/sqlalchemy_async.py @@ -99,13 +99,12 @@ class AsyncSelectPagination(Pagination): select = select_q.limit(self.per_page).offset(self._query_offset) session: AsyncSession = self._query_args["session"] out = (await session.execute(select)).scalars() - return out async def _query_count(self) -> int: select_q: Select = self._query_args["select"] sub = select_q.options(lazyload("*")).order_by(None).subquery() session: AsyncSession = self._query_args["session"] - out = (await session.execute(select(func.count()).select_from(sub))).scalar() + out = await session.execute(select(func.count()).select_from(sub)) return out def __init__(self, @@ -139,14 +138,16 @@ class AsyncSelectPagination(Pagination): self.error_out = error_out self.has_count = count - async def __aiter__(self): + async def __await__(self): self.items = await self._query_items() - if self.items is None: - raise RuntimeError('query returned None') if not self.items and self.page != 1 and self.error_out: raise self.error_out if self.has_count: self.total = await self._query_count() + return self + + async def __aiter__(self): + await self for i in self.items: yield i