Add wait parameter to pool.open()

diff --git a/docs/api/pool.rst b/docs/api/pool.rst
index 0ae7244b7cb1b11a21f65419650dc7281d62fec9..eb6aeb3ca5955f4018e5f2c7bc607fc7cd9251b4 100644 (file)
--- a/docs/api/pool.rst
+++ b/docs/api/pool.rst
@@ -150,7 +150,6 @@ The `!ConnectionPool` class
         you don't use the pool as a context manager) you might want to specify
         this parameter explicitly.
 
-   .. automethod:: wait
    .. automethod:: connection
 
       .. code:: python
@@ -177,6 +176,8 @@ The `!ConnectionPool` class
           with ConnectionPool(...) as pool:
               # code using the pool
 
+   .. automethod:: wait
+
    .. attribute:: name
       :type: str
 
@@ -244,7 +245,6 @@ listed here.
                  the pool.
    :type reset: `async Callable[[AsyncConnection], None]`
 
-   .. automethod:: wait
    .. automethod:: connection
 
       .. code:: python
@@ -255,7 +255,6 @@ listed here.
           # the connection is now back in the pool
 
    .. automethod:: open
-
    .. automethod:: close
 
    .. note::
@@ -269,6 +268,7 @@ listed here.
           async with AsyncConnectionPool(...) as pool:
               # code using the pool
 
+   .. automethod:: wait
    .. automethod:: resize
    .. automethod:: check
    .. automethod:: getconn

diff --git a/psycopg_pool/psycopg_pool/pool.py b/psycopg_pool/psycopg_pool/pool.py
index 7e0b88bcec8f20b1822b80422f27baf6d77e3e5e..4e90a5f6dc6360633b9be9a234bd069391ca7db4 100644 (file)
--- a/psycopg_pool/psycopg_pool/pool.py
+++ b/psycopg_pool/psycopg_pool/pool.py
@@ -218,24 +218,36 @@ class ConnectionPool(BasePool[Connection[Any]]):
         else:
             self._return_connection(conn)
 
-    def open(self) -> None:
+    def open(self, wait: bool = False, timeout: float = 30.0) -> None:
         """Open the pool by starting connecting and and accepting clients.
 
-        The method is no-op if the pool is already opened (because the method
-        was already called, or because the pool context was entered, or because
-        the pool was initialized with ``open=true``.
+        If *wait* is `!False`, return immediately and let the background worker
+        fill the pool if `min_size` > 0. Otherwise wait up to *timeout* seconds
+        for the requested number of connections to be ready (see `wait()` for
+        details).
+
+        It is safe to call `!open()` again on a pool already open (because the
+        method was already called, or because the pool context was entered, or
+        because the pool was initialized with *open* = `!True`) but you cannot
+        currently re-open a closed pool.
         """
         with self._lock:
-            if not self._closed:
-                return
+            self._open()
+
+        if wait:
+            self.wait(timeout=timeout)
+
+    def _open(self) -> None:
+        if not self._closed:
+            return
 
-            self._check_open()
+        self._check_open()
 
-            self._start_workers()
-            self._start_initial_tasks()
+        self._start_workers()
+        self._start_initial_tasks()
 
-            self._closed = False
-            self._opened = True
+        self._closed = False
+        self._opened = True
 
     def _start_workers(self) -> None:
         self._sched_runner = threading.Thread(

diff --git a/psycopg_pool/psycopg_pool/pool_async.py b/psycopg_pool/psycopg_pool/pool_async.py
index b74bc5c5e47896a8dfe6c69f662aa74ec97ee878..8a4391ceb5403827579de16e2e10e945d98daee0 100644 (file)
--- a/psycopg_pool/psycopg_pool/pool_async.py
+++ b/psycopg_pool/psycopg_pool/pool_async.py
@@ -179,10 +179,13 @@ class AsyncConnectionPool(BasePool[AsyncConnection[Any]]):
         else:
             await self._return_connection(conn)
 
-    async def open(self) -> None:
+    async def open(self, wait: bool = False, timeout: float = 30.0) -> None:
         async with self._lock:
             self._open()
 
+        if wait:
+            await self.wait(timeout=timeout)
+
     def _open(self) -> None:
         if not self._closed:
             return

diff --git a/tests/pool/test_pool.py b/tests/pool/test_pool.py
index f30d8513c4398b8661f5226ad8955f8cb41316eb..c24759a31d45e5fc451a10d96c84264de91e206e 100644 (file)
--- a/tests/pool/test_pool.py
+++ b/tests/pool/test_pool.py
@@ -744,6 +744,36 @@ def test_open_no_op(dsn):
         p.close()
 
 
+@pytest.mark.slow
+@pytest.mark.timing
+def test_open_wait(dsn, monkeypatch):
+    delay_connection(monkeypatch, 0.1)
+    with pytest.raises(pool.PoolTimeout):
+        p = pool.ConnectionPool(dsn, min_size=4, num_workers=1, open=False)
+        try:
+            p.open(wait=True, timeout=0.3)
+        finally:
+            p.close()
+
+    p = pool.ConnectionPool(dsn, min_size=4, num_workers=1, open=False)
+    try:
+        p.open(wait=True, timeout=0.5)
+    finally:
+        p.close()
+
+
+@pytest.mark.slow
+@pytest.mark.timing
+def test_open_as_wait(dsn, monkeypatch):
+    delay_connection(monkeypatch, 0.1)
+    with pytest.raises(pool.PoolTimeout):
+        with pool.ConnectionPool(dsn, min_size=4, num_workers=1) as p:
+            p.open(wait=True, timeout=0.3)
+
+    with pool.ConnectionPool(dsn, min_size=4, num_workers=1) as p:
+        p.open(wait=True, timeout=0.5)
+
+
 def test_reopen(dsn):
     p = pool.ConnectionPool(dsn)
     with p.connection() as conn:

diff --git a/tests/pool/test_pool_async.py b/tests/pool/test_pool_async.py
index d4aa5db8dde276b638ceb0a75c6260256280dde2..9618af1db9ad8feaa57e030028ac739d4e52e66c 100644 (file)
--- a/tests/pool/test_pool_async.py
+++ b/tests/pool/test_pool_async.py
@@ -719,6 +719,40 @@ async def test_open_no_op(dsn):
         await p.close()
 
 
+@pytest.mark.slow
+@pytest.mark.timing
+async def test_open_wait(dsn, monkeypatch):
+    delay_connection(monkeypatch, 0.1)
+    with pytest.raises(pool.PoolTimeout):
+        p = pool.AsyncConnectionPool(
+            dsn, min_size=4, num_workers=1, open=False
+        )
+        try:
+            await p.open(wait=True, timeout=0.3)
+        finally:
+            await p.close()
+
+    p = pool.AsyncConnectionPool(dsn, min_size=4, num_workers=1, open=False)
+    try:
+        await p.open(wait=True, timeout=0.5)
+    finally:
+        await p.close()
+
+
+@pytest.mark.slow
+@pytest.mark.timing
+async def test_open_as_wait(dsn, monkeypatch):
+    delay_connection(monkeypatch, 0.1)
+    with pytest.raises(pool.PoolTimeout):
+        async with pool.AsyncConnectionPool(
+            dsn, min_size=4, num_workers=1
+        ) as p:
+            await p.open(wait=True, timeout=0.3)
+
+    async with pool.AsyncConnectionPool(dsn, min_size=4, num_workers=1) as p:
+        await p.open(wait=True, timeout=0.5)
+
+
 async def test_reopen(dsn):
     p = pool.AsyncConnectionPool(dsn)
     async with p.connection() as conn: