1+ # Copyright The IETF Trust 2020, All Rights Reserved
2+ # -*- coding: utf-8 -*-
3+ from __future__ import absolute_import , print_function , unicode_literals
4+
5+ import collections
6+ import io
7+ import re
8+ import sys
9+ import warnings
10+
11+ from django .apps import apps
12+ from django .core import serializers
13+ from django .core .management .base import CommandError
14+ from django .core .management .commands .dumpdata import Command as DumpdataCommand
15+ from django .core .management .utils import parse_apps_and_model_labels
16+ from django .db import router
17+
18+ import debug # pyflakes:ignore
19+
20+ # ------------------------------------------------------------------------------
21+
22+ class Command (DumpdataCommand ):
23+ """
24+ Read 'INSERT INTO' lines from a (probably partial) SQL dump file, and
25+ extract table names and primary keys; then use these to do a data dump of
26+ the indicated records.
27+
28+ Only simpler variations on the full sql INSERT command are recognized.
29+
30+ The expected way to derive the input file is to do a diff between two sql
31+ dump files, and remove any diff line prefixes ('<' or '>' or '+' or -)
32+ from the diff, leaving only SQL "INSERT INTO" statements.
33+ """
34+ help = __doc__
35+
36+ def add_arguments (self , parser ):
37+ super (Command , self ).add_arguments (parser )
38+ # remove the usual positional args
39+ for i , a in enumerate (parser ._actions ):
40+ if a .dest == 'args' :
41+ break
42+ del parser ._actions [i ]
43+ parser .add_argument ('filenames' , nargs = '*' ,
44+ help = "One or more files to process" )
45+ parser .add_argument ('--pk-name' , default = 'id' , type = str ,
46+ help = "Use the specified name as the primary key filed name (default: '%(default)s')" )
47+
48+ def note (self , msg ):
49+ if self .verbosity > 1 :
50+ self .stderr .write ('%s\n ' % msg )
51+
52+ def warn (self , msg ):
53+ self .stderr .write ('Warning: %s\n ' % msg )
54+
55+ def err (self , msg ):
56+ self .stderr .write ('Error: %s\n ' % msg )
57+ sys .exit (1 )
58+
59+ def get_tables (self ):
60+ seen = set ([])
61+ tables = {}
62+ for name , appconf in apps .app_configs .items ():
63+ for model in appconf .get_models ():
64+ if not model in seen :
65+ seen .add (model )
66+ app_label = model ._meta .app_label
67+ tables [model ._meta .db_table ] = {
68+ 'app_config' : apps .get_app_config (app_label ),
69+ 'app_label' : app_label ,
70+ 'model' : model ,
71+ 'model_label' : model .__name__ ,
72+ 'pk' : model ._meta .pk .name ,
73+ }
74+ return tables
75+
76+ def get_pks (self , filenames , tables ):
77+ count = 0
78+ pks = {}
79+ for fn in filenames :
80+ prev = ''
81+ lc = 0
82+ with io .open (fn ) as f :
83+ for line in f :
84+ lc += 1
85+ line = line .strip ()
86+ if line [0 ] in ['<' , '>' ]:
87+ self .err ("Input file '%s' looks like a diff file. Please provide just the SQL 'INSERT' statements for the records to be dumped." % (fn , ))
88+ if prev :
89+ line = prev + line
90+ prev = None
91+ if not line .endswith (';' ):
92+ prev = line
93+ continue
94+ sql = line
95+ if not sql .upper ().startswith ('INSERT ' ):
96+ self .warn ("Skipping sql '%s...'" % sql [:64 ])
97+ else :
98+ sql = sql .replace ("\\ '" , "\\ x27" )
99+ match = re .match (r"INSERT( +(LOW_PRIORITY|DELAYED|HIGH_PRIORITY))*( +IGNORE)?( +INTO)?"
100+ r" +(?P<table>\S+)"
101+ r" +\((?P<fields>([^ ,]+)(, [^ ,]+)*)\)"
102+ r" +(VALUES|VALUE)"
103+ r" +\((?P<values>(\d+|'[^']*'|NULL)(,(\d+|'[^']*'|NULL))*)\)"
104+ r" *;"
105+ , sql )
106+ if not match :
107+ self .warn ("Unrecognized sql command: '%s'" % sql )
108+ else :
109+ table = match .group ('table' ).strip ('`' )
110+ if not table in pks :
111+ pks [table ] = []
112+ fields = match .group ('fields' )
113+ fields = [ f .strip ("`" ) for f in re .split (r"(`[^`]+`)" , fields ) if f and not re .match (r'\s*,\s*' , f )]
114+ values = match .group ('values' )
115+ values = [ v .strip ("'" ) for v in re .split (r"(\d+|'[^']*'|NULL)" , values ) if v and not re .match (r'\s*,\s*' , v ) ]
116+ try :
117+ pk_name = tables [table ]['pk' ]
118+ ididx = fields .index (pk_name )
119+ pk = values [ididx ]
120+ pks [table ].append (pk )
121+ count += 1
122+ except (KeyError , ValueError ) as e :
123+ pass
124+ return pks , count
125+
126+ def get_objects (self , app_list , pks , count_only = False ):
127+ """
128+ Collate the objects to be serialized. If count_only is True, just
129+ count the number of objects to be serialized.
130+ """
131+ models = serializers .sort_dependencies (app_list .items ())
132+ excluded_models , __ = parse_apps_and_model_labels (self .excludes )
133+ for model in models :
134+ if model in excluded_models :
135+ continue
136+ if not model ._meta .proxy and router .allow_migrate_model (self .using , model ):
137+ if self .use_base_manager :
138+ objects = model ._base_manager
139+ else :
140+ objects = model ._default_manager
141+
142+ queryset = objects .using (self .using ).order_by (model ._meta .pk .name )
143+ primary_keys = pks [model ._meta .db_table ] if model ._meta .db_table in pks else []
144+ if primary_keys :
145+ queryset = queryset .filter (pk__in = primary_keys )
146+ #self.stderr.write('+ %s: %s\n' % (model._meta.db_table, queryset.count() ))
147+ else :
148+ continue
149+ if count_only :
150+ yield queryset .order_by ().count ()
151+ else :
152+ for obj in queryset .iterator ():
153+ yield obj
154+
155+
156+ def handle (self , filenames = [], ** options ):
157+ self .verbosity = int (options .get ('verbosity' ))
158+ format = options ['format' ]
159+ indent = options ['indent' ]
160+ self .using = options ['database' ]
161+ self .excludes = options ['exclude' ]
162+ output = options ['output' ]
163+ show_traceback = options ['traceback' ]
164+ use_natural_foreign_keys = options ['use_natural_foreign_keys' ]
165+ use_natural_primary_keys = options ['use_natural_primary_keys' ]
166+ self .use_base_manager = options ['use_base_manager' ]
167+ pks = options ['primary_keys' ]
168+
169+ # Check that the serialization format exists; this is a shortcut to
170+ # avoid collating all the objects and _then_ failing.
171+ if format not in serializers .get_public_serializer_formats ():
172+ try :
173+ serializers .get_serializer (format )
174+ except serializers .SerializerDoesNotExist :
175+ pass
176+
177+ raise CommandError ("Unknown serialization format: %s" % format )
178+
179+ tables = self .get_tables ()
180+ pks , count = self .get_pks (filenames , tables )
181+ sys .stdout .write ("Found %s SQL records.\n " % count )
182+
183+ app_list = collections .OrderedDict ()
184+
185+ for t in tables :
186+ #print("%32s\t%s" % (t, ','.join(pks[t])))
187+ app_config = tables [t ]['app_config' ]
188+ app_list .setdefault (app_config , [])
189+ app_list [app_config ].append (tables [t ]['model' ])
190+
191+ #debug.pprint('app_list')
192+
193+ try :
194+ self .stdout .ending = None
195+ progress_output = None
196+ object_count = 0
197+ # If dumpdata is outputting to stdout, there is no way to display progress
198+ if (output and self .stdout .isatty () and options ['verbosity' ] > 0 ):
199+ progress_output = self .stdout
200+ object_count = sum (self .get_objects (app_list , pks , count_only = True ))
201+ stream = open (output , 'w' ) if output else None
202+ try :
203+ serializers .serialize (
204+ format , self .get_objects (app_list , pks ), indent = indent ,
205+ use_natural_foreign_keys = use_natural_foreign_keys ,
206+ use_natural_primary_keys = use_natural_primary_keys ,
207+ stream = stream or self .stdout , progress_output = progress_output ,
208+ object_count = object_count ,
209+ )
210+ sys .stdout .write ("Dumped %s objects.\n " % object_count )
211+ finally :
212+ if stream :
213+ stream .close ()
214+ except Exception as e :
215+ if show_traceback :
216+ raise
217+ raise CommandError ("Unable to serialize database: %s" % e )
218+
0 commit comments