creation.py 15 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381
  1. import os
  2. import sys
  3. from io import StringIO
  4. from django.apps import apps
  5. from django.conf import settings
  6. from django.core import serializers
  7. from django.db import router
  8. from django.db.transaction import atomic
  9. from django.utils.module_loading import import_string
  10. # The prefix to put on the default database name when creating
  11. # the test database.
  12. TEST_DATABASE_PREFIX = "test_"
  13. class BaseDatabaseCreation:
  14. """
  15. Encapsulate backend-specific differences pertaining to creation and
  16. destruction of the test database.
  17. """
  18. def __init__(self, connection):
  19. self.connection = connection
  20. def _nodb_cursor(self):
  21. return self.connection._nodb_cursor()
  22. def log(self, msg):
  23. sys.stderr.write(msg + os.linesep)
  24. def create_test_db(
  25. self, verbosity=1, autoclobber=False, serialize=True, keepdb=False
  26. ):
  27. """
  28. Create a test database, prompting the user for confirmation if the
  29. database already exists. Return the name of the test database created.
  30. """
  31. # Don't import django.core.management if it isn't needed.
  32. from django.core.management import call_command
  33. test_database_name = self._get_test_db_name()
  34. if verbosity >= 1:
  35. action = "Creating"
  36. if keepdb:
  37. action = "Using existing"
  38. self.log(
  39. "%s test database for alias %s..."
  40. % (
  41. action,
  42. self._get_database_display_str(verbosity, test_database_name),
  43. )
  44. )
  45. # We could skip this call if keepdb is True, but we instead
  46. # give it the keepdb param. This is to handle the case
  47. # where the test DB doesn't exist, in which case we need to
  48. # create it, then just not destroy it. If we instead skip
  49. # this, we will get an exception.
  50. self._create_test_db(verbosity, autoclobber, keepdb)
  51. self.connection.close()
  52. settings.DATABASES[self.connection.alias]["NAME"] = test_database_name
  53. self.connection.settings_dict["NAME"] = test_database_name
  54. try:
  55. if self.connection.settings_dict["TEST"]["MIGRATE"] is False:
  56. # Disable migrations for all apps.
  57. old_migration_modules = settings.MIGRATION_MODULES
  58. settings.MIGRATION_MODULES = {
  59. app.label: None for app in apps.get_app_configs()
  60. }
  61. # We report migrate messages at one level lower than that
  62. # requested. This ensures we don't get flooded with messages during
  63. # testing (unless you really ask to be flooded).
  64. call_command(
  65. "migrate",
  66. verbosity=max(verbosity - 1, 0),
  67. interactive=False,
  68. database=self.connection.alias,
  69. run_syncdb=True,
  70. )
  71. finally:
  72. if self.connection.settings_dict["TEST"]["MIGRATE"] is False:
  73. settings.MIGRATION_MODULES = old_migration_modules
  74. # We then serialize the current state of the database into a string
  75. # and store it on the connection. This slightly horrific process is so people
  76. # who are testing on databases without transactions or who are using
  77. # a TransactionTestCase still get a clean database on every test run.
  78. if serialize:
  79. self.connection._test_serialized_contents = self.serialize_db_to_string()
  80. call_command("createcachetable", database=self.connection.alias)
  81. # Ensure a connection for the side effect of initializing the test database.
  82. self.connection.ensure_connection()
  83. if os.environ.get("RUNNING_DJANGOS_TEST_SUITE") == "true":
  84. self.mark_expected_failures_and_skips()
  85. return test_database_name
  86. def set_as_test_mirror(self, primary_settings_dict):
  87. """
  88. Set this database up to be used in testing as a mirror of a primary
  89. database whose settings are given.
  90. """
  91. self.connection.settings_dict["NAME"] = primary_settings_dict["NAME"]
  92. def serialize_db_to_string(self):
  93. """
  94. Serialize all data in the database into a JSON string.
  95. Designed only for test runner usage; will not handle large
  96. amounts of data.
  97. """
  98. # Iteratively return every object for all models to serialize.
  99. def get_objects():
  100. from django.db.migrations.loader import MigrationLoader
  101. loader = MigrationLoader(self.connection)
  102. for app_config in apps.get_app_configs():
  103. if (
  104. app_config.models_module is not None
  105. and app_config.label in loader.migrated_apps
  106. and app_config.name not in settings.TEST_NON_SERIALIZED_APPS
  107. ):
  108. for model in app_config.get_models():
  109. if model._meta.can_migrate(
  110. self.connection
  111. ) and router.allow_migrate_model(self.connection.alias, model):
  112. queryset = model._base_manager.using(
  113. self.connection.alias,
  114. ).order_by(model._meta.pk.name)
  115. yield from queryset.iterator()
  116. # Serialize to a string
  117. out = StringIO()
  118. serializers.serialize("json", get_objects(), indent=None, stream=out)
  119. return out.getvalue()
  120. def deserialize_db_from_string(self, data):
  121. """
  122. Reload the database with data from a string generated by
  123. the serialize_db_to_string() method.
  124. """
  125. data = StringIO(data)
  126. table_names = set()
  127. # Load data in a transaction to handle forward references and cycles.
  128. with atomic(using=self.connection.alias):
  129. # Disable constraint checks, because some databases (MySQL) doesn't
  130. # support deferred checks.
  131. with self.connection.constraint_checks_disabled():
  132. for obj in serializers.deserialize(
  133. "json", data, using=self.connection.alias
  134. ):
  135. obj.save()
  136. table_names.add(obj.object.__class__._meta.db_table)
  137. # Manually check for any invalid keys that might have been added,
  138. # because constraint checks were disabled.
  139. self.connection.check_constraints(table_names=table_names)
  140. def _get_database_display_str(self, verbosity, database_name):
  141. """
  142. Return display string for a database for use in various actions.
  143. """
  144. return "'%s'%s" % (
  145. self.connection.alias,
  146. (" ('%s')" % database_name) if verbosity >= 2 else "",
  147. )
  148. def _get_test_db_name(self):
  149. """
  150. Internal implementation - return the name of the test DB that will be
  151. created. Only useful when called from create_test_db() and
  152. _create_test_db() and when no external munging is done with the 'NAME'
  153. settings.
  154. """
  155. if self.connection.settings_dict["TEST"]["NAME"]:
  156. return self.connection.settings_dict["TEST"]["NAME"]
  157. return TEST_DATABASE_PREFIX + self.connection.settings_dict["NAME"]
  158. def _execute_create_test_db(self, cursor, parameters, keepdb=False):
  159. cursor.execute("CREATE DATABASE %(dbname)s %(suffix)s" % parameters)
  160. def _create_test_db(self, verbosity, autoclobber, keepdb=False):
  161. """
  162. Internal implementation - create the test db tables.
  163. """
  164. test_database_name = self._get_test_db_name()
  165. test_db_params = {
  166. "dbname": self.connection.ops.quote_name(test_database_name),
  167. "suffix": self.sql_table_creation_suffix(),
  168. }
  169. # Create the test database and connect to it.
  170. with self._nodb_cursor() as cursor:
  171. try:
  172. self._execute_create_test_db(cursor, test_db_params, keepdb)
  173. except Exception as e:
  174. # if we want to keep the db, then no need to do any of the below,
  175. # just return and skip it all.
  176. if keepdb:
  177. return test_database_name
  178. self.log("Got an error creating the test database: %s" % e)
  179. if not autoclobber:
  180. confirm = input(
  181. "Type 'yes' if you would like to try deleting the test "
  182. "database '%s', or 'no' to cancel: " % test_database_name
  183. )
  184. if autoclobber or confirm == "yes":
  185. try:
  186. if verbosity >= 1:
  187. self.log(
  188. "Destroying old test database for alias %s..."
  189. % (
  190. self._get_database_display_str(
  191. verbosity, test_database_name
  192. ),
  193. )
  194. )
  195. cursor.execute("DROP DATABASE %(dbname)s" % test_db_params)
  196. self._execute_create_test_db(cursor, test_db_params, keepdb)
  197. except Exception as e:
  198. self.log("Got an error recreating the test database: %s" % e)
  199. sys.exit(2)
  200. else:
  201. self.log("Tests cancelled.")
  202. sys.exit(1)
  203. return test_database_name
  204. def clone_test_db(self, suffix, verbosity=1, autoclobber=False, keepdb=False):
  205. """
  206. Clone a test database.
  207. """
  208. source_database_name = self.connection.settings_dict["NAME"]
  209. if verbosity >= 1:
  210. action = "Cloning test database"
  211. if keepdb:
  212. action = "Using existing clone"
  213. self.log(
  214. "%s for alias %s..."
  215. % (
  216. action,
  217. self._get_database_display_str(verbosity, source_database_name),
  218. )
  219. )
  220. # We could skip this call if keepdb is True, but we instead
  221. # give it the keepdb param. See create_test_db for details.
  222. self._clone_test_db(suffix, verbosity, keepdb)
  223. def get_test_db_clone_settings(self, suffix):
  224. """
  225. Return a modified connection settings dict for the n-th clone of a DB.
  226. """
  227. # When this function is called, the test database has been created
  228. # already and its name has been copied to settings_dict['NAME'] so
  229. # we don't need to call _get_test_db_name.
  230. orig_settings_dict = self.connection.settings_dict
  231. return {
  232. **orig_settings_dict,
  233. "NAME": "{}_{}".format(orig_settings_dict["NAME"], suffix),
  234. }
  235. def _clone_test_db(self, suffix, verbosity, keepdb=False):
  236. """
  237. Internal implementation - duplicate the test db tables.
  238. """
  239. raise NotImplementedError(
  240. "The database backend doesn't support cloning databases. "
  241. "Disable the option to run tests in parallel processes."
  242. )
  243. def destroy_test_db(
  244. self, old_database_name=None, verbosity=1, keepdb=False, suffix=None
  245. ):
  246. """
  247. Destroy a test database, prompting the user for confirmation if the
  248. database already exists.
  249. """
  250. self.connection.close()
  251. if suffix is None:
  252. test_database_name = self.connection.settings_dict["NAME"]
  253. else:
  254. test_database_name = self.get_test_db_clone_settings(suffix)["NAME"]
  255. if verbosity >= 1:
  256. action = "Destroying"
  257. if keepdb:
  258. action = "Preserving"
  259. self.log(
  260. "%s test database for alias %s..."
  261. % (
  262. action,
  263. self._get_database_display_str(verbosity, test_database_name),
  264. )
  265. )
  266. # if we want to preserve the database
  267. # skip the actual destroying piece.
  268. if not keepdb:
  269. self._destroy_test_db(test_database_name, verbosity)
  270. # Restore the original database name
  271. if old_database_name is not None:
  272. settings.DATABASES[self.connection.alias]["NAME"] = old_database_name
  273. self.connection.settings_dict["NAME"] = old_database_name
  274. def _destroy_test_db(self, test_database_name, verbosity):
  275. """
  276. Internal implementation - remove the test db tables.
  277. """
  278. # Remove the test database to clean up after
  279. # ourselves. Connect to the previous database (not the test database)
  280. # to do so, because it's not allowed to delete a database while being
  281. # connected to it.
  282. with self._nodb_cursor() as cursor:
  283. cursor.execute(
  284. "DROP DATABASE %s" % self.connection.ops.quote_name(test_database_name)
  285. )
  286. def mark_expected_failures_and_skips(self):
  287. """
  288. Mark tests in Django's test suite which are expected failures on this
  289. database and test which should be skipped on this database.
  290. """
  291. # Only load unittest if we're actually testing.
  292. from unittest import expectedFailure, skip
  293. for test_name in self.connection.features.django_test_expected_failures:
  294. test_case_name, _, test_method_name = test_name.rpartition(".")
  295. test_app = test_name.split(".")[0]
  296. # Importing a test app that isn't installed raises RuntimeError.
  297. if test_app in settings.INSTALLED_APPS:
  298. test_case = import_string(test_case_name)
  299. test_method = getattr(test_case, test_method_name)
  300. setattr(test_case, test_method_name, expectedFailure(test_method))
  301. for reason, tests in self.connection.features.django_test_skips.items():
  302. for test_name in tests:
  303. test_case_name, _, test_method_name = test_name.rpartition(".")
  304. test_app = test_name.split(".")[0]
  305. # Importing a test app that isn't installed raises RuntimeError.
  306. if test_app in settings.INSTALLED_APPS:
  307. test_case = import_string(test_case_name)
  308. test_method = getattr(test_case, test_method_name)
  309. setattr(test_case, test_method_name, skip(reason)(test_method))
  310. def sql_table_creation_suffix(self):
  311. """
  312. SQL to append to the end of the test table creation statements.
  313. """
  314. return ""
  315. def test_db_signature(self):
  316. """
  317. Return a tuple with elements of self.connection.settings_dict (a
  318. DATABASES setting value) that uniquely identify a database
  319. accordingly to the RDBMS particularities.
  320. """
  321. settings_dict = self.connection.settings_dict
  322. return (
  323. settings_dict["HOST"],
  324. settings_dict["PORT"],
  325. settings_dict["ENGINE"],
  326. self._get_test_db_name(),
  327. )
  328. def setup_worker_connection(self, _worker_id):
  329. settings_dict = self.get_test_db_clone_settings(str(_worker_id))
  330. # connection.settings_dict must be updated in place for changes to be
  331. # reflected in django.db.connections. If the following line assigned
  332. # connection.settings_dict = settings_dict, new threads would connect
  333. # to the default database instead of the appropriate clone.
  334. self.connection.settings_dict.update(settings_dict)
  335. self.connection.close()