Function exec_sql

backend/apps/db/db.py:585–716

(ds: CoreDatasource | AssistantOutDsSchema, sql: str, origin_column=False)

def exec_sql(ds: CoreDatasource | AssistantOutDsSchema, sql: str, origin_column=False):
    while sql.endswith(';'):
        sql = sql[:-1]
    # check execute sql only contain read operations
    if not check_sql_read(sql, ds):
        raise ValueError(f"SQL can only contain read operations")

    db = DB.get_db(ds.type)
    if db.connect_type == ConnectType.sqlalchemy:
        with get_session(ds) as session:
            with session.execute(text(sql)) as result:
                try:
                    columns = result.keys()._keys if origin_column else [item.lower() for item in result.keys()._keys]
                    res = result.fetchall()
                    result_list = [
                        {str(columns[i]): convert_value(value) for i, value in enumerate(tuple_item)} for tuple_item in
                        res
                    ]
                    return {"fields": columns, "data": result_list,
                            "sql": bytes.decode(base64.b64encode(bytes(sql, 'utf-8')))}
                except Exception as ex:
                    raise ParseSQLResultError(str(ex))
    else:
        conf = DatasourceConf(**json.loads(aes_decrypt(ds.configuration)))
        extra_config_dict = get_extra_config(conf)
        if equals_ignore_case(ds.type, 'dm'):
            with dmPython.connect(user=conf.username, password=conf.password, server=conf.host,
                                  port=conf.port, **extra_config_dict) as conn, conn.cursor() as cursor:
                try:
                    cursor.execute(sql, timeout=conf.timeout)
                    res = cursor.fetchall()
                    columns = [field[0] for field in cursor.description] if origin_column else [field[0].lower() for
                                                                                                field in
                                                                                                cursor.description]
                    result_list = [
                        {str(columns[i]): convert_value(value) for i, value in enumerate(tuple_item)} for tuple_item in
                        res
                    ]
                    return {"fields": columns, "data": result_list,
                            "sql": bytes.decode(base64.b64encode(bytes(sql, 'utf-8')))}
                except Exception as ex:
                    raise ParseSQLResultError(str(ex))
        elif equals_ignore_case(ds.type, 'doris', 'starrocks'):
            ssl_args = {'ssl': {'ssl_mode': 'REQUIRE'}} if conf.ssl else {}
            with pymysql.connect(user=conf.username, passwd=conf.password, host=conf.host,
                                 port=conf.port, db=conf.database, connect_timeout=conf.timeout,
                                 read_timeout=conf.timeout, **extra_config_dict,
                                 **ssl_args) as conn, conn.cursor() as cursor:
                try:
                    cursor.execute(sql)
                    res = cursor.fetchall()
                    columns = [field[0] for field in cursor.description] if origin_column else [field[0].lower() for
                                                                                                field in
                                                                                                cursor.description]
                    result_list = [
                        {str(columns[i]): convert_value(value) for i, value in enumerate(tuple_item)} for tuple_item in
                        res
                    ]