From 78efd236ed7685fa111e8a33ba18ff01a7a502d4 Mon Sep 17 00:00:00 2001 From: Simha G <149056463+yaga-simha@users.noreply.github.com> Date: Wed, 15 Oct 2025 12:20:36 +0530 Subject: [PATCH] feat: Introduce query optimization recommendation service to background jobs (#8791) ### **User description** # PR Type Enhancement # Description This PR introduces a configurable query recommendation service / back ground job that will take org wise usage data from metadata and generate recommendations for the same and reingests this data into _meta org's query_recommendation stream. ___ ### **PR Type** Enhancement ___ ### **Description** - Add enterprise query recommendation background job - Expose config for schedule and limits - Map meta stream to enterprise type - Wire job init under enterprise feature ___ ### Diagram Walkthrough ```mermaid flowchart LR CFG["Config limits: recommendation params"] JOB["Background job: QueryRecommendationService"] CTX["OptimizerContext (engine impl)"] ORGS["List org IDs"] STREAMS["Get stream settings"] SEARCH["Run usage search"] INGEST["Ingest recommendations to _meta/query_recommendations"] CFG -- interval/duration/top_k --> JOB JOB -- uses --> CTX CTX -- fetch --> ORGS CTX -- fetch --> STREAMS CTX -- analyze --> SEARCH CTX -- publish --> INGEST ```

File Walkthrough
Relevant files
Enhancement
stream.rs
Add enterprise Stream conversion for recommendations         
src/common/meta/stream.rs
  • Gate enterprise import behind feature flag
  • Implement From to enterprise meta Stream
 +14/-0   
mod.rs
Wire query recommendation job into init                                   
src/job/mod.rs
  • Add enterprise-only module include
  • Spawn recommendation job at init
 +5/-0     
query_optimization_recommendation.rs
Implement enterprise query recommendation service job       
src/job/query_optimization_recommendation.rs
  • Add enterprise-only background job implementation
  • Implement engine trait with org/streams/search
  • Ingest recommendations into meta stream
  • Run service using configurable schedule
 +124/-0 
Configuration changes
config.rs
Configure query recommendation scheduling and limits         
src/config/src/config.rs
  • Add recommendation duration config
  • Add analysis interval config
  • Add top-k limit config
  • Set sensible defaults via env vars
 +7/-0     
___ Co-authored-by: OpenObserve --- src/common/meta/stream.rs | 14 +++ src/config/src/config.rs | 7 ++ src/job/mod.rs | 5 + src/job/query_optimization_recommendation.rs | 124 +++++++++++++++++++ 4 files changed, 150 insertions(+) create mode 100644 src/job/query_optimization_recommendation.rs diff --git a/src/common/meta/stream.rs b/src/common/meta/stream.rs index bfc5ae82ea..b9d74ce1e5 100644 --- a/src/common/meta/stream.rs +++ b/src/common/meta/stream.rs @@ -24,6 +24,8 @@ use config::{ utils::json, }; use datafusion::arrow::datatypes::Schema; +#[cfg(feature = "enterprise")] +use o2_enterprise::enterprise; use serde::{Deserialize, Serialize}; use utoipa::ToSchema; @@ -46,6 +48,18 @@ pub struct Stream { pub is_derived: Option, } +#[cfg(feature = "enterprise")] +impl From for enterprise::recommendations::meta::Stream { + fn from(value: Stream) -> Self { + Self { + name: value.name, + stream_type: value.stream_type, + stats: value.stats, + settings: value.settings, + } + } +} + #[derive(Clone, Debug, Serialize, Deserialize, ToSchema)] pub struct StreamCreate { pub fields: Vec, diff --git a/src/config/src/config.rs b/src/config/src/config.rs index 594fcb3643..8ad1c0708e 100644 --- a/src/config/src/config.rs +++ b/src/config/src/config.rs @@ -1278,6 +1278,13 @@ pub struct Limit { pub query_partition_by_secs: usize, #[env_config(name = "ZO_QUERY_GROUP_BASE_SPEED", default = 768)] // MB/s/core pub query_group_base_speed: usize, + // Default Config: Run Query Recommendation Analysis for last one hour for every hour + #[env_config(name = "ZO_QUERY_RECOMMENDATION_DURATION", default = 3600000000)] + pub query_recommendation_duration: i64, + #[env_config(name = "ZO_QUERY_RECOMMENDATION_INTERVAL", default = 3600000000)] + pub query_recommendation_analysis_interval: i64, + #[env_config(name = "ZO_QUERY_RECOMMENDATION_TOP_K", default = 128)] + pub query_recommendation_top_k: usize, #[env_config(name = "ZO_INGEST_ALLOWED_UPTO", default = 5)] // in hours - in past pub ingest_allowed_upto: i64, #[env_config(name = "ZO_INGEST_ALLOWED_IN_FUTURE", default = 24)] // in hours - in future diff --git a/src/job/mod.rs b/src/job/mod.rs index ef0e7440a9..8c2e5a2f2b 100644 --- a/src/job/mod.rs +++ b/src/job/mod.rs @@ -46,6 +46,8 @@ pub(crate) mod pipeline; mod pipeline_error_cleanup; mod promql; mod promql_self_consume; +#[cfg(feature = "enterprise")] +mod query_optimization_recommendation; mod stats; pub use file_downloader::{download_from_node, queue_download}; @@ -324,6 +326,9 @@ pub async fn init() -> Result<(), anyhow::Error> { } } + #[cfg(feature = "enterprise")] + tokio::task::spawn(async move { query_optimization_recommendation::run().await }); + // Shouldn't serve request until initialization finishes log::info!("Job initialization complete"); diff --git a/src/job/query_optimization_recommendation.rs b/src/job/query_optimization_recommendation.rs new file mode 100644 index 0000000000..baf7efb4bd --- /dev/null +++ b/src/job/query_optimization_recommendation.rs @@ -0,0 +1,124 @@ +// 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 . + +#![cfg(feature = "enterprise")] + +use std::{pin::Pin, sync::Arc}; + +use config::{META_ORG_ID, get_config, meta::stream::StreamType}; +use o2_enterprise::enterprise::recommendations::{ + meta::{OptimiserRecommendation, Stream}, + service::{QueryRecommendationEngine, QueryRecommendationService}, +}; +use proto::cluster_rpc::{ + IngestionData, IngestionRequest, IngestionResponse, IngestionType, ingest_server::Ingest, +}; + +use crate::{ + handler::grpc::request::ingest::Ingester, + service::{db::organization, search::search, stream::get_streams}, +}; + +#[derive(Clone)] +struct OptimizerContext; + +impl QueryRecommendationEngine for OptimizerContext { + fn get_all_org_ids( + &self, + limit: Option, + ) -> Pin, anyhow::Error>> + Send>> { + Box::pin(async move { + Ok(organization::list(limit) + .await? + .into_iter() + .map(|org| org.name) + .collect()) + }) + } + + fn get_all_stream_settings( + &self, + org_id: String, + ) -> Pin, anyhow::Error>> + Send>> { + Box::pin(async move { + Ok(get_streams(&org_id, None, true, None) + .await + .into_iter() + .map(Into::into) + .collect()) + }) + } + + fn search( + &self, + org_id: String, + request: config::meta::search::Request, + ) -> Pin> + Send>> + { + Box::pin(async move { + search("", &org_id, StreamType::Logs, None, &request) + .await + .map_err(Into::into) + }) + } + + fn ingest_recommendations( + &self, + recommendations: Vec, + ) -> Pin> + Send>> { + Box::pin(async move { + let ingester = Ingester::default(); + let request = IngestionRequest { + org_id: META_ORG_ID.to_string(), + stream_type: StreamType::Logs.to_string(), + stream_name: "query_recommendations".to_string(), + data: Some(IngestionData { + data: serde_json::to_vec_pretty(&recommendations).map_err(|e| { + anyhow::anyhow!("Recommendation serialization failed. Error={:?}", e) + })?, + }), + ingestion_type: Some(IngestionType::Json as i32), + metadata: None, + }; + Ok(ingester + .ingest(tonic::Request::new(request)) + .await? + .into_inner() + .into()) + }) + } +} + +pub async fn run() { + let cfg = get_config(); + let query_recommendation_service = QueryRecommendationService { + ctx: Arc::new(OptimizerContext), + query_recommendation_analysis_interval: cfg.limit.query_recommendation_analysis_interval, + query_recommendation_duration: cfg.limit.query_recommendation_duration, + query_recommendation_top_k: cfg.limit.query_recommendation_top_k, + }; + let _ = query_recommendation_service + .run() + .await + .inspect_err(|e| { + log::error!( + "Recommendation service stopped with an error: Error={:?}", + e + ); + }) + .inspect(|_| { + log::warn!("Recommendation service quietly ended without an error!"); + }); +}