Skip to content

Commit e0c6f3a

Browse files
levkowetzlevkowetz
committed
Added a management command to merge data from one of two divergent django databases to the other, taking care to insert or update records as appropriate, and update foreign keys and many-to-many keys appropriately. Accepts the dump formats generated by the dumpdata command.
 - Legacy-Id: 17299
1 parent 58d8c2f commit e0c6f3a

1 file changed

Lines changed: 286 additions & 0 deletions

File tree

‎ietf/utils/management/commands/mergedata.py‎

Lines changed: 286 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,286 @@
1+
# Copyright The IETF Trust 2020, All Rights Reserved
2+
# -*- coding: utf-8 -*-
3+
4+
from __future__ import absolute_import, print_function, unicode_literals
5+
6+
import glob
7+
import gzip
8+
import os
9+
import warnings
10+
import zipfile
11+
12+
from collections import defaultdict
13+
from itertools import product, chain
14+
15+
import debug # pyflakes:ignore
16+
17+
from django.apps import apps
18+
from django.conf import settings
19+
from django.contrib.admin.utils import NestedObjects
20+
from django.core import serializers
21+
from django.core.exceptions import ImproperlyConfigured
22+
from django.core.management.base import BaseCommand, CommandError
23+
from django.core.management.color import no_style
24+
from django.core.management.commands.loaddata import Command as LoadCommand, SingleZipReader, humanize
25+
from django.core.management.utils import parse_apps_and_model_labels
26+
from django.db import DEFAULT_DB_ALIAS, DatabaseError, IntegrityError, connections, router, transaction
27+
from django.db.models import ManyToManyField
28+
from django.utils import lru_cache
29+
from django.utils._os import upath
30+
from django.utils.encoding import force_text
31+
from django.utils.functional import cached_property
32+
from django.utils.glob import glob_escape
33+
34+
from ietf.utils.models import ForeignKey
35+
36+
37+
try:
38+
import bz2
39+
has_bz2 = True
40+
except ImportError:
41+
has_bz2 = False
42+
43+
44+
def flatten(l):
45+
if isinstance(l, list):
46+
for el in l:
47+
if isinstance(el, list):
48+
for sub in flatten(el):
49+
yield sub
50+
else:
51+
yield el
52+
else:
53+
yield l
54+
55+
56+
57+
class Command(LoadCommand):
58+
help = 'Merges the named fixture(s) into the database.'
59+
60+
def add_arguments(self, parser):
61+
# parser.add_argument(
62+
# '-e', '--exclude', dest='exclude', action='append', default=[],
63+
# help='An app_label or app_label.ModelName to exclude. Can be used multiple times.',
64+
# )
65+
super(Command, self).add_arguments(parser)
66+
67+
def load_label(self, fixture_label):
68+
"""
69+
Loads fixtures files for a given label.
70+
"""
71+
# def update_objects(objects, model, old_pk, new_pk):
72+
# debug.show('old_pk, new_pk')
73+
# for o in objects:
74+
# for f in o._meta.fields:
75+
# if type(f) == ForeignKey:
76+
# if f.pk in [old_pk, new_pk]:
77+
# #debug.show('f')
78+
# #debug.show('f.pk')
79+
# elif type(f) == ManyToManyKey:
80+
# pass
81+
82+
def obj_to_dict(obj):
83+
opts = obj._meta
84+
data = {}
85+
for f in chain(opts.concrete_fields, opts.private_fields):
86+
data[f.name] = f.value_from_object(obj)
87+
return data
88+
89+
def get_unique_field(obj):
90+
unique = {}
91+
for f in obj._meta.get_fields():
92+
if hasattr(f, 'primary_key') and f.primary_key:
93+
continue
94+
if hasattr(f, 'unique') and f.unique:
95+
unique[f.name] = f.value_from_object(obj)
96+
return unique
97+
98+
show_progress = self.verbosity >= 3
99+
if hasattr(settings, 'SERVER_MODE'):
100+
# Try to avoid sending mail during repair
101+
settings.SERVER_MODE = 'repair'
102+
for fixture_file, fixture_dir, fixture_name in self.find_fixtures(fixture_label):
103+
_, ser_fmt, cmp_fmt = self.parse_name(os.path.basename(fixture_file))
104+
open_method, mode = self.compression_formats[cmp_fmt]
105+
fixture = open_method(fixture_file, mode)
106+
try:
107+
self.fixture_count += 1
108+
objects_in_fixture = 0
109+
loaded_objects_in_fixture = 0
110+
if self.verbosity >= 2:
111+
self.stdout.write(
112+
"Installing %s fixture '%s' from %s."
113+
% (ser_fmt, fixture_name, humanize(fixture_dir))
114+
)
115+
116+
objects = list(serializers.deserialize(
117+
ser_fmt, fixture, using=self.using, ignorenonexistent=self.ignore,
118+
))
119+
120+
# Prime fkrefs and m2mrefs with our deserialized objects
121+
fkrefs = {}
122+
m2mrefs = {}
123+
for obj in objects:
124+
o = obj.object
125+
oname = o._meta.app_label + '.' + o.__class__.__name__
126+
fkrefs[(oname, o.pk)] = []
127+
m2mrefs[(oname, o.pk)] = []
128+
129+
# Tabulate all FKs and M2M fields
130+
for obj in objects:
131+
o = obj.object
132+
for f in o._meta.get_fields():
133+
if type(f) == ForeignKey:
134+
fobj = getattr(o, f.name)
135+
if fobj:
136+
fmod = f.related_model
137+
fname = fmod._meta.app_label + '.' + fmod.__name__
138+
key = (fname, fobj.pk)
139+
if key in fkrefs:
140+
fkrefs[key].append((o, f.name))
141+
elif type(f) == ManyToManyField:
142+
fobjs = getattr(o, f.name).all()
143+
if fobjs:
144+
fmod = f.related_model
145+
fname = fmod._meta.app_label + '.' + fmod.__name__
146+
for fobj in fobjs:
147+
key = (fname, fobj.pk)
148+
if key in m2mrefs:
149+
m2mrefs[key].append((o, f.name))
150+
else:
151+
pass
152+
#debug.type('f')
153+
154+
for obj in objects:
155+
objects_in_fixture += 1
156+
if (obj.object._meta.app_config in self.excluded_apps or
157+
type(obj.object) in self.excluded_models):
158+
continue
159+
if router.allow_migrate_model(self.using, obj.object.__class__):
160+
loaded_objects_in_fixture += 1
161+
self.models.add(obj.object.__class__)
162+
try:
163+
model = obj.object.__class__
164+
mname = model._meta.app_label + '.' + model.__name__
165+
old_pk = obj.object.pk
166+
unique = get_unique_field(obj.object)
167+
if unique:
168+
match = model.objects.filter(**unique)
169+
if not match.exists():
170+
obj.object.pk = None
171+
try:
172+
with transaction.atomic(using=self.using):
173+
obj.save(using=self.using)
174+
setattr(obj.object, '_saved', True)
175+
except IntegrityError as e:
176+
this_pk = match.first().pk if match else None
177+
new_dict = obj_to_dict(obj.object)
178+
self.stderr.write("\nSaving an updated object failed, possibly due to manually inserted conflicting data.\n"
179+
" Found PK: %s: %s\n"
180+
" Unique key: %s\n"
181+
" New record: %s\n"
182+
" Exception: %s" % (mname, this_pk, unique, new_dict, e))
183+
else:
184+
match = model.objects.filter(pk=obj.object.pk)
185+
if match.exists():
186+
#debug.say('PK exists: %s' % obj.object.pk)
187+
prev = match.first()
188+
prev_dict = obj_to_dict(prev)
189+
obj_dict = obj_to_dict(obj.object)
190+
if prev_dict == obj_dict:
191+
pass # nothing to do, the object is already there
192+
else:
193+
del obj_dict[obj.object._meta.pk.name]
194+
match = model.objects.filter(**obj_dict)
195+
if not match:
196+
try:
197+
obj.object.pk = None
198+
with transaction.atomic(using=self.using):
199+
obj.save(using=self.using)
200+
setattr(obj.object, '_saved', True)
201+
except IntegrityError as e:
202+
new_dict = obj_to_dict(obj.object)
203+
self.stderr.write("\nSaving an object failed, possibly due to manually inserted conflicting data.\n"
204+
" Object type: %s\n"
205+
" New record: %s\n"
206+
" Exception: %s" % (mname, new_dict, e))
207+
else:
208+
obj.object.pk = match.first().pk
209+
#debug.say("Found matching object with new pk: %s" % (obj.object.pk, ))
210+
new_pk = obj.object.pk
211+
if new_pk != old_pk:
212+
# Update other objects refering to this
213+
# object to use the new pk
214+
#debug.show('old_pk, new_pk')
215+
mname = model._meta.app_label + '.' + model.__name__
216+
key = (mname, old_pk)
217+
if fkrefs[key] or m2mrefs[key]:
218+
#debug.pprint('fkrefs[key]')
219+
for o, f in fkrefs[key]:
220+
setattr(o, f+'_id', new_pk)
221+
if getattr(o, '_saved', False) == True:
222+
try:
223+
with transaction.atomic(using=self.using):
224+
o.save()
225+
except IntegrityError as e:
226+
self.stderr.write("\nSaving an object failed, possibly due to manually inserted conflicting data.\n"
227+
" Object type: %s\n"
228+
" New record: %s\n"
229+
" Exception: %s" % (mname, obj_to_dict(o), e))
230+
231+
#debug.pprint('m2mrefs[key]')
232+
for o, f in m2mrefs[key]:
233+
m2mfield = getattr(o, f)
234+
if m2mfield.through:
235+
through = m2mfield.instance
236+
for ff in through._meta.fields:
237+
if type(ff) == ForeignKey:
238+
old_value = getattr(through, ff.name+'_id')
239+
if ff.related_model == model and old_value == old_pk:
240+
setattr(through, ff.name+'_id', new_pk)
241+
through.save()
242+
else:
243+
m2mfield.remove(old_pk)
244+
m2mfield.add(new_pk)
245+
246+
247+
if show_progress:
248+
self.stdout.write(
249+
'\rProcessed %i object(s).' % loaded_objects_in_fixture,
250+
ending=''
251+
)
252+
except (DatabaseError, IntegrityError) as e:
253+
e.args = ("Could not load %(app_label)s.%(object_name)s(pk=%(pk)s): %(error_msg)s\n %(data)s\n" % {
254+
'app_label': obj.object._meta.app_label,
255+
'object_name': obj.object._meta.object_name,
256+
'pk': obj.object.pk,
257+
'data': obj_to_dict(obj.object),
258+
'error_msg': force_text(e)
259+
},)
260+
raise
261+
if objects and show_progress:
262+
self.stdout.write('') # add a newline after progress indicator
263+
self.loaded_object_count += loaded_objects_in_fixture
264+
self.fixture_object_count += objects_in_fixture
265+
except Exception as e:
266+
if not isinstance(e, CommandError):
267+
e.args = ("Problem installing fixture '%s': %s" % (fixture_file, e),)
268+
raise
269+
finally:
270+
fixture.close()
271+
272+
# Warn if the fixture we loaded contains 0 objects.
273+
if objects_in_fixture == 0:
274+
warnings.warn(
275+
"No fixture data found for '%s'. (File format may be "
276+
"invalid.)" % fixture_name,
277+
RuntimeWarning
278+
)
279+
280+
def find_fixtures(self, fixture_label):
281+
fixture_files = []
282+
if os.path.exists(fixture_label):
283+
fixture_files = [ (fixture_label, os.path.dirname(fixture_label) or os.getcwd(), os.path.basename(fixture_label)) ]
284+
else:
285+
fixture_files = super(Command, self).find_fixtures(fixture_label)
286+
return fixture_files

0 commit comments

Comments
 (0)