Fix use of database router for write queries and remove Conf.HAS_REPLICA (#61)

Co-authored-by: Marc Sabatier

django_q/cluster.py

@@ -20,7 +20,6 @@ except core.exceptions.AppRegistryNotReady:
 
     django.setup()
 
-from django.conf import settings
 from django.utils import timezone
 from django.utils.translation import gettext_lazy as _
 
@@ -541,12 +540,7 @@ def save_task(task, broker: Broker):
                 value = get_func_repr(value)
             filters[Conf.SAVE_LIMIT_PER] = value
 
-        database_to_use = (
-            {"using": Conf.ORM if Conf.ORM else Schedule.objects.db}
-            if not Conf.HAS_REPLICA
-            else {}
-        )
-        with db.transaction.atomic(**database_to_use):
+        with db.transaction.atomic(using=db.router.db_for_write(Success)):
             last = Success.objects.filter(**filters).select_for_update().last()
             if (
                 task["success"]
@@ -652,12 +646,7 @@ def scheduler(broker: Broker = None):
         broker = get_broker()
     close_old_django_connections()
     try:
-        database_to_use = (
-            {"using": Conf.ORM if Conf.ORM else Schedule.objects.db}
-            if not Conf.HAS_REPLICA
-            else {}
-        )
-        with db.transaction.atomic(**database_to_use):
+        with db.transaction.atomic(using=db.router.db_for_write(Schedule)):
             for s in (
                 Schedule.objects.select_for_update()
                 .exclude(repeats=0)

django_q/conf.py

@@ -54,9 +54,6 @@ class Conf:
     # ORM broker
     ORM = conf.get("orm", None)
 
-    # ORM support for read/write replicas
-    HAS_REPLICA = conf.get("has_replica", False)
-
     # Custom broker class
     BROKER_CLASS = conf.get("broker_class", None)
 

django_q/tests/test_scheduler.py

@@ -50,25 +50,11 @@ def orm_broker(monkeypatch) -> None:
     monkeypatch.setattr(Conf, "ORM", "default")
 
 
-@pytest.fixture
-def orm_no_replica_broker(orm_broker, monkeypatch) -> Broker:
-    """Generates a Broker with a disabled read replica database configuration."""
-    monkeypatch.setattr(Conf, "HAS_REPLICA", False)
-    return get_broker(list_key="scheduler_test:q")
-
-
-@pytest.fixture
-def orm_replica_broker(orm_broker, monkeypatch) -> Broker:
-    """Generates a Broker with read replica database configuration."""
-    monkeypatch.setattr(Conf, "HAS_REPLICA", True)
-    return get_broker(list_key="scheduler_test:q")
-
-
 REPLICA_DATABASE_ROUTERS = [
     f"{TestingReplicaDatabaseRouter.__module__}.{TestingReplicaDatabaseRouter.__name__}"
 ]
 REPLICA_DATABASES = {
-    "default": {
+    "writable": {
         "ENGINE": "django.db.backends.sqlite3",
         "NAME": os.path.join(BASE_DIR, "db.sqlite3"),
     },
@@ -453,39 +439,18 @@ def test_intended_schedule_kwarg(broker, monkeypatch):
     DATABASE_ROUTERS=REPLICA_DATABASE_ROUTERS, DATABASES=REPLICA_DATABASES
 )
 @pytest.mark.django_db
-def test_scheduler_atomic_transaction_must_specify_a_database_when_no_replicas_are_used(
-    orm_no_replica_broker: Broker,
-):
-    """
-    GIVEN a environment without a read replica database
-    WHEN the scheduler is called
-    THEN the transaction atomic must be called using the configured database in the
-    Conf.ORM settings.
-    """
-    broker = orm_no_replica_broker
-    with mock.patch("django_q.cluster.db") as mocked_db:
-        scheduler(broker=broker)
-        # The router should correctly set the database to use!
-        mocked_db.transaction.atomic.assert_called_with(using=broker.connection.db)
-
-
-@override_settings(
-    DATABASE_ROUTERS=REPLICA_DATABASE_ROUTERS, DATABASES=REPLICA_DATABASES
-)
-@pytest.mark.django_db
-def test_scheduler_atomic_must_specify_no_db_when_read_write_replicas_are_used(
-    orm_replica_broker: Broker,
+def test_scheduler_atomic_must_specify_the_write_db(
+    orm_broker: Broker,
 ):
     """
     GIVEN a environment with a read/write configured replica database
     WHEN the scheduler is called
-    THEN the transaction must be called without a specific database, thus letting the
-    database router pick.
+    THEN the transaction must be called with the write database.
     """
-    with mock.patch("django_q.cluster.db") as mocked_db:
-        scheduler(broker=orm_replica_broker)
-        # No specific databases should be set here, this is the job of the router!
-        mocked_db.transaction.atomic.assert_called_with()
+    broker = get_broker(list_key="scheduler_test:q")
+    with mock.patch("django_q.cluster.db.transaction") as mocked_db:
+        scheduler(broker=broker)
+        mocked_db.atomic.assert_called_with(using="writable")
 
 
 @override_settings(
@@ -493,20 +458,17 @@ def test_scheduler_atomic_must_specify_no_db_when_read_write_replicas_are_used(
 )
 @pytest.mark.django_db
 def test_scheduler_atomic_must_specify_the_database_based_on_router_redirection(
-    orm_no_replica_broker: Broker,
+    orm_broker: Broker,
 ):
     """
     GIVEN a environment without a read replica database
     WHEN the scheduler is called
-    THEN the transaction atomic must be called using the configured database in the
-    Conf.ORM settings.
+    THEN the transaction atomic must be called using the default connection.
     """
-    broker = orm_no_replica_broker
-    with mock.patch("django_q.cluster.db") as mocked_db:
+    broker = get_broker(list_key="scheduler_test:q")
+    with mock.patch("django_q.cluster.db.transaction") as mocked_db:
         scheduler(broker=broker)
-        # The router should correctly set the database to use!
-        assert broker.connection.db == "default"
-        mocked_db.transaction.atomic.assert_called_with(using=broker.connection.db)
+        mocked_db.atomic.assert_called_with(using="default")
 
 
 def test_localtime():

django_q/tests/testing_utilities/multiple_database_routers.py

@@ -12,9 +12,9 @@ class TestingReplicaDatabaseRouter:
 
     def db_for_write(self, model, **hints):
         """
-        Always write to DEFAULT database
+        Always write to WRITABLE database
         """
-        return "default"
+        return "writable"
 
 
 class TestingMultipleAppsDatabaseRouter: