Merge pull request #185 from peopledoc/ele_pool

Create aiopg pool lazily

Author: Éric Lemoine

docs/reference.rst

@@ -20,11 +20,8 @@ are accessible through :py:attr:`App.builtin_tasks`:
 Connectors
 ----------
 
-.. This does not indicate that create_with_pool* is an async classmethod because of
-   https://github.com/sphinx-doc/sphinx/issues/7189
-
 .. autoclass:: procrastinate.PostgresConnector
-    :members: create_with_pool, create_with_pool_async, close, close_async
+    :members: set_pool, close, close_async
 
 .. autoclass:: procrastinate.testing.InMemoryConnector
     :members: reset

procrastinate/aiopg_connector.py

@@ -7,7 +7,7 @@
 import psycopg2.sql
 from psycopg2.extras import Json, RealDictCursor
 
-from procrastinate import connector, sql, utils
+from procrastinate import connector, exceptions, sql, utils
 
 logger = logging.getLogger(__name__)
 
@@ -18,70 +18,33 @@
 class PostgresConnector(connector.BaseConnector):
     def __init__(
         self,
-        pool: aiopg.Pool,
+        *,
         json_dumps: Optional[Callable] = None,
         json_loads: Optional[Callable] = None,
+        **kwargs: Any,
     ):
         """
-        The pool connections are expected to have jsonb adapters.
+        Create a PostgreSQL connector. The connector uses an :py:func:`aiopg.Pool`,
+        which is either created automatically upon first use, or set through the
+        py:func:`PostgresConnector.set_pool` method.
 
-        See parameter details in :py:func:`PostgresConnector.create_with_pool`.
+        All other arguments than ``json_dumps`` and ``json_loads`` are passed to
+        ``aiopg.create_pool`` (see aiopg documentation__), with default values that
+        may differ from those of ``aiopg`` (see the list of parameters below).
 
-        Parameters
-        ----------
-        pool:
-            An aiopg pool, either externally configured or passed by
-            :py:func:`PostgresConnector.create_with_pool`.
-
-        """
-        self._pool = pool
-        self.json_dumps = json_dumps
-        self.json_loads = json_loads
-
-    async def close_async(self) -> None:
-        """
-        Closes the pool and awaits all connections to be released.
-        """
-        self._pool.close()
-        await self._pool.wait_closed()
-
-    def _wrap_json(self, arguments: Dict[str, Any]):
-        return {
-            key: Json(value, dumps=self.json_dumps)
-            if isinstance(value, dict)
-            else value
-            for key, value in arguments.items()
-        }
-
-    @classmethod
-    async def create_with_pool_async(
-        cls,
-        json_dumps: Optional[Callable] = None,
-        json_loads: Optional[Callable] = None,
-        **kwargs,
-    ) -> aiopg.Pool:
-        """
-        Creates a connector, and its connection pool, using the provided parameters.
-        All additional parameters will be used to create a
-        :py:func:`aiopg.Pool` (see the documentation__), sometimes with a different
-        default value.
-
-        When using this method, you explicitely take the responsibility for opening the
-        pool. It's your responsibility to call
-        :py:func:`procrastinate.PostgresConnector.close` or
-        :py:func:`procrastinate.PostgresConnector.close_async` to close connections
-        when your process ends.
-
-        .. __: https://aiopg.readthedocs.io/en/stable/core.html#aiopg.create_pool
         .. _psycopg2 doc: https://www.psycopg.org/docs/extras.html#json-adaptation
+        .. __: https://aiopg.readthedocs.io/en/stable/core.html#aiopg.create_pool
 
+        Parameters
+        ----------
         json_dumps:
             The JSON dumps function to use for serializing job arguments. Defaults to
             the function used by psycopg2. See the `psycopg2 doc`_.
         json_loads:
             The JSON loads function to use for deserializing job arguments. Defaults
-            to the function used by psycopg2. See the `psycopg2 doc`_. Unused if pool
-            is passed.
+            to the function used by psycopg2. See the `psycopg2 doc`_. Unused if the
+            pool is externally created and set into the connector through the
+            :py:func:`PostgresConnector.set_pool` method.
         dsn (Optional[str]):
             Passed to aiopg. Default is "" instead of None, which means if no argument
             is passed, it will connect to localhost:5432 instead of a Unix-domain
@@ -102,32 +65,82 @@ async def create_with_pool_async(
         maxsize (int):
             Passed to aiopg. Cannot be lower than 2, otherwise worker won't be
             functionning normally (one connection for listen/notify, one for executing
-            tasks)
+            tasks).
+        minsize (int):
+            Passed to aiopg. Initial connections are not opened when the connector
+            is created, but at first use of the pool.
         """
