"""
ETL generic database utility functions.

This is an abstraction for basic database related calls, for any system
compliant with the DB API v2.
"""

from contextlib import contextmanager


@contextmanager
def context(connection, *args, **kwargs):
    """Convenience context manager to provide a cursor and connection.

    This also auto commits in Python runtime. Use the connection object for any
    intermediate commits or rollbacks.

    Usage example:

        with database.context() as (cursor, connection):
            cursor.execute('DROP TABLE super_important_stuffs')
            connection.rollback()

    Args:
        connection (Connection): DB API v2 compatible connection object.
        *args: Additional argument for cursor constructor.
        **kwargs: Additional keyword arguments for cursor constructor.

    Yields:
        tuple(Cursor, Connection): initialized cursor object and connection
            object as per DB API v2.
    """
    try:
        cursor = connection.cursor(*args, **kwargs)
        yield cursor, connection
        connection.commit()
    except:
        connection.rollback()
        raise
    finally:
        cursor.close()
        connection.close()


def execute(connection, sql, params=None, *args, **kwargs):
    """Convenience to use the database contextmanager to run SQL.

    This is for one-off style executions that do not need finer control with
    features like rollbacks or fetching from cursors.

    Args:
        connection (Connection): DB API v2 compatible connection object.
        sql (str): Query to run.
        params (dict): Params matching the sql.
        *args: Additional argument for cursor constructor.
        **kwargs: Additional keyword arguments for cursor constructor.

    Returns:
        mixed: query specific, or None
    """
    with context(connection, *args, **kwargs) as (cursor, connection):
        cursor.execute(sql, params)
        return cursor


def executemany(connection, sql, sequence, *args, **kwargs):
    """Convenience to use the database contextmanager to run SQL.

    This is for one-off style executionmany calls that do not need finer
    control with features like rollbacks or fetching from cursors.

    Args:
        connection (Connection): DB API v2 compatible connection object.
        sql (str): Query to run.
        sequence (iterable): rows to run sql on.
        *args: Additional argument for cursor constructor.
        **kwargs: Additional keyword arguments for cursor constructor.

    Returns:
        mixed: query specific, or None
    """
    with context(connection, *args, **kwargs) as (cursor, connection):
        cursor.executemany(sql, sequence)
        return cursor


def query(connection, sql, params=None, *args, **kwargs):
    """Convenience to execute a query and get a DB cursor for fetching.

    Unlike the execute function, this does not use a context manager that would
    automatically close the cursor.

    Args:
        connection (Connection): DB API v2 compatible connection object.
        sql (str): Query to run.
        params (dict): Params matching the sql.
        *args: Additional argument for cursor constructor.
        **kwargs: Additional keyword arguments for cursor constructor.

    Returns:
        Cursor: database cursor after a query is executed.
    """
    cur = connection.cursor()
    cur.execute(sql, params, *args, **kwargs)
    return cur