Skip to content

Commit

Permalink
move Command.get_connector to DjangoConnector.get_worker_connector
Browse files Browse the repository at this point in the history
  • Loading branch information
ewjoachim committed Jan 30, 2024
1 parent 8b7eecf commit 3347834
Show file tree
Hide file tree
Showing 2 changed files with 31 additions and 27 deletions.
29 changes: 27 additions & 2 deletions procrastinate/contrib/django/django_connector.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
"""
Expand Down Expand Up @@ -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))
29 changes: 4 additions & 25 deletions procrastinate/contrib/django/management/commands/procrastinate.py
Original file line number Diff line number Diff line change
@@ -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):
Expand All @@ -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))

0 comments on commit 3347834

Please # to comment.