# Copyright The IETF Trust 2020, All Rights Reserved # -*- coding: utf-8 -*- import os import warnings from itertools import chain import debug # pyflakes:ignore from django.conf import settings from django.core import serializers from django.core.management.base import CommandError from django.core.management.commands.loaddata import Command as LoadCommand, humanize from django.db import DatabaseError, IntegrityError, router, transaction from django.db.models import ManyToManyField from django.utils.encoding import force_str from ietf.utils.models import ForeignKey def flatten(l): if isinstance(l, list): for el in l: if isinstance(el, list): for sub in flatten(el): yield sub else: yield el else: yield l class Command(LoadCommand): help = 'Merges the named fixture(s) into the database.' def add_arguments(self, parser): # parser.add_argument( # '-e', '--exclude', dest='exclude', action='append', default=[], # help='An app_label or app_label.ModelName to exclude. Can be used multiple times.', # ) super(Command, self).add_arguments(parser) def load_label(self, fixture_label): """ Loads fixtures files for a given label. """ # def update_objects(objects, model, old_pk, new_pk): # debug.show('old_pk, new_pk') # for o in objects: # for f in o._meta.fields: # if type(f) == ForeignKey: # if f.pk in [old_pk, new_pk]: # #debug.show('f') # #debug.show('f.pk') # elif type(f) == ManyToManyKey: # pass def obj_to_dict(obj): opts = obj._meta data = {} for f in chain(opts.concrete_fields, opts.private_fields): data[f.name] = f.value_from_object(obj) return data def get_unique_field(obj): unique = {} for f in obj._meta.get_fields(): if hasattr(f, 'primary_key') and f.primary_key: continue if hasattr(f, 'unique') and f.unique: unique[f.name] = f.value_from_object(obj) return unique show_progress = self.verbosity >= 3 if hasattr(settings, 'SERVER_MODE'): # Try to avoid sending mail during repair settings.SERVER_MODE = 'repair' for fixture_file, fixture_dir, fixture_name in self.find_fixtures(fixture_label): _, ser_fmt, cmp_fmt = self.parse_name(os.path.basename(fixture_file)) open_method, mode = self.compression_formats[cmp_fmt] fixture = open_method(fixture_file, mode) try: self.fixture_count += 1 objects_in_fixture = 0 loaded_objects_in_fixture = 0 if self.verbosity >= 2: self.stdout.write( "Installing %s fixture '%s' from %s." % (ser_fmt, fixture_name, humanize(fixture_dir)) ) objects = list(serializers.deserialize( ser_fmt, fixture, using=self.using, ignorenonexistent=self.ignore, )) # Prime fkrefs and m2mrefs with our deserialized objects fkrefs = {} m2mrefs = {} for obj in objects: o = obj.object oname = o._meta.app_label + '.' + o.__class__.__name__ fkrefs[(oname, o.pk)] = [] m2mrefs[(oname, o.pk)] = [] # Tabulate all FKs and M2M fields for obj in objects: o = obj.object for f in o._meta.get_fields(): if type(f) == ForeignKey: fobj = getattr(o, f.name) if fobj: fmod = f.related_model fname = fmod._meta.app_label + '.' + fmod.__name__ key = (fname, fobj.pk) if key in fkrefs: fkrefs[key].append((o, f.name)) elif type(f) == ManyToManyField: fobjs = getattr(o, f.name).all() if fobjs: fmod = f.related_model fname = fmod._meta.app_label + '.' + fmod.__name__ for fobj in fobjs: key = (fname, fobj.pk) if key in m2mrefs: m2mrefs[key].append((o, f.name)) else: pass #debug.type('f') for obj in objects: objects_in_fixture += 1 if (obj.object._meta.app_config in self.excluded_apps or type(obj.object) in self.excluded_models): continue if router.allow_migrate_model(self.using, obj.object.__class__): loaded_objects_in_fixture += 1 self.models.add(obj.object.__class__) try: model = obj.object.__class__ mname = model._meta.app_label + '.' + model.__name__ old_pk = obj.object.pk unique = get_unique_field(obj.object) if unique: match = model.objects.filter(**unique) if not match.exists(): obj.object.pk = None try: with transaction.atomic(using=self.using): obj.save(using=self.using) setattr(obj.object, '_saved', True) except IntegrityError as e: this_pk = match.first().pk if match else None new_dict = obj_to_dict(obj.object) self.stderr.write("\nSaving an updated object failed, possibly due to manually inserted conflicting data.\n" " Found PK: %s: %s\n" " Unique key: %s\n" " New record: %s\n" " Exception: %s" % (mname, this_pk, unique, new_dict, e)) else: match = model.objects.filter(pk=obj.object.pk) if match.exists(): #debug.say('PK exists: %s' % obj.object.pk) prev = match.first() prev_dict = obj_to_dict(prev) obj_dict = obj_to_dict(obj.object) if prev_dict == obj_dict: pass # nothing to do, the object is already there else: del obj_dict[obj.object._meta.pk.name] match = model.objects.filter(**obj_dict) if not match: try: obj.object.pk = None with transaction.atomic(using=self.using): obj.save(using=self.using) setattr(obj.object, '_saved', True) except IntegrityError as e: new_dict = obj_to_dict(obj.object) self.stderr.write("\nSaving an object failed, possibly due to manually inserted conflicting data.\n" " Object type: %s\n" " New record: %s\n" " Exception: %s" % (mname, new_dict, e)) else: obj.object.pk = match.first().pk #debug.say("Found matching object with new pk: %s" % (obj.object.pk, )) new_pk = obj.object.pk if new_pk != old_pk: # Update other objects referring to this # object to use the new pk #debug.show('old_pk, new_pk') mname = model._meta.app_label + '.' + model.__name__ key = (mname, old_pk) if fkrefs[key] or m2mrefs[key]: #debug.pprint('fkrefs[key]') for o, f in fkrefs[key]: setattr(o, f+'_id', new_pk) if getattr(o, '_saved', False) == True: try: with transaction.atomic(using=self.using): o.save() except IntegrityError as e: self.stderr.write("\nSaving an object failed, possibly due to manually inserted conflicting data.\n" " Object type: %s\n" " New record: %s\n" " Exception: %s" % (mname, obj_to_dict(o), e)) #debug.pprint('m2mrefs[key]') for o, f in m2mrefs[key]: m2mfield = getattr(o, f) if m2mfield.through: through = m2mfield.instance for ff in through._meta.fields: if type(ff) == ForeignKey: old_value = getattr(through, ff.name+'_id') if ff.related_model == model and old_value == old_pk: setattr(through, ff.name+'_id', new_pk) through.save() else: m2mfield.remove(old_pk) m2mfield.add(new_pk) if show_progress: self.stdout.write( '\rProcessed %i object(s).' % loaded_objects_in_fixture, ending='' ) except (DatabaseError, IntegrityError) as e: e.args = ("Could not load %(app_label)s.%(object_name)s(pk=%(pk)s): %(error_msg)s\n %(data)s\n" % { 'app_label': obj.object._meta.app_label, 'object_name': obj.object._meta.object_name, 'pk': obj.object.pk, 'data': obj_to_dict(obj.object), 'error_msg': force_str(e) },) raise if objects and show_progress: self.stdout.write('') # add a newline after progress indicator self.loaded_object_count += loaded_objects_in_fixture self.fixture_object_count += objects_in_fixture except Exception as e: if not isinstance(e, CommandError): e.args = ("Problem installing fixture '%s': %s" % (fixture_file, e),) raise finally: fixture.close() # Warn if the fixture we loaded contains 0 objects. if objects_in_fixture == 0: warnings.warn( "No fixture data found for '%s'. (File format may be " "invalid.)" % fixture_name, RuntimeWarning ) def find_fixtures(self, fixture_label): fixture_files = [] if os.path.exists(fixture_label): fixture_files = [ (fixture_label, os.path.dirname(fixture_label) or os.getcwd(), os.path.basename(fixture_label)) ] else: fixture_files = super(Command, self).find_fixtures(fixture_label) return fixture_files