Merge pull request #946 from procrastinate-org/django-connector-uses-…

…psycopg3

.readthedocs.yml

@@ -5,15 +5,16 @@
 # Required
 version: 2
 
-sphinx:
-  fail_on_warning: true
-
 build:
-  os: ubuntu-20.04
+  os: "ubuntu-22.04"
   tools:
     python: "3.10"
   jobs:
+    post_create_environment:
+      - python -m pip install poetry
     post_install:
-      - pip install -U poetry
-      - poetry config virtualenvs.create false
-      - poetry install --with docs
+      - VIRTUAL_ENV=$READTHEDOCS_VIRTUALENV_PATH python -m poetry install --with docs
+
+sphinx:
+  configuration: docs/conf.py
+  fail_on_warning: true

docs/howto/django.md

@@ -87,8 +87,10 @@ itself, with some commands removed and the app is configured for you.
 You can also use other subcommands such as `./manage.py procrastinate defer`.
 
 :::{note}
-Procrastinate generates an app for you using a Psycopg or Aiopg connector
-depending on your Django setup, and connects using the `DATABASES` settings.
+Procrastinate generates an app for you using a `Psycopg` (by default) or
+`Aiopg` connector depending on whether `psycopg3` or `aiopg` is
+installed, and connects using the `DATABASES` settings. If neither library is
+installed, an error will be raised.
 :::
 
 ## Deferring jobs
@@ -189,9 +191,13 @@ if __name__ == "__main__":
     main()
 ```
 :::{note}
-The ``.get_worker_connector()`` method is only available on `DjangoConnector`
+The `.get_worker_connector()` method is only available on `DjangoConnector`
 and the API isn't guaranteed to be stable.
 :::
+:::{note}
+As mentionned above, either `psycopg` or `aiopg` need to be installed.
+`psycopg` will be used by default.
+:::
 
 ## Alternatives
 

procrastinate/contrib/django/django_connector.py

@@ -15,16 +15,12 @@
 
 if TYPE_CHECKING:
     from psycopg.types.json import Jsonb
-
-    is_psycopg3 = True
 else:
     try:
-        from django.db.backends.postgresql.psycopg_any import Jsonb, is_psycopg3
+        from django.db.backends.postgresql.psycopg_any import Jsonb
     except ImportError:
         from psycopg2.extras import Json as Jsonb
 
-        is_psycopg3 = False
-
 
 class DjangoConnector(connector.BaseAsyncConnector):
     """
@@ -122,6 +118,8 @@ def get_worker_connector(self) -> connector.BaseAsyncConnector:
         """
         The default DjangoConnector is not suitable for workers. This function
         returns a connector that uses the same database and is suitable for workers.
+        The type of connector returned is a `PsycopgConnector` if psycopg3 is installed,
+        otherwise an `AiopgConnector`.
 
         Returns
         -------
@@ -130,9 +128,15 @@ def get_worker_connector(self) -> connector.BaseAsyncConnector:
         """
         alias = utils.get_setting("DATABASE_ALIAS", default="default")
 
-        if is_psycopg3:
+        if utils.package_is_installed("psycopg3"):
             return psycopg_connector.PsycopgConnector(
                 kwargs=utils.connector_params(alias)
             )
-        else:
+        if utils.package_is_installed("aiopg"):
             return aiopg_connector.AiopgConnector(**utils.connector_params(alias))
+
+        raise django_exceptions.ImproperlyConfigured(
+            "You must install either psycopg(3) or aiopg to use "
+            "``./manage.py procrastinate`` or "
+            "``app.connector.get_worker_connector()``."
+        )

procrastinate/contrib/django/utils.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import importlib.util
 from typing import Any
 
 from django.conf import settings
@@ -31,3 +32,7 @@ def connector_params(alias: str = "default") -> dict[str, Any]:
 
 def get_setting(name: str, *, default) -> Any:
     return getattr(settings, f"PROCRASTINATE_{name}", default)
+
+
+def package_is_installed(name: str) -> bool:
+    return bool(importlib.util.find_spec(name))

tests/integration/contrib/django/test_django_connector.py

@@ -3,6 +3,8 @@
 import pytest
 from django.core import exceptions
 
+from procrastinate import psycopg_connector
+from procrastinate.contrib.aiopg import aiopg_connector
 from procrastinate.contrib.django import django_connector as django_connector_module
 
 
@@ -75,3 +77,40 @@ def test_execute_query_sync(django_connector):
         "SELECT obj_description('public.procrastinate_jobs'::regclass)"
     )
     assert result == [{"obj_description": "foo"}]
+
+
+@pytest.mark.parametrize(
+    "installed, type",
+    [
+        ({"psycopg3"}, psycopg_connector.PsycopgConnector),
+        ({"aiopg"}, aiopg_connector.AiopgConnector),
+    ],
+)
+async def test_get_worker_connector(installed, type, django_connector, mocker):
+    def package_is_installed(name):
+        return name in installed
+
+    mocker.patch(
+        "procrastinate.contrib.django.utils.package_is_installed",
+        side_effect=package_is_installed,
+    )
+
+    connector = django_connector.get_worker_connector()
+
+    assert isinstance(connector, type)
+
+    await connector.open_async()
+    try:
+        assert await connector.execute_query_one_async("SELECT 1 as x") == {"x": 1}
+    finally:
+        await connector.close_async()
+
+
+async def test_get_worker_connector__error(django_connector, mocker):
+    mocker.patch(
+        "procrastinate.contrib.django.utils.package_is_installed",
+        return_value=False,
+    )
+
+    with pytest.raises(exceptions.ImproperlyConfigured):
+        django_connector.get_worker_connector()

tests/unit/contrib/django/test_utils.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+import pytest
+
 from procrastinate.contrib.django import utils
 
 
@@ -14,3 +16,14 @@ def test_get_settings(settings):
 
 def test_get_settings_default():
     assert utils.get_setting("FOO", default="baz") == "baz"
+
+
+@pytest.mark.parametrize(
+    "package_name, expected",
+    [
+        ("foo" * 30, False),
+        ("pytest", True),
+    ],
+)
+def test_package_is_installed(package_name, expected):
+    assert utils.package_is_installed(package_name) is expected