import logging
import time
from util import config
import pandas as pd


class Postgres(object):
    def __init__(self, env, dbnm):
        self.config = config.init(env)
        self.conn = None
        self.dbnm = dbnm
        self.__open__()

    def __open__(self):
        import psycopg2
        try:
            if self.conn:
                self.__close__()
            del self.conn
        except Exception as e:
            pass

        try:
            self.conn = psycopg2.connect(f"""dbname={self.dbnm}
                                         user={self.config.get('postgres_user')}
                                         password={self.config.get('postgres_password')}
                                         host={self.config.get('postgres_host')}
                                         port={self.config.get('postgres_port')}""")
        except Exception as e:
            logging.error(
                "Fail to connect %s " %
                self.config.get("pgdb_host") + str(e))

    def __close__(self):
        try:
            self.conn.close()
        except Exception as e:
            logging.error(
                "Fail to close %s " %
                self.config.get("pgdb_host") + str(e))

    def select(self, sql):
        for _ in range(10):
            try:
                return pd.read_sql(sql, self.conn)
            except Exception as e:
                logging.error(f"Fail to exec {sql}  " + str(e))
                self.__open__()
                time.sleep(1)

        return pd.DataFrame()

    def update(self, sql):
        for _ in range(10):
            try:
                cur = self.conn.cursor()
                cur.execute(sql)
                updated_rows = cur.rowcount
                self.conn.commit()
                cur.close()
                return updated_rows
            except Exception as e:
                logging.error(f"Fail to exec {sql}  " + str(e))
                self.__open__()
                time.sleep(1)
        return 0


if __name__ == "__main__":
    Postgres("infiniflow", "docgpt")