Skip to content

Commit 6ae900d

Browse files
ulleoulleo
committed
fix: check connection in assistant
1 parent b8b0825 commit 6ae900d

1 file changed

Lines changed: 56 additions & 44 deletions

File tree

‎backend/apps/db/db.py‎

Lines changed: 56 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -102,10 +102,62 @@ def get_session(ds: CoreDatasource | AssistantOutDsSchema):
102102
return session
103103

104104

105-
def check_connection(trans: Optional[Trans], ds: CoreDatasource, is_raise: bool = False):
106-
db = DB.get_db(ds.type)
107-
if db.connect_type == ConnectType.sqlalchemy:
108-
conn = get_engine(ds, 10)
105+
def check_connection(trans: Optional[Trans], ds: CoreDatasource | AssistantOutDsSchema, is_raise: bool = False):
106+
if isinstance(ds, CoreDatasource):
107+
db = DB.get_db(ds.type)
108+
if db.connect_type == ConnectType.sqlalchemy:
109+
conn = get_engine(ds, 10)
110+
try:
111+
with conn.connect() as connection:
112+
SQLBotLogUtil.info("success")
113+
return True
114+
except Exception as e:
115+
SQLBotLogUtil.error(f"Datasource {ds.id} connection failed: {e}")
116+
if is_raise:
117+
raise HTTPException(status_code=500, detail=trans('i18n_ds_invalid') + f': {e.args}')
118+
return False
119+
else:
120+
conf = DatasourceConf(**json.loads(aes_decrypt(ds.configuration)))
121+
if ds.type == 'dm':
122+
with dmPython.connect(user=conf.username, password=conf.password, server=conf.host,
123+
port=conf.port) as conn, conn.cursor() as cursor:
124+
try:
125+
cursor.execute('select 1', timeout=10).fetchall()
126+
SQLBotLogUtil.info("success")
127+
return True
128+
except Exception as e:
129+
SQLBotLogUtil.error(f"Datasource {ds.id} connection failed: {e}")
130+
if is_raise:
131+
raise HTTPException(status_code=500, detail=trans('i18n_ds_invalid') + f': {e.args}')
132+
return False
133+
elif ds.type == 'doris':
134+
with pymysql.connect(user=conf.username, passwd=conf.password, host=conf.host,
135+
port=conf.port, db=conf.database, connect_timeout=10,
136+
read_timeout=10) as conn, conn.cursor() as cursor:
137+
try:
138+
cursor.execute('select 1')
139+
SQLBotLogUtil.info("success")
140+
return True
141+
except Exception as e:
142+
SQLBotLogUtil.error(f"Datasource {ds.id} connection failed: {e}")
143+
if is_raise:
144+
raise HTTPException(status_code=500, detail=trans('i18n_ds_invalid') + f': {e.args}')
145+
return False
146+
elif ds.type == 'redshift':
147+
with redshift_connector.connect(host=conf.host, port=conf.port, database=conf.database, user=conf.username,
148+
password=conf.password,
149+
timeout=10) as conn, conn.cursor() as cursor:
150+
try:
151+
cursor.execute('select 1')
152+
SQLBotLogUtil.info("success")
153+
return True
154+
except Exception as e:
155+
SQLBotLogUtil.error(f"Datasource {ds.id} connection failed: {e}")
156+
if is_raise:
157+
raise HTTPException(status_code=500, detail=trans('i18n_ds_invalid') + f': {e.args}')
158+
return False
159+
else:
160+
conn = get_ds_engine(ds)
109161
try:
110162
with conn.connect() as connection:
111163
SQLBotLogUtil.info("success")
@@ -115,46 +167,6 @@ def check_connection(trans: Optional[Trans], ds: CoreDatasource, is_raise: bool
115167
if is_raise:
116168
raise HTTPException(status_code=500, detail=trans('i18n_ds_invalid') + f': {e.args}')
117169
return False
118-
else:
119-
conf = DatasourceConf(**json.loads(aes_decrypt(ds.configuration)))
120-
if ds.type == 'dm':
121-
with dmPython.connect(user=conf.username, password=conf.password, server=conf.host,
122-
port=conf.port) as conn, conn.cursor() as cursor:
123-
try:
124-
cursor.execute('select 1', timeout=10).fetchall()
125-
SQLBotLogUtil.info("success")
126-
return True
127-
except Exception as e:
128-
SQLBotLogUtil.error(f"Datasource {ds.id} connection failed: {e}")
129-
if is_raise:
130-
raise HTTPException(status_code=500, detail=trans('i18n_ds_invalid') + f': {e.args}')
131-
return False
132-
elif ds.type == 'doris':
133-
with pymysql.connect(user=conf.username, passwd=conf.password, host=conf.host,
134-
port=conf.port, db=conf.database, connect_timeout=10,
135-
read_timeout=10) as conn, conn.cursor() as cursor:
136-
try:
137-
cursor.execute('select 1')
138-
SQLBotLogUtil.info("success")
139-
return True
140-
except Exception as e:
141-
SQLBotLogUtil.error(f"Datasource {ds.id} connection failed: {e}")
142-
if is_raise:
143-
raise HTTPException(status_code=500, detail=trans('i18n_ds_invalid') + f': {e.args}')
144-
return False
145-
elif ds.type == 'redshift':
146-
with redshift_connector.connect(host=conf.host, port=conf.port, database=conf.database, user=conf.username,
147-
password=conf.password,
148-
timeout=10) as conn, conn.cursor() as cursor:
149-
try:
150-
cursor.execute('select 1')
151-
SQLBotLogUtil.info("success")
152-
return True
153-
except Exception as e:
154-
SQLBotLogUtil.error(f"Datasource {ds.id} connection failed: {e}")
155-
if is_raise:
156-
raise HTTPException(status_code=500, detail=trans('i18n_ds_invalid') + f': {e.args}')
157-
return False
158170

159171

160172
def get_version(ds: CoreDatasource | AssistantOutDsSchema):

0 commit comments

Comments
 (0)