from __future__ import division

# -----------------------------------------------------------------------------
# Copyright (c) 2014--, The American Gut Development Team.
#
# Distributed under the terms of the BSD 3-clause License.
#
# The full license is in the file LICENSE, distributed with this software.
# -----------------------------------------------------------------------------

from contextlib import contextmanager
from psycopg2 import connect, Error as PostgresError
from psycopg2.extras import DictCursor

from amgut.lib.config_manager import AMGUT_CONFIG


class SQLConnectionHandler(object):
    """Encapsulates the DB connection with the Postgres DB"""
    def __init__(self, con=None):
        if not con:
            self._connection = connect(user=AMGUT_CONFIG.user,
                                       password=AMGUT_CONFIG.password,
                                       database=AMGUT_CONFIG.database,
                                       host=AMGUT_CONFIG.host,
                                       port=AMGUT_CONFIG.port)
        else:
            self._connection = con

        self.execute('set search_path to ag, public')

    def __del__(self):
        self._connection.close()

    @contextmanager
    def get_postgres_cursor(self):
        """ Returns a Postgres cursor, commits on close

        Returns
        -------
        pgcursor : psycopg2.cursor
        """
        try:
            with self._connection.cursor(cursor_factory=DictCursor) as cur:
                yield cur
        except:
            self._connection.rollback()
            raise
        else:
            self._connection.commit()

    def _check_sql_args(self, sql_args):
        """ Checks that sql_args have the correct type

        Inputs:
            sql_args: SQL arguments

        Raises a TypeError if sql_args does not have the correct type,
            otherwise it just returns the execution to the caller
        """
        # Check that sql arguments have the correct type
        if sql_args and type(sql_args) not in [tuple, list, dict]:
            raise TypeError("sql_args should be tuple, list or dict. Found %s "
                            % type(sql_args))

    @contextmanager
    def _sql_executor(self, sql, sql_args=None, many=False):
        """Executes an SQL query

        Parameters
        ----------
        sql: str
            The SQL query
        sql_args: tuple or list, optional
            The arguments for the SQL query
        many: bool, optional
            If true, performs an execute many call

        Returns
        -------
        pgcursor : psycopg2.cursor
            The cursor in which the SQL query was executed

        Raises
        ------
        ValueError
            If there is some error executing the SQL query
        """
        # Check that sql arguments have the correct type
        if many:
            for args in sql_args:
                self._check_sql_args(args)
        else:
            self._check_sql_args(sql_args)

        # Execute the query
        with self.get_postgres_cursor() as cur:
            try:
                if many:
                    cur.executemany(sql, sql_args)
                else:
                    cur.execute(sql, sql_args)
                yield cur
                self._connection.commit()
            except PostgresError as e:
                self._connection.rollback()
                try:
                    err_sql = cur.mogrify(sql, sql_args)
                except ValueError:
                    err_sql = cur.mogrify(sql, sql_args[0])
                raise ValueError(("\nError running SQL query: %s"
                                  "\nError: %s" % (err_sql, e)))

    def execute_fetchall(self, sql, sql_args=None):
        """ Executes a fetchall SQL query

        Parameters
        ----------
        sql: str
            The SQL query
        sql_args: tuple or list, optional
            The arguments for the SQL query

        Returns
        ------
        list of tuples
            The results of the fetchall query

        Note: from psycopg2 documentation, only variable values should be bound
            via sql_args, it shouldn't be used to set table or field names. For
            those elements, ordinary string formatting should be used before
            running execute.
        """
        with self._sql_executor(sql, sql_args) as pgcursor:
            result = pgcursor.fetchall()
        return result

    def execute_fetchone(self, sql, sql_args=None):
        """ Executes a fetchone SQL query

        Parameters
        ----------
        sql: str
            The SQL query
        sql_args: tuple or list, optional
            The arguments for the SQL query

        Returns
        -------
        Tuple
            The results of the fetchone query

        Notes
        -----
        from psycopg2 documentation, only variable values should be bound
            via sql_args, it shouldn't be used to set table or field names. For
            those elements, ordinary string formatting should be used before
            running execute.
        """
        with self._sql_executor(sql, sql_args) as pgcursor:
            result = pgcursor.fetchone()
        return result

    def execute(self, sql, sql_args=None):
        """ Executes an SQL query with no results

        Parameters
        ----------
        sql: str
            The SQL query
        sql_args: tuple or list, optional
            The arguments for the SQL query

        Notes
        -----
        from psycopg2 documentation, only variable values should be bound
            via sql_args, it shouldn't be used to set table or field names. For
            those elements, ordinary string formatting should be used before
            running execute.
        """
        with self._sql_executor(sql, sql_args):
            pass

    def executemany(self, sql, sql_args_list):
        """ Executes an executemany SQL query with no results

        Parameters
        ----------
        sql: str
            The SQL query
        sql_args: list of tuples
            The arguments for the SQL query

        Note: from psycopg2 documentation, only variable values should be bound
            via sql_args, it shouldn't be used to set table or field names. For
            those elements, ordinary string formatting should be used before
            running execute.
        """
        with self._sql_executor(sql, sql_args_list, True):
            pass

    def execute_proc_return_cursor(self, procname, proc_args):
        """Executes a stored procedure and returns a cursor

        Parameters
        ----------
        procname: str
            the name of the stored procedure
        proc_args: list
            arguments sent to the stored procedure
        """
        proc_args.append('cur2')
        cur = self._connection.cursor()
        cur.callproc(procname, proc_args)
        cur.close()
        return self._connection.cursor('cur2')

@@ -513,16 +513,54 @@
 ctCursor
+%0Afrom psycopg2 import OperationalError
 %0A%0Afrom a
@@ -757,16 +757,288 @@
 ot con:%0A
+            self._open_connection()%0A        else:%0A            self._connection = con%0A%0A        self.execute('set search_path to ag, public')%0A%0A    def __del__(self):%0A        self._connection.close()%0A        del self._connection%0A%0A    def _open_connection(self):%0A        try:%0A
         
@@ -1368,54 +1368,8 @@
     
-    else:%0A            self._connection = con%0A%0A
     
@@ -1422,63 +1422,188 @@
 c')%0A
-%0A    def __del__(self):%0A        self._connection.close(
+        except Exception as e:%0A            # catch any error from the database and raise for site to catch%0A            raise RuntimeError(%22Cannot connect to database!%5Cn%25s%22 %25 str(e)
 )%0A%0A 
@@ -1784,32 +1784,105 @@
 sor%0A        %22%22%22%0A
+        if self._connection.closed:%0A            self._open_connection()%0A%0A
         try:%0A   
@@ -2103,24 +2103,94 @@
 n.commit()%0A%0A
+    def __del__(self):%0A        self.con.close()%0A        del self.con%0A%0A
     def _che

add reconnection capability