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

src/common/utils/auth.rs

@@ -505,20 +505,18 @@ pub fn extract_auth_str(req: &HttpRequest) -> String {
             // If cookie was set but access token is still empty
             // we check auth_ext cookie to get the token.
             auth_ext_cookie(req)
-        } else {
-            if access_token.starts_with("Basic") || access_token.starts_with("Bearer") {
-                access_token
-            } else if access_token.starts_with("session") {
-                let session_key = access_token.strip_prefix("session ").unwrap().to_string();
-                match USER_SESSIONS.get(&session_key) {
-                    Some(token) => {
-                        format!("Bearer {}", *token)
-                    }
-                    None => access_token,
+        } else if access_token.starts_with("Basic") || access_token.starts_with("Bearer") {
+            access_token
+        } else if access_token.starts_with("session") {
+            let session_key = access_token.strip_prefix("session ").unwrap().to_string();
+            match USER_SESSIONS.get(&session_key) {
+                Some(token) => {
+                    format!("Bearer {}", *token)
                 }
-            } else {
-                format!("Bearer {}", access_token)
+                None => access_token,
             }
+        } else {
+            format!("Bearer {}", access_token)
         }
     } else if let Some(cookie) = req.cookie("auth_ext") {
         cookie.value().to_string()

src/config/src/meta/search.rs

@@ -465,6 +465,7 @@ impl From<Request> for cluster_rpc::SearchRequest {
             stream_type: "".to_string(),
             timeout: req.timeout,
             work_group: "".to_string(),
+            user_id: None,
         }
     }
 }

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)

src/infra/src/schema/mod.rs

@@ -20,7 +20,7 @@ use config::{
     get_config,
     meta::stream::{PartitionTimeLevel, StreamSettings, StreamType},
     utils::{json, schema_ext::SchemaExt},
-    RwAHashMap,
+    RwAHashMap, BLOOM_FILTER_DEFAULT_FIELDS, SQL_FULL_TEXT_SEARCH_FIELDS,
 };
 use datafusion::arrow::datatypes::{DataType, Field, Schema};
 use futures::{StreamExt, TryStreamExt};
@@ -266,17 +266,31 @@ pub fn unwrap_partition_time_level(
     }
 }
 
