diff --git a/src/common/utils/auth.rs b/src/common/utils/auth.rs index 100c7d2306..12dcbf3aee 100644 --- a/src/common/utils/auth.rs +++ b/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() diff --git a/src/config/src/meta/search.rs b/src/config/src/meta/search.rs index df32e9f15d..3244d982ab 100644 --- a/src/config/src/meta/search.rs +++ b/src/config/src/meta/search.rs @@ -465,6 +465,7 @@ impl From for cluster_rpc::SearchRequest { stream_type: "".to_string(), timeout: req.timeout, work_group: "".to_string(), + user_id: None, } } } diff --git a/src/handler/http/auth/validator.rs b/src/handler/http/auth/validator.rs index 34a6c7caf2..48ee91f191 100644 --- a/src/handler/http/auth/validator.rs +++ b/src/handler/http/auth/validator.rs @@ -589,7 +589,7 @@ pub async fn get_user_email_from_auth_str(auth_str: &str) -> Option { 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 { diff --git a/src/handler/http/request/search/job.rs b/src/handler/http/request/search/job.rs index a52c5ce485..b9766c84ce 100644 --- a/src/handler/http/request/search/job.rs +++ b/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 { 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::>(); + 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 Result { + let trace_ids: Vec = 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::>(); + 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 { + Ok(HttpResponse::Forbidden().json("Not Supported")) +} + #[cfg(feature = "enterprise")] #[get("/{org_id}/query_manager/status")] pub async fn query_status(_params: web::Path) -> Result { 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) -> Result) -> Result { Ok(HttpResponse::Forbidden().json("Not Supported")) } + +#[cfg(feature = "enterprise")] +async fn cancel_query_inner(trace_ids: &[&str]) -> Result { + 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)) +} diff --git a/src/handler/http/request/search/mod.rs b/src/handler/http/request/search/mod.rs index e228bf1d90..63218a451c 100644 --- a/src/handler/http/request/search/mod.rs +++ b/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, diff --git a/src/handler/http/router/mod.rs b/src/handler/http/router/mod.rs index a6d226f565..f2152f5c8e 100644 --- a/src/handler/http/router/mod.rs +++ b/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) diff --git a/src/infra/src/schema/mod.rs b/src/infra/src/schema/mod.rs index f05628aed8..d6c8c6de5f 100644 --- a/src/infra/src/schema/mod.rs +++ b/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> { +pub fn get_stream_setting_fts_fields(schema: &Schema) -> Vec { + 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> { +pub fn get_stream_setting_bloom_filter_fields(schema: &Schema) -> Vec { + 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()); } } diff --git a/src/ingester/src/partition.rs b/src/ingester/src/partition.rs index a4131d1e04..6b0b4412e9 100644 --- a/src/ingester/src/partition.rs +++ b/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![]) diff --git a/src/proto/proto/cluster/search.proto b/src/proto/proto/cluster/search.proto index acbfc6033b..843051f837 100644 --- a/src/proto/proto/cluster/search.proto +++ b/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 { diff --git a/src/service/compact/merge.rs b/src/service/compact/merge.rs index a8e6ef1323..84913af097 100644 --- a/src/service/compact/merge.rs +++ b/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 diff --git a/src/service/search/cluster/mod.rs b/src/service/search/cluster/mod.rs index 8d759f7cfd..833242fa9c 100644 --- a/src/service/search/cluster/mod.rs +++ b/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 = None; // 1. get work group #[cfg(feature = "enterprise")] - let work_group: Option = - Some(o2_enterprise::enterprise::search::queue::predict_work_group(&nodes, &file_list)); + let work_group: Option = 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)) } diff --git a/src/service/search/mod.rs b/src/service/search/mod.rs index fbacbdbeb5..f2bf066d1f 100644 --- a/src/service/search/mod.rs +++ b/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(); diff --git a/src/service/search/sql.rs b/src/service/search/sql.rs index 7376c353c5..4dd70716c6 100644 --- a/src/service/search/sql.rs +++ b/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::>() - }; + 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::>(); - 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();