Add 'init' parameter to create_pool(). Addresses issue #80.

Author: 1st1

asyncpg/connection.py

@@ -134,7 +134,8 @@ def transaction(self, *, isolation='read_committed', readonly=False,
         :param deferrable: Specifies whether or not this transaction is
                            deferrable.
 
-        .. _`PostgreSQL documentation`: https://www.postgresql.org/docs/current/static/sql-set-transaction.html
+        .. _`PostgreSQL documentation`: https://www.postgresql.org/docs/\
+                                        current/static/sql-set-transaction.html
         """
         return transaction.Transaction(self, isolation, readonly, deferrable)
 

asyncpg/pool.py

@@ -26,13 +26,14 @@ class Pool:
                  '_connect_args', '_connect_kwargs',
                  '_working_addr', '_working_opts',
                  '_con_count', '_max_queries', '_connections',
-                 '_initialized', '_closed', '_setup')
+                 '_initialized', '_closed', '_setup', '_init')
 
     def __init__(self, *connect_args,
                  min_size,
                  max_size,
                  max_queries,
                  setup,
+                 init,
                  loop,
                  **connect_kwargs):
 
@@ -57,6 +58,7 @@ def __init__(self, *connect_args,
         self._max_queries = max_queries
 
         self._setup = setup
+        self._init = init
 
         self._connect_args = connect_args
         self._connect_kwargs = connect_kwargs
@@ -88,10 +90,13 @@ async def _new_connection(self):
                                            loop=self._loop,
                                            **self._working_opts)
 
+        if self._init is not None:
+            await self._init(con)
+
         self._connections.add(con)
         return con
 
-    async def _init(self):
+    async def _initialize(self):
         if self._initialized:
             return
         if self._closed:
@@ -219,10 +224,10 @@ def _reset(self):
         self._queue = asyncio.Queue(maxsize=self._maxsize, loop=self._loop)
 
     def __await__(self):
-        return self._init().__await__()
+        return self._initialize().__await__()
 
     async def __aenter__(self):
-        await self._init()
+        await self._initialize()
         return self
 
     async def __aexit__(self, *exc):
@@ -261,6 +266,7 @@ def create_pool(dsn=None, *,
                 max_size=10,
                 max_queries=50000,
                 setup=None,
+                init=None,
                 loop=None,
                 **connect_kwargs):
     r"""Create a connection pool.
@@ -296,16 +302,22 @@ def create_pool(dsn=None, *,
     :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 coroutine setup: A coroutine to initialize a connection right before
+    :param coroutine setup: A coroutine to prepare a connection right before
                             it is returned from :meth:`~pool.Pool.acquire`.
                             An example use case would be to automatically
                             set up notifications listeners for all connections
                             of a pool.
+    :param coroutine init:  A coroutine to initialize a connection when it
+                            is created. An example use case would be to setup
+                            type codecs with
+                            :meth:`~asyncpg.connection.Connection.\
+                            set_builtin_type_codec` or :meth:`~asyncpg.\
+                            connection.Connection.set_type_codec`.
     :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, setup=setup,
+                max_queries=max_queries, loop=loop, setup=setup, init=init,
                 **connect_kwargs)

tests/test_pool.py

@@ -102,6 +102,33 @@ async def setup(con):
 
         self.assertIs(con, await fut)
 
+    async def test_pool_07(self):
+        cons = set()
+
+        async def setup(con):
+            if con not in cons:
+                raise RuntimeError('init was not called before setup')
+
+        async def init(con):
+            if con in cons:
+                raise RuntimeError('init was called more than once')
+            cons.add(con)
+
+        async def user(pool):
+            async with pool.acquire() as con:
+                if con not in cons:
+                    raise RuntimeError('init was not called')
+
+        async with self.create_pool(database='postgres',
+                                    min_size=2, max_size=5,
+                                    init=init,
+                                    setup=setup) as pool:
+            users = asyncio.gather(*[user(pool) for _ in range(10)],
+                                   loop=self.loop)
+            await users
+
+        self.assertEqual(len(cons), 5)
+
     async def test_pool_auth(self):
         if not self.cluster.is_managed():
             self.skipTest('unmanaged cluster')