From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Sat, 6 Mar 2021 02:38:45 +0000 (+0100)
Subject: Add pool.check()
X-Git-Tag: 3.0.dev0~87^2~30
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=cc647c1ea3609622eb127bae3ef8eaa4363fddbe;p=thirdparty%2Fpsycopg.git

Add pool.check()
---

diff --git a/psycopg3/psycopg3/pool/pool.py b/psycopg3/psycopg3/pool/pool.py
index 5154f2c84..bbc89776a 100644
--- a/psycopg3/psycopg3/pool/pool.py
+++ b/psycopg3/psycopg3/pool/pool.py
@@ -246,6 +246,26 @@ class ConnectionPool(BasePool[Connection]):
         for i in range(ngrow):
             self.run_task(tasks.AddConnection(self))
 
+    def check(self) -> None:
+        """Verify the state of the connections currently in the pool.
+
+        Test each connection: if it works return it to the pool, otherwise
+        dispose of it and create a new one.
+        """
+        with self._lock:
+            conns = list(self._pool)
+            self._pool.clear()
+
+        while conns:
+            conn = conns.pop()
+            try:
+                conn.execute("select 1")
+            except Exception:
+                logger.warning("discarding broken connection: %s", conn)
+                self.run_task(tasks.AddConnection(self))
+            else:
+                self._add_to_pool(conn)
+
     def configure(self, conn: Connection) -> None:
         """Configure a connection after creation."""
         self._configure(conn)
diff --git a/tests/pool/test_pool.py b/tests/pool/test_pool.py
index b043b86c1..6c400ac22 100644
--- a/tests/pool/test_pool.py
+++ b/tests/pool/test_pool.py
@@ -659,6 +659,27 @@ def test_max_lifetime(dsn, retries):
             assert pids[0] == pids[1] != pids[2] == pids[3] != pids[4], pids
 
 
+def test_check(dsn, caplog):
+    caplog.set_level(logging.WARNING, logger="psycopg3.pool")
+    with pool.ConnectionPool(dsn, minconn=4) as p:
+        p.wait_ready(1.0)
+        with p.connection() as conn:
+            pid = conn.pgconn.backend_pid
+
+        p.wait_ready(1.0)
+        pids = set(conn.pgconn.backend_pid for conn in p._pool)
+        assert pid in pids
+        conn.close()
+
+        assert len(caplog.records) == 0
+        p.check()
+        assert len(caplog.records) == 1
+        p.wait_ready(1.0)
+        pids2 = set(conn.pgconn.backend_pid for conn in p._pool)
+        assert len(pids & pids2) == 3
+        assert pid not in pids2
+
+
 def delay_connection(monkeypatch, sec):
     """
     Return a _connect_gen function delayed by the amount of seconds