mysql2csv

#!/usr/bin/env python
# mysql2csv (part of ossobv/vcutil) // wdoekes/2011-2016 // Public Domain
#
# This mysql2csv dump script is a convenient way to dump and/or prune
# large tables. Instead of selecting all data at once, it takes chunks
# of data based on a datetime column supplied by the caller.
#
# For example, if you run this:
#
#     mysql2csv ... mydatabase large_table created
#
# Then the dump script will check the MIN(created) on large_table, and
# from that point, start BETWEEN queries between that time and the next
# hour.
#
#     SELECT * FROM mydatabase.large_table
#     WHERE created >= '2011-01-21 03:00:00' AND
#           created  < '2011-01-21 04:00:00'
#     ORDER BY created
#
# If you have proper indexes on the created column, these
# "hour-selecting" queries should be light enough to handle without
# locking up your production database.
#
# See mysql2csv -h for more info.
#
# Created by Walter Doekes, OSSO B.V.
# - 2011-01-21: Initial from http://wjd.nu/files/2011/01/mysql2csv.py
# - 2016-05-23: PEP cleanup, add into vcutil repository.
# - 2016-05-23: Only write CSVs if there is data for that period. End
#               before until_date if the max date is lower.
# - 2016-05-23: Use argparse for the options. Read ini file as well.
# - 2016-05-23: Add --quiet option to silence the output.
# - 2016-05-23: Append to CSV instead clobbering it. Safer for data.
# - 2016-05-23: Allow sleep-time to be set.
# - 2016-05-23: Prune during dump, using the same date filters.
# - 2016-05-23: Fix problem with TIMESTAMP column and DST changes.
# - 2016-05-25: Add --extra-where parameter to limit the query results.
#
from __future__ import print_function
import argparse  # pip install argparse on old pre-2.7 python
import getpass
import sys
import time

from MySQLdb import connect
from datetime import date, datetime, timedelta
from decimal import Decimal

try:
    from configparser import RawConfigParser
except ImportError:  # python2-
    from ConfigParser import RawConfigParser


DEFAULT_FILENAME = '%(database)s.%(table)s.%%Y%%m.csv'
DEFAULT_UNTIL_DATE = datetime(2035, 1, 1)
DEFAULT_SLEEP_TIME = 0.005  # sleep N secs every iteration, easing server load


def dump2csv_and_prune(db_connection, kwargs):
    # Get relevant args.
    table_name, date_column, until, template, quiet, sleep, prune = tuple(
        kwargs[i] for i in (
            'table', 'date_column', 'until', 'template', 'quiet',
            'sleep', 'prune'))

    extra_where = (kwargs.get('extra_where') or '').strip()
    if extra_where:
        extra_where = 'AND (%s)' % (extra_where.replace('%', '%%'),)

    rows = describe_table(db_connection, table_name)
    begin_date = min_column_value(
        db_connection, table_name, date_column, extra_where)
    if not begin_date or begin_date > until:
        return
    begin_date = begin_date.replace(minute=0, second=0)
    last_date = max_column_value(
        db_connection, table_name, date_column, until, extra_where)
    last_date = (last_date.replace(minute=0, second=0) + timedelta(hours=1))

    # Create cursor. Increment in 1 hour steps.
    cursor = db_connection.cursor()
    period = timedelta(seconds=3600)
    end_date = begin_date + period
    dest_file = None
    dest_file_name = ''

    rotate = 0
    while begin_date < last_date:
        # Check if we need to open a new file.
        if dest_file and begin_date.strftime(template) != dest_file_name:
            if dest_file:
                dest_file.close()
                dest_file = None

        # For certain columns, string/time comparisons can fail during
        # DST changes, for example:
        #   TABLE auth (ts timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP)
        #   mysql> SELECT MIN(ts) FROM auth\G
        #   min(ts): 2014-10-26 02:00:02
        #   mysql> SELECT COUNT(*) FROM auth WHERE '2014-10-26 02:00:00' <= ts
        #            AND ts < '2014-10-26 03:00:00' ORDER BY ts\G
        #   COUNT(*): 0
        #   mysql> SELECT COUNT(*) FROM auth WHERE '2014-10-26 01:59:59' <= ts
        #            AND ts < '2014-10-26 03:00:00' ORDER BY ts\G
        #   COUNT(*): 256
        # For DATETIME columns, the problem does not occur.
        # As a workaround we'll have to do an extra check to see if we've just
        # crossed a DST line and add more time to the query.
        extra_seconds = has_crossed_dst(db_connection, begin_date)
        begin_date_hack = (begin_date - timedelta(seconds=extra_seconds))

        # Query a bit of data and write to file.
        cursor.execute(
            "SELECT * FROM %s WHERE '%s' <= %s AND %s < '%s' %s "
            "ORDER BY %s" % (
                table_name,
                begin_date_hack.strftime('%Y-%m-%d %H:%M:%S'), date_column,
                date_column, end_date.strftime('%Y-%m-%d %H:%M:%S'),
                extra_where, date_column))

        row = None
        for row in cursor:
            if dest_file is None:
                dest_file_name = begin_date.strftime(template)
                dest_file = open(dest_file_name, 'a')
                dest_file.write('%s\n' % (','.join(rows),))

                if not quiet:
                    sys.stdout.write('+')
                    sys.stdout.flush()

            dest_file.write('%s\n' % (','.join(to_string(i) for i in row),))

        # Prune immediately in equal sized chunks.
        if prune and row is not None:
            cursor.execute(
                "DELETE FROM %s WHERE '%s' <= %s AND %s < '%s' %s" % (
                    table_name,
                    begin_date_hack.strftime('%Y-%m-%d %H:%M:%S'), date_column,
                    date_column, end_date.strftime('%Y-%m-%d %H:%M:%S'),
                    extra_where))
            db_connection.commit()
            if not quiet:
                sys.stdout.write('!')
                sys.stdout.flush()

        # Increase dates.
        begin_date, end_date = end_date, end_date + period
        if not quiet:
            sys.stdout.write('%s\b' % (r'\|/-'[rotate],))
            sys.stdout.flush()
            rotate = (rotate + 1) % 4
        if sleep:
            time.sleep(sleep)

    if dest_file:
        dest_file.close()
        dest_file = None


