creation.py 15 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384
  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. chunk_size = (
  116. 2000 if queryset._prefetch_related_lookups else None
  117. )
  118. yield from queryset.iterator(chunk_size=chunk_size)
  119. # Serialize to a string
  120. out = StringIO()
  121. serializers.serialize("json", get_objects(), indent=None, stream=out)
  122. return out.getvalue()
  123. def deserialize_db_from_string(self, data):
  124. """
  125. Reload the database with data from a string generated by
  126. the serialize_db_to_string() method.
  127. """
  128. data = StringIO(data)
  129. table_names = set()
  130. # Load data in a transaction to handle forward references and cycles.
  131. with atomic(using=self.connection.alias):
  132. # Disable constraint checks, because some databases (MySQL) doesn't
  133. # support deferred checks.
  134. with self.connection.constraint_checks_disabled():
  135. for obj in serializers.deserialize(
  136. "json", data, using=self.connection.alias
  137. ):
  138. obj.save()
  139. table_names.add(obj.object.__class__._meta.db_table)
  140. # Manually check for any invalid keys that might have been added,
  141. # because constraint checks were disabled.
  142. self.connection.check_constraints(table_names=table_names)
  143. def _get_database_display_str(self, verbosity, database_name):
  144. """
  145. Return display string for a database for use in various actions.
  146. """
  147. return "'%s'%s" % (
  148. self.connection.alias,
  149. (" ('%s')" % database_name) if verbosity >= 2 else "",
  150. )
  151. def _get_test_db_name(self):
  152. """
  153. Internal implementation - return the name of the test DB that will be
  154. created. Only useful when called from create_test_db() and
  155. _create_test_db() and when no external munging is done with the 'NAME'
  156. settings.
  157. """
  158. if self.connection.settings_dict["TEST"]["NAME"]:
  159. return self.connection.settings_dict["TEST"]["NAME"]
  160. return TEST_DATABASE_PREFIX + self.connection.settings_dict["NAME"]
  161. def _execute_create_test_db(self, cursor, parameters, keepdb=False):
  162. cursor.execute("CREATE DATABASE %(dbname)s %(suffix)s" % parameters)
  163. def _create_test_db(self, verbosity, autoclobber, keepdb=False):
  164. """
  165. Internal implementation - create the test db tables.
  166. """
  167. test_database_name = self._get_test_db_name()
  168. test_db_params = {
  169. "dbname": self.connection.ops.quote_name(test_database_name),
  170. "suffix": self.sql_table_creation_suffix(),
  171. }
  172. # Create the test database and connect to it.
  173. with self._nodb_cursor() as cursor:
  174. try:
  175. self._execute_create_test_db(cursor, test_db_params, keepdb)
  176. except Exception as e:
  177. # if we want to keep the db, then no need to do any of the below,
  178. # just return and skip it all.
  179. if keepdb:
  180. return test_database_name
  181. self.log("Got an error creating the test database: %s" % e)
  182. if not autoclobber:
  183. confirm = input(
  184. "Type 'yes' if you would like to try deleting the test "
  185. "database '%s', or 'no' to cancel: " % test_database_name
  186. )
  187. if autoclobber or confirm == "yes":
  188. try:
  189. if verbosity >= 1:
  190. self.log(
  191. "Destroying old test database for alias %s..."
  192. % (
  193. self._get_database_display_str(
  194. verbosity, test_database_name
  195. ),
  196. )
  197. )
  198. cursor.execute("DROP DATABASE %(dbname)s" % test_db_params)
  199. self._execute_create_test_db(cursor, test_db_params, keepdb)
  200. except Exception as e:
  201. self.log("Got an error recreating the test database: %s" % e)
  202. sys.exit(2)
  203. else:
  204. self.log("Tests cancelled.")
  205. sys.exit(1)
  206. return test_database_name
  207. def clone_test_db(self, suffix, verbosity=1, autoclobber=False, keepdb=False):
  208. """
  209. Clone a test database.
  210. """
  211. source_database_name = self.connection.settings_dict["NAME"]
  212. if verbosity >= 1:
  213. action = "Cloning test database"
  214. if keepdb:
  215. action = "Using existing clone"
  216. self.log(
  217. "%s for alias %s..."
  218. % (
  219. action,
  220. self._get_database_display_str(verbosity, source_database_name),
  221. )
  222. )
  223. # We could skip this call if keepdb is True, but we instead
  224. # give it the keepdb param. See create_test_db for details.
  225. self._clone_test_db(suffix, verbosity, keepdb)
  226. def get_test_db_clone_settings(self, suffix):
  227. """
  228. Return a modified connection settings dict for the n-th clone of a DB.
  229. """
  230. # When this function is called, the test database has been created
  231. # already and its name has been copied to settings_dict['NAME'] so
  232. # we don't need to call _get_test_db_name.
  233. orig_settings_dict = self.connection.settings_dict
  234. return {
  235. **orig_settings_dict,
  236. "NAME": "{}_{}".format(orig_settings_dict["NAME"], suffix),
  237. }
  238. def _clone_test_db(self, suffix, verbosity, keepdb=False):
  239. """
  240. Internal implementation - duplicate the test db tables.
  241. """
  242. raise NotImplementedError(
  243. "The database backend doesn't support cloning databases. "
  244. "Disable the option to run tests in parallel processes."
  245. )
  246. def destroy_test_db(
  247. self, old_database_name=None, verbosity=1, keepdb=False, suffix=None
  248. ):
  249. """
  250. Destroy a test database, prompting the user for confirmation if the
  251. database already exists.
  252. """
  253. self.connection.close()
  254. if suffix is None:
  255. test_database_name = self.connection.settings_dict["NAME"]
  256. else:
  257. test_database_name = self.get_test_db_clone_settings(suffix)["NAME"]
  258. if verbosity >= 1:
  259. action = "Destroying"
  260. if keepdb:
  261. action = "Preserving"
  262. self.log(
  263. "%s test database for alias %s..."
  264. % (
  265. action,
  266. self._get_database_display_str(verbosity, test_database_name),
  267. )
  268. )
  269. # if we want to preserve the database
  270. # skip the actual destroying piece.
  271. if not keepdb:
  272. self._destroy_test_db(test_database_name, verbosity)
  273. # Restore the original database name
  274. if old_database_name is not None:
  275. settings.DATABASES[self.connection.alias]["NAME"] = old_database_name
  276. self.connection.settings_dict["NAME"] = old_database_name
  277. def _destroy_test_db(self, test_database_name, verbosity):
  278. """
  279. Internal implementation - remove the test db tables.
  280. """
  281. # Remove the test database to clean up after
  282. # ourselves. Connect to the previous database (not the test database)
  283. # to do so, because it's not allowed to delete a database while being
  284. # connected to it.
  285. with self._nodb_cursor() as cursor:
  286. cursor.execute(
  287. "DROP DATABASE %s" % self.connection.ops.quote_name(test_database_name)
  288. )
  289. def mark_expected_failures_and_skips(self):
  290. """
  291. Mark tests in Django's test suite which are expected failures on this
  292. database and test which should be skipped on this database.
  293. """
  294. # Only load unittest if we're actually testing.
  295. from unittest import expectedFailure, skip
  296. for test_name in self.connection.features.django_test_expected_failures:
  297. test_case_name, _, test_method_name = test_name.rpartition(".")
  298. test_app = test_name.split(".")[0]
  299. # Importing a test app that isn't installed raises RuntimeError.
  300. if test_app in settings.INSTALLED_APPS:
  301. test_case = import_string(test_case_name)
  302. test_method = getattr(test_case, test_method_name)
  303. setattr(test_case, test_method_name, expectedFailure(test_method))
  304. for reason, tests in self.connection.features.django_test_skips.items():
  305. for test_name in tests:
  306. test_case_name, _, test_method_name = test_name.rpartition(".")
  307. test_app = test_name.split(".")[0]
  308. # Importing a test app that isn't installed raises RuntimeError.
  309. if test_app in settings.INSTALLED_APPS:
  310. test_case = import_string(test_case_name)
  311. test_method = getattr(test_case, test_method_name)
  312. setattr(test_case, test_method_name, skip(reason)(test_method))
  313. def sql_table_creation_suffix(self):
  314. """
  315. SQL to append to the end of the test table creation statements.
  316. """
  317. return ""
  318. def test_db_signature(self):
  319. """
  320. Return a tuple with elements of self.connection.settings_dict (a
  321. DATABASES setting value) that uniquely identify a database
  322. accordingly to the RDBMS particularities.
  323. """
  324. settings_dict = self.connection.settings_dict
  325. return (
  326. settings_dict["HOST"],
  327. settings_dict["PORT"],
  328. settings_dict["ENGINE"],
  329. self._get_test_db_name(),
  330. )
  331. def setup_worker_connection(self, _worker_id):
  332. settings_dict = self.get_test_db_clone_settings(str(_worker_id))
  333. # connection.settings_dict must be updated in place for changes to be
  334. # reflected in django.db.connections. If the following line assigned
  335. # connection.settings_dict = settings_dict, new threads would connect
  336. # to the default database instead of the appropriate clone.
  337. self.connection.settings_dict.update(settings_dict)
  338. self.connection.close()