mirror of
https://github.com/infiniflow/ragflow.git
synced 2025-12-08 20:42:30 +08:00
@ -3,6 +3,7 @@ import time
|
||||
from util import config
|
||||
import pandas as pd
|
||||
|
||||
|
||||
class Postgres(object):
|
||||
def __init__(self, env, dbnm):
|
||||
self.config = config.init(env)
|
||||
@ -13,36 +14,42 @@ class Postgres(object):
|
||||
def __open__(self):
|
||||
import psycopg2
|
||||
try:
|
||||
if self.conn:self.__close__()
|
||||
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('pgdb_usr')} password={self.config.get('pgdb_pwd')} host={self.config.get('pgdb_host')} port={self.config.get('pgdb_port')}")
|
||||
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))
|
||||
|
||||
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))
|
||||
|
||||
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))
|
||||
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:
|
||||
@ -53,11 +60,11 @@ class Postgres(object):
|
||||
cur.close()
|
||||
return updated_rows
|
||||
except Exception as e:
|
||||
logging.error(f"Fail to exec {sql} "+str(e))
|
||||
logging.error(f"Fail to exec {sql} " + str(e))
|
||||
self.__open__()
|
||||
time.sleep(1)
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
Postgres("infiniflow", "docgpt")
|
||||
|
||||
|
||||
Reference in New Issue
Block a user