import psycopg2
import psycopg2.extensions
import psycopg2.extras


class DatabaseError(Exception):
    def __init__(self, msg, query=None):
        self.query = query
        Exception.__init__(self, msg)


class BadConnectionError(Exception):
    pass


class Row(dict):
    def __getattr__(self, key):
        return self[key]


class Database:
    def __init__(self):
        # pool of database connections
        # indexed by conninfo, items are lists of connections
        self.pool = {}
        # number of unused connections per conninfo to keep open
        self.pool_keep_open = 1


    def __del__(self):
        for conninfo in self.pool.keys():
            for conn in self.pool[conninfo]:
                conn.close()


    def connect(self, conninfo):
        try:
            conn = psycopg2.connect(conninfo, async=1)
            psycopg2.extras.wait_select(conn)
        except psycopg2.DatabaseError as e:
            raise DatabaseError(str(e))
        return conn


    def get_conn(self, conninfo):
        if not conninfo in self.pool:
            self.pool[conninfo] = []
            return self.connect(conninfo)
        else:
            conn = None
            while len(self.pool[conninfo]) and conn is None:
                conn = self.pool[conninfo].pop()
                if conn.closed:
                    conn = None
            if conn is None:
                return self.connect(conninfo)
        return conn


    def put_conn(self, conninfo, conn):
        if len(self.pool[conninfo]) >= self.pool_keep_open:
            conn.close()
        else:
            self.pool[conninfo].append(conn)


    def execute(self, q, args=[]):
        conn = self.get_conn()
        try:
            curs = conn.cursor()
            curs.execute(q, args)
            psycopg2.extras.wait_select(curs.connection)
#            conn.commit()
        except psycopg2.OperationalError as e:
            # disconnected?
#            conn.rollback()
            conn.close()
            raise BadConnectionError(str(e))
        except psycopg2.DatabaseError as e:
#            conn.rollback()
            raise DatabaseError(str(e), curs.query)
        return curs


    def finish(self, curs):
        self.put_conn(curs.connection)


    def row(self, curs, row):
        return Row(zip([x[0] for x in curs.description], row))


    def fetchone(self, curs):
        return self.row(curs, curs.fetchone())


    def fetchall(self, curs):
        rows = curs.fetchall()
        return [self.row(curs, row) for row in rows]