fix database name, fix imports (direct -> utils) to catch situation w…

…here psycopg is not installed
05bit · Oct 10, 2024 · 4e3ca11 · 4e3ca11
1 parent 9212a5a
commit 4e3ca11
Show file tree

Hide file tree

Showing 5 changed files with 8 additions and 9 deletions.
diff --git a/peewee_async/__init__.py b/peewee_async/__init__.py
@@ -22,7 +22,7 @@
     PooledPostgresqlDatabase,
     PooledPostgresqlExtDatabase,
     PooledMySQLDatabase,
-    PooledPsycopg3PostgresqlDatabase,
+    Psycopg3Database,
 )
 from .pool import PostgresqlPoolBackend, MysqlPoolBackend
 from .transactions import Transaction
@@ -44,6 +44,6 @@
 
 register_database(PooledPostgresqlDatabase, 'postgres+pool+async', 'postgresql+pool+async')
 register_database(PooledPostgresqlExtDatabase, 'postgresext+pool+async', 'postgresqlext+pool+async')
-register_database(PooledPsycopg3PostgresqlDatabase, 'postgres+psycopg+pool+async', 'postgres+psycopg+pool+async')
+register_database(Psycopg3Database, 'postgres+psycopg+pool+async', 'postgres+psycopg+pool+async')
 register_database(PooledMySQLDatabase, 'mysql+pool+async')
 
diff --git a/peewee_async/databases.py b/peewee_async/databases.py
@@ -197,7 +197,7 @@ def init(self, database: Optional[str], **kwargs: Any) -> None:
         super().init(database, **kwargs)
 
 
-class PooledPsycopg3PostgresqlDatabase(AioDatabase, peewee.PostgresqlDatabase):
+class Psycopg3Database(AioDatabase, peewee.PostgresqlDatabase):
     """Extension for `peewee.PostgresqlDatabase` providing extra methods
     for managing async connection based on psycopg3 pool backend.
 

diff --git a/peewee_async/pool.py b/peewee_async/pool.py
@@ -2,10 +2,7 @@
 import asyncio
 from typing import Any, Optional, cast
 
-import psycopg_pool
-from psycopg import AsyncClientCursor
-
-from .utils import aiopg, aiomysql, PoolProtocol, ConnectionProtocol, format_dsn
+from .utils import aiopg, aiomysql, PoolProtocol, ConnectionProtocol, format_dsn, psycopg, psycopg_pool
 
 
 class PoolBackend(metaclass=abc.ABCMeta):
@@ -110,7 +107,7 @@ async def create(self) -> None:
             max_lifetime=self.connect_params.get('pool_recycle', 60 * 60.0),
             open=False,
             kwargs={
-                'cursor_factory': AsyncClientCursor,
+                'cursor_factory': psycopg.AsyncClientCursor,
                 'autocommit': True,
             }
         )

diff --git a/peewee_async/utils.py b/peewee_async/utils.py
@@ -10,8 +10,10 @@
 
 try:
     import psycopg
+    import psycopg_pool
 except ImportError:
     psycopg = None  # type: ignore
+    psycopg_pool = None  # type: ignore
 
 try:
     import aiomysql

diff --git a/tests/db_config.py b/tests/db_config.py
@@ -34,6 +34,6 @@
 DB_CLASSES = {
     'postgres-pool': peewee_async.PooledPostgresqlDatabase,
     'postgres-pool-ext': peewee_async.PooledPostgresqlExtDatabase,
-    'psycopg-pool': peewee_async.PooledPsycopg3PostgresqlDatabase,
+    'psycopg-pool': peewee_async.Psycopg3Database,
     'mysql-pool': peewee_async.PooledMySQLDatabase
 }