"""Helper functions for sql2sf workflow."""

import pymysql


# TODO: Make it more generic and move it to contrib-mysql
def execute_with_mysql(
        sql, command, host, port, user, password, cursor_type=None):
    """Execute SQL query and return results.

    Args:
        sql (str): The SQL query to execute.
        command (str): 'fetchall' or 'fetchone'.
        host (str): The database host to connect.
        port (int): The database port to connect.
        user (str): The database username to connect.
        password (str): The database password to connect.
        cursor_type (str): 'DictCursor' if we want to execute with DictCursor.

    Returns:
        dict or tuple: Dict if DictCursor was used, tuple if regular one.
    """
    conn = pymysql.connect(
        host=host,
        port=port,
        user=user,
        password=password,
        charset='utf8mb4')
    if cursor_type == 'DictCursor':
        cursor = conn.cursor(pymysql.cursors.DictCursor)
    else:
        cursor = conn.cursor()

    cursor.execute(sql)

    if command == 'fetchall':
        results = cursor.fetchall()
    elif command == 'fetchone':
        results = cursor.fetchone()
    else:
        raise ValueError(
            'Please specify command "fetchall" or "fetchone" for '
            'the execute_with_mysql function')

    cursor.close()
    conn.close()
    return results