blob: 2639a9b82d7ac883766022437d375537f4453935 [file] [log] [blame]
# Copyright (c) 2012 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""Module to upload a MySQL dump file to Cloud SQL.
Usage: [-h] [--resume NUM] [--user USER] [--passwd PASSWD] FILE
Uploads MySQL dump file to a MySQL database or Cloud SQL. With no optional
arguments will connect to localhost as root with an empty password.
positional arguments:
FILE text dump file containing MySQL commands
REMOTE Cloud SQL instance name or MySQL hostname
optional arguments:
-h, --help show this help message and exit
--resume NUM resume dump at command NUM
--user USER user (ignored for CloudSQL)
--passwd PASSWD passwd (ignored for CloudSQL)
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import argparse
import collections
import datetime
import os
import re
import sys
import time
import six
BYTES_PER_GB = 2**30
class MySQLConnectionManager(object):
"""Manages connections to a MySQL database.
factory: A *ConnectionFactory.
connected: Whether we currently hold a live DB connection.
cmd_num: The number of commands executed.
def __init__(self, connection_factory):
self.factory = connection_factory
self.connected = False
self.cmd_num = 0
def write(self, data, execute_cmd=True, increment_cmd=False):
"""Buffers writes to command boundaries.
data: A line of data from the MySQL dump.
execute_cmd: Whether to execute the command, defaults to True.
increment_cmd: Whether to increment cmd_num, defaults to False.
if not data or not data.strip() or data == '\n' or data[:2] == '--':
self._cmd += data[:-1] if data[-1] == '\n' else data
if self._cmd[-1] != ';':
# Execute command.
if execute_cmd:
self._cursor.execute(six.ensure_text(self._cmd, 'utf-8'))
self._cmd = ''
if increment_cmd:
self.cmd_num += 1
def disconnect(self):
"""Closes the current database connection."""
if self.connected:
self.connected = False
def connect(self):
"""Creates a new database connection."""
self._db = self.factory.connect()
self.connected = True
self._cursor = self._db.cursor()
self._cmd = ''
class CloudSQLConnectionFactory(object):
"""Creates Cloud SQL database connections."""
def __init__(self, cloudsql_instance):
self._instance = cloudsql_instance
def connect(self):
"""Connects to the Cloud SQL database and returns the connection.
A MySQLdb compatible database connection to the Cloud SQL instance.
print('Connecting to Cloud SQL instance %s.' % self._instance)
from import rdbms_googleapi
except ImportError:
sys.exit('Unable to import rdbms_googleapi. Add the AppEngine SDK '
'directory to your PYTHONPATH. Download the SDK from: '
return rdbms_googleapi.connect(None, instance=self._instance)
class LocalSQLConnectionFactory(object):
"""Creates local MySQL database connections."""
def __init__(self, host=None, user='root', passwd=''):
if not host:
host = 'localhost'
self._host = host
self._user = user
self._passwd = passwd
def connect(self):
"""Connects to the local MySQL database and returns the connection.
A MySQLdb database connection to the local MySQL database.
print('Connecting to mysql at localhost as %s.' % self._user)
import MySQLdb
except ImportError:
sys.exit('Unable to import MySQLdb. To install on Ubuntu: '
'apt-get install python-mysqldb')
return MySQLdb.connect(host=self._host, user=self._user,
class MySQLState(object):
"""Maintains the MySQL global state.
This is a hack that keeps record of all MySQL lines that set global state.
These are needed to reconstruct the MySQL state on resume.
_set_regex = re.compile('\S*\s*SET(.*)[\s=]')
def __init__(self):
self._db_line = ''
self._table_lock = []
self._sets = collections.OrderedDict()
def process(self, line):
"""Check and save lines that affect the global state.
line: A line from the MySQL dump file.
# Most recent USE line.
if line[:3] == 'USE':
self._db_line = line
# SET variables.
m = self._set_regex.match(line)
if m:
self._sets[] = line
# Maintain LOCK TABLES
if (line[:11] == 'LOCK TABLES' or
('ALTER TABLE' in line and 'DISABLE KEYS' in line)):
if (line[:14] == 'UNLOCK TABLES;'):
self._table_lock = []
def write(self, out):
"""Print lines to recreate the saved state.
out: A File-like object to write out saved state.
for v in six.itervalues(self._sets):
for l in self._table_lock:
def breakpoint(self, line):
"""Returns true if we can handle breaking after this line.
line: A line from the MySQL dump file.
Boolean indicating whether we can break after |line|.
return (line[:28] == '-- Table structure for table' or
line[:11] == 'INSERT INTO')
def dump_to_cloudsql(dumpfile, manager, cmd_offset=0):
"""Dumps a MySQL dump file to a database through a MySQLConnectionManager.
dumpfile: Path to a file from which to read the MySQL dump.
manager: An instance of MySQLConnectionManager.
cmd_offset: No commands will be executed on the database before this count
is reached. Used to continue an uncompleted dump. Defaults to 0.
state = MySQLState()
total = os.path.getsize(dumpfile)
start_time = time.time()
line_num = 0
with open(dumpfile, 'r') as dump:
for line in dump:
line_num += 1
if not manager.connected:
# Construct commands from lines and execute them.
if manager.cmd_num == cmd_offset and cmd_offset != 0:
print('\nRecreating state at line: %d' % line_num)
manager.write(line, manager.cmd_num >= cmd_offset, True)
# Print status.
'\rstatus: %.3f%% %0.2f GB %d commands ' %
(100 * dump.tell() / total, dump.tell() / BYTES_PER_GB,
# Handle interrupts and connection failures.
except KeyboardInterrupt:
print('\nInterrupted while executing command: %d' %
print('\nFailed while executing command: %d' % manager.cmd_num)
delta = int(time.time() - start_time)
print('Total time: %s' % str(datetime.timedelta(seconds=delta)))
if state.breakpoint(line):
# Attempt to resume.
print('Execution can resume from here (line = %d)' %
manager.cmd_num += 1
cmd_offset = manager.cmd_num
print('Will now attempt to auto-resume at command: %d' %
print('Execution may fail to resume correctly from here.')
print('Use --resume=%d to attempt to resume the dump.' %
if __name__ == '__main__':
"""Imports a MySQL database from a dump file.
Interprets command line arguments and calls dump_to_cloudsql appropriately.
description = """Uploads MySQL dump file to a MySQL database or Cloud SQL.
With no optional arguments will connect to localhost as root
with an empty password."""
parser = argparse.ArgumentParser(description=description)
parser.add_argument('mysqldump', metavar='FILE',
help='text dump file containing MySQL commands')
parser.add_argument('remote', default=None, nargs='?', metavar='REMOTE',
help='either a Cloud SQL account:instance or a hostname')
parser.add_argument('--resume', default=0, type=int, metavar='NUM',
help='resume dump at command NUM')
parser.add_argument('--user', default='root', metavar='USER',
help='user (ignored for Cloud SQL)')
parser.add_argument('--passwd', default='', metavar='PASSWD',
help='passwd (ignored for Cloud SQL)')
args = parser.parse_args()
if args.remote and ':' in args.remote:
connection = CloudSQLConnectionFactory(args.remote)
connection = LocalSQLConnectionFactory(args.remote, args.user,
if args.resume:
print('Resuming execution at command: %d' % options.resume)
dump_to_cloudsql(args.mysqldump, MySQLConnectionManager(connection),