General MySQL database utility functions.

mysql_db_utils.py

# coding: utf-8

from sh import cat, mysql, mysqladmin, ErrorReturnCode as ShellReturnCodeError
from pipes import quote as shell_quote
import re

def _drop_create_common_args(database_name, action='create', user=None, passwd=None, host='localhost', verbose=False):
    args = []

    if user:
        args.append('-u%s' % (shell_quote(user)))
    if passwd:
        args.append('-p%s' % (shell_quote(passwd)))
    if verbose:
        args.append('-v')

    args.extend([
        '-h%s' % (shell_quote(host)),
        '-f',
        action,
        database_name,
    ])

    return args

def drop_database(database_name, user=None, passwd=None, host='localhost', verbose=False):
    args = _drop_create_common_args(database_name, action='drop', user=user, passwd=passwd, host=host, verbose=verbose)

    try:
        ret = mysqladmin(args)
    except ShellReturnCodeError as e:
        if 'database doesn\'t exist' in e.stderr:
            return True

        raise e

    return ret


def create_database(database_name, user=None, passwd=None, host='localhost', verbose=False):
    args = _drop_create_common_args(database_name, user=user, passwd=passwd, host=host, verbose=verbose)

    return mysqladmin(args)


def import_from_file(sql_file, username, db_name, host='localhost', password=None, verbose=False):
    command_line = ['-u', username, '-h', host, db_name]

    if password:
        command_line.append('-p%s' % (password))

    if verbose:
        # Don't expose password
        args = re.sub('\-p[^\s]+', '-pxxxx', ' '.join(command_line))
        print('mysql %s < %s' % (args, sql_file))

    return mysql(cat(sql_file), command_line)


def user_exists(connection, user_name, host='localhost'):
    c = connection.cursor()

    c.execute('SELECT EXISTS(SELECT 1 FROM mysql.user WHERE User = %s AND Host = %s)', args=(user_name, host,))

    return c.fetchone()[0] == 1


def create_user_for_db(connection, user_name, database, password=None, host='localhost'):
    if user_name == 'root':
        return

    c = connection.cursor()
    args = [user_name, host]
    sql = 'CREATE USER %s@%s'

    if mysql_user_exists(connection, user_name, host=host):
        c.execute('DROP USER %s@%s', args=(user_name, host,))

    if password:
        sql += ' IDENTIFIED BY %s'
        args.append(password)

    c.execute(sql, args=args)

    if not mysql_user_exists(connection, user_name, host=host):
        raise Exception('Failed to create user %s@%s' % (user_name, host))

    # Can't use built-in escaper for this
    #c.execute('GRANT ALL PRIVILEGES ON %s.* TO %s@%s WITH GRANT OPTION', args=(database, user_name, host,))
    c.execute('GRANT ALL PRIVILEGES ON %s.* TO \'%s\'@\'%s\'' % (database, user_name, host,))