pool: Make pool.acquire() compatible with 'async with'

Author: 1st1

asyncpg/pool.py

@@ -12,18 +12,27 @@
 
 
 class Pool:
+    """A connection pool.
+
+    Connection pool can be used to manage a set of connections to the database.
+    Connections are first acquired from the pool, then used, and then released
+    back to the pool.  Once a connection is released, it's reset to close all
+    open cursors and other resources *except* prepared statements.
+
+    Pools are created by calling :func:`~asyncpg.pool.create_pool`.
+    """
 
     __slots__ = ('_queue', '_loop', '_minsize', '_maxsize',
                  '_connect_args', '_connect_kwargs',
                  '_working_addr', '_working_opts',
                  '_con_count', '_max_queries', '_connections',
-                 '_initialized')
+                 '_initialized', '_closed')
 
     def __init__(self, *connect_args,
-                 min_size=10,
-                 max_size=10,
-                 max_queries=50000,
-                 loop=None,
+                 min_size,
+                 max_size,
+                 max_queries,
+                 loop,
                  **connect_kwargs):
 
         if loop is None:
@@ -54,6 +63,8 @@ def __init__(self, *connect_args,
 
         self._reset()
 
+        self._closed = False
+
     async def _new_connection(self, timeout=None):
         if self._working_addr is None:
             con = await connection.connect(*self._connect_args,
@@ -82,6 +93,8 @@ async def _new_connection(self, timeout=None):
     async def _init(self):
         if self._initialized:
             return
+        if self._closed:
+            raise exceptions.InterfaceError('pool is closed')
 
         for _ in range(self._minsize):
             self._con_count += 1
@@ -95,7 +108,32 @@ async def _init(self):
         self._initialized = True
         return self
 
-    async def acquire(self, *, timeout=None):
+    def acquire(self, *, timeout=None):
+        """Acquire a database connection from the pool.
+
+        :param float timeout: A timeout for acquiring a Connection.
+        :return: An instance of :class:`~asyncpg.connection.Connection`.
+
+        Can be used in an ``await`` expression or with an ``async with`` block.
+
+        .. code-block:: python
+
+            async with pool.acquire() as con:
+                await con.execute(...)
+
+        Or:
+
+        .. code-block:: python
+
+            con = await pool.acquire()
+            try:
+                await con.execute(...)
+            finally:
+                await pool.release(con)
+        """
+        return PoolAcquireContext(self, timeout)
+
+    async def _acquire(self, timeout):
         self._check_init()
 
         try:
@@ -119,6 +157,7 @@ async def acquire(self, *, timeout=None):
                                           loop=self._loop)
 
     async def release(self, connection):
+        """Release a database connection back to the pool."""
         self._check_init()
         if connection.is_closed():
             self._con_count -= 1
@@ -132,22 +171,32 @@ async def release(self, connection):
             self._queue.put_nowait(connection)
 
     async def close(self):
+        """Gracefully close all connections in the pool."""
+        if self._closed:
+            return
         self._check_init()
+        self._closed = True
         coros = []
         for con in self._connections:
             coros.append(con.close())
         await asyncio.gather(*coros, loop=self._loop)
         self._reset()
 
     def terminate(self):
+        """Terminate all connections in the pool."""
+        if self._closed:
+            return
         self._check_init()
+        self._closed = True
         for con in self._connections:
             con.terminate()
         self._reset()
 
     def _check_init(self):
         if not self._initialized:
             raise exceptions.InterfaceError('pool is not initialized')
+        if self._closed:
+            raise exceptions.InterfaceError('pool is closed')
 
     def _reset(self):
         self._connections = set()
@@ -166,5 +215,77 @@ async def __aexit__(self, *exc):
         await self.close()
 
 
-def create_pool(*args, **kwargs):
-    return Pool(*args, **kwargs)
+class PoolAcquireContext:
+
+    __slots__ = ('timeout', 'connection', 'done', 'pool')
+
+    def __init__(self, pool, timeout):
+        self.pool = pool
+        self.timeout = timeout
+        self.connection = None
+        self.done = False
+
+    async def __aenter__(self):
+        if self.connection is not None or self.done:
+            raise exceptions.InterfaceError('a connection is already acquired')
+        self.connection = await self.pool._acquire(self.timeout)
+        return self.connection
+
+    async def __aexit__(self, *exc):
+        self.done = True
+        con = self.connection
+        self.connection = None
+        await self.pool.release(con)
+
+    def __await__(self):
+        self.done = True
+        return self.pool._acquire(self.timeout).__await__()
+
+
+def create_pool(dsn=None, *,
+                min_size=10,
+                max_size=10,
+                max_queries=50000,
+                loop=None,
+                **connect_kwargs):
+    r"""Create a connection pool.
+
+    Can be used either with an ``async with`` block:
+
+    .. code-block:: python
+
+        async with asyncpg.create_pool(user='postgres',
+                                       command_timeout=60) as pool:
+            async with poll.acquire() as con:
+                await con.fetch('SELECT 1')
+
+    Or directly with ``await``:
+
+    .. code-block:: python
+
+        pool = await asyncpg.create_pool(user='postgres', command_timeout=60)
+        con = await poll.acquire()
+        try:
+            await con.fetch('SELECT 1')
+        finally:
+            await pool.release(con)
+
+    :param str dsn: Connection arguments specified using as a single string in
+                    the following format:
+                    ``postgres://user:pass@host:port/database?option=value``.
+
+    :param \*\*connect_kwargs: Keyword arguments for the
+                               :func:`~asyncpg.connection.connect` function.
+    :param int min_size: Number of connection the pool will be initialized
+                         with.
+    :param int max_size: Max number of connections in the pool.
+    :param int max_queries: Number of queries after a connection is closed
+                            and replaced with a new connection.
+    :param loop: An asyncio event loop instance.  If ``None``, the default
+                 event loop will be used.
+    :return: An instance of :class:`~asyncpg.pool.Pool`.
+    """
+    return Pool(dsn,
+                min_size=min_size, max_size=max_size,
+                max_queries=max_queries, loop=loop,
+                **connect_kwargs)

tests/test_pool.py

@@ -73,7 +73,27 @@ async def test_pool_04(self):
         con.terminate()
         await pool.release(con)
 
+        async with pool.acquire(timeout=0.1):
+            con.terminate()
+
         con = await pool.acquire(timeout=0.1)
         self.assertEqual(await con.fetchval('SELECT 1'), 1)
 
         await pool.close()
+
+    async def test_pool_05(self):
+        for n in {1, 3, 5, 10, 20, 100}:
+            with self.subTest(tasksnum=n):
+                addr = self.cluster.get_connection_addr()
+                pool = await asyncpg.create_pool(host=addr[0], port=addr[1],
+                                                 database='postgres',
+                                                 loop=self.loop, min_size=5,
+                                                 max_size=10)
+
+                async def worker():
+                    async with pool.acquire() as con:
+                        self.assertEqual(await con.fetchval('SELECT 1'), 1)
+
+                tasks = [worker() for _ in range(n)]
+                await asyncio.gather(*tasks, loop=self.loop)
+                await pool.close()