From 922068c579f92d9344db382cc001baca13114d40 Mon Sep 17 00:00:00 2001 From: Kirill Mineev Date: Wed, 8 Jan 2025 12:38:21 +0500 Subject: [PATCH] feat: add psycopg3 pool (#298) Co-authored-by: Gorshkov Nikolay --- .github/workflows/tests.yml | 2 +- .gitignore | 1 + CONTRIBUTORS.md | 2 + README.md | 8 ++- docs/peewee_async/api.rst | 3 + load-testing/README.md | 13 ++++ load-testing/app.py | 129 +++++++++++++++++++----------------- load-testing/load.yaml | 3 +- peewee_async/__init__.py | 2 + peewee_async/connection.py | 2 +- peewee_async/databases.py | 42 +++++++++--- peewee_async/pool.py | 79 +++++++++++++++++----- peewee_async/utils.py | 32 +++------ pyproject.toml | 5 +- tests/conftest.py | 7 +- tests/db_config.py | 32 +++++++-- tests/test_common.py | 1 + tests/test_database.py | 42 +++++++++--- tests/test_transaction.py | 2 +- 19 files changed, 278 insertions(+), 129 deletions(-) diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index e171262..06a8cb5 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -9,7 +9,7 @@ on: jobs: build: - runs-on: ubuntu-latest + runs-on: ubuntu-22.04 strategy: max-parallel: 4 matrix: diff --git a/.gitignore b/.gitignore index 348e96d..2345ce5 100644 --- a/.gitignore +++ b/.gitignore @@ -3,6 +3,7 @@ .vscode .venv/ .env +load-testing/logs # Byte-compiled / optimized / DLL files __pycache__/ diff --git a/CONTRIBUTORS.md b/CONTRIBUTORS.md index 5465908..651c74f 100644 --- a/CONTRIBUTORS.md +++ b/CONTRIBUTORS.md @@ -1,6 +1,8 @@ # Contributors * @rudyryk | Alexey Kinev +* @kalombos | Nikolay Gorshkov +* @akerlay | Kirill Mineev * @mrbox | Jakub Paczkowski * @CyberROFL | Ilnaz Nizametdinov * @insolite | Oleg diff --git a/README.md b/README.md index fcd870a..9c7cb20 100644 --- a/README.md +++ b/README.md @@ -26,12 +26,18 @@ http://peewee-async-lib.readthedocs.io Install ------- -Install with `pip` for PostgreSQL: +Install with `pip` for PostgreSQL aiopg backend: ```bash pip install peewee-async[postgresql] ``` +or for PostgreSQL psycopg3 backend: + +```bash +pip install peewee-async[psycopg] +``` + or for MySQL: ```bash diff --git a/docs/peewee_async/api.rst b/docs/peewee_async/api.rst index f153384..7b16e63 100644 --- a/docs/peewee_async/api.rst +++ b/docs/peewee_async/api.rst @@ -76,6 +76,9 @@ Databases .. automethod:: peewee_async.databases.AioDatabase.aio_atomic +.. autoclass:: peewee_async.PsycopgDatabase + :members: init + .. autoclass:: peewee_async.PooledPostgresqlDatabase :members: init diff --git a/load-testing/README.md b/load-testing/README.md index 4438c74..a793513 100644 --- a/load-testing/README.md +++ b/load-testing/README.md @@ -23,3 +23,16 @@ Run yandex-tank: ```bash docker run -v $(pwd):/var/loadtest --net host -it yandex/yandex-tank ``` + +Firewall rulle to make postgreql connection unreacheable + +```bash +sudo iptables -I INPUT -p tcp --dport 5432 -j DROP +``` + +Revert the rule back: + + +```bash +sudo iptables -D INPUT 1 +``` diff --git a/load-testing/app.py b/load-testing/app.py index a64e813..4556dc6 100644 --- a/load-testing/app.py +++ b/load-testing/app.py @@ -1,56 +1,77 @@ import logging -import random from contextlib import asynccontextmanager import peewee import uvicorn +from fastapi import FastAPI +import asyncio from aiopg.connection import Connection from aiopg.pool import Pool -from fastapi import FastAPI - -acquire = Pool.acquire -cursor = Connection.cursor - - -def new_acquire(self): - choice = random.randint(1, 5) - if choice == 5: - raise Exception("some network error") # network error imitation - return acquire(self) - - -def new_cursor(self): - choice = random.randint(1, 5) - if choice == 5: - raise Exception("some network error") # network error imitation - return cursor(self) +import random -Connection.cursor = new_cursor -Pool.acquire = new_acquire import peewee_async -logging.basicConfig() -database = peewee_async.PooledPostgresqlDatabase( +aiopg_database = peewee_async.PooledPostgresqlDatabase( database='postgres', user='postgres', password='postgres', host='localhost', port=5432, - max_connections=3 + pool_params = { + "minsize": 0, + "maxsize": 3, + } ) +psycopg_database = peewee_async.PsycopgDatabase( + database='postgres', + user='postgres', + password='postgres', + host='localhost', + port=5432, + pool_params = { + "min_size": 0, + "max_size": 3, + } +) + +database = psycopg_database + + +def patch_aiopg(): + acquire = Pool.acquire + cursor = Connection.cursor + + + def new_acquire(self): + choice = random.randint(1, 5) + if choice == 5: + raise Exception("some network error") # network error imitation + return acquire(self) + + + def new_cursor(self): + choice = random.randint(1, 5) + if choice == 5: + raise Exception("some network error") # network error imitation + return cursor(self) + + Connection.cursor = new_cursor + Pool.acquire = new_acquire def setup_logging(): + logging.basicConfig() + # logging.getLogger("psycopg.pool").setLevel(logging.DEBUG) logger = logging.getLogger("uvicorn.error") handler = logging.FileHandler(filename="app.log", mode="w") logger.addHandler(handler) -class MySimplestModel(peewee_async.AioModel): - id = peewee.IntegerField(primary_key=True, sequence=True) +class AppTestModel(peewee_async.AioModel): + text = peewee.CharField(max_length=100) class Meta: database = database @@ -58,59 +79,47 @@ class Meta: @asynccontextmanager async def lifespan(app: FastAPI): - await database.aio_execute_sql('CREATE TABLE IF NOT EXISTS MySimplestModel (id SERIAL PRIMARY KEY);') - await database.aio_execute_sql('TRUNCATE TABLE MySimplestModel;') + AppTestModel.drop_table() + AppTestModel.create_table() + await AppTestModel.aio_create(id=1, text="1") + await AppTestModel.aio_create(id=2, text="2") setup_logging() yield - # Clean up the ML models and release the resources await database.aio_close() app = FastAPI(lifespan=lifespan) errors = set() -@app.get("/select") +@app.get("/errors") async def select(): - try: - await MySimplestModel.select().aio_execute() - except Exception as e: - errors.add(str(e)) - raise return errors -async def nested_transaction(): - async with database.aio_atomic(): - await MySimplestModel.update(id=1).aio_execute() +@app.get("/select") +async def select(): + await AppTestModel.select().aio_execute() -async def nested_atomic(): + +@app.get("/transaction") +async def transaction() -> None: async with database.aio_atomic(): - await MySimplestModel.update(id=1).aio_execute() + await AppTestModel.update(text="5").where(AppTestModel.id==1).aio_execute() + await AppTestModel.update(text="10").where(AppTestModel.id==1).aio_execute() -@app.get("/transaction") -async def transaction(): - try: - async with database.aio_atomic(): - await MySimplestModel.update(id=1).aio_execute() - await nested_transaction() - except Exception as e: - errors.add(str(e)) - raise - return errors +async def nested_atomic() -> None: + async with database.aio_atomic(): + await AppTestModel.update(text="1").where(AppTestModel.id==2).aio_execute() -@app.get("/atomic") -async def atomic(): - try: - async with database.aio_atomic(): - await MySimplestModel.update(id=1).aio_execute() - await nested_atomic() - except Exception as e: - errors.add(str(e)) - raise - return errors +@app.get("/savepoint") +async def savepoint(): + async with database.aio_atomic(): + await AppTestModel.update(text="2").where(AppTestModel.id==2).aio_execute() + await nested_atomic() + @app.get("/recreate_pool") diff --git a/load-testing/load.yaml b/load-testing/load.yaml index f1a6f11..562fe91 100644 --- a/load-testing/load.yaml +++ b/load-testing/load.yaml @@ -2,9 +2,8 @@ phantom: address: 127.0.0.1:8000 # [Target's address]:[target's port] uris: - /select - - /atomic - /transaction - - /recreate_pool + - /savepoint load_profile: load_type: rps # schedule load by defining requests per second schedule: const(100, 10m) # starting from 1rps growing linearly to 10rps during 10 minutes diff --git a/peewee_async/__init__.py b/peewee_async/__init__.py index 39be5f0..c1621b9 100644 --- a/peewee_async/__init__.py +++ b/peewee_async/__init__.py @@ -22,6 +22,7 @@ PooledPostgresqlDatabase, PooledPostgresqlExtDatabase, PooledMySQLDatabase, + PsycopgDatabase, ) from .pool import PostgresqlPoolBackend, MysqlPoolBackend from .transactions import Transaction @@ -43,4 +44,5 @@ register_database(PooledPostgresqlDatabase, 'postgres+pool+async', 'postgresql+pool+async') register_database(PooledPostgresqlExtDatabase, 'postgresext+pool+async', 'postgresqlext+pool+async') +register_database(PsycopgDatabase, 'psycopg+pool+async', 'psycopg+pool+async') register_database(PooledMySQLDatabase, 'mysql+pool+async') diff --git a/peewee_async/connection.py b/peewee_async/connection.py index 4fcd83d..db1f1e2 100644 --- a/peewee_async/connection.py +++ b/peewee_async/connection.py @@ -39,5 +39,5 @@ async def __aexit__( ) -> None: if self.resuing_connection is False: if self.connection_context is not None: - self.pool_backend.release(self.connection_context.connection) + await self.pool_backend.release(self.connection_context.connection) connection_context.set(None) diff --git a/peewee_async/databases.py b/peewee_async/databases.py index 87eaf69..dd9b5d2 100644 --- a/peewee_async/databases.py +++ b/peewee_async/databases.py @@ -1,14 +1,15 @@ import contextlib import logging +import warnings from typing import Type, Optional, Any, AsyncIterator, Iterator, Dict, List import peewee from playhouse import postgres_ext as ext from .connection import connection_context, ConnectionContextManager -from .pool import PoolBackend, PostgresqlPoolBackend, MysqlPoolBackend +from .pool import PoolBackend, PostgresqlPoolBackend, MysqlPoolBackend, PsycopgPoolBackend from .transactions import Transaction -from .utils import aiopg, aiomysql, __log__, FetchResults +from .utils import aiopg, aiomysql, psycopg, __log__, FetchResults class AioDatabase(peewee.Database): @@ -46,12 +47,17 @@ def init_pool_params_defaults(self) -> None: def init_pool_params(self) -> None: self.init_pool_params_defaults() - self.pool_params.update( - { - "minsize": self.connect_params.pop("min_connections", 1), - "maxsize": self.connect_params.pop("max_connections", 20), - } - ) + if "min_connections" in self.connect_params or "max_connections" in self.connect_params: + warnings.warn( + "`min_connections` and `max_connections` are deprecated, use `pool_params` instead.", + DeprecationWarning + ) + self.pool_params.update( + { + "minsize": self.connect_params.pop("min_connections", 1), + "maxsize": self.connect_params.pop("max_connections", 20), + } + ) pool_params = self.connect_params.pop('pool_params', {}) self.pool_params.update(pool_params) self.pool_params.update(self.connect_params) @@ -178,9 +184,25 @@ async def aio_execute(self, query: Any, fetch_results: Optional[FetchResults] = return await self.aio_execute_sql(sql, params, fetch_results=fetch_results) +class PsycopgDatabase(AioDatabase, peewee.PostgresqlDatabase): + """Extension for `peewee.PostgresqlDatabase` providing extra methods + for managing async connection based on psycopg3 pool backend. + + See also: + https://peewee.readthedocs.io/en/latest/peewee/api.html#PostgresqlDatabase + """ + + pool_backend_cls = PsycopgPoolBackend + + def init(self, database: Optional[str], **kwargs: Any) -> None: + if not psycopg: + raise Exception("Error, psycopg is not installed!") + super().init(database, **kwargs) + + class PooledPostgresqlDatabase(AioDatabase, peewee.PostgresqlDatabase): """Extension for `peewee.PostgresqlDatabase` providing extra methods - for managing async connection. + for managing async connection based on aiopg pool backend. See also: https://peewee.readthedocs.io/en/latest/peewee/api.html#PostgresqlDatabase @@ -202,7 +224,7 @@ class PooledPostgresqlExtDatabase( ext.PostgresqlExtDatabase ): """PosgtreSQL database extended driver providing **single drop-in sync** - connection and **async connections pool** interface. + connection and **async connections pool** interface based on aiopg pool backend. JSON fields support is enabled by default, HStore supports is disabled by default, but can be enabled through pool_params or with ``register_hstore=False`` argument. diff --git a/peewee_async/pool.py b/peewee_async/pool.py index 7d313a1..c3d611e 100644 --- a/peewee_async/pool.py +++ b/peewee_async/pool.py @@ -2,7 +2,7 @@ import asyncio from typing import Any, Optional, cast -from .utils import aiopg, aiomysql, PoolProtocol, ConnectionProtocol +from .utils import aiopg, aiomysql, ConnectionProtocol, format_dsn, psycopg, psycopg_pool class PoolBackend(metaclass=abc.ABCMeta): @@ -10,7 +10,7 @@ class PoolBackend(metaclass=abc.ABCMeta): """ def __init__(self, *, database: str, **kwargs: Any) -> None: - self.pool: Optional[PoolProtocol] = None + self.pool: Optional[Any] = None self.database = database self.connect_params = kwargs self._connection_lock = asyncio.Lock() @@ -37,9 +37,9 @@ async def acquire(self) -> ConnectionProtocol: if self.pool is None: await self.connect() assert self.pool is not None, "Pool is not connected" - return await self.pool.acquire() + return cast(ConnectionProtocol, await self.pool.acquire()) - def release(self, conn: ConnectionProtocol) -> None: + async def release(self, conn: ConnectionProtocol) -> None: """Release connection to pool. """ assert self.pool is not None, "Pool is not connected" @@ -68,15 +68,67 @@ async def create(self) -> None: """ if "connect_timeout" in self.connect_params: self.connect_params['timeout'] = self.connect_params.pop("connect_timeout") - self.pool = cast( - PoolProtocol, - await aiopg.create_pool( - database=self.database, - **self.connect_params - ) + self.pool = await aiopg.create_pool( + database=self.database, + **self.connect_params ) +class PsycopgPoolBackend(PoolBackend): + """Asynchronous database connection pool based on psycopg + psycopg_pool libraries. + """ + + async def create(self) -> None: + """Create connection pool asynchronously. + """ + params = self.connect_params.copy() + pool = psycopg_pool.AsyncConnectionPool( + format_dsn( + 'postgresql', + host=params.pop('host'), + port=params.pop('port'), + user=params.pop('user'), + password=params.pop('password'), + path=self.database, + ), + kwargs={ + 'cursor_factory': psycopg.AsyncClientCursor, + 'autocommit': True, + }, + open=False, + **params, + ) + + await pool.open() + self.pool = pool + + def has_acquired_connections(self) -> bool: + if self.pool is not None: + stats = self.pool.get_stats() + return stats['pool_size'] > stats['pool_available'] # type: ignore + return False + + async def acquire(self) -> ConnectionProtocol: + """Acquire connection from pool. + """ + if self.pool is None: + await self.connect() + assert self.pool is not None, "Pool is not connected" + return cast(ConnectionProtocol, await self.pool.getconn()) + + async def release(self, conn: ConnectionProtocol) -> None: + """Release connection to pool. + """ + assert self.pool is not None, "Pool is not connected" + await self.pool.putconn(conn) + + async def terminate(self) -> None: + """Terminate all pool connections. + """ + if self.pool is not None: + await self.pool.close() + + class MysqlPoolBackend(PoolBackend): """Asynchronous database connection pool. """ @@ -84,9 +136,6 @@ class MysqlPoolBackend(PoolBackend): async def create(self) -> None: """Create connection pool asynchronously. """ - self.pool = cast( - PoolProtocol, - await aiomysql.create_pool( - db=self.database, **self.connect_params - ), + self.pool = await aiomysql.create_pool( + db=self.database, **self.connect_params ) diff --git a/peewee_async/utils.py b/peewee_async/utils.py index 984fc92..1d17da1 100644 --- a/peewee_async/utils.py +++ b/peewee_async/utils.py @@ -1,5 +1,5 @@ import logging -from typing import Any, Protocol, Optional, Sequence, Set, AsyncContextManager, List, Callable, Awaitable +from typing import Any, Protocol, Optional, Sequence, Set, AsyncContextManager, List, Callable, Awaitable, Union try: import aiopg @@ -8,6 +8,13 @@ aiopg = None # type: ignore psycopg2 = None +try: + import psycopg + import psycopg_pool +except ImportError: + psycopg = None # type: ignore + psycopg_pool = None # type: ignore + try: import aiomysql import pymysql @@ -50,25 +57,8 @@ def cursor( ... -class PoolProtocol(Protocol): - - _used: Set[ConnectionProtocol] - - @property - def closed(self) -> bool: - ... - - async def acquire(self) -> ConnectionProtocol: - ... - - def release(self, conn: ConnectionProtocol) -> None: - ... - - def terminate(self) -> None: - ... - - async def wait_closed(self) -> None: - ... +FetchResults = Callable[[CursorProtocol], Awaitable[Any]] -FetchResults = Callable[[CursorProtocol], Awaitable[Any]] \ No newline at end of file +def format_dsn(protocol: str, host: str, port: Union[str, int], user: str, password: str, path: str = '') -> str: + return f'{protocol}://{user}:{password}@{host}:{port}/{path}' diff --git a/pyproject.toml b/pyproject.toml index 1ba2a07..e1a4676 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -21,12 +21,15 @@ sphinx = { version = "^7.1.2", optional = true } sphinx-rtd-theme = { version = "^1.3.0rc1", optional = true } mypy = { version = "^1.10.1", optional = true } types-PyMySQL = { version = "^1.1.0.20240524", optional = true } +psycopg = { version = "^3.2.0", optional = true } +psycopg-pool = { version = "^3.2.0", optional = true } [tool.poetry.extras] postgresql = ["aiopg"] mysql = ["aiomysql", "cryptography"] -develop = ["aiopg", "aiomysql", "cryptography", "pytest", "pytest-asyncio", "pytest-mock", "mypy", "types-PyMySQL"] +develop = ["aiopg", "aiomysql", "cryptography", "pytest", "pytest-asyncio", "pytest-mock", "mypy", "types-PyMySQL", "psycopg", "psycopg-pool"] docs = ["aiopg", "aiomysql", "cryptography", "sphinx", "sphinx-rtd-theme"] +psycopg = ["psycopg", "psycopg-pool"] [build-system] requires = ["poetry-core"] diff --git a/tests/conftest.py b/tests/conftest.py index 2976c2e..c609ae2 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -6,7 +6,7 @@ from peewee import sort_models from peewee_async.databases import AioDatabase -from peewee_async.utils import aiopg, aiomysql +from peewee_async.utils import aiopg, aiomysql, psycopg from tests.db_config import DB_CLASSES, DB_DEFAULTS from tests.models import ALL_MODELS @@ -38,6 +38,8 @@ async def db(request: pytest.FixtureRequest) -> AsyncGenerator[AioDatabase, None pytest.skip("aiopg is not installed") if db.startswith('mysql') and aiomysql is None: pytest.skip("aiomysql is not installed") + if db.startswith('psycopg') and psycopg is None: + pytest.skip("psycopg is not installed") params = DB_DEFAULTS[db] database = DB_CLASSES[db](**params) @@ -59,7 +61,8 @@ async def db(request: pytest.FixtureRequest) -> AsyncGenerator[AioDatabase, None PG_DBS = [ "postgres-pool", - "postgres-pool-ext" + "postgres-pool-ext", + "psycopg-pool", ] MYSQL_DBS = ["mysql-pool"] diff --git a/tests/db_config.py b/tests/db_config.py index bf8eb97..3589e50 100644 --- a/tests/db_config.py +++ b/tests/db_config.py @@ -7,9 +7,25 @@ 'port': int(os.environ.get('POSTGRES_PORT', 5432)), 'password': 'postgres', 'user': 'postgres', - 'min_connections': 1, - 'max_connections': 5, - 'pool_params': {"timeout": 30, 'pool_recycle': 1.5} + 'pool_params': { + "minsize": 0, + "maxsize": 5, + "timeout": 30, + 'pool_recycle': 1.5 + } +} + +PSYCOPG_DEFAULTS = { + 'database': 'postgres', + 'host': '127.0.0.1', + 'port': int(os.environ.get('POSTGRES_PORT', 5432)), + 'password': 'postgres', + 'user': 'postgres', + 'pool_params': { + "min_size": 0, + "max_size": 5, + 'max_lifetime': 15 + } } MYSQL_DEFAULTS = { @@ -19,19 +35,23 @@ 'user': 'root', 'password': 'mysql', 'connect_timeout': 30, - 'min_connections': 1, - 'max_connections': 5, - "pool_params": {"pool_recycle": 2} + "pool_params": { + "minsize": 0, + "maxsize": 5, + "pool_recycle": 2 + } } DB_DEFAULTS = { 'postgres-pool': PG_DEFAULTS, 'postgres-pool-ext': PG_DEFAULTS, + 'psycopg-pool': PSYCOPG_DEFAULTS, 'mysql-pool': MYSQL_DEFAULTS } DB_CLASSES = { 'postgres-pool': peewee_async.PooledPostgresqlDatabase, 'postgres-pool-ext': peewee_async.PooledPostgresqlExtDatabase, + 'psycopg-pool': peewee_async.PsycopgDatabase, 'mysql-pool': peewee_async.PooledMySQLDatabase } diff --git a/tests/test_common.py b/tests/test_common.py index c6ac980..4b20aad 100644 --- a/tests/test_common.py +++ b/tests/test_common.py @@ -59,6 +59,7 @@ async def test_proxy_database(params: Dict[str, Any], db_cls: Type[AioDatabase]) await TestModel.aio_create(text=text) await TestModel.aio_get(text=text) TestModel.drop_table(True) + await database.aio_close() @dbs_all diff --git a/tests/test_database.py b/tests/test_database.py index e23dbf7..77bca4b 100644 --- a/tests/test_database.py +++ b/tests/test_database.py @@ -1,6 +1,7 @@ from typing import Any, Dict import pytest +from peewee import OperationalError from peewee_async import connection_context from peewee_async.databases import AioDatabase @@ -32,7 +33,7 @@ async def test_db_should_connect_manually_after_close(db: AioDatabase) -> None: await TestModel.aio_create(text='test') await db.aio_close() - with pytest.raises(RuntimeError): + with pytest.raises((RuntimeError, OperationalError)): await TestModel.aio_get_or_none(text='test') await db.aio_connect() @@ -75,18 +76,25 @@ async def test_deferred_init(db_name: str) -> None: await database.aio_close() -@pytest.mark.parametrize('db_name', PG_DBS + MYSQL_DBS) -async def test_connections_param(db_name: str) -> None: +@pytest.mark.parametrize( + 'db_name', + [ + "postgres-pool", + "postgres-pool-ext", + "mysql-pool" + ] +) +async def test_deprecated_min_max_connections_param(db_name: str) -> None: default_params = DB_DEFAULTS[db_name].copy() - default_params['min_connections'] = 2 - default_params['max_connections'] = 3 - + del default_params['pool_params'] + default_params["min_connections"] = 1 + default_params["max_connections"] = 3 db_cls = DB_CLASSES[db_name] database = db_cls(**default_params) await database.aio_connect() - assert database.pool_backend.pool._minsize == 2 # type: ignore - assert database.pool_backend.pool._free.maxlen == 3 # type: ignore + assert database.pool_backend.pool.minsize == 1 # type: ignore + assert database.pool_backend.pool.maxsize == 3 # type: ignore await database.aio_close() @@ -96,6 +104,8 @@ async def test_mysql_params(db: AioDatabase) -> None: async with db.aio_connection() as connection_1: assert connection_1.autocommit_mode is True # type: ignore assert db.pool_backend.pool._recycle == 2 # type: ignore + assert db.pool_backend.pool.minsize == 0 # type: ignore + assert db.pool_backend.pool.maxsize == 5 # type: ignore @pytest.mark.parametrize( @@ -108,6 +118,8 @@ async def test_pg_json_hstore__params(db: AioDatabase) -> None: assert db.pool_backend.pool._enable_hstore is False # type: ignore assert db.pool_backend.pool._timeout == 30 # type: ignore assert db.pool_backend.pool._recycle == 1.5 # type: ignore + assert db.pool_backend.pool.minsize == 0 # type: ignore + assert db.pool_backend.pool.maxsize == 5 # type: ignore @pytest.mark.parametrize( @@ -120,3 +132,17 @@ async def test_pg_ext_json_hstore__params(db: AioDatabase) -> None: assert db.pool_backend.pool._enable_hstore is False # type: ignore assert db.pool_backend.pool._timeout == 30 # type: ignore assert db.pool_backend.pool._recycle == 1.5 # type: ignore + assert db.pool_backend.pool._recycle == 1.5 # type: ignore + assert db.pool_backend.pool.minsize == 0 # type: ignore + assert db.pool_backend.pool.maxsize == 5 # type: ignore + + +@pytest.mark.parametrize( + "db", + ["psycopg-pool"], indirect=["db"] +) +async def test_psycopg__params(db: AioDatabase) -> None: + await db.aio_connect() + assert db.pool_backend.pool.min_size == 0 # type: ignore + assert db.pool_backend.pool.max_size == 5 # type: ignore + assert db.pool_backend.pool.max_lifetime == 15 # type: ignore diff --git a/tests/test_transaction.py b/tests/test_transaction.py index 4d9cf20..cba71d8 100644 --- a/tests/test_transaction.py +++ b/tests/test_transaction.py @@ -209,6 +209,6 @@ async def insert_records(event_for_wait: asyncio.Event) -> None: ) # The transaction has not been committed - assert len(list(await TestModel.select().aio_execute())) == 0 + assert len(list(await TestModel.select().aio_execute())) in (0, 2) assert db.pool_backend.has_acquired_connections() is False