|
| 1 | +from typing import Any, Type, Optional, AsyncIterator, TypeVar |
| 2 | + |
| 3 | +import asyncpg |
| 4 | +import pydantic |
| 5 | + |
| 6 | +from sqlc_runtime import AsyncCursor, AsyncConnection |
| 7 | + |
| 8 | +T = TypeVar("T", bound=pydantic.BaseModel) |
| 9 | + |
| 10 | + |
| 11 | +def build_asyncpg_connection(conn: asyncpg.Connection) -> AsyncConnection: |
| 12 | + return AsyncpgConnection(conn) |
| 13 | + |
| 14 | + |
| 15 | +class AsyncpgConnection: |
| 16 | + def __init__(self, conn: asyncpg.Connection): |
| 17 | + self._conn = conn |
| 18 | + |
| 19 | + async def execute(self, query: str, *params: Any) -> AsyncCursor: |
| 20 | + return await self._conn.cursor(query, *params) |
| 21 | + |
| 22 | + async def execute_none(self, query: str, *params: Any) -> None: |
| 23 | + await self._conn.execute(query, *params) |
| 24 | + |
| 25 | + async def execute_rowcount(self, query: str, *params: Any) -> int: |
| 26 | + status = await self._conn.execute(query, *params) |
| 27 | + return int(status.split(" ")[-1]) |
| 28 | + |
| 29 | + async def execute_one(self, query: str, *params: Any) -> Any: |
| 30 | + row = await self._conn.fetchrow(query, *params) |
| 31 | + return row[0] if row is not None else None |
| 32 | + |
| 33 | + async def execute_one_model( |
| 34 | + self, model: Type[T], query: str, *params: Any |
| 35 | + ) -> Optional[T]: |
| 36 | + row = await self._conn.fetchrow(query, *params) |
| 37 | + if row is None: |
| 38 | + return None |
| 39 | + return model.parse_obj(row) |
| 40 | + |
| 41 | + async def execute_many(self, query: str, *params: Any) -> AsyncIterator[Any]: |
| 42 | + async with self._conn.transaction(): |
| 43 | + async for row in self._conn.cursor(query, *params): |
| 44 | + yield row[0] |
| 45 | + |
| 46 | + async def execute_many_model( |
| 47 | + self, model: Type[T], query: str, *params: Any |
| 48 | + ) -> AsyncIterator[T]: |
| 49 | + async with self._conn.transaction(): |
| 50 | + async for row in self._conn.cursor(query, *params): |
| 51 | + yield model.parse_obj(row) |
0 commit comments