You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

loaddata.py 14KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351
  1. import functools
  2. import glob
  3. import gzip
  4. import os
  5. import sys
  6. import warnings
  7. import zipfile
  8. from itertools import product
  9. from django.apps import apps
  10. from django.conf import settings
  11. from django.core import serializers
  12. from django.core.exceptions import ImproperlyConfigured
  13. from django.core.management.base import BaseCommand, CommandError
  14. from django.core.management.color import no_style
  15. from django.core.management.utils import parse_apps_and_model_labels
  16. from django.db import (
  17. DEFAULT_DB_ALIAS, DatabaseError, IntegrityError, connections, router,
  18. transaction,
  19. )
  20. from django.utils.functional import cached_property
  21. try:
  22. import bz2
  23. has_bz2 = True
  24. except ImportError:
  25. has_bz2 = False
  26. READ_STDIN = '-'
  27. class Command(BaseCommand):
  28. help = 'Installs the named fixture(s) in the database.'
  29. missing_args_message = (
  30. "No database fixture specified. Please provide the path of at least "
  31. "one fixture in the command line."
  32. )
  33. def add_arguments(self, parser):
  34. parser.add_argument('args', metavar='fixture', nargs='+', help='Fixture labels.')
  35. parser.add_argument(
  36. '--database', default=DEFAULT_DB_ALIAS,
  37. help='Nominates a specific database to load fixtures into. Defaults to the "default" database.',
  38. )
  39. parser.add_argument(
  40. '--app', dest='app_label',
  41. help='Only look for fixtures in the specified app.',
  42. )
  43. parser.add_argument(
  44. '--ignorenonexistent', '-i', action='store_true', dest='ignore',
  45. help='Ignores entries in the serialized data for fields that do not '
  46. 'currently exist on the model.',
  47. )
  48. parser.add_argument(
  49. '-e', '--exclude', action='append', default=[],
  50. help='An app_label or app_label.ModelName to exclude. Can be used multiple times.',
  51. )
  52. parser.add_argument(
  53. '--format',
  54. help='Format of serialized data when reading from stdin.',
  55. )
  56. def handle(self, *fixture_labels, **options):
  57. self.ignore = options['ignore']
  58. self.using = options['database']
  59. self.app_label = options['app_label']
  60. self.verbosity = options['verbosity']
  61. self.excluded_models, self.excluded_apps = parse_apps_and_model_labels(options['exclude'])
  62. self.format = options['format']
  63. with transaction.atomic(using=self.using):
  64. self.loaddata(fixture_labels)
  65. # Close the DB connection -- unless we're still in a transaction. This
  66. # is required as a workaround for an edge case in MySQL: if the same
  67. # connection is used to create tables, load data, and query, the query
  68. # can return incorrect results. See Django #7572, MySQL #37735.
  69. if transaction.get_autocommit(self.using):
  70. connections[self.using].close()
  71. def loaddata(self, fixture_labels):
  72. connection = connections[self.using]
  73. # Keep a count of the installed objects and fixtures
  74. self.fixture_count = 0
  75. self.loaded_object_count = 0
  76. self.fixture_object_count = 0
  77. self.models = set()
  78. self.serialization_formats = serializers.get_public_serializer_formats()
  79. # Forcing binary mode may be revisited after dropping Python 2 support (see #22399)
  80. self.compression_formats = {
  81. None: (open, 'rb'),
  82. 'gz': (gzip.GzipFile, 'rb'),
  83. 'zip': (SingleZipReader, 'r'),
  84. 'stdin': (lambda *args: sys.stdin, None),
  85. }
  86. if has_bz2:
  87. self.compression_formats['bz2'] = (bz2.BZ2File, 'r')
  88. # Django's test suite repeatedly tries to load initial_data fixtures
  89. # from apps that don't have any fixtures. Because disabling constraint
  90. # checks can be expensive on some database (especially MSSQL), bail
  91. # out early if no fixtures are found.
  92. for fixture_label in fixture_labels:
  93. if self.find_fixtures(fixture_label):
  94. break
  95. else:
  96. return
  97. with connection.constraint_checks_disabled():
  98. self.objs_with_deferred_fields = []
  99. for fixture_label in fixture_labels:
  100. self.load_label(fixture_label)
  101. for obj in self.objs_with_deferred_fields:
  102. obj.save_deferred_fields(using=self.using)
  103. # Since we disabled constraint checks, we must manually check for
  104. # any invalid keys that might have been added
  105. table_names = [model._meta.db_table for model in self.models]
  106. try:
  107. connection.check_constraints(table_names=table_names)
  108. except Exception as e:
  109. e.args = ("Problem installing fixtures: %s" % e,)
  110. raise
  111. # If we found even one object in a fixture, we need to reset the
  112. # database sequences.
  113. if self.loaded_object_count > 0:
  114. sequence_sql = connection.ops.sequence_reset_sql(no_style(), self.models)
  115. if sequence_sql:
  116. if self.verbosity >= 2:
  117. self.stdout.write("Resetting sequences\n")
  118. with connection.cursor() as cursor:
  119. for line in sequence_sql:
  120. cursor.execute(line)
  121. if self.verbosity >= 1:
  122. if self.fixture_object_count == self.loaded_object_count:
  123. self.stdout.write(
  124. "Installed %d object(s) from %d fixture(s)"
  125. % (self.loaded_object_count, self.fixture_count)
  126. )
  127. else:
  128. self.stdout.write(
  129. "Installed %d object(s) (of %d) from %d fixture(s)"
  130. % (self.loaded_object_count, self.fixture_object_count, self.fixture_count)
  131. )
  132. def load_label(self, fixture_label):
  133. """Load fixtures files for a given label."""
  134. show_progress = self.verbosity >= 3
  135. for fixture_file, fixture_dir, fixture_name in self.find_fixtures(fixture_label):
  136. _, ser_fmt, cmp_fmt = self.parse_name(os.path.basename(fixture_file))
  137. open_method, mode = self.compression_formats[cmp_fmt]
  138. fixture = open_method(fixture_file, mode)
  139. try:
  140. self.fixture_count += 1
  141. objects_in_fixture = 0
  142. loaded_objects_in_fixture = 0
  143. if self.verbosity >= 2:
  144. self.stdout.write(
  145. "Installing %s fixture '%s' from %s."
  146. % (ser_fmt, fixture_name, humanize(fixture_dir))
  147. )
  148. objects = serializers.deserialize(
  149. ser_fmt, fixture, using=self.using, ignorenonexistent=self.ignore,
  150. handle_forward_references=True,
  151. )
  152. for obj in objects:
  153. objects_in_fixture += 1
  154. if (obj.object._meta.app_config in self.excluded_apps or
  155. type(obj.object) in self.excluded_models):
  156. continue
  157. if router.allow_migrate_model(self.using, obj.object.__class__):
  158. loaded_objects_in_fixture += 1
  159. self.models.add(obj.object.__class__)
  160. try:
  161. obj.save(using=self.using)
  162. if show_progress:
  163. self.stdout.write(
  164. '\rProcessed %i object(s).' % loaded_objects_in_fixture,
  165. ending=''
  166. )
  167. # psycopg2 raises ValueError if data contains NUL chars.
  168. except (DatabaseError, IntegrityError, ValueError) as e:
  169. e.args = ("Could not load %(app_label)s.%(object_name)s(pk=%(pk)s): %(error_msg)s" % {
  170. 'app_label': obj.object._meta.app_label,
  171. 'object_name': obj.object._meta.object_name,
  172. 'pk': obj.object.pk,
  173. 'error_msg': e,
  174. },)
  175. raise
  176. if obj.deferred_fields:
  177. self.objs_with_deferred_fields.append(obj)
  178. if objects and show_progress:
  179. self.stdout.write('') # add a newline after progress indicator
  180. self.loaded_object_count += loaded_objects_in_fixture
  181. self.fixture_object_count += objects_in_fixture
  182. except Exception as e:
  183. if not isinstance(e, CommandError):
  184. e.args = ("Problem installing fixture '%s': %s" % (fixture_file, e),)
  185. raise
  186. finally:
  187. fixture.close()
  188. # Warn if the fixture we loaded contains 0 objects.
  189. if objects_in_fixture == 0:
  190. warnings.warn(
  191. "No fixture data found for '%s'. (File format may be "
  192. "invalid.)" % fixture_name,
  193. RuntimeWarning
  194. )
  195. @functools.lru_cache(maxsize=None)
  196. def find_fixtures(self, fixture_label):
  197. """Find fixture files for a given label."""
  198. if fixture_label == READ_STDIN:
  199. return [(READ_STDIN, None, READ_STDIN)]
  200. fixture_name, ser_fmt, cmp_fmt = self.parse_name(fixture_label)
  201. databases = [self.using, None]
  202. cmp_fmts = list(self.compression_formats) if cmp_fmt is None else [cmp_fmt]
  203. ser_fmts = serializers.get_public_serializer_formats() if ser_fmt is None else [ser_fmt]
  204. if self.verbosity >= 2:
  205. self.stdout.write("Loading '%s' fixtures..." % fixture_name)
  206. if os.path.isabs(fixture_name):
  207. fixture_dirs = [os.path.dirname(fixture_name)]
  208. fixture_name = os.path.basename(fixture_name)
  209. else:
  210. fixture_dirs = self.fixture_dirs
  211. if os.path.sep in os.path.normpath(fixture_name):
  212. fixture_dirs = [os.path.join(dir_, os.path.dirname(fixture_name))
  213. for dir_ in fixture_dirs]
  214. fixture_name = os.path.basename(fixture_name)
  215. suffixes = (
  216. '.'.join(ext for ext in combo if ext)
  217. for combo in product(databases, ser_fmts, cmp_fmts)
  218. )
  219. targets = {'.'.join((fixture_name, suffix)) for suffix in suffixes}
  220. fixture_files = []
  221. for fixture_dir in fixture_dirs:
  222. if self.verbosity >= 2:
  223. self.stdout.write("Checking %s for fixtures..." % humanize(fixture_dir))
  224. fixture_files_in_dir = []
  225. path = os.path.join(fixture_dir, fixture_name)
  226. for candidate in glob.iglob(glob.escape(path) + '*'):
  227. if os.path.basename(candidate) in targets:
  228. # Save the fixture_dir and fixture_name for future error messages.
  229. fixture_files_in_dir.append((candidate, fixture_dir, fixture_name))
  230. if self.verbosity >= 2 and not fixture_files_in_dir:
  231. self.stdout.write("No fixture '%s' in %s." %
  232. (fixture_name, humanize(fixture_dir)))
  233. # Check kept for backwards-compatibility; it isn't clear why
  234. # duplicates are only allowed in different directories.
  235. if len(fixture_files_in_dir) > 1:
  236. raise CommandError(
  237. "Multiple fixtures named '%s' in %s. Aborting." %
  238. (fixture_name, humanize(fixture_dir)))
  239. fixture_files.extend(fixture_files_in_dir)
  240. if not fixture_files:
  241. raise CommandError("No fixture named '%s' found." % fixture_name)
  242. return fixture_files
  243. @cached_property
  244. def fixture_dirs(self):
  245. """
  246. Return a list of fixture directories.
  247. The list contains the 'fixtures' subdirectory of each installed
  248. application, if it exists, the directories in FIXTURE_DIRS, and the
  249. current directory.
  250. """
  251. dirs = []
  252. fixture_dirs = settings.FIXTURE_DIRS
  253. if len(fixture_dirs) != len(set(fixture_dirs)):
  254. raise ImproperlyConfigured("settings.FIXTURE_DIRS contains duplicates.")
  255. for app_config in apps.get_app_configs():
  256. app_label = app_config.label
  257. app_dir = os.path.join(app_config.path, 'fixtures')
  258. if app_dir in fixture_dirs:
  259. raise ImproperlyConfigured(
  260. "'%s' is a default fixture directory for the '%s' app "
  261. "and cannot be listed in settings.FIXTURE_DIRS." % (app_dir, app_label)
  262. )
  263. if self.app_label and app_label != self.app_label:
  264. continue
  265. if os.path.isdir(app_dir):
  266. dirs.append(app_dir)
  267. dirs.extend(fixture_dirs)
  268. dirs.append('')
  269. dirs = [os.path.abspath(os.path.realpath(d)) for d in dirs]
  270. return dirs
  271. def parse_name(self, fixture_name):
  272. """
  273. Split fixture name in name, serialization format, compression format.
  274. """
  275. if fixture_name == READ_STDIN:
  276. if not self.format:
  277. raise CommandError('--format must be specified when reading from stdin.')
  278. return READ_STDIN, self.format, 'stdin'
  279. parts = fixture_name.rsplit('.', 2)
  280. if len(parts) > 1 and parts[-1] in self.compression_formats:
  281. cmp_fmt = parts[-1]
  282. parts = parts[:-1]
  283. else:
  284. cmp_fmt = None
  285. if len(parts) > 1:
  286. if parts[-1] in self.serialization_formats:
  287. ser_fmt = parts[-1]
  288. parts = parts[:-1]
  289. else:
  290. raise CommandError(
  291. "Problem installing fixture '%s': %s is not a known "
  292. "serialization format." % ('.'.join(parts[:-1]), parts[-1]))
  293. else:
  294. ser_fmt = None
  295. name = '.'.join(parts)
  296. return name, ser_fmt, cmp_fmt
  297. class SingleZipReader(zipfile.ZipFile):
  298. def __init__(self, *args, **kwargs):
  299. super().__init__(*args, **kwargs)
  300. if len(self.namelist()) != 1:
  301. raise ValueError("Zip-compressed fixtures must contain one file.")
  302. def read(self):
  303. return zipfile.ZipFile.read(self, self.namelist()[0])
  304. def humanize(dirname):
  305. return "'%s'" % dirname if dirname else 'absolute path'