· 8 years ago · Sep 25, 2017, 12:48 AM
1#!/usr/bin/env python
2
3"""my2pg.py: MySQL to PostgreSQL database conversion
4
5Copyright (c) 2010 Matrix Group International.
6Licensed under the MIT license; see LICENSE file for terms.
7
8"""
9
10import sys, os
11import optparse, logging, traceback
12import re, collections, pickle
13
14import MySQLdb, psycopg2
15from MySQLdb.cursors import DictCursor
16
17# We commit data rows every so often.
18COMMIT_AFTER_ROWS = 10000
19ROWS_PER_BLOCK = 100000
20
21def pg_execute(pg_conn, options, sql, args=()):
22 """(Connection, Options, str, tuple)
23
24 Log and execute a SQL command on the PostgreSQL connection.
25 """
26 #print sql
27 if not options.dry_run:
28 pg_cur = pg_conn.cursor()
29 pg_cur.execute(sql, args)
30
31# XXX need to expand this set of words.
32_reserved_words = set("""end user""".split())
33
34def is_reserved_word(word):
35 """(str): bool
36
37 Returns true if this word is a PostgreSQL reserved-word.
38 """
39 return word in _reserved_words
40
41def fix_reserved_word(S):
42 """(str): str
43
44 Takes a MySQL name, and adds an underscore if it's a PostgreSQL
45 reserved word.
46 """
47 if is_reserved_word(S.lower()):
48 S += '_'
49 return S
50
51def convert_type(typ):
52 """(str): str
53
54 Parses a MySQL type declaration and returns the corresponding PostgreSQL
55 type.
56 """
57 if re.match('tinyint([(]\d+[)])?', typ):
58 # MySQL tinyint is 1 byte, -128 to 127; we'll use the 2-byte int.
59 return 'smallint'
60 elif re.match('smallint([(]\d+[)])?', typ):
61 return 'smallint'
62 elif re.match('mediumint([(]\d+[)])?', typ):
63 # MySQL medium int is 3 bytes; we'll use the 4-byte int.
64 return 'integer'
65 elif re.match('bigint([(]\d+[)])?', typ):
66 # XXX use the parametrized number?
67 # XXX 'bigint NOT NULL auto_increment' -> bigserial
68 return 'bigint'
69 elif re.match('integer([(]\d+[)])?', typ):
70 return 'integer'
71 elif re.match('int([(]\d+[)])?', typ):
72 return 'integer'
73 elif typ == 'float':
74 return 'real'
75 elif re.match('double([(]\d+,\d+[)])?', typ):
76 return 'double precision'
77 elif typ == 'datetime':
78 return 'timestamp'
79 elif typ in ('tinytext', 'text', 'mediumtext', 'longtext'):
80 return 'text'
81 elif typ in ('tinyblob', 'blob', 'mediumblob', 'longblob'):
82 return 'bytea'
83 elif typ.startswith('enum('):
84 # For enums, we'll return a varchar long enough to hold
85 # the longest possible value.
86 # XXX this parsing is very dumb.
87 values = typ[5:-1] # Chop off enum( and ).
88 values = eval(values)
89 longest = max(len(v) for v in values)
90 return 'varchar(%i)' % longest
91
92 # Give up and just return the input type.
93 return typ
94
95def convert_data(col, data):
96 """(Column, any) : any
97
98 Convert a Python value retrieved from MySQL into a PostgreSQL value.
99 """
100 if isinstance(data, str):
101 data = data.decode('latin-1')
102 data = data.encode('utf-8')
103
104 if col.type in ('tinyblob', 'blob', 'mediumblob', 'longblob') and data:
105 # Convert to a BYTEA literal. We just use octal escapes for
106 # everything.
107 L = [('\\%03o') % ord(ch) for ch in data]
108 data = ''.join(L)
109 return data
110
111class Column:
112 """
113 Represents a column.
114
115 Instance attributes:
116 name : str
117 type : str
118 position : int
119 default : str
120 is_nullable : bool
121
122 """
123
124 def __init__(self, **kw):
125 for k,v in kw.items():
126 setattr(self, k, v)
127
128 def pg_decl(self):
129 """(): str
130
131 Return the PostgreSQL declaration syntax for this column.
132 """
133 typ = convert_type(self.type)
134 decl_typ = typ
135 decl = ' %s %s' % (fix_reserved_word(self.name), decl_typ)
136 if self.default:
137 default = self.get_default()
138 decl += ' DEFAULT %s' % default
139 if not self.is_nullable:
140 decl += ' NOT NULL'
141 return decl
142
143 def get_default(self):
144 if (self.type in ('datetime', 'timestamp') and
145 self.default == "0000-00-00 00:00:00"):
146 return 'NULL'
147
148 typ = convert_type(self.type)
149 if typ.startswith(('char', 'varchar')):
150 return "'" + self.default + "'"
151
152 return self.default
153
154class Index:
155 """
156 Represents an index.
157
158 Instance attributes:
159 name : str
160 table : str
161 type : str
162 column_names : [str]
163 non_unique : bool
164 nullable : bool
165
166 """
167
168 def __init__(self, **kw):
169 self.column_names = []
170 for k,v in kw.items():
171 setattr(self, k, v)
172
173 def pg_decl(self):
174 """(): str
175
176 Return the PostgreSQL declaration syntax for this index.
177 """
178 # We'll ignore the MySQL index name, and invent a new name.
179 name = 'idx_' + '_'.join([self.table] + self.column_names)
180 sql = 'CREATE INDEX %s ON %s(%s)' % (fix_reserved_word(name),
181 fix_reserved_word(self.table),
182 ','.join(self.column_names))
183 if self.type:
184 # XXX convert index_type:
185 # BTREE, etc.
186 pass
187 return sql
188
189def read_mysql_tables(mysql_cur, mysql_db, options):
190 """(Cursor):
191 """
192 logging.info('Reading structure of MySQL database')
193 mysql_cur.execute("""
194SELECT * FROM information_schema.tables WHERE table_schema = %s
195""", mysql_db)
196 rows = mysql_cur.fetchall()
197 tables = sorted(row['TABLE_NAME'] for row in rows)
198 if options.starting_table:
199 tables = [t for t in tables if options.starting_table <= t]
200
201 # Convert tables
202 table_cols = {}
203 table_indexes = {}
204 for table in tables:
205 logging.debug('Reading table %s', table)
206 mysql_cur.execute("""
207SELECT * FROM information_schema.columns
208WHERE table_schema = %s and table_name = %s
209""", (mysql_db, table))
210 cols = table_cols[table] = []
211 for row in mysql_cur.fetchall():
212 c = Column()
213 cols.append(c)
214 c.name = row['COLUMN_NAME']
215 c.type = row['COLUMN_TYPE'] #
216 c.position = row['ORDINAL_POSITION']
217 c.default = row['COLUMN_DEFAULT']
218 c.is_nullable = bool(row['IS_NULLABLE'] == 'YES')
219 # XXX character set?
220
221 # Sort columns into left-to-right order.
222 cols.sort(key=lambda c: c.position)
223
224 # Convert indexes
225 mysql_cur.execute("""
226SELECT * FROM information_schema.statistics
227WHERE table_schema = %s AND table_name = %s
228""", (mysql_db, table))
229 d = collections.defaultdict(Index)
230 for row in mysql_cur.fetchall():
231 index_name = row['INDEX_NAME']
232 i = d[index_name]
233 i.table = table
234 i.name = index_name
235 i.column_names.append(row['COLUMN_NAME'])
236 i.type = row['INDEX_TYPE']
237 i.non_unique = bool(row['NON_UNIQUE'])
238 i.nullable = bool(row['NULLABLE'] == 'YES')
239 table_indexes[table] = d.values()
240
241 return tables, table_cols, table_indexes
242
243
244def main ():
245 parser = optparse.OptionParser(
246 '%prog [options] mysql-host mysql-db pg-host pg-db')
247 parser.add_option('--data-only',
248 action="store_true", default=False,
249 dest="data_only",
250 help="Assume the tables already exist, and only convert data")
251 parser.add_option('--drop-tables',
252 action="store_true", default=False,
253 dest="drop_tables",
254 help="Drop existing PostgreSQL tables (if any) before creating")
255 parser.add_option('-n', '--dry-run',
256 action="store_true", default=False,
257 dest="dry_run",
258 help="Make no changes to PostgreSQL database")
259 parser.add_option('--mysql-user',
260 action="store",
261 dest="mysql_user",
262 help="User for login if not current user.")
263 parser.add_option('--mysql-password',
264 action="store",
265 dest="mysql_password",
266 help="Password to use when connecting to server.")
267 parser.add_option('--pg-user',
268 action="store",
269 dest="pg_user",
270 help="User for login if not current user.")
271 parser.add_option('--pg-password',
272 action="store", default='',
273 dest="pg_password",
274 help="Password to use when connecting to server.")
275 parser.add_option('--pickle=',
276 action="store", default='',
277 dest="pickle",
278 help="File for storing a pickled version of the MySQL table structure.")
279 parser.add_option('--starting-table',
280 action="store", default=None,
281 dest="starting_table",
282 help="Name of table to start conversion with")
283 parser.add_option('-v', '--verbose',
284 action="count", default=0,
285 dest="verbose",
286 help="Display more output as the script runs")
287
288 options, args = parser.parse_args()
289 if len(args) != 4:
290 parser.print_help()
291 sys.exit(1)
292
293 mysql_host, mysql_db, pg_host, pg_db = args
294
295 # Set logging level.
296 if options.verbose:
297 logging.basicConfig(level=logging.INFO)
298 if options.verbose > 1:
299 logging.basicConfig(level=logging.DEBUG)
300
301 # Set up connections
302 logging.info('Connecting to databases')
303
304 mysql_conn = MySQLdb.Connection(
305 user=options.mysql_user,
306 passwd=options.mysql_password,
307 db=mysql_db,
308 host=mysql_host,
309 )
310 pg_conn = psycopg2.connect(
311 database=pg_db,
312 host=pg_host,
313 user=options.pg_user,
314 password=options.pg_password,
315 )
316 mysql_cur = mysql_conn.cursor(cursorclass=DictCursor)
317
318 # Make list of tables to process.
319 if options.pickle and os.path.exists(options.pickle):
320 f = open(options.pickle, 'rb')
321 tables, table_cols, table_indexes = pickle.load(f)
322 f.close()
323
324 # Discard tables that we don't need to process.
325 if options.starting_table:
326 tables = [t for t in tables if options.starting_table <= t]
327
328 else:
329 tables, table_cols, table_indexes = read_mysql_tables(mysql_cur,
330 mysql_db,
331 options)
332 if options.pickle and not options.starting_table:
333 f = open(options.pickle, 'wb')
334 t = (tables, table_cols, table_indexes)
335 pickle.dump(t, f)
336 f.close()
337
338 #
339 # Convert the table structure.
340 #
341 if not options.data_only:
342 for table in tables:
343 cols = table_cols[table]
344 indexes = table_indexes[table]
345
346 # Drop table if necessary.
347 if options.drop_tables:
348 sql = "DROP TABLE IF EXISTS %s" % fix_reserved_word(table)
349 pg_execute(pg_conn, options, sql)
350
351 # Assemble into a PGSQL declaration
352 pg_table = fix_reserved_word(table)
353 sql = "CREATE TABLE %s (\n" % pg_table
354 sql += ',\n'.join(c.pg_decl() for c in cols) + '\n'
355
356 # Look for index named PRIMARY, and add PRIMARY KEY if found.
357 primary_L = [i for i in indexes if i.name == 'PRIMARY']
358 if len(primary_L):
359 if len(primary_L) > 1:
360 logging.warn('%s: Multiple PRIMARY indexes on table',
361 table)
362 else:
363 primary = primary_L.pop()
364 sql = sql.rstrip() + ',\n'
365 sql += ' PRIMARY KEY (%s)' % ','.join(primary.column_names)
366
367 sql += ');'
368 pg_execute(pg_conn, options, sql)
369
370 pg_conn.commit()
371
372 # Create indexes
373 for i in indexes:
374 if i.name == 'PRIMARY':
375 continue
376
377 sql = i.pg_decl()
378 try:
379 pg_execute(pg_conn, options, sql)
380 except Exception:
381 logging.error('Failure creating index on table %s', table,
382 exc_info=True)
383
384 pg_conn.commit()
385
386 #
387 # Convert data.
388 #
389 logging.info('Converting data')
390 for table in tables:
391 # Convert data.
392 logging.info('Converting data in table %s', table)
393 pg_table = fix_reserved_word(table)
394 cols = table_cols[table]
395
396 # Assemble the INSERT statement once.
397 ins_sql = ('INSERT INTO %s (%s) VALUES (%s);' %
398 (pg_table,
399 ', '.join(fix_reserved_word(c.name) for c in cols),
400 ','.join(['%s'] * len(cols))))
401
402 # Ensure the table is empty.
403 pg_execute(pg_conn, options, "DELETE FROM %s" % pg_table)
404
405
406 # We don't do a fetchall() since the table contents are
407 # very likely to not fit into memory.
408 row_count = 0
409 errors = 0
410 rowsHere=True
411 while rowsHere:
412 logging.info('executing query from offset %d'%row_count)
413 mysql_cur.execute("SELECT * FROM %s limit %d offset %d" % (table,ROWS_PER_BLOCK,row_count))
414 rowsHere=False
415 while True:
416 row = mysql_cur.fetchone()
417 if row is None:
418 break
419 rowsHere=True
420
421 # Assemble a list of the output data that we'll subsequently
422 # convert to a tuple.
423 output_L = []
424 for c in cols:
425 data = row[c.name]
426 newdata = convert_data(c, data)
427 output_L.append(newdata)
428
429 try:
430 pg_execute(pg_conn, options, ins_sql, tuple(output_L))
431 except KeyboardInterrupt:
432 raise
433 except:
434 logging.error('Failure inserting row into table %s', table,
435 exc_info=True)
436 errors += 1
437 else:
438 row_count += 1
439 if (row_count % COMMIT_AFTER_ROWS) == 0:
440 logging.debug('Committing transaction after %i rows', COMMIT_AFTER_ROWS)
441 pg_conn.commit()
442
443 logging.info("Table %s: %i rows converted (%i errors)",
444 table, row_count, errors)
445 pg_conn.commit()
446
447 # Close connections
448 logging.info('Closing database connections')
449 mysql_conn.close()
450 pg_conn.close()
451
452
453if __name__ == '__main__':
454 main()