@@ -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,8 @@ 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
170+
171+ return False
158172
159173
160174def get_version (ds : CoreDatasource | AssistantOutDsSchema ):
0 commit comments