allow atomic on external db (Koed00#14)

Co-authored-by: Rajeev KL <[email protected]>
msabatier · Oct 15, 2022 · ce98aaa · ce98aaa
1 parent 00fe611
commit ce98aaa
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/django_q/cluster.py b/django_q/cluster.py
@@ -496,7 +496,9 @@ def save_task(task, broker: Broker):
             if Conf.SAVE_LIMIT_PER == "func":
                 value = get_func_repr(value)
             filters[Conf.SAVE_LIMIT_PER] = value
-        with db.transaction.atomic():
+
+        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):
             last = Success.objects.filter(**filters).select_for_update().last()
             if task["success"] and 0 < Conf.SAVE_LIMIT <= Success.objects.filter(**filters).count():
                 last.delete()
@@ -599,7 +601,7 @@ def scheduler(broker: Broker = None):
         broker = get_broker()
     close_old_django_connections()
     try:
-        database_to_use = {"using": Conf.ORM} if not Conf.HAS_REPLICA else {}
+        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):
             for s in (
                 Schedule.objects.select_for_update()