"""MySQL logic.""" import glob import pymysql import sqlparse from pymysql import DatabaseError from pymysql.constants import ER from pymysql.cursors import DictCursor from lambdacommon.common_config import logger SELECT_USERS_SQL = ( 'select user, host, authentication_string from mysql.user ' "where user not in ('rdsadmin', 'rdsrepladmin', 'mysql.sys', %s)" ) SHOW_GRANTS_SQL = 'show grants for %s@%s' CREATE_USER_SQL = 'create user %s@%s identified by password %s' DROP_USER_SQL = 'drop user %s@%s' def run_scripts(db_credentials, scripts_dir): """Run a set of SQL scripts from a directory. Args: db_credentials: credentials which identify the database to run the scripts against scripts_dir: the directory containing the scripts Returns: The number of scripts run """ connection = _get_connection(db_credentials) with connection: with connection.cursor() as cursor: scripts = glob.glob(f'{scripts_dir}/*.sql') logger.info(f'Found {len(scripts)} scripts in {scripts_dir}') for script in sorted(scripts): logger.info(f'Running SQL script {script}') _run_script(cursor, script) return len(scripts) def _run_script(cursor, script_path): """Run a single SQL script against a database.""" with open(script_path) as script: script_body = script.read() commands = sqlparse.split(script_body) for command in commands: logger.info(f'Running command {command}') cursor.execute(command) def copy_users(source_db_credentials, target_db_credentials, role_passwords=None): """Copy MySQL users from one database to another. role_passwords is accepted for parity with the PostgreSQL logic's copy_users interface but is unused: MySQL authentication strings are readable and copied directly, so no externally supplied passwords are needed. """ users = _get_user_details(source_db_credentials) logger.info( f"Copying {len(users)} users from {source_db_credentials['host']} " f"to {target_db_credentials['host']}") return _restore_users(target_db_credentials, users) def _get_user_details(db_credentials): """Retrieve details of users and their grants. Returns all users in the database identified by the given credentials, except for the user identified by the credentials and RDS system users. """ connection = _get_connection(db_credentials) with connection: with connection.cursor(DictCursor) as cursor: users = _get_users(cursor, db_credentials['username']) with connection.cursor() as cursor: for user in users: user['grants'] = _get_grants(cursor, user) return users def _restore_users(creds, users): """Restore users and their grants to a database.""" connection = _get_connection(creds) with connection: with connection.cursor(DictCursor) as cursor: existing_users = _get_users(cursor, creds['username']) logger.info( f'Dropping {len(existing_users)} existing users ' f"from {creds['host']}") for existing_user in existing_users: _drop_user(cursor, existing_user) logger.info( f"Successfully dropped existing users from {creds['host']}") logger.info(f"Creating {len(users)} users in {creds['host']}") for user in users: _create_user(cursor, user) logger.info(f"Successfully created users in {creds['host']}") return len(users) def _get_connection(db_credentials): """Return a database connection from the given credentials.""" return pymysql.connect( host=db_credentials['host'], user=db_credentials['username'], password=db_credentials['password'], autocommit=True, ) def _get_users(cursor, user_to_exclude): """Retrieve users from a MySQL database. Returns all users except for RDS system users and the user specified by the user_to_exclude parameter. """ cursor.execute( ( 'select user, host, plugin, authentication_string from mysql.user ' ' where user not in (' " 'rdsadmin'," " 'rdsrepladmin'," " 'rdsrepladmin_priv_checks_user'," " 'rds_superuser_role'," " 'mysql.infoschema'," " 'mysql.session'," " 'mysql.sys'," " 'AWS_BEDROCK_ACCESS'," " 'AWS_COMPREHEND_ACCESS'," " 'AWS_LAMBDA_ACCESS'," " 'AWS_LOAD_S3_ACCESS'," " 'AWS_SAGEMAKER_ACCESS'," " 'AWS_SELECT_S3_ACCESS'," ' %s' ' )' ), user_to_exclude) return cursor.fetchall() def _get_grants(cursor, user): """Retrieve the grants for the given user.""" cursor.execute('show grants for %s@%s', (user['user'], user['host'])) rows = cursor.fetchall() grants = [] for row in rows: grants.append(row[0]) return grants def _drop_user(cursor, user): """Drop the given user from the database.""" logger.debug(f'Dropping user {user["user"]}@{user["host"]}') cursor.execute('drop user %s@%s', (user['user'], user['host'])) def _create_user(cursor, user): """Create a user and their grants.""" logger.debug(f'Creating user {user["user"]}@{user["host"]}') try: cursor.execute( 'create user %s@%s identified with %s as %s', ( user['user'], user['host'], user['plugin'], user['authentication_string'] ) ) except DatabaseError as e: logger.error( f'Error creating user: {user["user"]}@{user["host"]}. ' ) raise e for grant in user['grants']: try: cursor.execute(grant) except DatabaseError as e: code, message = e.args if code in [ ER.NO_SUCH_TABLE, ER.BAD_FIELD_ERROR, ER.SP_DOES_NOT_EXIST ]: logger.info( f'Ignoring create grant failure as object does not exist. ' f'Grant: {grant}' ) else: logger.error( f'Error creating grant: {grant}. Error: {message}' ) raise e