def describe_table(db_connection, table_name):
    cursor = db_connection.cursor()
    cursor.execute('DESCRIBE %s;' % (table_name,))
    ret = []
    for row in cursor.fetchall():
        ret.append(row[0])
    return ret


def min_column_value(db_connection, table_name, column, extra_where):
    cursor = db_connection.cursor()
    # We use 2 jan because 1 jan fails if the timezone is positive.
    cursor.execute(
        "SELECT MIN(%s) FROM %s WHERE %s >= '1970-01-02' %s" % (
            column, table_name, column, extra_where))
    ret = cursor.fetchall()
    return ret[0][0]


def max_column_value(db_connection, table_name, column, max_date,
                     extra_where):
    cursor = db_connection.cursor()
    cursor.execute(
        "SELECT MAX(%s) FROM %s WHERE %s < '%s' %s" % (
            column, table_name, column, max_date.strftime('%Y-%m-%d'),
            extra_where))
    ret = cursor.fetchall()
    return ret[0][0]


def has_crossed_dst(db_connection, begin_date):
    cursor = db_connection.cursor()
    cursor.execute(
        "SELECT (UNIX_TIMESTAMP('%s') - UNIX_TIMESTAMP('%s')) AS v" % (
            begin_date.strftime('%Y-%m-%d %H:%M:%S'),
            (begin_date - timedelta(seconds=1)).strftime('%Y-%m-%d %H:%M:%S')))
    ret = cursor.fetchall()
    if ret[0][0] == 1:
        return 0
    return 1  # ret[0][0] == 3601?


def to_string(data):
    if isinstance(data, str):
        return '"%s"' % (data.replace('"', '""'),)
    if isinstance(data, unicode):
        return '"%s"' % (data.encode('utf-8').replace('"', '""'),)
    if isinstance(data, int) or isinstance(data, long):
        return str(data)
    if isinstance(data, datetime):
        return '"%s"' % (data.strftime('%Y-%m-%d %H:%M:%S'),)
    if isinstance(data, date):  # must come after datetime
        return '"%s"' % (data.strftime('%Y-%m-%d'),)
    if isinstance(data, Decimal):
        return str(data)
    if data is None:
        return 'NULL'
    raise TypeError('Unexpected type %s for data %r: %s' % (
        type(data), data, data))


def parse_defaults_file(filename):
    config = RawConfigParser()
    with open(filename, 'r') as fp:
        config.readfp(fp)

    sections = config.sections()
    for possible_section in ('mysql2csv', 'client'):
        if possible_section in sections:
            break
    else:
        return {}

    return dict(config.items(possible_section))


def parse_date(date_):
    return datetime.fromtimestamp(
        time.mktime(time.strptime(date_, '%Y-%m-%d')))


