S01E02

Author: ansipunk

databases/backends/dialects/psycopg.py

@@ -58,8 +58,12 @@ def get_dialect() -> Dialect:
     return dialect
 
 
-def compile_query(query: ClauseElement, dialect: Dialect) -> typing.Tuple[str, list, tuple]:
-    compiled = query.compile(dialect=dialect, compile_kwargs={"render_postcompile": True})
+def compile_query(
+    query: ClauseElement, dialect: Dialect
+) -> typing.Tuple[str, list, tuple]:
+    compiled = query.compile(
+        dialect=dialect, compile_kwargs={"render_postcompile": True}
+    )
 
     if not isinstance(query, DDLElement):
         compiled_params = sorted(compiled.params.items())

databases/backends/psycopg.py

@@ -37,7 +37,8 @@ async def connect(self) -> None:
             return
 
         self._pool = psycopg_pool.AsyncConnectionPool(
-            self._database_url.url, open=False, **self._options)
+            self._database_url._url, open=False, **self._options
+        )
 
         # TODO: Add configurable timeouts
         await self._pool.open()
@@ -86,13 +87,33 @@ async def fetch_all(self, query: ClauseElement) -> typing.List[RecordInterface]:
             raise RuntimeError("Connection is not acquired")
 
         query_str, args, result_columns = compile_query(query, self._dialect)
-        rows = await self._connection.fetch(query_str, *args)
+
+        async with self._connection.cursor() as cursor:
+            await cursor.execute(query_str, args)
+            rows = await cursor.fetchall()
+
         column_maps = create_column_maps(result_columns)
         return [Record(row, result_columns, self._dialect, column_maps) for row in rows]
-        raise NotImplementedError()  # pragma: no cover
 
     async def fetch_one(self, query: ClauseElement) -> typing.Optional[RecordInterface]:
-        raise NotImplementedError()  # pragma: no cover
+        if self._connection is None:
+            raise RuntimeError("Connection is not acquired")
+
+        query_str, args, result_columns = compile_query(query, self._dialect)
+
+        async with self._connection.cursor() as cursor:
+            await cursor.execute(query_str, args)
+            row = await cursor.fetchone()
+
+        if row is None:
+            return None
+
+        return Record(
+            row,
+            result_columns,
+            self._dialect,
+            create_column_maps(result_columns),
+        )
 
     async def fetch_val(
         self, query: ClauseElement, column: typing.Any = 0
@@ -101,25 +122,47 @@ async def fetch_val(
         return None if row is None else row[column]
 
     async def execute(self, query: ClauseElement) -> typing.Any:
-        raise NotImplementedError()  # pragma: no cover
+        if self._connection is None:
+            raise RuntimeError("Connection is not acquired")
+
+        query_str, args, _ = compile_query(query, self._dialect)
+
+        async with self._connection.cursor() as cursor:
+            await cursor.execute(query_str, args)
 
     async def execute_many(self, queries: typing.List[ClauseElement]) -> None:
-        raise NotImplementedError()  # pragma: no cover
+        # TODO: Find a way to use psycopg's executemany
+        for query in queries:
+            await self.execute(query)
 
     async def iterate(
         self, query: ClauseElement
     ) -> typing.AsyncGenerator[typing.Mapping, None]:
-        raise NotImplementedError()  # pragma: no cover
-        # mypy needs async iterators to contain a `yield`
-        # https://github.com/python/mypy/issues/5385#issuecomment-407281656
-        yield True  # pragma: no cover
+        if self._connection is None:
+            raise RuntimeError("Connection is not acquired")
+
+        query_str, args, result_columns = compile_query(query, self._dialect)
+        column_maps = create_column_maps(result_columns)
+
+        async with self._connection.cursor() as cursor:
+            await cursor.execute(query_str, args)
+
+            while True:
+                row = await cursor.fetchone()
+
+                if row is None:
+                    break
+
+                yield Record(row, result_columns, self._dialect, column_maps)
 
     def transaction(self) -> "TransactionBackend":
         raise NotImplementedError()  # pragma: no cover
 
     @property
     def raw_connection(self) -> typing.Any:
-        raise NotImplementedError()  # pragma: no cover
+        if self._connection is None:
+            raise RuntimeError("Connection is not acquired")
+        return self._connection
 
 
 class PsycopgTransaction(TransactionBackend):

requirements.txt

@@ -8,6 +8,7 @@ aiosqlite==0.20.0
 asyncpg==0.29.0
 psycopg==3.1.18
 psycopg-binary==3.1.18
+psycopg-pool==3.2.1
 
 # Sync database drivers for standard tooling around setup/teardown/migrations.
 psycopg==3.1.18

setup.py

@@ -55,7 +55,7 @@ def get_packages(package):
         "postgresql": ["asyncpg"],
         "aiopg": ["aiopg"],
         "asyncpg": ["asyncpg"],
-        "psycopg3": ["psycopg"],
+        "psycopg3": ["psycopg", "psycopg-pool"],
         "sqlite": ["aiosqlite"],
         "aiosqlite": ["aiosqlite"],
     },