-pub fn get_stream_setting_fts_fields(schema: &Schema) -> Result<Vec<String>> {
+pub fn get_stream_setting_fts_fields(schema: &Schema) -> Vec<String> {
+    let default_fields = SQL_FULL_TEXT_SEARCH_FIELDS.clone();
     match unwrap_stream_settings(schema) {
-        Some(setting) => Ok(setting.full_text_search_keys),
-        None => Ok(vec![]),
+        Some(setting) => {
+            let mut fields = setting.full_text_search_keys;
+            fields.extend(default_fields);
+            fields.sort();
+            fields.dedup();
+            fields
+        }
+        None => default_fields,
     }
 }
 
-pub fn get_stream_setting_bloom_filter_fields(schema: &Schema) -> Result<Vec<String>> {
+pub fn get_stream_setting_bloom_filter_fields(schema: &Schema) -> Vec<String> {
+    let default_fields = BLOOM_FILTER_DEFAULT_FIELDS.clone();
     match unwrap_stream_settings(schema) {
-        Some(setting) => Ok(setting.bloom_filter_fields),
-        None => Ok(vec![]),
+        Some(setting) => {
+            let mut fields = setting.bloom_filter_fields;
+            fields.extend(default_fields);
+            fields.sort();
+            fields.dedup();
+            fields
+        }
+        None => default_fields,
     }
 }
 
@@ -737,6 +751,6 @@ mod tests {
     fn test_get_stream_setting_fts_fields() {
         let sch = Schema::new(vec![Field::new("f.c", DataType::Int32, false)]);
         let res = get_stream_setting_fts_fields(&sch);
-        assert!(res.is_ok());
+        assert!(!res.is_empty());
     }
 }

src/ingester/src/partition.rs

@@ -110,10 +110,9 @@ impl Partition {
             let (bloom_filter_fields, full_text_search_fields) =
                 if self.schema.fields().len() >= cfg.limit.file_move_fields_limit {
                     let bloom_filter_fields =
-                        infra::schema::get_stream_setting_bloom_filter_fields(self.schema.as_ref())
-                            .unwrap();
+                        infra::schema::get_stream_setting_bloom_filter_fields(self.schema.as_ref());
                     let full_text_search_fields =
-                        infra::schema::get_stream_setting_fts_fields(self.schema.as_ref()).unwrap();
+                        infra::schema::get_stream_setting_fts_fields(self.schema.as_ref());
                     (bloom_filter_fields, full_text_search_fields)
                 } else {
                     (vec![], vec![])

src/proto/proto/cluster/search.proto

@@ -44,6 +44,7 @@ message SearchRequest {
     repeated SearchAggRequest aggs = 7;
     int64                  timeout = 8;
     string              work_group = 9;
+    optional string       user_id = 10;
 }
 
 message SearchResponse {

src/service/compact/merge.rs

@@ -631,8 +631,8 @@ pub async fn merge_files(
         infra::schema::get_versions(org_id, stream_name, stream_type, Some((min_ts, max_ts)))
             .await?;
     let schema_latest_id = schema_versions.len() - 1;
-    let bloom_filter_fields = get_stream_setting_bloom_filter_fields(&schema_latest).unwrap();
-    let full_text_search_fields = get_stream_setting_fts_fields(&schema_latest).unwrap();
+    let bloom_filter_fields = get_stream_setting_bloom_filter_fields(&schema_latest);
+    let full_text_search_fields = get_stream_setting_fts_fields(&schema_latest);
     if cfg.common.widening_schema_evolution && schema_versions.len() > 1 {
         for file in new_file_list.iter() {
             // get the schema version of the file

src/service/search/cluster/mod.rs

@@ -69,8 +69,16 @@ pub async fn search(
     let cfg = get_config();
     // if the request is a super cluster request, then forward it to the super cluster service
     let is_final_phase = req.stype != cluster_rpc::SearchType::SuperCluster as i32;
-
     let stream_type = StreamType::from(req.stream_type.as_str());
+    if req.timeout == 0 {
+        req.timeout = cfg.limit.query_timeout as i64;
+    }
+
+    // get user_id
+    #[cfg(feature = "enterprise")]
+    let user_id = req.user_id.clone();
+    #[cfg(feature = "enterprise")]
+    let user_id = user_id.as_deref();
 
     // get nodes from cluster
     let mut nodes = infra_cluster::get_cached_online_query_nodes()
@@ -256,8 +264,9 @@ pub async fn search(
     let work_group: Option<String> = None;
     // 1. get work group
     #[cfg(feature = "enterprise")]
-    let work_group: Option<o2_enterprise::enterprise::search::WorkGroup> =
-        Some(o2_enterprise::enterprise::search::queue::predict_work_group(&nodes, &file_list));
+    let work_group: Option<o2_enterprise::enterprise::search::WorkGroup> = Some(
+        o2_enterprise::enterprise::search::work_group::predict(&nodes, &file_list),
+    );
     // 2. check concurrency
     let work_group_str = if let Some(wg) = &work_group {
         wg.to_string()
@@ -265,36 +274,78 @@ pub async fn search(
         "global".to_string()
     };
 
-    let locker_key = "/search/cluster_queue/".to_string() + work_group_str.as_str();
+    let locker_key = format!("/search/cluster_queue/{}", work_group_str);
     // get a cluster search queue lock
     let locker = if cfg.common.local_mode || !cfg.common.feature_query_queue_enabled {
         None
     } else {
-        dist_lock::lock(&locker_key, 0).await?
+        dist_lock::lock(&locker_key, req.timeout as u64).await?
     };
     #[cfg(feature = "enterprise")]
-    loop {
-        match work_group.as_ref().unwrap().need_wait().await {
-            Ok(true) => {
-                tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
+    {
+        // check user concurrency
+        if user_id.is_some() {
+            loop {
+                if start.elapsed().as_millis() as u64 >= req.timeout as u64 * 1000 {
+                    dist_lock::unlock(&locker).await?;
+                    return Err(Error::Message(format!(
+                        "[trace_id {trace_id}] search: timeout in queue"
+                    )));
+                }
+                match work_group.as_ref().unwrap().need_wait(user_id).await {
+                    Ok(true) => {
+                        tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
+                    }
+                    Ok(false) => {
+                        break;
+                    }
+                    Err(e) => {
+                        dist_lock::unlock(&locker).await?;
+                        return Err(Error::Message(e.to_string()));
+                    }
+                }
             }
-            Ok(false) => {
-                break;
-            }
-            Err(e) => {
+            // process the search in the work group
+            if let Err(e) = work_group
+                .as_ref()
+                .unwrap()
+                .process(trace_id, user_id)
+                .await
+            {
                 dist_lock::unlock(&locker).await?;
                 return Err(Error::Message(e.to_string()));
             }
         }
-    }
-    // 3. process the search in the work group
-    #[cfg(feature = "enterprise")]
-    if let Err(e) = work_group.as_ref().unwrap().process(trace_id).await {
+        // check global concurrency
+        loop {
+            if start.elapsed().as_millis() as u64 >= req.timeout as u64 * 1000 {
+                dist_lock::unlock(&locker).await?;
+                return Err(Error::Message(format!(
+                    "[trace_id {trace_id}] search: timeout in queue"
+                )));
+            }
+            match work_group.as_ref().unwrap().need_wait(None).await {
+                Ok(true) => {
+                    tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
+                }
+                Ok(false) => {
+                    break;
+                }
+                Err(e) => {
+                    dist_lock::unlock(&locker).await?;
+                    return Err(Error::Message(e.to_string()));
+                }
+            }
+        }
+        // 3. process the search in the work group
+        if let Err(e) = work_group.as_ref().unwrap().process(trace_id, None).await {
+            dist_lock::unlock(&locker).await?;
+            return Err(Error::Message(e.to_string()));
+        }
         dist_lock::unlock(&locker).await?;
-        return Err(Error::Message(e.to_string()));
     }
-    #[cfg(feature = "enterprise")]
-    dist_lock::unlock(&locker).await?;
+
+    // done in the queue
     let took_wait = start.elapsed().as_millis() as usize - file_list_took;
     log::info!(
         "[trace_id {trace_id}] search: wait in queue took: {} ms",
@@ -430,9 +481,17 @@ pub async fn search(
             work_group
                 .as_ref()
                 .unwrap()
-                .done(&trace_id)
+                .done(&trace_id, None)
                 .await
                 .map_err(|e| Error::Message(e.to_string()))?;
+            if user_id.is_some() {
+                work_group
+                    .as_ref()
+                    .unwrap()
+                    .done(&trace_id, user_id)
+                    .await
+                    .map_err(|e| Error::Message(e.to_string()))?;
+            }
             return Err(Error::ErrorCode(ErrorCodes::SearchCancelQuery(format!(
                 "[trace_id {trace_id}] search->grpc: search canceled before call search->grpc"
             ))));
@@ -557,12 +616,22 @@ pub async fn search(
                 #[cfg(not(feature = "enterprise"))]
                 dist_lock::unlock(&locker).await?;
                 #[cfg(feature = "enterprise")]
-                work_group
-                    .as_ref()
-                    .unwrap()
-                    .done(trace_id)
-                    .await
-                    .map_err(|e| Error::Message(e.to_string()))?;
+                {
+                    work_group
+                        .as_ref()
+                        .unwrap()
+                        .done(trace_id, None)
+                        .await
+                        .map_err(|e| Error::Message(e.to_string()))?;
+                    if user_id.is_some() {
+                        work_group
+                            .as_ref()
+                            .unwrap()
+                            .done(trace_id, user_id)
+                            .await
+                            .map_err(|e| Error::Message(e.to_string()))?;
+                    }
+                }
                 return Err(Error::ErrorCode(ErrorCodes::ServerInternalError(
                     e.to_string(),
                 )));
@@ -583,12 +652,22 @@ pub async fn search(
                 #[cfg(not(feature = "enterprise"))]
                 dist_lock::unlock(&locker).await?;
                 #[cfg(feature = "enterprise")]
-                work_group
-                    .as_ref()
-                    .unwrap()
-                    .done(trace_id)
-                    .await
-                    .map_err(|e| Error::Message(e.to_string()))?;
+                {
+                    work_group
+                        .as_ref()
+                        .unwrap()
+                        .done(trace_id, None)
+                        .await
+                        .map_err(|e| Error::Message(e.to_string()))?;
+                    if user_id.is_some() {
+                        work_group
+                            .as_ref()
+                            .unwrap()
+                            .done(trace_id, user_id)
+                            .await
+                            .map_err(|e| Error::Message(e.to_string()))?;
+                    }
+                }
                 return Err(e);
             }
         };
@@ -598,12 +677,22 @@ pub async fn search(
     #[cfg(not(feature = "enterprise"))]
     dist_lock::unlock(&locker).await?;
     #[cfg(feature = "enterprise")]
-    work_group
-        .as_ref()
-        .unwrap()
-        .done(trace_id)
-        .await
-        .map_err(|e| Error::Message(e.to_string()))?;
+    {
+        work_group
+            .as_ref()
+            .unwrap()
+            .done(trace_id, None)
+            .await
+            .map_err(|e| Error::Message(e.to_string()))?;
+        if user_id.is_some() {
+            work_group
+                .as_ref()
+                .unwrap()
+                .done(trace_id, user_id)
+                .await
+                .map_err(|e| Error::Message(e.to_string()))?;
+        }
+    }
 
     Ok((merge_batches, scan_stats, took_wait, is_partial))
 }

src/service/search/mod.rs

@@ -123,6 +123,7 @@ pub async fn search(
     req.org_id = org_id.to_string();
     req.stype = cluster_rpc::SearchType::Cluster as _;
     req.stream_type = stream_type.to_string();
+    req.user_id = user_id.clone();
 
     let req_query = req.clone().query.unwrap();
 

src/service/search/sql.rs

@@ -25,7 +25,7 @@ use config::{
         sql::{Sql as MetaSql, SqlOperator},
         stream::{FileKey, StreamPartition, StreamType},
     },
-    QUICK_MODEL_FIELDS, SQL_FULL_TEXT_SEARCH_FIELDS,
+    QUICK_MODEL_FIELDS,
 };
 use datafusion::arrow::datatypes::{DataType, Schema};
 use hashbrown::HashSet;
@@ -389,15 +389,7 @@ impl Sql {
 
         // fetch fts fields
         let mut fts_terms = HashSet::new();
-        let fts_fields = get_stream_setting_fts_fields(&schema).unwrap();
-        let match_all_fields = if !fts_fields.is_empty() {
-            fts_fields.iter().map(|v| v.to_lowercase()).collect()
-        } else {
-            SQL_FULL_TEXT_SEARCH_FIELDS
-                .iter()
-                .map(|v| v.to_string())
-                .collect::<Vec<String>>()
-        };
+        let fts_fields = get_stream_setting_fts_fields(&schema);
 
         // Hack for quick_mode
         // replace `select *` to `select f1,f2,f3`
@@ -415,7 +407,7 @@ impl Sql {
             } else {
                 None
             };
-            let fields = generate_quick_mode_fields(&schema, cached_fields, &match_all_fields);
+            let fields = generate_quick_mode_fields(&schema, cached_fields, &fts_fields);
             let select_fields = "SELECT ".to_string() + &fields.join(",");
             origin_sql = RE_ONLY_SELECT
                 .replace(origin_sql.as_str(), &select_fields)
@@ -462,7 +454,7 @@ impl Sql {
 
         // use full text search instead if inverted index feature is not enabled but it's an
         // inverted index search
-        let rerouted = if !cfg.common.inverted_index_enabled & fulltext.is_empty() {
+        let ignore_case = if !cfg.common.inverted_index_enabled && fulltext.is_empty() {
             fulltext = std::mem::take(&mut indexed_text);
             true
         } else {
@@ -473,17 +465,15 @@ impl Sql {
         for item in indexed_text.iter() {
             let mut indexed_search = Vec::new();
             for field in &schema_fields {
-                if !match_all_fields.contains(&field.name().to_lowercase()) {
+                if !fts_fields.contains(&field.name().to_lowercase()) {
                     continue;
                 }
                 if !field.data_type().eq(&DataType::Utf8) || field.name().starts_with('@') {
                     continue;
                 }
                 indexed_search.push(format!("\"{}\" LIKE '%{}%'", field.name(), item.1));
-
                 // add full text field to meta fields
                 meta.fields.push(field.name().to_string());
-
                 fts_terms.insert(item.1.clone());
             }
             if indexed_search.is_empty() {
@@ -496,14 +486,14 @@ impl Sql {
         for item in fulltext.iter() {
             let mut fulltext_search = Vec::new();
             for field in &schema_fields {
-                if !match_all_fields.contains(&field.name().to_lowercase()) {
+                if !fts_fields.contains(&field.name().to_lowercase()) {
                     continue;
                 }
                 if !field.data_type().eq(&DataType::Utf8) || field.name().starts_with('@') {
                     continue;
                 }
                 let mut func = "LIKE";
-                if rerouted || item.0.to_lowercase().contains("_ignore_case") {
+                if ignore_case || item.0.to_lowercase().contains("_ignore_case") {
                     func = "ILIKE";
                 }
                 fulltext_search.push(format!("\"{}\" {} '%{}%'", field.name(), func, item.1));
@@ -517,44 +507,6 @@ impl Sql {
             origin_sql = origin_sql.replace(item.0.as_str(), &fulltext_search);
         }
 
-        // Hack: str_match
-        for key in [
-            "match",
-            "match_ignore_case",
-            // "str_match"             // use UDF will get better result
-            // "str_match_ignore_case" // use UDF will get better result
-        ] {
-            let re_str_match = Regex::new(&format!(r"(?i)\b{key}\b\(([^\)]*)\)")).unwrap();
-            let re_fn = if key == "match" || key == "str_match" {
-                "LIKE"
-            } else {
-                "ILIKE"
-            };
-            for token in &where_tokens {
-                if !token.to_lowercase().starts_with("match")
-                    && !token.to_lowercase().starts_with("str_match")
-                {
-                    continue;
-                }
-                for cap in re_str_match.captures_iter(token.as_str()) {
-                    let attrs = cap
-                        .get(1)
-                        .unwrap()
-                        .as_str()
-                        .splitn(2, ',')
-                        .map(|v| v.trim().trim_matches(|v| v == '\'' || v == '"'))
-                        .collect::<Vec<&str>>();
-                    let field = attrs.first().unwrap();
-                    let value = attrs.last().unwrap();
-                    origin_sql = origin_sql.replace(
-                        cap.get(0).unwrap().as_str(),
-                        &format!("\"{field}\" {re_fn} '%{value}%'"),
-                    );
-                    fts_terms.insert(value.to_string());
-                }
-            }
-        }
-
         // Hack for histogram
         let mut histogram_interval = None;
         let from_pos = origin_sql.to_lowercase().find(" from ").unwrap();