diff --git a/peewee_async/__init__.py b/peewee_async/__init__.py index adc4f77..bcf4af4 100644 --- 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 index b2fff04..7cf3322 100644 --- 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 index f7fd3d1..e8586bb 100644 --- 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 index ec53476..bcf3d87 100644 --- 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 index 22ec425..5f39391 100644 --- 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 }