(trans: Optional[Trans], ds: CoreDatasource | AssistantOutDsSchema, is_raise: bool = False)
| 178 | |
| 179 | |
| 180 | def check_connection(trans: Optional[Trans], ds: CoreDatasource | AssistantOutDsSchema, is_raise: bool = False): |
| 181 | if isinstance(ds, AssistantOutDsSchema): |
| 182 | out_conf = get_out_ds_conf(ds, 10) |
| 183 | ds.configuration = out_conf |
| 184 | |
| 185 | db = DB.get_db(ds.type) |
| 186 | if db.connect_type == ConnectType.sqlalchemy: |
| 187 | conn = get_engine(ds, 10) |
| 188 | try: |
| 189 | with conn.connect() as connection: |
| 190 | SQLBotLogUtil.info("success") |
| 191 | return True |
| 192 | except Exception as e: |
| 193 | SQLBotLogUtil.error(f"Datasource {ds.id} connection failed: {e}") |
| 194 | if is_raise: |
| 195 | raise HTTPException(status_code=500, detail=trans('i18n_ds_invalid') + f': {e.args}') |
| 196 | return False |
| 197 | else: |
| 198 | conf = DatasourceConf(**json.loads(aes_decrypt(ds.configuration))) |
| 199 | extra_config_dict = get_extra_config(conf) |
| 200 | if equals_ignore_case(ds.type, 'dm'): |
| 201 | with dmPython.connect(user=conf.username, password=conf.password, server=conf.host, |
| 202 | port=conf.port, **extra_config_dict) as conn, conn.cursor() as cursor: |
| 203 | try: |
| 204 | cursor.execute('select 1', timeout=10).fetchall() |
| 205 | SQLBotLogUtil.info("success") |
| 206 | return True |
| 207 | except Exception as e: |
| 208 | SQLBotLogUtil.error(f"Datasource {ds.id} connection failed: {e}") |
| 209 | if is_raise: |
| 210 | raise HTTPException(status_code=500, detail=trans('i18n_ds_invalid') + f': {e.args}') |
| 211 | return False |
| 212 | elif equals_ignore_case(ds.type, 'doris', 'starrocks'): |
| 213 | ssl_args = {'ssl': {'ssl_mode': 'REQUIRE'}} if conf.ssl else {} |
| 214 | with pymysql.connect(user=conf.username, passwd=conf.password, host=conf.host, |
| 215 | port=conf.port, db=conf.database, connect_timeout=10, |
| 216 | read_timeout=10, **extra_config_dict, **ssl_args) as conn, conn.cursor() as cursor: |
| 217 | try: |
| 218 | cursor.execute('select 1') |
| 219 | SQLBotLogUtil.info("success") |
| 220 | return True |
| 221 | except Exception as e: |
| 222 | SQLBotLogUtil.error(f"Datasource {ds.id} connection failed: {e}") |
| 223 | if is_raise: |
| 224 | raise HTTPException(status_code=500, detail=trans('i18n_ds_invalid') + f': {e.args}') |
| 225 | return False |
| 226 | elif equals_ignore_case(ds.type, 'redshift'): |
| 227 | with redshift_connector.connect(host=conf.host, port=conf.port, database=conf.database, |
| 228 | user=conf.username, |
| 229 | password=conf.password, |
| 230 | timeout=10, **extra_config_dict) as conn, conn.cursor() as cursor: |
| 231 | try: |
| 232 | cursor.execute('select 1') |
| 233 | SQLBotLogUtil.info("success") |
| 234 | return True |
| 235 | except Exception as e: |
| 236 | SQLBotLogUtil.error(f"Datasource {ds.id} connection failed: {e}") |
| 237 | if is_raise: |
no test coverage detected