Fix set_rollback for non_atomic_requests views

rest_framework/deprecation.py

@@ -21,6 +21,10 @@ class RemovedInDRF320Warning(PendingDeprecationWarning):
     pass
 
 
+class RemovedInDRF321Warning(PendingDeprecationWarning):
+    pass
+
+
 # Aliases that always track the current release cycle, so that projects can
 # filter on them without editing their configuration on every DRF release.
 RemovedInNextDRFVersionWarning = RemovedInDRF319Warning

rest_framework/views.py

@@ -1,6 +1,8 @@
 """
 Provides an APIView class that is the base of all views in REST framework.
 """
+import warnings
+
 from django.conf import settings
 from django.core.exceptions import PermissionDenied
 from django.db import connections, models
@@ -13,6 +15,7 @@
 
 from rest_framework import exceptions, status
 from rest_framework.compat import split_header_value
+from rest_framework.deprecation import RemovedInDRF321Warning
 from rest_framework.request import Request
 from rest_framework.response import Response
 from rest_framework.schemas import DefaultSchema
@@ -63,10 +66,45 @@ def get_view_description(view, html=False):
     return description
 
 
-def set_rollback():
+def set_rollback(request=None):
+    """
+    Mark databases for rollback that were atomically wrapped by
+    ATOMIC_REQUESTS.
+
+    If `request` is provided, we mirror Django's
+    ``BaseHandler.make_view_atomic`` logic and skip databases where the
+    resolved view is decorated with ``@transaction.non_atomic_requests``.
+    This prevents DRF from accidentally marking outer transactions (for
+    example, the one Django's ``TestCase`` wraps each test in) for
+    rollback.
+
+    Passing ``request`` is recommended. Calling without it is deprecated.
+    """
+    if request is None:
+        warnings.warn(
+            "Calling `set_rollback()` without a `request` argument is "
+            "deprecated. `set_rollback(request)` will be required in a "
+            "future release.",
+            RemovedInDRF321Warning,
+            stacklevel=2,
+        )
+        non_atomic_requests = None
+    else:
+        # Best-effort: RequestFactory may not set `resolver_match`.
+        view_func = getattr(
+            getattr(request, 'resolver_match', None), 'func', None
+        )
+        non_atomic_requests = getattr(view_func, '_non_atomic_requests', None)
+
     for db in connections.all(initialized_only=True):
-        if db.settings_dict['ATOMIC_REQUESTS'] and db.in_atomic_block:
-            db.set_rollback(True)
+        if not db.settings_dict['ATOMIC_REQUESTS']:
+            continue
+        if not db.in_atomic_block:
+            continue
+        # Mirror Django's make_view_atomic.
+        if non_atomic_requests is not None and db.alias in non_atomic_requests:
+            continue
+        db.set_rollback(True)
 
 
 def exception_handler(exc, context):
@@ -96,7 +134,7 @@ def exception_handler(exc, context):
         else:
             data = {'detail': exc.detail}
 
-        set_rollback()
+        set_rollback(context.get('request'))
         return Response(data, status=exc.status_code, headers=headers)
 
     return None
@@ -230,7 +268,7 @@ def get_exception_handler_context(self):
             'view': self,
             'args': getattr(self, 'args', ()),
             'kwargs': getattr(self, 'kwargs', {}),
-            'request': getattr(self, 'request', None)
+            'request': getattr(self, 'request', None),
         }
 
     def get_view_name(self):

tests/test_atomic_requests.py

@@ -1,11 +1,13 @@
 import unittest
+import warnings
 
 from django.db import connection, connections, transaction
 from django.http import Http404
 from django.test import TestCase, TransactionTestCase, override_settings
 from django.urls import path
 
 from rest_framework import status
+from rest_framework.deprecation import RemovedInDRF321Warning
 from rest_framework.exceptions import APIException
 from rest_framework.response import Response
 from rest_framework.test import APIRequestFactory
@@ -39,11 +41,12 @@ def dispatch(self, *args, **kwargs):
         return super().dispatch(*args, **kwargs)
 
     def get(self, request, *args, **kwargs):
-        BasicModel.objects.all()
+        list(BasicModel.objects.all())
         raise Http404
 
 
 urlpatterns = (
+    path('non-atomic-exception', NonAtomicAPIExceptionView.as_view()),
     path('', NonAtomicAPIExceptionView.as_view()),
 )
 
@@ -95,7 +98,8 @@ def test_generic_exception_delegate_transaction_management(self):
             # 2 - insert
             # 3 - release savepoint
             with transaction.atomic():
-                self.assertRaises(Exception, self.view, request)
+                with self.assertRaises(Exception):
+                    self.view(request)
                 assert not transaction.get_rollback()
         assert BasicModel.objects.count() == 1
 
@@ -174,15 +178,18 @@ class NonAtomicDBTransactionAPIExceptionTests(TransactionTestCase):
     def setUp(self):
         connections.databases['default']['ATOMIC_REQUESTS'] = True
 
-    def tearDown(self):
-        connections.databases['default']['ATOMIC_REQUESTS'] = False
+        @self.addCleanup
+        def restore_atomic_requests():
+            connections.databases['default']['ATOMIC_REQUESTS'] = False
 
     def test_api_exception_rollback_transaction_non_atomic_view(self):
-        response = self.client.get('/')
+        response = self.client.get('/non-atomic-exception')
 
-        # without checking connection.in_atomic_block view raises 500
-        # due attempt to rollback without transaction
+        # without check for db.in_atomic_block, would raise 500 due to attempt
+        # to rollback without transaction
         assert response.status_code == status.HTTP_404_NOT_FOUND
+        # Check we can still perform DB queries
+        list(BasicModel.objects.all())
 
 
 @unittest.skipUnless(
@@ -197,8 +204,9 @@ def tearDown(self):
         connections.databases['default']['ATOMIC_REQUESTS'] = False
 
     def test_marks_initialized_atomic_connection_for_rollback(self):
+        request = factory.post('/')
         with transaction.atomic():
-            set_rollback()
+            set_rollback(request)
             assert transaction.get_rollback()
 
 
@@ -233,7 +241,8 @@ def tearDown(self):
         super().tearDown()
 
     def test_does_not_initialize_unused_connections(self):
-        set_rollback()
+        request = factory.post('/')
+        set_rollback(request)
         assert not hasattr(connections._connections, 'secondary')
 
 
@@ -252,3 +261,71 @@ def test_api_exception_leaves_unused_connection_uninitialized(self):
         response = self.view(request)
         assert response.status_code == status.HTTP_500_INTERNAL_SERVER_ERROR
         assert not hasattr(connections._connections, 'secondary')
+
+
+@override_settings(ROOT_URLCONF='tests.test_atomic_requests')
+class NonAtomicViewUnderTestCaseTests(TestCase):
+    """
+    Regression test for #6921.
+
+    A view decorated with ``@transaction.non_atomic_requests`` must not
+    poison the outer transaction that Django's ``TestCase`` wraps each
+    test in. Before the fix, ``set_rollback`` would see
+    ``connection.in_atomic_block == True`` and mark the outer transaction
+    for rollback, breaking any subsequent queries in the test.
+    """
+    def setUp(self):
+        connections.databases['default']['ATOMIC_REQUESTS'] = True
+
+    def tearDown(self):
+        connections.databases['default']['ATOMIC_REQUESTS'] = False
+
+    def test_non_atomic_view_does_not_mark_outer_transaction_for_rollback(self):
+        # TestCase has already opened an atomic block around this test.
+        assert transaction.get_rollback() is False
+
+        # Sending the request through the test client populates
+        # ``request.resolver_match``, so ``set_rollback`` can see the
+        # view's ``_non_atomic_requests`` attribute and skip the database.
+        response = self.client.get('/')
+        assert response.status_code == status.HTTP_404_NOT_FOUND
+
+        # The outer transaction must remain usable. Without the fix,
+        # ``set_rollback`` would have marked it for rollback and the
+        # following query would raise.
+        assert transaction.get_rollback() is False
+        BasicModel.objects.create()
+
+
+class SetRollbackDeprecationTests(TestCase):
+    """
+    Calling ``set_rollback()`` without a ``request`` argument should emit
+    a ``RemovedInDRF321Warning``.
+    """
+    def test_calling_without_request_emits_warning(self):
+        with warnings.catch_warnings(record=True) as caught:
+            warnings.simplefilter('always')
+            set_rollback()
+        assert len(caught) == 1
+        assert issubclass(caught[0].category, RemovedInDRF321Warning)
+        assert "without a `request` argument is deprecated" in str(caught[0].message)
+
+
+class SetRollbackWithoutResolverMatchTests(TestCase):
+    """
+    ``APIRequestFactory`` produces requests without a ``resolver_match``.
+    ``set_rollback`` must fall back to the previous
+    ``ATOMIC_REQUESTS + in_atomic_block`` behavior in that case.
+    """
+    def setUp(self):
+        connections.databases['default']['ATOMIC_REQUESTS'] = True
+
+    def tearDown(self):
+        connections.databases['default']['ATOMIC_REQUESTS'] = False
+
+    def test_falls_back_when_resolver_match_missing(self):
+        request = factory.post('/')
+        assert getattr(request, 'resolver_match', None) is None
+        with transaction.atomic():
+            set_rollback(request)
+            assert transaction.get_rollback()