diff --git a/procrastinate/contrib/django/django_connector.py b/procrastinate/contrib/django/django_connector.py index cae12871a..3de066cff 100644 --- a/procrastinate/contrib/django/django_connector.py +++ b/procrastinate/contrib/django/django_connector.py @@ -9,16 +9,22 @@ from django.db.backends.base.base import BaseDatabaseWrapper from typing_extensions import LiteralString -from procrastinate import connector +from procrastinate import connector, psycopg_connector +from procrastinate.contrib.aiopg import aiopg_connector +from procrastinate.contrib.django import utils if TYPE_CHECKING: from psycopg.types.json import Jsonb + + is_psycopg3 = True else: try: - from django.db.backends.postgresql.psycopg_any import Jsonb + from django.db.backends.postgresql.psycopg_any import Jsonb, is_psycopg3 except ImportError: from psycopg2.extras import Json as Jsonb + is_psycopg3 = False + class DjangoConnector(connector.BaseAsyncConnector): """ @@ -111,3 +117,22 @@ async def listen_notify( raise NotImplementedError( "listen/notify is not supported with Django connector" ) + + 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. + + Returns + ------- + ``procrastinate.contrib.aiopg.AiopgConnector`` or ``procrastinate.contrib.psycopg3.PsycopgConnector`` + A connector that can be used in a worker + """ + alias = utils.get_setting("DATABASE_ALIAS", default="default") + + if is_psycopg3: + return psycopg_connector.PsycopgConnector( + kwargs=utils.connector_params(alias) + ) + else: + return aiopg_connector.AiopgConnector(**utils.connector_params(alias)) diff --git a/procrastinate/contrib/django/management/commands/procrastinate.py b/procrastinate/contrib/django/management/commands/procrastinate.py index 42596f212..e177b0c11 100644 --- a/procrastinate/contrib/django/management/commands/procrastinate.py +++ b/procrastinate/contrib/django/management/commands/procrastinate.py @@ -1,21 +1,11 @@ from __future__ import annotations import asyncio -from typing import TYPE_CHECKING from django.core.management.base import BaseCommand -from procrastinate import cli, psycopg_connector -from procrastinate.contrib.aiopg import aiopg_connector -from procrastinate.contrib.django import app, apps, utils - -if TYPE_CHECKING: - is_psycopg3 = True -else: - try: - from django.db.backends.postgresql.psycopg_any import is_psycopg3 - except ImportError: - is_psycopg3 = False +from procrastinate import cli +from procrastinate.contrib.django import app, django_connector class Command(BaseCommand): @@ -32,17 +22,6 @@ def add_arguments(self, parser): def handle(self, *args, **kwargs): kwargs = {k: v for k, v in kwargs.items() if k not in self._django_options} - kwargs["app"] = app.with_connector(self.get_connector()) + if isinstance(app.connector, django_connector.DjangoConnector): + kwargs["app"] = app.with_connector(app.connector.get_worker_connector()) asyncio.run(cli.execute_command(kwargs)) - - def get_connector(self): - # It's not possible to use the Django connector for the worker, so - # it's easier to just a classic Procrastinate connector. - alias = apps.get_setting("DATABASE_ALIAS", default="default") - - if is_psycopg3: - return psycopg_connector.PsycopgConnector( - kwargs=utils.connector_params(alias) - ) - else: - return aiopg_connector.AiopgConnector(**utils.connector_params(alias))