Add pool implementation and test fixes for psycopg3.

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:

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: