tests.py 10 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270
  1. from __future__ import absolute_import
  2. import sys
  3. import time
  4. import unittest
  5. from django.conf import settings
  6. from django.db import transaction, connection
  7. from django.db.utils import ConnectionHandler, DEFAULT_DB_ALIAS, DatabaseError
  8. from django.test import (TransactionTestCase, skipIfDBFeature,
  9. skipUnlessDBFeature)
  10. from .models import Person
  11. # Some tests require threading, which might not be available. So create a
  12. # skip-test decorator for those test functions.
  13. try:
  14. import threading
  15. except ImportError:
  16. threading = None
  17. requires_threading = unittest.skipUnless(threading, 'requires threading')
  18. class SelectForUpdateTests(TransactionTestCase):
  19. available_apps = ['select_for_update']
  20. def setUp(self):
  21. transaction.enter_transaction_management()
  22. self.person = Person.objects.create(name='Reinhardt')
  23. # We have to commit here so that code in run_select_for_update can
  24. # see this data.
  25. transaction.commit()
  26. # We need another database connection to test that one connection
  27. # issuing a SELECT ... FOR UPDATE will block.
  28. new_connections = ConnectionHandler(settings.DATABASES)
  29. self.new_connection = new_connections[DEFAULT_DB_ALIAS]
  30. self.new_connection.enter_transaction_management()
  31. # We need to set settings.DEBUG to True so we can capture
  32. # the output SQL to examine.
  33. self._old_debug = settings.DEBUG
  34. settings.DEBUG = True
  35. def tearDown(self):
  36. try:
  37. # We don't really care if this fails - some of the tests will set
  38. # this in the course of their run.
  39. transaction.abort()
  40. self.new_connection.abort()
  41. except transaction.TransactionManagementError:
  42. pass
  43. self.new_connection.close()
  44. settings.DEBUG = self._old_debug
  45. try:
  46. self.end_blocking_transaction()
  47. except (DatabaseError, AttributeError):
  48. pass
  49. def start_blocking_transaction(self):
  50. # Start a blocking transaction. At some point,
  51. # end_blocking_transaction() should be called.
  52. self.cursor = self.new_connection.cursor()
  53. sql = 'SELECT * FROM %(db_table)s %(for_update)s;' % {
  54. 'db_table': Person._meta.db_table,
  55. 'for_update': self.new_connection.ops.for_update_sql(),
  56. }
  57. self.cursor.execute(sql, ())
  58. self.cursor.fetchone()
  59. def end_blocking_transaction(self):
  60. # Roll back the blocking transaction.
  61. self.new_connection.rollback()
  62. def has_for_update_sql(self, tested_connection, nowait=False):
  63. # Examine the SQL that was executed to determine whether it
  64. # contains the 'SELECT..FOR UPDATE' stanza.
  65. for_update_sql = tested_connection.ops.for_update_sql(nowait)
  66. sql = tested_connection.queries[-1]['sql']
  67. return bool(sql.find(for_update_sql) > -1)
  68. @skipUnlessDBFeature('has_select_for_update')
  69. def test_for_update_sql_generated(self):
  70. """
  71. Test that the backend's FOR UPDATE variant appears in
  72. generated SQL when select_for_update is invoked.
  73. """
  74. list(Person.objects.all().select_for_update())
  75. self.assertTrue(self.has_for_update_sql(connection))
  76. @skipUnlessDBFeature('has_select_for_update_nowait')
  77. def test_for_update_sql_generated_nowait(self):
  78. """
  79. Test that the backend's FOR UPDATE NOWAIT variant appears in
  80. generated SQL when select_for_update is invoked.
  81. """
  82. list(Person.objects.all().select_for_update(nowait=True))
  83. self.assertTrue(self.has_for_update_sql(connection, nowait=True))
  84. # In Python 2.6 beta and some final releases, exceptions raised in __len__
  85. # are swallowed (Python issue 1242657), so these cases return an empty
  86. # list, rather than raising an exception. Not a lot we can do about that,
  87. # unfortunately, due to the way Python handles list() calls internally.
  88. # Python 2.6.1 is the "in the wild" version affected by this, so we skip
  89. # the test for that version.
  90. @requires_threading
  91. @skipUnlessDBFeature('has_select_for_update_nowait')
  92. @unittest.skipIf(sys.version_info[:3] == (2, 6, 1), "Python version is 2.6.1")
  93. def test_nowait_raises_error_on_block(self):
  94. """
  95. If nowait is specified, we expect an error to be raised rather
  96. than blocking.
  97. """
  98. self.start_blocking_transaction()
  99. status = []
  100. thread = threading.Thread(
  101. target=self.run_select_for_update,
  102. args=(status,),
  103. kwargs={'nowait': True},
  104. )
  105. thread.start()
  106. time.sleep(1)
  107. thread.join()
  108. self.end_blocking_transaction()
  109. self.assertIsInstance(status[-1], DatabaseError)
  110. # In Python 2.6 beta and some final releases, exceptions raised in __len__
  111. # are swallowed (Python issue 1242657), so these cases return an empty
  112. # list, rather than raising an exception. Not a lot we can do about that,
  113. # unfortunately, due to the way Python handles list() calls internally.
  114. # Python 2.6.1 is the "in the wild" version affected by this, so we skip
  115. # the test for that version.
  116. @skipIfDBFeature('has_select_for_update_nowait')
  117. @skipUnlessDBFeature('has_select_for_update')
  118. @unittest.skipIf(sys.version_info[:3] == (2, 6, 1), "Python version is 2.6.1")
  119. def test_unsupported_nowait_raises_error(self):
  120. """
  121. If a SELECT...FOR UPDATE NOWAIT is run on a database backend
  122. that supports FOR UPDATE but not NOWAIT, then we should find
  123. that a DatabaseError is raised.
  124. """
  125. self.assertRaises(
  126. DatabaseError,
  127. list,
  128. Person.objects.all().select_for_update(nowait=True)
  129. )
  130. def run_select_for_update(self, status, nowait=False):
  131. """
  132. Utility method that runs a SELECT FOR UPDATE against all
  133. Person instances. After the select_for_update, it attempts
  134. to update the name of the only record, save, and commit.
  135. This function expects to run in a separate thread.
  136. """
  137. status.append('started')
  138. try:
  139. # We need to enter transaction management again, as this is done on
  140. # per-thread basis
  141. transaction.enter_transaction_management()
  142. people = list(
  143. Person.objects.all().select_for_update(nowait=nowait)
  144. )
  145. people[0].name = 'Fred'
  146. people[0].save()
  147. transaction.commit()
  148. except DatabaseError as e:
  149. status.append(e)
  150. finally:
  151. # This method is run in a separate thread. It uses its own
  152. # database connection. Close it without waiting for the GC.
  153. transaction.abort()
  154. connection.close()
  155. @requires_threading
  156. @skipUnlessDBFeature('has_select_for_update')
  157. @skipUnlessDBFeature('supports_transactions')
  158. def test_block(self):
  159. """
  160. Check that a thread running a select_for_update that
  161. accesses rows being touched by a similar operation
  162. on another connection blocks correctly.
  163. """
  164. # First, let's start the transaction in our thread.
  165. self.start_blocking_transaction()
  166. # Now, try it again using the ORM's select_for_update
  167. # facility. Do this in a separate thread.
  168. status = []
  169. thread = threading.Thread(
  170. target=self.run_select_for_update, args=(status,)
  171. )
  172. # The thread should immediately block, but we'll sleep
  173. # for a bit to make sure.
  174. thread.start()
  175. sanity_count = 0
  176. while len(status) != 1 and sanity_count < 10:
  177. sanity_count += 1
  178. time.sleep(1)
  179. if sanity_count >= 10:
  180. raise ValueError('Thread did not run and block')
  181. # Check the person hasn't been updated. Since this isn't
  182. # using FOR UPDATE, it won't block.
  183. p = Person.objects.get(pk=self.person.pk)
  184. self.assertEqual('Reinhardt', p.name)
  185. # When we end our blocking transaction, our thread should
  186. # be able to continue.
  187. self.end_blocking_transaction()
  188. thread.join(5.0)
  189. # Check the thread has finished. Assuming it has, we should
  190. # find that it has updated the person's name.
  191. self.assertFalse(thread.isAlive())
  192. # We must commit the transaction to ensure that MySQL gets a fresh read,
  193. # since by default it runs in REPEATABLE READ mode
  194. transaction.commit()
  195. p = Person.objects.get(pk=self.person.pk)
  196. self.assertEqual('Fred', p.name)
  197. @requires_threading
  198. @skipUnlessDBFeature('has_select_for_update')
  199. def test_raw_lock_not_available(self):
  200. """
  201. Check that running a raw query which can't obtain a FOR UPDATE lock
  202. raises the correct exception
  203. """
  204. self.start_blocking_transaction()
  205. def raw(status):
  206. try:
  207. list(
  208. Person.objects.raw(
  209. 'SELECT * FROM %s %s' % (
  210. Person._meta.db_table,
  211. connection.ops.for_update_sql(nowait=True)
  212. )
  213. )
  214. )
  215. except DatabaseError as e:
  216. status.append(e)
  217. finally:
  218. # This method is run in a separate thread. It uses its own
  219. # database connection. Close it without waiting for the GC.
  220. connection.close()
  221. status = []
  222. thread = threading.Thread(target=raw, kwargs={'status': status})
  223. thread.start()
  224. time.sleep(1)
  225. thread.join()
  226. self.end_blocking_transaction()
  227. self.assertIsInstance(status[-1], DatabaseError)
  228. @skipUnlessDBFeature('has_select_for_update')
  229. def test_transaction_dirty_managed(self):
  230. """ Check that a select_for_update sets the transaction to be
  231. dirty when executed under txn management. Setting the txn dirty
  232. means that it will be either committed or rolled back by Django,
  233. which will release any locks held by the SELECT FOR UPDATE.
  234. """
  235. people = list(Person.objects.select_for_update())
  236. self.assertTrue(transaction.is_dirty())