Forward connection pool parameters from django to psycopg

Author: ticosax

procrastinate/contrib/django/django_connector.py

@@ -166,18 +166,19 @@ def get_worker_connector(self) -> connector.BaseAsyncConnector:
         """
         alias = settings.settings.DATABASE_ALIAS
 
+        kwargs, pool_kwargs = utils.connector_params(alias)
         if utils.package_is_installed("psycopg") and utils.package_is_version(
             "psycopg", 3
         ):
             from procrastinate import psycopg_connector
 
             return psycopg_connector.PsycopgConnector(
-                kwargs=utils.connector_params(alias)
+                kwargs=kwargs, pool_kwargs=pool_kwargs
             )
         if utils.package_is_installed("aiopg"):
             from procrastinate.contrib.aiopg import aiopg_connector
 
-            return aiopg_connector.AiopgConnector(**utils.connector_params(alias))
+            return aiopg_connector.AiopgConnector(**kwargs)
 
         raise django_exceptions.ImproperlyConfigured(
             "You must install either psycopg(3) or aiopg to use "

procrastinate/contrib/django/utils.py

@@ -7,7 +7,7 @@
 from django.db import connections
 
 
-def connector_params(alias: str = "default") -> dict[str, Any]:
+def connector_params(alias: str = "default") -> tuple[dict[str, Any], dict[str, Any]]:
     """
     Returns parameters for in a format that is suitable to be passed to a
     connector constructor (see `howto/django`).
@@ -27,7 +27,7 @@ def connector_params(alias: str = "default") -> dict[str, Any]:
     params = wrapper.get_connection_params()
     params.pop("cursor_factory", None)
     params.pop("context", None)
-    return params
+    return params, wrapper.settings_dict.get("OPTIONS", {}).get("pool", {})
 
 
 def package_is_installed(name: str) -> bool:

procrastinate/psycopg_connector.py

@@ -49,6 +49,7 @@ def __init__(
         pool_factory: Callable[
             ..., psycopg_pool.AsyncConnectionPool
         ] = psycopg_pool.AsyncConnectionPool,
+        pool_kwargs: dict[str, Any] | None = None,
         **kwargs: Any,
     ):
         """
@@ -93,6 +94,7 @@ def __init__(
         self._json_loads = json_loads
         self._json_dumps = json_dumps
         self._pool_args = kwargs
+        self._pool_kwargs = pool_kwargs
         self._sync_connector: connector.BaseConnector | None = None
 
     def get_sync_connector(self) -> connector.BaseConnector:
@@ -140,24 +142,27 @@ async def open_async(
             self._pool_externally_set = True
             self._async_pool = pool
         else:
-            self._async_pool = await self._create_pool(self._pool_args)
+            self._async_pool = await self._create_pool(
+                self._pool_args, self._pool_kwargs
+            )
 
             await self._async_pool.open(wait=True)  # type: ignore
 
     @wrap_exceptions()
     async def _create_pool(
-        self,
-        pool_args: dict[str, Any],
+        self, pool_args: dict[str, Any], pool_kwargs: dict[str, Any] | None
     ) -> psycopg_pool.AsyncConnectionPool:
+        pool_kwargs = pool_kwargs if pool_kwargs is not None else {}
         return self._pool_factory(
-            **pool_args,
+            kwargs=pool_args.get("kwargs", {}),
             # Not specifying open=False raises a warning and will be deprecated.
             # It makes sense, as we can't really make async I/Os in a constructor.
             open=False,
             # Enables a check that will ensure the connections returned when
             # using the pool are still alive. If they have been closed by the
             # database, they will be seamlessly replaced by a new connection.
             check=psycopg_pool.AsyncConnectionPool.check_connection,
+            **pool_kwargs,
         )
 
     @wrap_exceptions()

tests/acceptance/django_settings.py

@@ -8,6 +8,7 @@
         "ENGINE": "django.db.backends.postgresql",
         "NAME": os.environ.get("PGDATABASE", "procrastinate"),
         "TEST": {"NAME": "procrastinate_django_test"},
+        "OPTIONS": {"pool": {"min_size": 1, "max_size": 4}},
     },
 }
 INSTALLED_APPS = [

tests/unit/contrib/django/test_django_connector.py

@@ -6,6 +6,7 @@
 
 from procrastinate import exceptions
 from procrastinate.contrib.django import django_connector as django_connector_module
+from procrastinate.psycopg_connector import PsycopgConnector
 
 
 def test_wrap_exceptions__no_cause():
@@ -18,3 +19,10 @@ def test_wrap_exceptions__with_cause():
     with pytest.raises(exceptions.ConnectorException):
         with django_connector_module.wrap_exceptions():
             raise django_db.DatabaseError from psycopg_errors.Error()
+
+
+def test_get_worker_connector():
+    django_connector = django_connector_module.DjangoConnector()
+    worker_connector = django_connector.get_worker_connector()
+    assert isinstance(worker_connector, PsycopgConnector)
+    assert worker_connector._pool_kwargs == {"min_size": 1, "max_size": 4}

tests/unit/contrib/django/test_utils.py

@@ -6,7 +6,9 @@
 
 
 def test_connector_params():
-    assert "dbname" in utils.connector_params()
+    params, pool_params = utils.connector_params()
+    assert "dbname" in params
+    assert pool_params == {"min_size": 1, "max_size": 4}
 
 
 @pytest.mark.parametrize(

tests/unit/test_psycopg_connector.py

@@ -50,7 +50,7 @@ async def test_open_async_no_pool_specified(mocker, connector):
 
     await connector.open_async()
 
-    assert connector._create_pool.call_count == 1
+    connector._create_pool.assert_called_once_with({}, None)
     assert connector._async_pool.open.await_count == 1