"""Util module for using database.

It helps with loading SQL queries/templates from files and executing them.
Glossary:
SQL query         - plain SQL code.
SQL template      - SQL code that contains format placeholders such as
                    {keyword} or %(keyword)s or $s.
Parameters        - keyword arguments used for formatting {keyword}
                    placeholders. Use for table names, schemas, etc.
Escape parameters - keyword arguments used for filling sql while execution with
                    DB API v2. It's save way for parametrise queries with
                    'external' data such as dates passed from the context.
USAGE EXAMPLE:
count_for_date_range.sql:
    /* Count rows in date range.
    Parameters:
        table_name (str): name of the table.
    Escaped parameters:
        start_date (date): start date.
        end_date (date): end date.
    */
    SELECT count(*)
    FROM {table_name}
    WHERE date BETWEEN %(start_date)s AND %(end_date)s;

config.py:
    import os

    sql = SQLLoader(os.path.realpath(os.path.dirname(__file__)) + '/queries')
    redshift = {
        'driver': Driver.REDSHIFT,
        'host': '***',
        'port': '***',
        'user': '***',
        'password': '***',
        'db': '***'}

tasks.py:
    import config
    from config import sql

    def count_for_date_range(table_name, start_date, end_date):
        sql['count_for_date_range'].fill(
            table_name=table_name
        ).fill_escaped(
            start_date=start_date,
            end_date=end_date
        ).execute(
            config.redshift)

Author: Daniil Omelchenko (domelchenko@theorchard.com).
"""

from contextlib import contextmanager

import psycopg2
import pymysql


class Query:
    """Construct and execute SQL queries."""

    def __init__(self, sql_template):
        """Create query object from the sql_template.

        Args:
            sql_template (str): SQL template - SQL code with
            {params} and %(escape_params)s.
        """
        self.sql = sql_template
        self.escape_parameters = {}

    def execute(self, db_config):
        """Execute query with the following config and parameters.

        Args:
            db_config (dict): Database credentials for DB API V2 connection.
        """
        with db_context(connect(db_config)) as (cursor, connection):
            cursor.execute(self.sql, self.escape_parameters)

    def fetchall(self, db_config):
        """Execute query and fetch all results.

        Args:
            db_config (dict): Database credentials for DB API V2 connection.
        Returns:
            list(tuple): List of rows with query result.
        """
        with db_context(connect(db_config)) as (cursor, connection):
            cursor.execute(self.sql, self.escape_parameters)
            return cursor.fetchall()

    def fetchone(self, db_config):
        """Execute query and fetch first row of a query result set.

        Args:
            db_config (dict): Database credentials for DB API V2 connection.
        Returns:
            tuple: Row with query result.
        """
        with db_context(connect(db_config)) as (cursor, connection):
            cursor.execute(self.sql, self.escape_parameters)
            return cursor.fetchone()

    def fetchmany(self, db_config, size):
        """Generator fetches set of n rows of a query result (n=size).

        Stops when no more results available.

        Args:
            db_config (dict): Database credentials for DB API V2 connection.
            size (int): Number of fetching rows.
        Yields:
            list(tuple): List of maximum n rows with query result (n = size).
        """
        with db_context(connect(db_config)) as (cursor, connection):
            cursor.execute(self.sql, self.escape_parameters)
            while True:
                batch = cursor.fetchmany(size)
                if not batch:
                    break
                yield batch

    def fill(self, **parameters):
        """Format SQL template with parameters.

        Args:
            parameters (**): Keyword parameters for SQL.
        """
        self.sql = self.sql.format(**parameters)
        return self

    def fill_escaped(self, *escape_parameters, **keyword_escape_parameters):
        """Prepare escaped before query execution.

        Args:
            escape_parameters (*): escape parameters for SQL.
            keyword_escape_parameters (**): Keyword escape parameters for SQL.
        Raises:
            ValueError: argument formats can't be mixed.
            Only one of escape_parameters and keyword_escape_parameters
            could be provided.
        """
        if escape_parameters and keyword_escape_parameters:
            raise ValueError("Argument formats can't be mixed.")
        self.escape_parameters = escape_parameters or keyword_escape_parameters
        return self


class SQLLoader:
    """Loads SQL queries/templates from files."""

    def __init__(self, sql_files_root):
        """Constructor for SQL loader.

        >>> # placed in config file:
        >>> import os
        >>> sql = SQLLoader(
        >>>     os.path.realpath(os.path.dirname(__file__)) + '/queries')
        Args:
             sql_files_root (str): absolute path to folder with queries.
        """
        self.sql_files_root = sql_files_root
        self.query_cash = {}

    def _load_query(self, query_name):
        """Load query from the disk.

        Args:
            query_name (str): name of the query (filename without extension).
        """
        path = '{root}/{query_name}.sql'.format(
            root=self.sql_files_root,
            query_name=query_name)
        with open(path, 'r') as q:
            query = q.read()
        return query

    def __getitem__(self, item):
        """Access queries by name.

        This method uses cache to reduce I/O operations.
        >>> sql = SQLLoader('query_path')
        >>> sql['query_name']
        Args:
            item (str): name of the query (filename without extension).
        Returns:
            Query: Query object initialized with loaded SQL query/template
        """
        if item not in self.query_cash:
            self.query_cash[item] = self._load_query(item)
        return Query(self.query_cash[item])


@contextmanager
def db_context(connection):
    """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 db_context() as (cursor, connection):
            cursor.execute('DROP TABLE super_important_stuffs')
            connection.rollback()
    Args:
        connection (Connection): DB API v2 compatible connection object.
    Yields:
        tuple(Cursor, Connection): initialized cursor object and connection
            object as per DB API v2.
    """
    try:
        cursor = connection.cursor()
        yield cursor, connection
        connection.commit()
    except:
        connection.rollback()
        raise
    finally:
        cursor.close()
        connection.close()


class Driver:
    """Enumeration of available SQL drivers."""

    MYSQL = 'mysql'
    REDSHIFT = 'redshift'  # postgres-like db.


class UnsupportedDriverException(Exception):
    """Exception rises for usage of unsupported DB driver."""

    pass


def connect(db_config):
    """Create connection to the DB by config parameters.

    Args:
        db_config (dict): Database credentials for DB API V2 connection.
    Returns:
        Connection: Connection that supports DB API v2 interface.
    Raises:
        UnsupportedDriverException: In case unsupported driver provided.
    """
    current_driver = db_config['driver']
    if current_driver == Driver.REDSHIFT:
        return psycopg2.connect(
            host=db_config['host'],
            port=db_config['port'],
            user=db_config['user'],
            password=db_config['password'],
            database=db_config['db'])
    if current_driver == Driver.MYSQL:
        return pymysql.connect(
            host=db_config['host'],
            user=db_config['user'],
            password=db_config['password'],
            db=db_config['db'])
    raise UnsupportedDriverException()


def escape_quotes(string):
    """Replace single quote characters with its escaped variant.

    Args:
        string (str): string to be escaped.
    Returns:
        str: string with escaped single quotes.
    """
    string = string.replace("'", "\\'")
    return string