-        base_on_connect = kwargs.pop("on_connect", None)
+        self._pool: Optional[aiopg.Pool] = None
+        self.json_dumps = json_dumps
+        self.json_loads = json_loads
+        self._pool_args = self._adapt_pool_args(kwargs, json_loads)
+        self._lock = asyncio.Lock()
+
+    @staticmethod
+    def _adapt_pool_args(
+        pool_args: Dict[str, Any], json_loads: Optional[Callable]
+    ) -> Dict[str, Any]:
+        """
+        Adapt the pool args for ``aiopg``, using sensible defaults for Procrastinate.
+        """
+        base_on_connect = pool_args.pop("on_connect", None)
 
         async def on_connect(connection):
             if base_on_connect:
                 await base_on_connect(connection)
             if json_loads:
                 psycopg2.extras.register_default_jsonb(connection.raw, loads=json_loads)
 
-        defaults = {
+        final_args = {
             "dsn": "",
             "enable_json": False,
             "enable_hstore": False,
             "enable_uuid": False,
             "on_connect": on_connect,
             "cursor_factory": RealDictCursor,
         }
-        if "maxsize" in kwargs:
-            kwargs["maxsize"] = max(2, kwargs["maxsize"])
+        if "maxsize" in pool_args:
+            pool_args["maxsize"] = max(2, pool_args["maxsize"])
 
-        defaults.update(kwargs)
+        final_args.update(pool_args)
+        return final_args
 
-        pool = await aiopg.create_pool(**defaults)
+    async def close_async(self) -> None:
+        """
+        Close the pool and awaits all connections to be released.
+        """
+        if self._pool:
+            self._pool.close()
+            await self._pool.wait_closed()
+            self._pool = None
 
-        return cls(pool=pool, json_dumps=json_dumps, json_loads=json_loads)
+    def _wrap_json(self, arguments: Dict[str, Any]):
+        return {
+            key: Json(value, dumps=self.json_dumps)
+            if isinstance(value, dict)
+            else value
+            for key, value in arguments.items()
+        }
+
+    @staticmethod
+    async def _create_pool(pool_args: Dict[str, Any]) -> aiopg.Pool:
+        return await aiopg.create_pool(**pool_args)
+
+    def set_pool(self, pool: aiopg.Pool) -> None:
+        """
+        Set the connection pool. Raises an exception if the pool is already set.
+        """
+        if self._pool:
+            raise exceptions.PoolAlreadySet
+        self._pool = pool
+
+    async def _get_pool(self) -> aiopg.Pool:
+        if self._pool:
+            return self._pool
+        async with self._lock:
+            if not self._pool:
+                self.set_pool(await self._create_pool(self._pool_args))
+        return self._pool
 
     # Pools and single connections do not exactly share their cursor API:
     # - connection.cursor() is an async context manager (async with)
@@ -136,7 +149,8 @@ async def on_connect(connection):
     # a pool or from a connection
 
     async def execute_query(self, query: str, **arguments: Any) -> None:
-        with await self._pool.cursor() as cursor:
+        pool = await self._get_pool()
+        with await pool.cursor() as cursor:
             await cursor.execute(query, self._wrap_json(arguments))
 
     async def _execute_query_connection(
@@ -146,16 +160,17 @@ async def _execute_query_connection(
             await cursor.execute(query, self._wrap_json(arguments))
 
     async def execute_query_one(self, query: str, **arguments: Any) -> Dict[str, Any]:
-        with await self._pool.cursor() as cursor:
+        pool = await self._get_pool()
+        with await pool.cursor() as cursor:
             await cursor.execute(query, self._wrap_json(arguments))
 
             return await cursor.fetchone()
 
     async def execute_query_all(
         self, query: str, **arguments: Any
     ) -> List[Dict[str, Any]]:
-
-        with await self._pool.cursor() as cursor:
+        pool = await self._get_pool()
+        with await pool.cursor() as cursor:
             await cursor.execute(query, self._wrap_json(arguments))
 
             return await cursor.fetchall()
@@ -171,10 +186,11 @@ def make_dynamic_query(self, query: str, **identifiers: str) -> str:
     async def listen_notify(
         self, event: asyncio.Event, channels: Iterable[str]
     ) -> NoReturn:
+        pool = await self._get_pool()
         # We need to acquire a dedicated connection, and use the listen
         # query
         while True:
-            async with self._pool.acquire() as connection:
+            async with pool.acquire() as connection:
                 for channel_name in channels:
                     await self._execute_query_connection(
                         connection=connection,
@@ -215,4 +231,4 @@ def PostgresJobStore(*args, **kwargs):
     )
     logger.warning(f"Deprecation Warning: {message}")
     warnings.warn(DeprecationWarning(message))
-    return PostgresConnector.create_with_pool(*args, **kwargs)
+    return PostgresConnector(**kwargs)

procrastinate/exceptions.py

@@ -20,3 +20,7 @@ class LoadFromPathError(ImportError, ProcrastinateException):
 class JobRetry(ProcrastinateException):
     def __init__(self, scheduled_at: datetime.datetime):
         self.scheduled_at = scheduled_at
+
+
+class PoolAlreadySet(ProcrastinateException):
+    pass

procrastinate_demo/app.py

@@ -1,6 +1,6 @@
 import procrastinate
 
 app = procrastinate.App(
-    connector=procrastinate.PostgresConnector.create_with_pool(),  # type: ignore
+    connector=procrastinate.PostgresConnector(),
     import_paths=["procrastinate_demo.tasks"],
 )

tests/acceptance/app.py

@@ -25,7 +25,7 @@ def decode(dct):
 json_loads = functools.partial(json.loads, object_hook=decode)
 
 app = procrastinate.App(
-    connector=procrastinate.PostgresConnector.create_with_pool(  # type: ignore
+    connector=procrastinate.PostgresConnector(
         json_dumps=json_dumps, json_loads=json_loads
     )
 )
@@ -38,6 +38,10 @@ def sum_task(a, b):
 
 @app.task(queue="default")
 def sum_task_param(p1: Param, p2: Param):
+    if not isinstance(p1, Param):
+        raise Exception("wrong type for p1")
+    if not isinstance(p2, Param):
+        raise Exception("wrong type for p2")
     print(p1 + p2)
 
 

tests/conftest.py

@@ -37,9 +37,7 @@ def setup_db():
             )
             _execute(cursor, "CREATE DATABASE {}", "procrastinate_test_template")
 
-    connector = aiopg_connector.PostgresConnector.create_with_pool(
-        dbname="procrastinate_test_template"
-    )
+    connector = aiopg_connector.PostgresConnector(dbname="procrastinate_test_template")
     schema_manager = schema.SchemaManager(connector=connector)
     schema_manager.apply_schema()
     # We need to close the psycopg2 underlying connection synchronously
@@ -84,9 +82,7 @@ async def connection(connection_params):
 
 @pytest.fixture
 async def pg_connector(connection_params):
-    connector = await aiopg_connector.PostgresConnector.create_with_pool_async(
-        **connection_params
-    )
+    connector = aiopg_connector.PostgresConnector(**connection_params)
     yield connector
     await connector.close_async()