diff --git a/src/suou/sqlalchemy_async.py b/src/suou/sqlalchemy_async.py index 99f0777..7a710bc 100644 --- a/src/suou/sqlalchemy_async.py +++ b/src/suou/sqlalchemy_async.py @@ -99,12 +99,13 @@ 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)) + out = (await session.execute(select(func.count()).select_from(sub))).scalar() return out def __init__(self, @@ -138,16 +139,14 @@ class AsyncSelectPagination(Pagination): self.error_out = error_out self.has_count = count - async def __await__(self): + async def __aiter__(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