From 00972e293eb69977bac42a74e1697fddfdbec457 Mon Sep 17 00:00:00 2001 From: udev Date: Wed, 16 Apr 2025 19:45:38 +0530 Subject: [PATCH] feat: Implement WS for Values API (#6463) Co-authored-by: omkark06 --- src/config/src/config.rs | 2 +- src/config/src/meta/search.rs | 36 +++ src/config/src/meta/websocket.rs | 27 +- src/config/src/utils/sql.rs | 12 + src/handler/http/request/search/mod.rs | 316 +++++++++++++++++++--- src/handler/http/request/ws_v2/session.rs | 228 ++++++++++++++-- src/service/search/mod.rs | 5 + src/service/search/partition.rs | 2 + src/service/websocket_events/mod.rs | 2 + src/service/websocket_events/search.rs | 115 +++++++- src/service/websocket_events/utils.rs | 7 +- src/service/websocket_events/values.rs | 275 +++++++++++++++++++ web/src/composables/useLogs.ts | 4 +- web/src/plugins/logs/IndexList.vue | 254 +++++++++++++++++ web/src/ts/interfaces/query.ts | 8 + 15 files changed, 1222 insertions(+), 71 deletions(-) create mode 100644 src/service/websocket_events/values.rs diff --git a/src/config/src/config.rs b/src/config/src/config.rs index 9255fb0b27..b2dcbd6506 100644 --- a/src/config/src/config.rs +++ b/src/config/src/config.rs @@ -738,7 +738,7 @@ pub struct Common { pub feature_query_without_index: bool, #[env_config(name = "ZO_FEATURE_QUERY_REMOVE_FILTER_WITH_INDEX", default = true)] pub feature_query_remove_filter_with_index: bool, - #[env_config(name = "ZO_FEATURE_QUERY_STREAMING_AGGS", default = false)] + #[env_config(name = "ZO_FEATURE_QUERY_STREAMING_AGGS", default = true)] pub feature_query_streaming_aggs: bool, #[env_config(name = "ZO_FEATURE_JOIN_MATCH_ONE_ENABLED", default = false)] pub feature_join_match_one_enabled: bool, diff --git a/src/config/src/meta/search.rs b/src/config/src/meta/search.rs index 5e8ccd5b9d..519c4ed846 100644 --- a/src/config/src/meta/search.rs +++ b/src/config/src/meta/search.rs @@ -16,6 +16,7 @@ use proto::cluster_rpc; use serde::{Deserialize, Deserializer, Serialize}; use utoipa::ToSchema; +use crate::meta::stream::StreamType; use crate::{ meta::sql::OrderBy, @@ -817,6 +818,14 @@ impl TryFrom<&str> for SearchEventType { } } +#[derive(Clone, Debug, PartialEq, Default, Serialize, Deserialize, ToSchema)] +#[serde(rename_all = "snake_case")] +pub struct ValuesEventContext { + pub top_k: Option, + pub no_count: bool, + pub field: String, +} + #[derive(Clone, Debug, PartialEq, Default, Serialize, Deserialize, ToSchema)] #[serde(rename_all = "snake_case")] pub struct SearchEventContext { @@ -1059,6 +1068,33 @@ pub struct PaginationQuery { pub size: Option, } +#[derive(Debug, Deserialize, Clone, Serialize)] +pub struct ValuesRequest { + pub fields: Vec, + #[serde(default)] + pub size: Option, + pub no_count: bool, + #[serde(default)] + pub regions: Vec, + #[serde(default)] + pub clusters: Vec, + #[serde(default)] + pub vrl_fn: Option, + #[serde(default)] + pub start_time: Option, + #[serde(default)] + pub end_time: Option, + #[serde(default)] + pub filter: Option, + #[serde(default)] + pub timeout: Option, + #[serde(default)] + pub use_cache: bool, + pub stream_name: String, + pub stream_type: StreamType, + pub sql: String, +} + #[cfg(test)] mod tests { use super::*; diff --git a/src/config/src/meta/websocket.rs b/src/config/src/meta/websocket.rs index ffedc7b785..3b85f2769d 100644 --- a/src/config/src/meta/websocket.rs +++ b/src/config/src/meta/websocket.rs @@ -15,7 +15,7 @@ use serde::{Deserialize, Serialize}; -use crate::meta::search::{Response, SearchEventContext}; +use crate::meta::search::{Response, SearchEventContext, ValuesEventContext}; pub const MAX_QUERY_RANGE_LIMIT_ERROR_MESSAGE: &str = "Reached Max query range limit."; @@ -40,6 +40,31 @@ pub struct SearchEventReq { pub org_id: String, // TODO: is it PII safe? pub user_id: Option, + /// Top K is used only in values search + #[serde(default)] + pub values_event_context: Option, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +pub struct ValuesEventReq { + pub trace_id: String, + pub payload: crate::meta::search::ValuesRequest, + pub stream_type: crate::meta::stream::StreamType, + pub use_cache: bool, + pub org_id: String, + pub user_id: Option, +} + +impl ValuesEventReq { + pub fn is_valid(&self) -> bool { + if self.trace_id.is_empty() { + return false; + } + if self.org_id.is_empty() { + return false; + } + true + } } impl SearchEventReq { diff --git a/src/config/src/utils/sql.rs b/src/config/src/utils/sql.rs index 2bc706ba12..3809c4d99e 100644 --- a/src/config/src/utils/sql.rs +++ b/src/config/src/utils/sql.rs @@ -194,3 +194,15 @@ impl Visitor for SubqueryVisitor { ControlFlow::Continue(()) } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn check_is_simple_aggregate() { + let query = r#"SELECT histogram(_timestamp) AS zo_sql_time, "kubernetes_docker_id" AS zo_sql_key, SUM(count) AS zo_sql_num FROM "distinct_values_logs_default22" GROUP BY zo_sql_time, zo_sql_key ORDER BY zo_sql_time ASC, zo_sql_num DESC"#; + let ab = is_aggregate_query(query); + print!("{:?}", ab); + } +} \ No newline at end of file diff --git a/src/handler/http/request/search/mod.rs b/src/handler/http/request/search/mod.rs index b2d8dded0e..9783d0e2d1 100644 --- a/src/handler/http/request/search/mod.rs +++ b/src/handler/http/request/search/mod.rs @@ -13,7 +13,7 @@ // You should have received a copy of the GNU Affero General Public License // along with this program. If not, see . -use std::io::Error; +use std::{cmp::Reverse, collections::BinaryHeap, io::Error}; use actix_web::{HttpRequest, HttpResponse, get, http::StatusCode, post, web}; use arrow_schema::Schema; @@ -651,6 +651,215 @@ pub async fn values( .await } +pub type FieldName = String; + +/// Builds a search request per field +/// +/// This function builds a search request per field based on the given request and parameters. +/// Search request can basically be of two types: +/// 1. Search on distinct values stream +/// 2. Search on original data stream +/// If the field is a distinct field, we will search of the distinct values stream. +/// Otherwise, we will search on the original data stream. +/// +/// The `use_result_cache` parameter is used to determine the projection of the SQL query. +/// This flag will toggle the resultant requests between streaming aggregations and result cache. +/// By default, this function will produce an SQL which utilizes `Streaming Aggregations` to send the results to the client. +/// The SQL will be a simple aggregation query in case streaming aggregations are used. +/// If `use_result_cache` is set to true, the SQL projectoin will include a histogram which will allow the use of result cache. +/// +/// Another parameter is `no_count` which is used to determine if the count is needed or not. +/// `no_count` is used when only distinct values (sorted in alphabetical order) are needed but not the frequency of the values. +/// For example, Dashboards, where we show the values listed in alphabetical order. +/// +/// Since values request can contain multiple fields, we return a vector of requests. +/// Each request is a tuple of `Request`, `StreamType`, and `FieldName`. +/// The `Request` contains the SQL query, from, size, start_time, end_time, etc. +/// The `StreamType` is the type of the stream to search on. +/// The `FieldName` is the name of the field to search on. +pub async fn build_search_request_per_field( + req: &config::meta::search::ValuesRequest, + org_id: &str, + stream_type: StreamType, + stream_name: &str, +) -> Result, Error> { + let query_fn = req.vrl_fn + .as_ref() + .and_then(|v| base64::decode_url(v.as_ref()).ok()) + .map(|vrl| { + if !vrl.trim().ends_with('.') { + format!("{} \n .", vrl) + } else { + vrl + } + }); + + let regions = req.regions.clone(); + let clusters = req.clusters.clone(); + + let schema = infra::schema::get(org_id, stream_name, stream_type) + .await + .unwrap_or(Schema::empty()); + let fields = req.fields + .iter() + .filter(|field| schema.field_with_name(field).is_ok()) + .cloned() + .collect::>(); + + + let no_count = req.no_count; + + let start_time = if stream_type.eq(&StreamType::EnrichmentTables) { + BASE_TIME.timestamp_micros() + } else { + req.start_time.unwrap_or(0) + }; + + if start_time == 0 { + return Err(Error::other("start_time is empty")); + } + + let end_time = if stream_type.eq(&StreamType::EnrichmentTables) { + now_micros() + } else { + req.end_time.unwrap_or(0) + }; + + if end_time == 0 { + return Err(Error::other("end_time is empty")); + } + + let (start_time, end_time) = if start_time == end_time { + (start_time - 1, end_time + 1) + } else { + (start_time, end_time) + }; + + let mut uses_fn = false; + let (sql_where, can_use_distinct_stream) = match req.filter.as_ref() { + None => { + if !req.sql.is_empty() { + if let Ok(sql) = base64::decode_url(&req.sql) { + uses_fn = functions::get_all_transform_keys(org_id) + .await + .iter() + .any(|fn_name| sql.contains(&format!("{}(", fn_name))); + + // pick up where clause from sql + let sql_where_from_query = match SearchService::sql::pickup_where(&sql, None) { + Ok(Some(v)) => format!("WHERE {}", v), + Ok(None) => "".to_string(), + Err(e) => { + return Err(Error::other(e)); + } + }; + let can_use_distinct_stream = can_use_distinct_stream( + org_id, + stream_name, + stream_type, + &fields, + &sql, + start_time, + ) + .await; + (sql_where_from_query, can_use_distinct_stream) + } else { + ("".to_string(), false) + } + } else { + ("".to_string(), false) + } + } + Some(v) => { + if v.is_empty() { + ("".to_string(), false) + } else { + let columns = v.splitn(2, '=').collect::>(); + if columns.len() < 2 { + return Err(Error::other("Invalid filter format")); + } + let vals = columns[1].split(',').collect::>().join("','"); + let sql_where = format!("WHERE {} IN ('{}')", columns[0], vals); + + // Define the default_sql here + let default_sql = format!("SELECT {} FROM \"{stream_name}\"", TIMESTAMP_COL_NAME); + + let can_use_distinct_stream = can_use_distinct_stream( + org_id, + stream_name, + stream_type, + &fields, + &format!("{} {}", default_sql, sql_where), + start_time, + ) + .await; + + (sql_where, can_use_distinct_stream) + } + } + }; + + let timeout = req.timeout.unwrap_or(0); + + let req = config::meta::search::Request { + query: config::meta::search::Query { + sql: String::new(), // Will be populated per field in the loop below + from: 0, + size: config::meta::sql::MAX_LIMIT, + start_time, + end_time, + uses_zo_fn: uses_fn, + query_fn: query_fn.clone(), + ..Default::default() + }, + encoding: config::meta::search::RequestEncoding::Empty, + regions, + clusters, + timeout, + search_type: Some(SearchEventType::Values), + search_event_context: None, + use_cache: Some(req.use_cache), + local_mode: None, + }; + + let distinct_prefix = if can_use_distinct_stream { + format!("{}_{}_", DISTINCT_STREAM_PREFIX, stream_type.as_str()) + } else { + "".to_string() + }; + + let count_fn = if can_use_distinct_stream { + "SUM(count)".to_string() + } else { + "COUNT(*)".to_string() + }; + + let actual_stream_type = if can_use_distinct_stream { + StreamType::Metadata + } else { + stream_type + }; + + let mut requests = Vec::new(); + for field in fields { + let sql = if no_count { + format!( + "SELECT \"{field}\" AS zo_sql_key FROM \"{distinct_prefix}{stream_name}\" {sql_where} GROUP BY zo_sql_key" + ) + } else { + format!( + "SELECT \"{field}\" AS zo_sql_key, {count_fn} AS zo_sql_num FROM \"{distinct_prefix}{stream_name}\" {sql_where} GROUP BY zo_sql_key" + ) + }; + + let mut req = req.clone(); + req.query.sql = sql; + requests.push((req, actual_stream_type, field)); + } + + Ok(requests) +} + // If all fields requested in the query AND fields from the // sql query in the query are stored in distinct stream, // this will search on the distinct stream, otherwise @@ -673,14 +882,16 @@ async fn values_v1( Some(v) => v.split(',').map(|s| s.to_string()).collect::>(), None => return Ok(MetaHttpResponse::bad_request("fields is empty")), }; - let mut query_fn = query + let query_fn = query .get("query_fn") - .and_then(|v| base64::decode_url(v).ok()); - if let Some(vrl_function) = &query_fn { - if !vrl_function.trim().ends_with('.') { - query_fn = Some(format!("{} \n .", vrl_function)); - } - } + .and_then(|v| base64::decode_url(v.as_ref()).ok()) + .map(|vrl_function| { + if !vrl_function.trim().ends_with('.') { + format!("{} \n .", vrl_function) + } else { + vrl_function + } + }); let default_sql = format!("SELECT {} FROM \"{stream_name}\"", TIMESTAMP_COL_NAME); let mut query_sql = match query.get("filter") { @@ -729,9 +940,6 @@ async fn values_v1( } }; - let size = query - .get("size") - .map_or(10, |v| v.parse::().unwrap_or(10)); // EnrichmentTable need query without time range let start_time = if stream_type.eq(&StreamType::EnrichmentTables) { BASE_TIME.timestamp_micros() @@ -892,6 +1100,12 @@ async fn values_v1( let mut resp = config::meta::search::Response::default(); let mut hit_values: Vec = Vec::new(); let mut work_group_set = Vec::with_capacity(query_results.len()); + + // Get the size from query parameter for limiting results + let size = query + .get("size") + .map_or(10, |v| v.parse::().unwrap_or(10)); + for (key, ret) in query_results { let mut top_hits: HashMap = HashMap::default(); for row in ret.hits { @@ -906,27 +1120,65 @@ async fn values_v1( let key_num = top_hits.entry(key).or_insert(0); *key_num += num; } - let mut top_hits = top_hits.into_iter().collect::>(); - if no_count { - top_hits.sort_by(|a, b| a.0.cmp(&b.0)); - } else { - top_hits.sort_by(|a, b| b.1.cmp(&a.1)); - } - let top_hits = top_hits - .into_iter() - .take(size as usize) - .map(|(k, v)| { - let mut item = json::Map::new(); - item.insert("zo_sql_key".to_string(), json::Value::String(k)); - item.insert("zo_sql_num".to_string(), json::Value::Number(v.into())); - json::Value::Object(item) - }) - .collect::>(); - let mut field_value: json::Map = json::Map::new(); - field_value.insert("field".to_string(), json::Value::String(key)); - field_value.insert("values".to_string(), json::Value::Array(top_hits)); - hit_values.push(json::Value::Object(field_value)); + // Use a min heap (BinaryHeap with Reverse) to find top k elements + let mut min_heap: BinaryHeap> = + BinaryHeap::with_capacity(size as usize); + + if no_count { + // For alphabetical sorting, collect all entries first + let mut all_entries: Vec<_> = top_hits.into_iter().collect(); + all_entries.sort_by(|a, b| a.0.cmp(&b.0)); + all_entries.truncate(size as usize); + + let top_hits = all_entries + .into_iter() + .map(|(k, v)| { + let mut item = json::Map::new(); + item.insert("zo_sql_key".to_string(), json::Value::String(k)); + item.insert("zo_sql_num".to_string(), json::Value::Number(v.into())); + json::Value::Object(item) + }) + .collect::>(); + + let mut field_value: json::Map = json::Map::new(); + field_value.insert("field".to_string(), json::Value::String(key)); + field_value.insert("values".to_string(), json::Value::Array(top_hits)); + hit_values.push(json::Value::Object(field_value)); + } else { + // For value-based sorting, use a min heap to get top k elements + for (k, v) in top_hits { + if min_heap.len() < size as usize { + // If heap not full, just add + min_heap.push(Reverse((v, k))); + } else if !min_heap.is_empty() && v > min_heap.peek().unwrap().0.0 { + // If current value is larger than smallest in heap, replace it + min_heap.pop(); + min_heap.push(Reverse((v, k))); + } + } + + // Convert heap to vector and sort in descending order + let mut top_elements: Vec<_> = + min_heap.into_iter().map(|Reverse((v, k))| (k, v)).collect(); + top_elements.sort_by(|a, b| b.1.cmp(&a.1)); + + let top_hits = top_elements + .into_iter() + .map(|(k, v)| { + let mut item = json::Map::new(); + item.insert("zo_sql_key".to_string(), json::Value::String(k)); + item.insert("zo_sql_num".to_string(), json::Value::Number(v.into())); + json::Value::Object(item) + }) + .collect::>(); + + let mut field_value: json::Map = json::Map::new(); + field_value.insert("field".to_string(), json::Value::String(key)); + field_value.insert("values".to_string(), json::Value::Array(top_hits)); + hit_values.push(json::Value::Object(field_value)); + } + resp.scan_size = std::cmp::max(resp.scan_size, ret.scan_size); resp.scan_records = std::cmp::max(resp.scan_records, ret.scan_records); resp.cached_ratio = std::cmp::max(resp.cached_ratio, ret.cached_ratio); @@ -1317,4 +1569,4 @@ pub async fn search_history( .await; Ok(HttpResponse::Ok().json(search_res)) -} +} \ No newline at end of file diff --git a/src/handler/http/request/ws_v2/session.rs b/src/handler/http/request/ws_v2/session.rs index ab5724d5cc..0a55c9c10f 100644 --- a/src/handler/http/request/ws_v2/session.rs +++ b/src/handler/http/request/ws_v2/session.rs @@ -19,7 +19,7 @@ use actix_http::ws::{CloseCode, CloseReason}; use actix_ws::{AggregatedMessageStream, Session}; use config::{ get_config, - meta::websocket::{SearchEventReq, SearchResultType}, + meta::websocket::{SearchEventReq, SearchResultType, ValuesEventReq}, utils::time::now_micros, }; use futures::StreamExt; @@ -45,7 +45,7 @@ use crate::{ }, // router::http::ws_v2::types::StreamMessage, service::websocket_events::{ - WsClientEvents, WsServerEvents, handle_search_request, search_registry_utils::SearchState, + WsClientEvents, WsServerEvents, handle_search_request, handle_values_request, search_registry_utils::SearchState, sessions_cache_utils, }, }; @@ -262,6 +262,41 @@ pub async fn run( cleanup_and_close_session(&req_id, close_reason).await; } +/// Resolves user ID based on the execution mode and request data +/// +/// # Parameters +/// * `default_user_id` - Default user ID from the HTTP request +/// * `user_id_from_event` - Optional user ID from the event +/// +/// # Returns +/// * `Option` - Resolved user ID if successful, None otherwise +#[cfg(feature = "enterprise")] +async fn resolve_enterprise_user_id( + default_user_id: &str, + user_id_from_event: Option<&String>, +) -> Option { + if get_config().common.local_mode { + // Single node mode, use user_id from the HTTP request header + Some(default_user_id.to_string()) + } else { + // Cluster mode, try to determine user ID + // First check if we're running without router nodes + let router_nodes = get_cached_online_router_nodes().await; + if let Some(nodes) = router_nodes { + if nodes.is_empty() { + // Single node enterprise deployment + return Some(default_user_id.to_string()); + } + } + + // Next, try to use user_id from the event + if let Some(id) = user_id_from_event { + return Some(id.clone()); + } + None + } +} + /// Handle the incoming text message /// Text message is parsed into `WsClientEvents` and processed accordingly /// Depending on each event type, audit must be done @@ -285,38 +320,21 @@ pub async fn handle_text_message(user_id: &str, req_id: &str, msg: String, path: match client_msg { WsClientEvents::Search(ref search_req) => { #[allow(unused_mut)] - let mut user_id = user_id; + let mut user_id = user_id.to_string(); // verify user_id for handling stream permissions #[cfg(feature = "enterprise")] { - user_id = if get_config().common.local_mode { - // single node mode, use user_id from the http request header - user_id - } else { - // cluster mode, use user_id from the ws event added by the router - let router_nodes = get_cached_online_router_nodes().await; - let user_id = router_nodes.and_then(|nodes| { - if nodes.is_empty() { - // if router nodes are not found, use the user_id from the - // search request - // since its a single node, running enterprise - Some(user_id) - } else { - None - } - }); - if let Some(user_id) = user_id { - user_id - } else if let Some(user_id) = &search_req.user_id { - user_id - } else { + user_id = match resolve_enterprise_user_id( + &user_id, + search_req.user_id.as_ref(), + ).await { + Some(id) => id, + None => { log::error!("[WS_HANDLER]: User id not found in search request"); let err_res = WsServerEvents::error_response( - errors::Error::Message( - "User id not found in search request".to_string(), - ), + errors::Error::Message("User id not found in search request".to_string()), Some(req_id.to_string()), - Some(search_req.trace_id.clone()), + Some(search_req.trace_id.to_string()), Default::default(), ); let _ = send_message(req_id, err_res.to_json()).await; @@ -327,7 +345,40 @@ pub async fn handle_text_message(user_id: &str, req_id: &str, msg: String, path: handle_search_event( search_req, &search_req.org_id, - user_id, + &user_id, + req_id, + path.clone(), + ) + .await; + } + WsClientEvents::Values(ref values_req) => { + #[allow(unused_mut)] + let mut user_id = user_id.to_string(); + // verify user_id for handling stream permissions + #[cfg(feature = "enterprise")] + { + user_id = match resolve_enterprise_user_id( + &user_id, + values_req.user_id.as_ref(), + ).await { + Some(id) => id, + None => { + log::error!("[WS_HANDLER]: User id not found in values request"); + let err_res = WsServerEvents::error_response( + errors::Error::Message("User id not found in values request".to_string()), + Some(req_id.to_string()), + Some(values_req.trace_id.to_string()), + Default::default(), + ); + let _ = send_message(req_id, err_res.to_json()).await; + return; + } + }; + } + handle_values_event( + values_req, + &values_req.org_id, + &user_id, req_id, path.clone(), ) @@ -720,3 +771,122 @@ async fn cleanup_search_resources(trace_id: &str) { drop(w); log::debug!("[WS_HANDLER]: trace_id: {}, Resources cleaned up", trace_id); } + +// Main values handler +async fn handle_values_event( + values_req: &ValuesEventReq, + org_id: &str, + user_id: &str, + req_id: &str, + #[allow(unused_variables)] path: String, +) { + let (cancel_tx, mut cancel_rx) = mpsc::channel(1); + + let org_id = org_id.to_string(); + let user_id = user_id.to_string(); + let req_id = req_id.to_string(); + let trace_id = values_req.trace_id.clone(); + let trace_id_for_task = trace_id.clone(); + let values_req = values_req.clone(); + + #[cfg(feature = "enterprise")] + let is_audit_enabled = get_o2_config().common.audit_enabled; + + #[cfg(feature = "enterprise")] + let client_msg = WsClientEvents::Values(Box::new(values_req.clone())); + + // Register running values search BEFORE spawning the values task + let mut w = WS_SEARCH_REGISTRY.write().await; + w.insert( + trace_id.clone(), + SearchState::Running { + cancel_tx, + req_id: req_id.to_string(), + }, + ); + drop(w); + + // Create a vector to accumulate results + let mut accumulated_results: Vec = Vec::new(); + + // Spawn the values task + tokio::spawn(async move { + // Handle the values request + // If values search is cancelled, the task will exit + // Otherwise, the task will complete and the results will be sent to the client + // The task will also update the values state to completed + // The task will also cleanup the values search resources + tokio::select! { + values_result = handle_values_request( + &org_id, + &user_id, + &req_id, + values_req.clone(), + &mut accumulated_results, + ) => { + match values_result { + Ok(_) => { + let mut w = WS_SEARCH_REGISTRY.write().await; + if let Some(state) = w.get_mut(&trace_id_for_task) { + *state = SearchState::Completed { + req_id: req_id.to_string(), + }; + } + drop(w); + + // Add audit before closing + #[cfg(feature = "enterprise")] + if is_audit_enabled { + audit(AuditMessage { + user_email: user_id, + org_id, + _timestamp: chrono::Utc::now().timestamp(), + protocol: Protocol::Ws(WsMeta { + path: path.clone(), + message_type: client_msg.get_type(), + content: client_msg.to_json(), + close_reason: "".to_string(), + }), + }) + .await; + } + + cleanup_search_resources(&trace_id_for_task).await; + } + Err(e) => { + // Convert anyhow::Error to our Error type + let error = Error::Message(e.to_string()); + let _ = handle_search_error(error, &req_id, &trace_id_for_task).await; + + // Add audit before closing + #[cfg(feature = "enterprise")] + if is_audit_enabled { + audit(AuditMessage { + user_email: user_id, + org_id, + _timestamp: chrono::Utc::now().timestamp(), + protocol: Protocol::Ws(WsMeta { + path: path.clone(), + message_type: client_msg.get_type(), + content: client_msg.to_json(), + close_reason: "".to_string(), + }), + }) + .await; + } + + // Even if the values search is cancelled, we need to cleanup the resources + cleanup_search_resources(&trace_id_for_task).await; + } + } + } + _ = cancel_rx.recv() => { + // if values search is cancelled, update the state + // the cancel handler will close the session + + // Just cleanup resources when cancelled + cleanup_search_resources(&trace_id_for_task).await; + } + } + }); +} diff --git a/src/service/search/mod.rs b/src/service/search/mod.rs index 125edb8497..f1c352d336 100644 --- a/src/service/search/mod.rs +++ b/src/service/search/mod.rs @@ -696,6 +696,7 @@ pub async fn search_partition( }); } } + log::info!("[trace_id {trace_id}] max_query_range: {}, max_query_range_in_hour: {}", max_query_range, max_query_range_in_hour); let file_list_took = start.elapsed().as_millis() as usize; log::info!( @@ -710,9 +711,11 @@ pub async fn search_partition( response.partitions.push([req.start_time, req.end_time]); response.max_query_range = max_query_range_in_hour; response.histogram_interval = sql.histogram_interval; + log::info!("[trace_id {trace_id}] search_partition: returning single partition"); return Ok(response); }; + log::info!("[trace_id {trace_id}] search_partition: getting nodes"); let nodes = infra_cluster::get_cached_online_querier_nodes(Some(RoleGroup::Interactive)) .await .unwrap_or_default(); @@ -725,6 +728,7 @@ pub async fn search_partition( let (records, original_size) = files.iter().fold((0, 0), |(records, original_size), f| { (records + f.records, original_size + f.original_size) }); + let mut resp = search::SearchPartitionResponse { trace_id: trace_id.to_string(), file_num: files.len(), @@ -803,6 +807,7 @@ pub async fn search_partition( }) .unwrap_or(OrderBy::Desc); + log::debug!("[trace_id {trace_id}] total_secs: {}, partition_num: {}, step: {}, min_step: {}, is_histogram: {}", total_secs, part_num, step, min_step, is_histogram); // Create a partition generator let generator = partition::PartitionGenerator::new( min_step, diff --git a/src/service/search/partition.rs b/src/service/search/partition.rs index 4c41355cac..4e51f954a2 100644 --- a/src/service/search/partition.rs +++ b/src/service/search/partition.rs @@ -84,6 +84,8 @@ impl PartitionGenerator { let mini_partition_size_microseconds = std::cmp::min(self.mini_partition_duration_secs * 1_000_000, step as u64) as i64; + log::info!("mini_partition_size_microseconds: {}, step: {}, end_time: {}, start_time: {}", mini_partition_size_microseconds, step, end_time, start_time); + // Different behavior based on order if order_by == OrderBy::Desc { // For DESC order, mini partition at the end time diff --git a/src/service/websocket_events/mod.rs b/src/service/websocket_events/mod.rs index 2d85bd3ce9..0d1a5892e8 100644 --- a/src/service/websocket_events/mod.rs +++ b/src/service/websocket_events/mod.rs @@ -16,6 +16,8 @@ pub mod search; pub mod sort; pub mod utils; +pub mod values; pub use search::*; pub use utils::*; +pub use values::*; diff --git a/src/service/websocket_events/search.rs b/src/service/websocket_events/search.rs index d6e84b3d5e..ca98922dfa 100644 --- a/src/service/websocket_events/search.rs +++ b/src/service/websocket_events/search.rs @@ -18,13 +18,14 @@ use config::{ meta::{ search::{ PARTIAL_ERROR_RESPONSE_MESSAGE, Response, SearchEventType, SearchPartitionRequest, - SearchPartitionResponse, + SearchPartitionResponse, ValuesEventContext, }, sql::{OrderBy, resolve_stream_names}, websocket::{MAX_QUERY_RANGE_LIMIT_ERROR_MESSAGE, SearchEventReq, SearchResultType}, }, }; use infra::errors::{Error, ErrorCodes}; +use std::time::Instant; use tracing::Instrument; use super::sort::order_search_results; @@ -50,6 +51,10 @@ use crate::{ }, }; +use std::cmp::Reverse; +use std::collections::BinaryHeap; +use config::utils::json::{Map, Value}; + #[cfg(feature = "enterprise")] pub async fn handle_cancel(trace_id: &str, org_id: &str) -> WsServerEvents { match crate::service::search::cancel_query(org_id, trace_id).await { @@ -343,7 +348,7 @@ fn handle_partial_response(mut res: Response) -> Response { } #[allow(clippy::too_many_arguments)] -async fn handle_cache_responses_and_deltas( +pub async fn handle_cache_responses_and_deltas( req_id: &str, req: &SearchEventReq, trace_id: String, @@ -616,6 +621,14 @@ async fn process_delta( // `result_cache_ratio` will be 0 for delta search let result_cache_ratio = search_res.result_cache_ratio; + if req.search_type == SearchEventType::Values && req.values_event_context.is_some() { + log::debug!("Getting top k values for partition {idx}"); + let top_k_values = tokio::task::spawn_blocking(move || { + get_top_k_values(&search_res.hits, &req.values_event_context.clone().unwrap()) + }).await.unwrap(); + search_res.hits = top_k_values?; + } + let ws_search_res = WsServerEvents::SearchResponse { trace_id: trace_id.clone(), results: Box::new(search_res.clone()), @@ -799,7 +812,7 @@ async fn send_cached_responses( // Do partitioned search without cache #[allow(clippy::too_many_arguments)] -async fn do_partitioned_search( +pub async fn do_partitioned_search( req_id: &str, req: &mut SearchEventReq, trace_id: &str, @@ -852,7 +865,7 @@ async fn do_partitioned_search( let mut curr_res_size = 0; log::info!( - "[WS_SEARCH] Found {} partitions for trace_id: {}, partitions: {:#?}", + "[WS_SEARCH] Found {} partitions for trace_id: {}, partitions: {:?}", partitions.len(), trace_id, &partitions @@ -911,6 +924,16 @@ async fn do_partitioned_search( accumulated_results.push(SearchResultType::Search(search_res.clone())); } + if req.search_type == SearchEventType::Values && req.values_event_context.is_some() { + let instant = Instant::now(); + let top_k_values = tokio::task::spawn_blocking(move || { + get_top_k_values(&search_res.hits, &req.values_event_context.clone().unwrap()) + }).await.unwrap(); + search_res.hits = top_k_values?; + let duration = instant.elapsed(); + log::debug!("Top k values for partition {idx} took {:?}", duration); + } + // Send the cached response let ws_search_res = WsServerEvents::SearchResponse { trace_id: trace_id.to_string(), @@ -1003,7 +1026,7 @@ async fn send_partial_search_resp( Ok(()) } -async fn write_results_to_cache( +pub async fn write_results_to_cache( c_resp: MultiCachedQueryResponse, start_time: i64, end_time: i64, @@ -1088,3 +1111,85 @@ async fn write_results_to_cache( Ok(()) } + +/// This function will compute top k values for values request +/// This is a com +pub fn get_top_k_values(hits: &Vec, ctx: &ValuesEventContext) -> Result, Error> { + let mut top_k_values: Vec = Vec::new(); + + if ctx.field.is_empty() { + log::error!("Field is empty for values search"); + return Err(Error::Message("field is empty".to_string())); + } + + let k_limit = ctx.top_k.unwrap_or(10); + let no_count = ctx.no_count; + + let mut search_result_hits = Vec::new(); + for hit in hits { + let key = hit.get("zo_sql_key") + .and_then(|v| v.as_str().map(|s| s.to_string())) + .unwrap_or_default(); + let num = hit.get("zo_sql_num") + .and_then(|v| v.as_i64()) + .unwrap_or_default(); + search_result_hits.push((key, num)); + } + + if no_count { + // For alphabetical sorting, collect all entries first + let mut all_entries: Vec<_> = search_result_hits; + all_entries.sort_by(|a, b| a.0.cmp(&b.0)); + all_entries.truncate(k_limit as usize); + + let top_hits = all_entries + .into_iter() + .map(|(k, v)| { + let mut item = Map::new(); + item.insert("zo_sql_key".to_string(), Value::String(k)); + item.insert("zo_sql_num".to_string(), Value::Number(v.into())); + Value::Object(item) + }) + .collect::>(); + + let mut field_value: Map = Map::new(); + field_value.insert("field".to_string(), Value::String(ctx.field.clone())); + field_value.insert("values".to_string(), Value::Array(top_hits)); + top_k_values.push(Value::Object(field_value)); + } else { + // For value-based sorting, use a min heap to get top k elements + let mut min_heap: BinaryHeap> = BinaryHeap::with_capacity(k_limit as usize); + for (k, v) in search_result_hits { + if min_heap.len() < k_limit as usize { + // If heap not full, just add + min_heap.push(Reverse((v, k))); + } else if !min_heap.is_empty() && v > min_heap.peek().unwrap().0.0 { + // If current value is larger than smallest in heap, replace it + min_heap.pop(); + min_heap.push(Reverse((v, k))); + } + } + + // Convert heap to vector and sort in descending order + let mut top_elements: Vec<_> = + min_heap.into_iter().map(|Reverse((v, k))| (k, v)).collect(); + top_elements.sort_by(|a, b| b.1.cmp(&a.1)); + + let top_hits = top_elements + .into_iter() + .map(|(k, v)| { + let mut item = Map::new(); + item.insert("zo_sql_key".to_string(), Value::String(k)); + item.insert("zo_sql_num".to_string(), Value::Number(v.into())); + Value::Object(item) + }) + .collect::>(); + + let mut field_value: Map = Map::new(); + field_value.insert("field".to_string(), Value::String(ctx.field.clone())); + field_value.insert("values".to_string(), Value::Array(top_hits)); + top_k_values.push(Value::Object(field_value)); + } + + Ok(top_k_values) +} \ No newline at end of file diff --git a/src/service/websocket_events/utils.rs b/src/service/websocket_events/utils.rs index c7ecf32616..67dc5efc5a 100644 --- a/src/service/websocket_events/utils.rs +++ b/src/service/websocket_events/utils.rs @@ -14,7 +14,7 @@ // along with this program. If not, see . use actix_web::http::StatusCode; -use config::meta::websocket::SearchEventReq; +use config::meta::websocket::{SearchEventReq, ValuesEventReq}; use infra::errors; use serde::{Deserialize, Serialize}; use tokio_tungstenite::tungstenite; @@ -321,6 +321,7 @@ pub mod search_registry_utils { )] pub enum WsClientEvents { Search(Box), + Values(Box), #[cfg(feature = "enterprise")] Cancel { trace_id: String, @@ -341,6 +342,7 @@ impl WsClientEvents { pub fn get_type(&self) -> String { match self { WsClientEvents::Search(_) => "search", + WsClientEvents::Values(_) => "values", #[cfg(feature = "enterprise")] WsClientEvents::Cancel { .. } => "cancel", WsClientEvents::Benchmark { .. } => "benchmark", @@ -356,6 +358,7 @@ impl WsClientEvents { pub fn get_trace_id(&self) -> String { match &self { Self::Search(req) => req.trace_id.clone(), + Self::Values(req) => req.trace_id.clone(), #[cfg(feature = "enterprise")] Self::Cancel { trace_id, .. } => trace_id.clone(), Self::Benchmark { id } => id.clone(), @@ -366,6 +369,7 @@ impl WsClientEvents { pub fn is_valid(&self) -> bool { match self { Self::Search(req) => req.is_valid(), + Self::Values(req) => req.is_valid(), #[cfg(feature = "enterprise")] Self::Cancel { trace_id, org_id, .. @@ -381,6 +385,7 @@ impl WsClientEvents { pub fn append_user_id(&mut self, user_id: Option) { match self { Self::Search(req) => req.user_id = user_id, + Self::Values(req) => req.user_id = user_id, Self::Cancel { user_id: uid, .. } => *uid = user_id, _ => {} } diff --git a/src/service/websocket_events/values.rs b/src/service/websocket_events/values.rs new file mode 100644 index 0000000000..3ed8ebadbc --- /dev/null +++ b/src/service/websocket_events/values.rs @@ -0,0 +1,275 @@ +// Copyright 2025 OpenObserve Inc. +// +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU Affero General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// This program is distributed in the hope that it will be useful +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU Affero General Public License for more details. +// +// You should have received a copy of the GNU Affero General Public License +// along with this program. If not, see . + +use config::meta::{ + search::{SearchEventType, ValuesEventContext}, + websocket::{ValuesEventReq, SearchResultType}, +}; +use crate::handler::http::request::search::build_search_request_per_field; +use config::get_config; +use chrono::DateTime; + +#[cfg(feature = "enterprise")] +use crate::service::websocket_events::enterprise_utils; +use crate::{ + common::utils::stream::get_max_query_range, + handler::http::request::ws_v2::session::send_message, + service::{ + search::{cache, sql::Sql}, + websocket_events::{WsServerEvents, search::{ + do_partitioned_search, write_results_to_cache, handle_cache_responses_and_deltas + }}, + }, +}; + +pub async fn handle_values_request( + org_id: &str, + user_id: &str, + request_id: &str, + req: ValuesEventReq, + accumulated_results: &mut Vec, +) -> Result<(), anyhow::Error> { + let cfg = get_config(); + let trace_id = req.trace_id.clone(); + let stream_type = req.stream_type; + let payload = req.payload; + let start_time = payload.start_time; + let end_time = payload.end_time; + let stream_name = payload.stream_name.clone(); + let no_count = payload.no_count; + + log::info!( + "[WS_VALUES] trace_id: {} Received values request, start_time: {}, end_time: {}", + trace_id, + DateTime::from_timestamp_micros(start_time.unwrap_or(0) / 1_000).map_or("None".to_string(), |dt| dt.to_string()), + DateTime::from_timestamp_micros(end_time.unwrap_or(0) / 1_000).map_or("None".to_string(), |dt| dt.to_string()) + ); + + // Check permissions for each stream + #[cfg(feature = "enterprise")] + { + if let Err(e) = + enterprise_utils::check_permissions(&stream_name, stream_type, user_id, org_id).await + { + let err_res = WsServerEvents::error_response( + infra::errors::Error::Message(e), + Some(request_id.to_string()), + Some(trace_id), + Default::default(), + ); + send_message(request_id, err_res.to_json()).await?; + return Ok(()); + } + } + + // handle search result size + // use this size as K in top K values logic + let top_k = payload.size.or(Some(cfg.limit.query_default_limit)); + + // get values req query + let reqs = build_search_request_per_field(&payload, org_id, stream_type, &stream_name).await?; + + for (search_req, stream_type, field) in reqs { + // Convert the request query to SearchQuery type if needed + // Here we're converting from the search_req.query to the expected type + let search_query = search_req.query.clone().into(); + + let sql = Sql::new(&search_query, org_id, stream_type, search_req.search_type).await?; + let order_by = sql.order_by.first().map(|v| v.1).unwrap_or_default(); + + if search_req.query.from == 0 { + let c_resp = + cache::check_cache_v2(&trace_id, org_id, stream_type, &search_req, search_req.use_cache.unwrap_or(false)) + .await?; + let local_c_resp = c_resp.clone(); + let cached_resp = local_c_resp.cached_response; + let mut deltas = local_c_resp.deltas; + deltas.sort(); + deltas.dedup(); + + let cached_hits = cached_resp + .iter() + .fold(0, |acc, c| acc + c.cached_response.hits.len()); + + let c_start_time = cached_resp + .first() + .map(|c| c.response_start_time) + .unwrap_or_default(); + + let c_end_time = cached_resp + .last() + .map(|c| c.response_end_time) + .unwrap_or_default(); + + log::info!( + "[WS_VALUES] trace_id: {}, found cache responses len:{}, with hits: {}, cache_start_time: {:#?}, cache_end_time: {:#?}", + trace_id, + cached_resp.len(), + cached_hits, + c_start_time, + c_end_time + ); + + // handle cache responses and deltas + if !cached_resp.is_empty() && cached_hits > 0 { + // `max_query_range` is used initialize `remaining_query_range` + // set max_query_range to i64::MAX if it is 0, to ensure unlimited query range + // for cache only search + let max_query_range = + get_max_query_range(&[stream_name.clone()], org_id, user_id, stream_type).await; // hours + let remaining_query_range = if max_query_range == 0 { + i64::MAX + } else { + max_query_range + }; // hours + + // Convert SearchRequest to SearchEventReq for compatibility with search functions + let search_event_req = config::meta::websocket::SearchEventReq { + org_id: org_id.to_string(), + stream_type, + fallback_order_by_col: None, + search_type: search_req.search_type.unwrap_or(SearchEventType::Values), + search_event_context: None, + trace_id: trace_id.clone(), + use_cache: search_req.use_cache.unwrap_or(false), + payload: search_req.clone(), + user_id: Some(user_id.to_string()), + time_offset: None, + values_event_context: Some(ValuesEventContext { + top_k, + no_count, + field, + }), + }; + + // Step 1(a): handle cache responses & query the deltas + handle_cache_responses_and_deltas( + request_id, + &search_event_req, + trace_id.clone(), + config::meta::sql::MAX_LIMIT, + cached_resp, + deltas, + accumulated_results, + user_id, + max_query_range, + remaining_query_range, + &order_by, + ) + .await?; + } else { + // Step 2: Search without cache + // no caches found process req directly + log::info!( + "[WS_VALUES] trace_id: {} No cache found, processing search request", + trace_id + ); + let max_query_range = + get_max_query_range(&[stream_name.clone()], org_id, user_id, stream_type).await; // hours + + // Convert SearchRequest to SearchEventReq for compatibility with search functions + let mut search_event_req = config::meta::websocket::SearchEventReq { + org_id: org_id.to_string(), + stream_type, + fallback_order_by_col: None, + search_type: search_req.search_type.unwrap_or(SearchEventType::Values), + search_event_context: None, + trace_id: trace_id.clone(), + use_cache: search_req.use_cache.unwrap_or(false), + payload: search_req.clone(), + user_id: Some(user_id.to_string()), + time_offset: None, + values_event_context: Some(ValuesEventContext { + top_k, + no_count, + field, + }), + }; + + do_partitioned_search( + request_id, + &mut search_event_req, + &trace_id, + config::meta::sql::MAX_LIMIT, + user_id, + accumulated_results, + max_query_range, + ) + .await?; + } + // Step 3: Write to results cache + // cache only if from is 0 and is not an aggregate_query + if search_req.query.from == 0 { + // Safely unwrap Option values with defaults + let safe_start_time = start_time.unwrap_or(0); + let safe_end_time = end_time.unwrap_or(0); + + write_results_to_cache(c_resp, safe_start_time, safe_end_time, accumulated_results) + .await + .map_err(|e| { + log::error!( + "[WS_VALUES] trace_id: {}, Error writing results to cache: {:?}", + trace_id, + e + ); + anyhow::anyhow!("Error writing results to cache: {}", e) + })?; + } + } else { + // Step 4: Search without cache for req with from > 0 + let max_query_range = + get_max_query_range(&[stream_name.clone()], org_id, user_id, stream_type).await; // hours + + // Convert SearchRequest to SearchEventReq for compatibility with search functions + let mut search_event_req = config::meta::websocket::SearchEventReq { + org_id: org_id.to_string(), + stream_type, + fallback_order_by_col: None, + search_type: search_req.search_type.unwrap_or(SearchEventType::Values), + search_event_context: None, + trace_id: trace_id.clone(), + use_cache: search_req.use_cache.unwrap_or(false), + payload: search_req.clone(), + user_id: Some(user_id.to_string()), + time_offset: None, + values_event_context: Some(ValuesEventContext { + top_k, + no_count, + field, + }), + }; + + do_partitioned_search( + request_id, + &mut search_event_req, + &trace_id, + config::meta::sql::MAX_LIMIT, + user_id, + accumulated_results, + max_query_range, + ) + .await?; + } + } + + // Once all searches are complete, write the accumulated results to a file + log::info!("[WS_VALUES] trace_id {} all searches completed", trace_id); + let end_res = WsServerEvents::End { + trace_id: Some(trace_id.clone()), + }; + send_message(request_id, end_res.to_json()).await?; + + Ok(()) +} diff --git a/web/src/composables/useLogs.ts b/web/src/composables/useLogs.ts index d5caba659d..7a4c1abfde 100644 --- a/web/src/composables/useLogs.ts +++ b/web/src/composables/useLogs.ts @@ -4899,7 +4899,7 @@ const useLogs = () => { const buildWebSocketPayload = ( queryReq: SearchRequestPayload, isPagination: boolean, - type: "search" | "histogram" | "pageCount", + type: "search" | "histogram" | "pageCount" | "values", meta?: any, ) => { const { traceId } = generateTraceContext(); @@ -4907,7 +4907,7 @@ const useLogs = () => { const payload: { queryReq: SearchRequestPayload; - type: "search" | "histogram" | "pageCount"; + type: "search" | "histogram" | "pageCount" | "values"; isPagination: boolean; traceId: string; org_id: string; diff --git a/web/src/plugins/logs/IndexList.vue b/web/src/plugins/logs/IndexList.vue index 5373798462..68ee7d3afb 100644 --- a/web/src/plugins/logs/IndexList.vue +++ b/web/src/plugins/logs/IndexList.vue @@ -254,6 +254,7 @@ along with this program. If not, see . @before-show=" (event: any) => openFilterCreator(event, props.row) " + @before-hide="(event: any) => cancelTraceId(props.row.name)" >