feat: support user level search work group (#3790)

src/handler/http/auth/validator.rs

@@ -589,7 +589,7 @@ pub async fn get_user_email_from_auth_str(auth_str: &str) -> Option<String> {
         super::token::get_user_name_from_token(auth_str).await
     } else if auth_str.starts_with("{\"auth_ext\":") {
         let auth_tokens: AuthTokensExt =
-            config::utils::json::from_str(&auth_str).unwrap_or_default();
+            config::utils::json::from_str(auth_str).unwrap_or_default();
         if chrono::Utc::now().timestamp() - auth_tokens.request_time > auth_tokens.expires_in {
             None
         } else {

src/handler/http/request/search/job.rs

@@ -15,17 +15,17 @@
 
 use std::io::Error;
 
-use actix_web::{delete, get, web, HttpResponse};
+use actix_web::{delete, get, put, web, HttpResponse};
+
+#[cfg(feature = "enterprise")]
+use crate::common::meta::http::HttpResponse as MetaHttpResponse;
 
 #[cfg(feature = "enterprise")]
 #[delete("/{org_id}/query_manager/{trace_id}")]
 pub async fn cancel_query(params: web::Path<(String, String)>) -> Result<HttpResponse, Error> {
     let (_, trace_id) = params.into_inner();
-    let res = crate::service::search::cancel_query(&trace_id).await;
-    match res {
-        Ok(status) => Ok(HttpResponse::Ok().json(status)),
-        Err(e) => Ok(HttpResponse::InternalServerError().body(format!("{:?}", e))),
-    }
+    let trace_ids = trace_id.split(',').collect::<Vec<&str>>();
+    cancel_query_inner(&trace_ids).await
 }
 
 #[cfg(not(feature = "enterprise"))]
@@ -34,13 +34,34 @@ pub async fn cancel_query(_params: web::Path<(String, String)>) -> Result<HttpRe
     Ok(HttpResponse::Forbidden().json("Not Supported"))
 }
 
+#[cfg(feature = "enterprise")]
+#[put("/{org_id}/query_manager/cancel")]
+pub async fn cancel_multiple_query(body: web::Bytes) -> Result<HttpResponse, Error> {
+    let trace_ids: Vec<String> = match config::utils::json::from_slice(&body) {
+        Ok(v) => v,
+        Err(e) => {
+            return Ok(MetaHttpResponse::bad_request(e));
+        }
+    };
+    let trace_ids = trace_ids.iter().map(|s| s.as_str()).collect::<Vec<&str>>();
+    cancel_query_inner(&trace_ids).await
+}
+
+#[cfg(not(feature = "enterprise"))]
+#[put("/{org_id}/query_manager/cancel")]
+pub async fn cancel_multiple_query(
+    _params: web::Path<(String, String)>,
+) -> Result<HttpResponse, Error> {
+    Ok(HttpResponse::Forbidden().json("Not Supported"))
+}
+
 #[cfg(feature = "enterprise")]
 #[get("/{org_id}/query_manager/status")]
 pub async fn query_status(_params: web::Path<String>) -> Result<HttpResponse, Error> {
     let res = crate::service::search::query_status().await;
     match res {
         Ok(query_status) => Ok(HttpResponse::Ok().json(query_status)),
-        Err(e) => Ok(HttpResponse::InternalServerError().body(format!("{:?}", e))),
+        Err(e) => Ok(MetaHttpResponse::bad_request(e)),
     }
 }
 
@@ -49,3 +70,18 @@ pub async fn query_status(_params: web::Path<String>) -> Result<HttpResponse, Er
 pub async fn query_status(_params: web::Path<String>) -> Result<HttpResponse, Error> {
     Ok(HttpResponse::Forbidden().json("Not Supported"))
 }
+
+#[cfg(feature = "enterprise")]
+async fn cancel_query_inner(trace_ids: &[&str]) -> Result<HttpResponse, Error> {
+    if trace_ids.is_empty() {
+        return Ok(HttpResponse::BadRequest().json("Invalid trace_id"));
+    }
+    let mut res = Vec::with_capacity(trace_ids.len());
+    for trace_id in trace_ids {
+        match crate::service::search::cancel_query(trace_id).await {
+            Ok(status) => res.push(status),
+            Err(e) => return Ok(MetaHttpResponse::bad_request(e)),
+        }
+    }
+    Ok(HttpResponse::Ok().json(res))
+}

src/handler/http/request/search/mod.rs

@@ -243,7 +243,7 @@ pub async fn search(
     let took_wait = start.elapsed().as_millis() as usize;
     #[cfg(feature = "enterprise")]
     let took_wait = 0;
-    log::info!("http search API wait in queue took: {} ms", took_wait);
+    log::info!("http search API wait in local queue took: {} ms", took_wait);
 
     let search_fut = SearchService::search(
         &trace_id,

src/handler/http/router/mod.rs

@@ -384,6 +384,7 @@ pub fn get_service_routes(cfg: &mut web::ServiceConfig) {
             .service(prom::format_query_post)
             .service(enrichment_table::save_enrichment_table)
             .service(search::search)
+            .service(search::job::cancel_multiple_query)
             .service(search::job::cancel_query)
             .service(search::job::query_status)
             .service(search::search_partition)