import psycopg2
from psycopg2.extras import DictCursor
import sys
import traceback

class Database(object):
    def __init__(self, **kwargs):
        self.args = kwargs
        self.conn = psycopg2.connect(keepalives_idle=60, **kwargs)
        self.cursor = self.conn.cursor(cursor_factory=DictCursor)
        self.rowcount = 0

    def fetchall(self):
        return self.cursor.fetchall()

    def fetchone(self):
        return self.cursor.fetchone()

    def execute(self, string, *args, commit=True, retry=True):
        if args and type(args[0]) != tuple:
            args = list(args)
            args[0] = [args[0]]
        try:
            self.cursor.execute(string, *tuple(args))
        except psycopg2.IntegrityError:
            self.rowcount = 0
            self.conn.rollback()
        except (psycopg2.OperationalError, psycopg2.InterfaceError):
            if retry:
                self.__init__(**self.args)
                return self.execute(string, *args, commit, retry=False)
        except:
            traceback.print_exc()
            self.rowcount = self.cursor.rowcount
            self.conn.rollback()
        else:
            self.rowcount = self.cursor.rowcount
            if commit:
                self.conn.commit()
        return self

    def commit(self):
        self.conn.commit()

    def close(self):
        self.conn.close()