if __name__ == '__main__':
    example_filename = DEFAULT_FILENAME % {
        'database': '<DATABASE>', 'table': '<TABLE>'}
    example_until = DEFAULT_UNTIL_DATE.strftime('%Y-%m-%d')

    parser = argparse.ArgumentParser(
        prog='mysql2csv',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        description=('''\
Dumps all data from the MySQL table to a CSV-formatted file on stdout IN A NON-
LOCKING FASHION. In most cases you can safely use this on a production database
to dump old records that you will be pruning later on.

Supply the necessary HOST/USER/DATABASE to connect to, the password is prompted
on stdin if required. The data from TABLE will be exported ordered by the date
from DATE_COLUMN in ascending order.

If strftime variables are used are used in the CSV filename template, they are
replaced with the values from the fetched period. The default filename TEMPLATE
is: '%(filename)s'

The default UNTIL time is %(until)s which should be far enough in the future
to get "all" records. (It is used as "exclusive" date, meaning that no records
from that date and onwards will be output/pruned.)

You can use the date(1) shell utility to specify a relative date for --until;
for example one year ago would be: date +%%Y-%%m-%%d -d '-1 year'

Supply --prune as argument to immediately prune the data after dumping it.

''' % {'filename': example_filename, 'until': example_until}))

    parser.add_argument(
        '-H', '--host',
        help='which database host to connect to')
    parser.add_argument(
        '-p', '--port', type=int,
        help='the database port if not 3306')
    parser.add_argument(
        '-u', '--user',
        help='the username if not the current user')
    parser.add_argument(
        '-P', action='store_true',
        help='ask for password')
    parser.add_argument(
        '-s', '--socket',
        help='unix socket to connect to instead of tcp host')
    parser.add_argument(
        '-t', '--template',
        help="the filename template to use instead of '%s'" % (
            example_filename.replace('%', '%%'),))
    parser.add_argument(  # not through defaults-file
        '--until', type=parse_date, default=DEFAULT_UNTIL_DATE,
        help='stop when arriving at this date')
    parser.add_argument(
        '-q', '--quiet', action='store_true',
        help='be quiet, apart from showing errors')
    parser.add_argument(  # not through defaults-file
        '--prune', action='store_true', default=False,
        help='prunes the data after writing the csv (DESTRUCTIVE!)')
    parser.add_argument(
        '--sleep', type=float, metavar='SECONDS',
        help=('sleep this long between every data fetch/prune; accepts '
              'fractional seconds'))
    parser.add_argument(
        '--extra-where',
        help=('additional filter to add to the SELECT and DELETE queries; '
              'for example: is_active = 0'))
    parser.add_argument(
        '-f', '--defaults-file',
        help=("a defaults-file in mysql-client fashion to set the username, "
              "password or any of the other long options; checks the "
              "[mysql2csv] section or [client] if the former doesn't exist"))
    parser.add_argument(
        'database',
        help='the database to connect to')
    parser.add_argument(
        'table',
        help='the table to dump and/or prune')
    parser.add_argument(
        'date_column',
        help='the column that holds the indexed date')

    # Extract args and merge with defaults. If a defaults-file is set,
    # use that and have the command line overrule it.
    args = parser.parse_args()
    # Copy defaults-file contents.
    kwargs = {'password': None}
    if args.defaults_file:
        kwargs.update(parse_defaults_file(args.defaults_file))
        if 'port' in kwargs:
            kwargs['port'] = int(kwargs['port'])
        if 'sleep' in kwargs:
            kwargs['sleep'] = float(kwargs['sleep'])
    # Copy command line args.
    for arg in ('host', 'port', 'user', 'socket', 'template',
                'database', 'table', 'date_column', 'until',
                'quiet', 'sleep', 'extra_where'):
        value = getattr(args, arg)
        if value is not None or arg not in kwargs:
            kwargs[arg] = value
    kwargs['prune'] = args.prune
    # Set defaults.
    if kwargs['template'] is None:
        kwargs['template'] = DEFAULT_FILENAME % {
            'database': kwargs['database'],
            'table': kwargs['table']}
    if kwargs['sleep'] is None:
        kwargs['sleep'] = DEFAULT_SLEEP_TIME

    # Get password?
    if args.P:
        kwargs['password'] = getpass.getpass('Enter password: ')

    # Extract mysql connect args.
    my_kwargs = {}
    for from_, to in (
            ('host', 'host'), ('port', 'port'), ('user', 'user'),
            ('password', 'passwd'), ('database', 'db'),
            ('socket', 'unix_socket')):
        if kwargs[from_] is not None:
            my_kwargs[to] = kwargs[from_]

    # Print a bit of info.
    if not kwargs['quiet']:
        print('MySQL server settings: %s@%s/%s' % (
            kwargs['user'] or '<nobody>',
            kwargs['socket'] or kwargs['host'] or '<local>',
            kwargs['database']))
        print('Table to dump: %s (ordered by %s, from begin to %s)' % (
            kwargs['table'], kwargs['date_column'],
            kwargs['until'].strftime('%Y-%m-%d')))
        print('Filename template: "%s"' % (kwargs['template'],))
        print('Processing %s.%s (ordered by %s)' % (
            my_kwargs['db'], kwargs['table'], kwargs['date_column']))

    # Connect and do work.
    conn = connect(**my_kwargs)
    dump2csv_and_prune(conn, kwargs)
    conn.close()

    if not kwargs['quiet']:
        print()

# vim: set ts=8 sw=4 sts=4 et ai: