diff --git a/playhouse/pool.py b/playhouse/pool.py index 2ee3b486f..1b4f843ec 100644 --- a/playhouse/pool.py +++ b/playhouse/pool.py @@ -289,6 +289,37 @@ class PooledPostgresqlExtDatabase(_PooledPostgresqlDatabase, PostgresqlExtDataba PooledPostgresqlExtDatabase = None +try: + from playhouse.psycopg3_ext import Psycopg3Database + + class PooledPsycopg3Database(PooledDatabase, Psycopg3Database): + def _is_closed(self, conn): + if conn.closed: + return True + + txn_status = conn.pgconn.transaction_status + if txn_status == conn.TransactionStatus.UNKNOWN: + return True + elif txn_status != conn.TransactionStatus.IDLE: + conn.rollback() + return False + + def _can_reuse(self, conn): + txn_status = conn.pgconn.transaction_status + # Do not return connection in an error state, as subsequent queries + # will all fail. If the status is unknown then we lost the connection + # to the server and the connection should not be re-used. + if txn_status == conn.TransactionStatus.UNKNOWN: + return False + elif txn_status == conn.TransactionStatus.INERROR: + conn.reset() + elif txn_status != conn.TransactionStatus.IDLE: + conn.rollback() + return True +except ImportError: + PooledPsycopg3Database = None + + class _PooledSqliteDatabase(PooledDatabase): def _is_closed(self, conn): try: diff --git a/tests/pool.py b/tests/pool.py index df93599dc..319afcdb0 100644 --- a/tests/pool.py +++ b/tests/pool.py @@ -8,6 +8,7 @@ from peewee import _transaction from playhouse.cockroachdb import PooledCockroachDatabase from playhouse.pool import * +from playhouse.psycopg3_ext import Psycopg3Database from .base import BACKEND from .base import BaseTestCase @@ -372,7 +373,10 @@ def setUp(self): if IS_MYSQL: db_class = PooledMySQLDatabase elif IS_POSTGRESQL: - db_class = PooledPostgresqlDatabase + if isinstance(self.database, Psycopg3Database): + db_class = PooledPsycopg3Database + else: + db_class = PooledPostgresqlDatabase elif IS_CRDB: db_class = PooledCockroachDatabase else: