Skip to content

Commit dc027d2

Browse files
jdetaeyeauvipybrowniebroke
authored
set_rollback should only rollback initialized connections (#9599)
* set_roll_back should only rollback initialized connections The exception handler call "connections.all()". This creates a database connection to all defined databases. As a small optimization you can use "connections.all(initialized_only=True)" to rollback only the database to which the current thread has open connections. (My application can have many databases defined, and this loop is identified as a source of many idle database connections) * Add regression test for the change --------- Co-authored-by: Asif Saif Uddin {"Auvi":"অভি"} <auvipy@gmail.com> Co-authored-by: Bruno Alla <alla.brunoo@gmail.com>
1 parent 094bb6d commit dc027d2

2 files changed

Lines changed: 71 additions & 2 deletions

File tree

rest_framework/views.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -64,7 +64,7 @@ def get_view_description(view, html=False):
6464

6565

6666
def set_rollback():
67-
for db in connections.all():
67+
for db in connections.all(initialized_only=True):
6868
if db.settings_dict['ATOMIC_REQUESTS'] and db.in_atomic_block:
6969
db.set_rollback(True)
7070

tests/test_atomic_requests.py

Lines changed: 70 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@
99
from rest_framework.exceptions import APIException
1010
from rest_framework.response import Response
1111
from rest_framework.test import APIRequestFactory
12-
from rest_framework.views import APIView
12+
from rest_framework.views import APIView, set_rollback
1313
from tests.models import BasicModel
1414

1515
factory = APIRequestFactory()
@@ -183,3 +183,72 @@ def test_api_exception_rollback_transaction_non_atomic_view(self):
183183
# without checking connection.in_atomic_block view raises 500
184184
# due attempt to rollback without transaction
185185
assert response.status_code == status.HTTP_404_NOT_FOUND
186+
187+
188+
@unittest.skipUnless(
189+
connection.features.uses_savepoints,
190+
"'atomic' requires transactions and savepoints."
191+
)
192+
class SetRollbackTests(TestCase):
193+
def setUp(self):
194+
connections.databases['default']['ATOMIC_REQUESTS'] = True
195+
196+
def tearDown(self):
197+
connections.databases['default']['ATOMIC_REQUESTS'] = False
198+
199+
def test_marks_initialized_atomic_connection_for_rollback(self):
200+
with transaction.atomic():
201+
set_rollback()
202+
assert transaction.get_rollback()
203+
204+
205+
class UninitializedSecondaryConnectionMixin:
206+
"""
207+
Remove the 'secondary' connection wrapper from the current thread for
208+
the duration of a test, restoring the original wrapper afterwards so
209+
Django's test-case connection patching still finds it at class cleanup.
210+
"""
211+
def setUp(self):
212+
super().setUp()
213+
self._saved_secondary = getattr(connections._connections, 'secondary', None)
214+
if self._saved_secondary is not None:
215+
delattr(connections._connections, 'secondary')
216+
217+
def tearDown(self):
218+
stray = getattr(connections._connections, 'secondary', None)
219+
if stray is not None and stray is not self._saved_secondary:
220+
stray.close()
221+
if self._saved_secondary is not None:
222+
setattr(connections._connections, 'secondary', self._saved_secondary)
223+
super().tearDown()
224+
225+
226+
class SetRollbackUninitializedConnectionTests(UninitializedSecondaryConnectionMixin, TestCase):
227+
def setUp(self):
228+
super().setUp()
229+
connections.databases['secondary']['ATOMIC_REQUESTS'] = True
230+
231+
def tearDown(self):
232+
connections.databases['secondary']['ATOMIC_REQUESTS'] = False
233+
super().tearDown()
234+
235+
def test_does_not_initialize_unused_connections(self):
236+
set_rollback()
237+
assert not hasattr(connections._connections, 'secondary')
238+
239+
240+
class MultiDBUnusedConnectionAPIExceptionTests(UninitializedSecondaryConnectionMixin, TestCase):
241+
def setUp(self):
242+
super().setUp()
243+
self.view = APIExceptionView.as_view()
244+
connections.databases['secondary']['ATOMIC_REQUESTS'] = True
245+
246+
def tearDown(self):
247+
connections.databases['secondary']['ATOMIC_REQUESTS'] = False
248+
super().tearDown()
249+
250+
def test_api_exception_leaves_unused_connection_uninitialized(self):
251+
request = factory.post('/')
252+
response = self.view(request)
253+
assert response.status_code == status.HTTP_500_INTERNAL_SERVER_ERROR
254+
assert not hasattr(connections._connections, 'secondary')

0 commit comments

Comments
 (0)