From 2b2e876edf686f3f4ddf5f36df9b978f42504082 Mon Sep 17 00:00:00 2001 From: Charles Leifer Date: Mon, 1 Jul 2024 11:26:41 -0500 Subject: [PATCH] Update pool as well. --- playhouse/pool.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/playhouse/pool.py b/playhouse/pool.py index d5a35eacf..030519b4d 100644 --- a/playhouse/pool.py +++ b/playhouse/pool.py @@ -47,6 +47,10 @@ class Meta: TRANSACTION_STATUS_IDLE = \ TRANSACTION_STATUS_INERROR = \ TRANSACTION_STATUS_UNKNOWN = None +try: + from psycopg.pq import TransactionStatus +except ImportError: + pass from peewee import MySQLDatabase from peewee import PostgresqlDatabase @@ -322,9 +326,9 @@ def _is_closed(self, conn): return True txn_status = conn.pgconn.transaction_status - if txn_status == conn.TransactionStatus.UNKNOWN: + if txn_status == TransactionStatus.UNKNOWN: return True - elif txn_status != conn.TransactionStatus.IDLE: + elif txn_status != TransactionStatus.IDLE: conn.rollback() return False @@ -333,11 +337,11 @@ def _can_reuse(self, conn): # 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: + if txn_status == TransactionStatus.UNKNOWN: return False - elif txn_status == conn.TransactionStatus.INERROR: + elif txn_status == TransactionStatus.INERROR: conn.reset() - elif txn_status != conn.TransactionStatus.IDLE: + elif txn_status != TransactionStatus.IDLE: conn.rollback() return True except ImportError: