From e92777025cfd56656cafd6ce7782cce115db45ec Mon Sep 17 00:00:00 2001
From: Ahmad Nofal <nofal.ahmad@outlook.com>
Date: Wed, 22 Jun 2022 12:50:59 +0400
Subject: [PATCH] Fix atomic transaction not routing to the the correct DB
 (#324)

* Fix atomic transaction not routing to the the correct DB

* Added test case for atomic transaction to ensure its using the correct db

* simplify the test

* lint
---
 django_celery_results/managers.py |  2 +-
 t/unit/test_models.py             | 29 +++++++++++++++++++++++++++++
 2 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/django_celery_results/managers.py b/django_celery_results/managers.py
index 3242e1d7..7df78e4e 100644
--- a/django_celery_results/managers.py
+++ b/django_celery_results/managers.py
@@ -87,7 +87,7 @@ def get_all_expired(self, expires):
 
     def delete_expired(self, expires):
         """Delete all expired results."""
-        with transaction.atomic():
+        with transaction.atomic(using=self.db):
             raw_delete(queryset=self.get_all_expired(expires))
 
 
diff --git a/t/unit/test_models.py b/t/unit/test_models.py
index 4c2cff62..e655f4ce 100644
--- a/t/unit/test_models.py
+++ b/t/unit/test_models.py
@@ -209,3 +209,32 @@ class TransactionError(Exception):
                 raise TransactionError()
         except TransactionError:
             pass
+
+
+@pytest.mark.usefixtures('depends_on_current_app')
+class test_ModelsWithoutDefaultDB(TransactionTestCase):
+    """
+    This class to ensure all operations are done on the
+    same db we use and dont leak accidentally into another
+    db. we dont include the default db in databases as by
+    default an incorrect behavior would route there and
+    would not be detectable.
+
+    The tests will fail with the below error incase we
+    try to interact from a db other than the one we have
+    specified.
+
+    `AssertionError: Database connections to 'default' are
+    not allowed in this test`
+    """
+
+    non_default_test_db = 'secondary'
+    databases = [non_default_test_db]
+
+    def test_operations_with_atomic_transactions(self):
+        TaskResult.objects.db_manager(
+            self.non_default_test_db
+        ).delete_expired(expires=10)
+        GroupResult.objects.db_manager(
+            self.non_default_test_db
+        ).delete_expired(expires=10)