mirror of
https://github.com/openobserve/openobserve.git
synced 2026-11-06 00:09:22 +00:00
refactor: make search engine sources crate-independent (#13234)
## Summary - add `src/vrl` as a leaf package named `openobserve-vrl` with Rust path `openobserve_vrl` - move `QUERY_FUNCTIONS`, enrichment-table ownership, GeoIP/MaxMind table registration, VRL compiler configuration, and shared runtime helpers into `openobserve-vrl` - let ingestion, jobs, core, and search use the shared VRL integration without depending on one another - keep DataFusion UDF registration direct through `Sql.org_id`; no `QueryExecutionContext`, `PreparedQueryTransform`, or `prepare_query_transforms` - move the cipher registry, query parameters, and file-cache helpers to `search` - move the generic file downloader implementation to `infra` - add feature forwarding for the `search` crate - keep compatibility re-exports for existing API, ingestion, gRPC, jobs, and Enterprise call sites Enterprise companion: https://github.com/openobserve/o2-enterprise/pull/2196
This commit is contained in:
Generated
+23
-3
@@ -6770,6 +6770,7 @@ dependencies = [
|
||||
"o2_enterprise",
|
||||
"o2_openfga",
|
||||
"object_store",
|
||||
"openobserve-vrl",
|
||||
"opentelemetry",
|
||||
"opentelemetry-proto",
|
||||
"parking_lot 0.12.5",
|
||||
@@ -6842,6 +6843,7 @@ dependencies = [
|
||||
"o2_enterprise",
|
||||
"o2_openfga",
|
||||
"openobserve-core",
|
||||
"openobserve-vrl",
|
||||
"prometheus",
|
||||
"proto",
|
||||
"regex",
|
||||
@@ -6853,6 +6855,18 @@ dependencies = [
|
||||
"wal",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "openobserve-vrl"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"async-trait",
|
||||
"config",
|
||||
"enrichment",
|
||||
"serde_json",
|
||||
"tokio",
|
||||
"vrl",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "openssl-probe"
|
||||
version = "0.2.1"
|
||||
@@ -9211,14 +9225,20 @@ version = "0.1.0"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"arrow-schema",
|
||||
"async-trait",
|
||||
"chrono",
|
||||
"config",
|
||||
"enrichment",
|
||||
"datafusion",
|
||||
"hashbrown 0.16.1",
|
||||
"infra",
|
||||
"log",
|
||||
"o2_enterprise",
|
||||
"openobserve-vrl",
|
||||
"parking_lot 0.12.5",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"tokio",
|
||||
"tracing",
|
||||
"utoipa",
|
||||
"vrl",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
||||
@@ -239,6 +239,7 @@ members = [
|
||||
"src/api_query",
|
||||
"src/api_common",
|
||||
"src/core",
|
||||
"src/vrl",
|
||||
"src/search",
|
||||
"src/jobs",
|
||||
"src/config",
|
||||
@@ -269,6 +270,7 @@ flight = { path = "src/flight" }
|
||||
tantivy_utils = { path = "src/tantivy_utils" }
|
||||
web = { path = "src/web" }
|
||||
openobserve-core = { path = "src/core" }
|
||||
openobserve-vrl = { path = "src/vrl" }
|
||||
search = { path = "src/search" }
|
||||
openobserve-api = { path = "src/api" }
|
||||
openobserve-api-management = { path = "src/api_management" }
|
||||
|
||||
@@ -28,6 +28,7 @@ exceptions = [
|
||||
{ name = "openobserve-api-query", allow = ["AGPL-3.0"] },
|
||||
{ name = "openobserve-core", allow = ["AGPL-3.0"] },
|
||||
{ name = "openobserve-jobs", allow = ["AGPL-3.0"] },
|
||||
{ name = "openobserve-vrl", allow = ["AGPL-3.0"] },
|
||||
{ name = "search", allow = ["AGPL-3.0"] },
|
||||
# Empty dependency-resolution stubs for the private enterprise crates.
|
||||
{ name = "o2_dex", allow = ["AGPL-3.0"] },
|
||||
|
||||
@@ -61,6 +61,7 @@ pub static COMMIT_HASH: &str = env!("GIT_COMMIT_HASH");
|
||||
pub static BUILD_DATE: &str = env!("GIT_BUILD_DATE");
|
||||
|
||||
pub const META_ORG_ID: &str = "_meta";
|
||||
pub const DEFAULT_ORG: &str = "default";
|
||||
|
||||
pub const MMDB_CITY_FILE_NAME: &str = "GeoLite2-City.mmdb";
|
||||
pub const MMDB_ASN_FILE_NAME: &str = "GeoLite2-ASN.mmdb";
|
||||
|
||||
+3
-1
@@ -20,15 +20,17 @@ enterprise = [
|
||||
"search/enterprise",
|
||||
]
|
||||
cloud = ["infra/cloud", "enterprise", "o2_enterprise/cloud"]
|
||||
vectorscan = ["o2_enterprise/vectorscan"]
|
||||
vectorscan = ["o2_enterprise/vectorscan", "search/vectorscan"]
|
||||
vortex = [
|
||||
"dep:vortex",
|
||||
"dep:vortex-datafusion",
|
||||
"config/vortex",
|
||||
"o2_enterprise?/vortex",
|
||||
"search/vortex",
|
||||
]
|
||||
|
||||
[dependencies]
|
||||
openobserve-vrl.workspace = true
|
||||
search.workspace = true
|
||||
axum.workspace = true
|
||||
http.workspace = true
|
||||
|
||||
@@ -13,12 +13,11 @@
|
||||
// You should have received a copy of the GNU Affero General Public License
|
||||
// along with this program. If not, see <http://www.gnu.org/licenses/>.
|
||||
|
||||
pub use ::search::cipher::registry;
|
||||
use o2_enterprise::enterprise::cipher::http_repr::HttpKey;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use utoipa::ToSchema;
|
||||
|
||||
pub mod registry;
|
||||
|
||||
#[derive(Clone, Debug, Deserialize, ToSchema)]
|
||||
pub struct KeyAddRequest {
|
||||
pub name: String,
|
||||
|
||||
@@ -21,7 +21,6 @@ use config::{
|
||||
alerts::alert::Alert,
|
||||
destinations::{Destination, Template},
|
||||
folder::Folder,
|
||||
function::Transform,
|
||||
pipeline::Pipeline,
|
||||
promql::ClusterLeader,
|
||||
ratelimit::CachedUserRoles,
|
||||
@@ -33,24 +32,18 @@ use config::{
|
||||
use dashmap::DashMap;
|
||||
use hashbrown::HashMap;
|
||||
use infra::table::short_urls::ShortUrlRecord;
|
||||
use parking_lot::RwLock;
|
||||
pub use search::enrichment::ENRICHMENT_TABLES;
|
||||
use vector_enrichment::TableRegistry;
|
||||
pub use openobserve_vrl::{QUERY_FUNCTIONS, enrichment::ENRICHMENT_TABLES};
|
||||
|
||||
use crate::{
|
||||
common::meta::{
|
||||
maxmind::MaxmindClient,
|
||||
organization::{OrgStatus, Organization, OrganizationSetting},
|
||||
},
|
||||
service::{
|
||||
db::scheduler as db_scheduler, enrichment_table::geoip::Geoip,
|
||||
pipeline::batch_execution::ExecutablePipeline,
|
||||
},
|
||||
service::{db::scheduler as db_scheduler, pipeline::batch_execution::ExecutablePipeline},
|
||||
};
|
||||
|
||||
// global cache variables
|
||||
pub static KVS: Lazy<RwHashMap<String, bytes::Bytes>> = Lazy::new(Default::default);
|
||||
pub static QUERY_FUNCTIONS: Lazy<RwHashMap<String, Transform>> = Lazy::new(DashMap::default);
|
||||
pub static USERS: Lazy<RwHashMap<String, infra::table::users::UserRecord>> =
|
||||
Lazy::new(DashMap::default);
|
||||
pub static ORG_USERS: Lazy<RwHashMap<String, infra::table::org_users::OrgUserRecord>> =
|
||||
@@ -74,22 +67,9 @@ pub static REALTIME_ALERT_TRIGGERS: Lazy<RwAHashMap<String, db_scheduler::Trigge
|
||||
Lazy::new(Default::default);
|
||||
pub static ALERTS_TEMPLATES: Lazy<RwHashMap<String, Template>> = Lazy::new(Default::default);
|
||||
pub static DESTINATIONS: Lazy<RwHashMap<String, Destination>> = Lazy::new(Default::default);
|
||||
pub static ENRICHMENT_REGISTRY: Lazy<Arc<TableRegistry>> =
|
||||
Lazy::new(|| Arc::new(TableRegistry::default()));
|
||||
|
||||
pub static MAXMIND_DB_CLIENT: Lazy<Arc<tokio::sync::RwLock<Option<MaxmindClient>>>> =
|
||||
Lazy::new(|| Arc::new(tokio::sync::RwLock::new(None)));
|
||||
|
||||
pub static GEOIP_CITY_TABLE: Lazy<Arc<RwLock<Option<Geoip>>>> =
|
||||
Lazy::new(|| Arc::new(RwLock::new(None)));
|
||||
|
||||
pub static GEOIP_ASN_TABLE: Lazy<Arc<RwLock<Option<Geoip>>>> =
|
||||
Lazy::new(|| Arc::new(RwLock::new(None)));
|
||||
|
||||
#[cfg(feature = "enterprise")]
|
||||
pub static GEOIP_ENT_TABLE: Lazy<Arc<RwLock<Option<Geoip>>>> =
|
||||
Lazy::new(|| Arc::new(RwLock::new(None)));
|
||||
|
||||
pub static STREAM_EXECUTABLE_PIPELINES: Lazy<RwAHashMap<StreamParams, Vec<ExecutablePipeline>>> =
|
||||
Lazy::new(Default::default);
|
||||
pub static PIPELINE_STREAM_MAPPING: Lazy<RwAHashMap<String, StreamParams>> =
|
||||
@@ -139,10 +119,6 @@ mod tests {
|
||||
// Test Arc<RwLock> variables
|
||||
let users_rum_token = USERS_RUM_TOKEN.clone();
|
||||
assert_eq!(users_rum_token.len(), 0);
|
||||
|
||||
// Test enrichment registry
|
||||
let registry = ENRICHMENT_REGISTRY.clone();
|
||||
assert!(!std::ptr::eq(registry.as_ref(), std::ptr::null()));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -258,29 +234,6 @@ mod tests {
|
||||
assert!(read_guard.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_geoip_tables() {
|
||||
// Test GEOIP_CITY_TABLE
|
||||
let city_table = GEOIP_CITY_TABLE.clone();
|
||||
let read_guard = city_table.read();
|
||||
assert!(read_guard.is_none());
|
||||
drop(read_guard);
|
||||
|
||||
// Test GEOIP_ASN_TABLE
|
||||
let asn_table = GEOIP_ASN_TABLE.clone();
|
||||
let read_guard = asn_table.read();
|
||||
assert!(read_guard.is_none());
|
||||
drop(read_guard);
|
||||
|
||||
#[cfg(feature = "enterprise")]
|
||||
{
|
||||
// Test GEOIP_ENT_TABLE (only available with enterprise feature)
|
||||
let ent_table = GEOIP_ENT_TABLE.clone();
|
||||
let read_guard = ent_table.read();
|
||||
assert!(read_guard.is_none());
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_dashmap_operations() {
|
||||
// Test basic DashMap operations on KVS
|
||||
@@ -303,6 +256,8 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_query_functions_cache() {
|
||||
use config::meta::function::Transform;
|
||||
|
||||
let test_key = "test_function".to_string();
|
||||
let test_transform = Transform {
|
||||
function: "test_function".to_string(),
|
||||
|
||||
@@ -13,14 +13,13 @@
|
||||
// You should have received a copy of the GNU Affero General Public License
|
||||
// along with this program. If not, see <http://www.gnu.org/licenses/>.
|
||||
|
||||
pub use config::DEFAULT_ORG;
|
||||
use config::{meta::user::UserRole, stats::MemorySize};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use utoipa::ToSchema;
|
||||
|
||||
#[cfg(feature = "cloud")]
|
||||
use super::user::InviteStatus;
|
||||
|
||||
pub const DEFAULT_ORG: &str = "default";
|
||||
pub const CUSTOM: &str = "custom";
|
||||
pub const USER_DEFAULT: &str = "user_default";
|
||||
pub const THRESHOLD: i64 = 9383939382;
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
// 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
|
||||
// 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.
|
||||
@@ -13,160 +13,6 @@
|
||||
// You should have received a copy of the GNU Affero General Public License
|
||||
// along with this program. If not, see <http://www.gnu.org/licenses/>.
|
||||
|
||||
use std::collections::HashMap;
|
||||
|
||||
use config::{
|
||||
GEO_IP_ASN_ENRICHMENT_TABLE, GEO_IP_CITY_ENRICHMENT_TABLE, meta::function::VRLCompilerConfig,
|
||||
pub use openobserve_vrl::{
|
||||
get_all_transform_keys, get_enrichment_tables, get_vrl_compiler_config, init_vrl_runtime,
|
||||
};
|
||||
use vector_enrichment::{Table, TableRegistry};
|
||||
|
||||
use crate::common::{
|
||||
infra::config::{ENRICHMENT_TABLES, GEOIP_ASN_TABLE, GEOIP_CITY_TABLE},
|
||||
meta::organization::DEFAULT_ORG,
|
||||
};
|
||||
|
||||
pub async fn get_all_transform_keys(org_id: &str) -> Vec<String> {
|
||||
let org_key = &format!("{org_id}/");
|
||||
|
||||
crate::common::infra::config::QUERY_FUNCTIONS
|
||||
.clone()
|
||||
.iter()
|
||||
.filter_map(|transform| {
|
||||
let key = transform.key();
|
||||
key.strip_prefix(org_key).map(|x| x.to_string())
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub fn init_vrl_runtime() -> vrl::compiler::runtime::Runtime {
|
||||
vrl::compiler::runtime::Runtime::new(vrl::prelude::state::RuntimeState::default())
|
||||
}
|
||||
|
||||
pub fn get_vrl_compiler_config(org_id: &str) -> VRLCompilerConfig {
|
||||
let en_tables = ENRICHMENT_TABLES.clone();
|
||||
let mut functions = vrl::stdlib::all();
|
||||
functions.append(&mut vector_enrichment::vrl_functions());
|
||||
let registry = TableRegistry::default();
|
||||
let mut tables: HashMap<String, Box<dyn Table + Send + Sync>> = HashMap::new();
|
||||
|
||||
for table in en_tables.iter() {
|
||||
if table.org_id == org_id || table.org_id == DEFAULT_ORG {
|
||||
tables.insert(
|
||||
table.stream_name.to_owned(),
|
||||
Box::new(table.value().clone()),
|
||||
);
|
||||
}
|
||||
}
|
||||
drop(en_tables);
|
||||
|
||||
if let Some(v) = GEOIP_CITY_TABLE.read().as_ref() {
|
||||
tables.insert(GEO_IP_CITY_ENRICHMENT_TABLE.to_owned(), Box::new(v.clone()));
|
||||
}
|
||||
|
||||
if let Some(v) = GEOIP_ASN_TABLE.read().as_ref() {
|
||||
tables.insert(GEO_IP_ASN_ENRICHMENT_TABLE.to_owned(), Box::new(v.clone()));
|
||||
}
|
||||
|
||||
#[cfg(feature = "enterprise")]
|
||||
if o2_enterprise::enterprise::common::config::get_config()
|
||||
.common
|
||||
.enable_enterprise_mmdb
|
||||
{
|
||||
let geoip_ent = crate::common::infra::config::GEOIP_ENT_TABLE.read();
|
||||
if let Some(table) = geoip_ent.as_ref() {
|
||||
tables.insert(
|
||||
o2_enterprise::enterprise::common::config::GEO_IP_ENTERPRISE_ENRICHMENT_TABLE
|
||||
.to_owned(),
|
||||
Box::new(table.clone()),
|
||||
);
|
||||
}
|
||||
};
|
||||
|
||||
registry.load(tables);
|
||||
let mut config = vrl::compiler::CompileConfig::default();
|
||||
config.set_custom(registry);
|
||||
VRLCompilerConfig { config, functions }
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_init_vrl_runtime() {
|
||||
let runtime = init_vrl_runtime();
|
||||
assert!(runtime.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_get_vrl_compiler_config() {
|
||||
let config = get_vrl_compiler_config("default");
|
||||
assert!(!config.functions.is_empty());
|
||||
assert!(config.config.get_custom::<TableRegistry>().is_some());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_get_all_transform_keys() {
|
||||
// Setup test data
|
||||
let test_org = "test_org";
|
||||
let test_key = format!("{test_org}/test_transform");
|
||||
crate::common::infra::config::QUERY_FUNCTIONS.insert(
|
||||
test_key.clone(),
|
||||
config::meta::function::Transform {
|
||||
name: test_key.clone(),
|
||||
function: ".".to_string(),
|
||||
params: "row".to_string(),
|
||||
num_args: 1,
|
||||
trans_type: Some(0),
|
||||
streams: Some(vec![config::meta::function::StreamOrder {
|
||||
stream: "test".to_string(),
|
||||
order: 1,
|
||||
stream_type: config::meta::stream::StreamType::Logs,
|
||||
is_removed: false,
|
||||
apply_before_flattening: false,
|
||||
}]),
|
||||
},
|
||||
);
|
||||
|
||||
// Test the function
|
||||
let keys = get_all_transform_keys(test_org).await;
|
||||
assert!(keys.contains(&"test_transform".to_string()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_get_vrl_compiler_config_with_enrichment_tables() {
|
||||
// Setup test data
|
||||
let test_org = "test_org";
|
||||
let test_stream = "test_stream";
|
||||
let table_name = "test_table";
|
||||
let en_table = crate::service::enrichment::StreamTable {
|
||||
org_id: test_org.to_string(),
|
||||
stream_name: test_stream.to_string(),
|
||||
data: Arc::new(vec![]),
|
||||
};
|
||||
ENRICHMENT_TABLES.insert(table_name.to_string(), en_table);
|
||||
|
||||
// Test the function
|
||||
let config = get_vrl_compiler_config(test_org);
|
||||
let registry = config.config.get_custom::<TableRegistry>().unwrap();
|
||||
|
||||
// Verify enrichment table is loaded
|
||||
assert!(registry.table_ids().contains(&test_stream.to_string()));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_get_all_transform_keys_empty() {
|
||||
let keys = get_all_transform_keys("nonexistent_org").await;
|
||||
assert!(keys.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_get_vrl_compiler_config_empty() {
|
||||
let config = get_vrl_compiler_config("nonexistent_org");
|
||||
let registry = config.config.get_custom::<TableRegistry>().unwrap();
|
||||
assert!(registry.table_ids().is_empty());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -13,720 +13,4 @@
|
||||
// You should have received a copy of the GNU Affero General Public License
|
||||
// along with this program. If not, see <http://www.gnu.org/licenses/>.
|
||||
|
||||
use std::{
|
||||
collections::VecDeque,
|
||||
sync::{Arc, LazyLock as Lazy},
|
||||
};
|
||||
|
||||
use config::{
|
||||
cluster::LOCAL_NODE,
|
||||
get_config,
|
||||
meta::cluster::{Role, RoleGroup, get_internal_grpc_token},
|
||||
metrics,
|
||||
utils::time::{day_micros, now_micros},
|
||||
};
|
||||
use futures_util::StreamExt;
|
||||
use hashbrown::{HashMap, HashSet};
|
||||
use infra::{cache::file_data, cluster};
|
||||
use proto::cluster_rpc::{SimpleFileList, event_client::EventClient};
|
||||
use tokio::sync::{
|
||||
Mutex,
|
||||
mpsc::{Receiver, Sender},
|
||||
};
|
||||
use tonic::{codec::CompressionEncoding, metadata::MetadataValue};
|
||||
|
||||
/// (trace_id, file_id, account, file, size, cache_type)
|
||||
type FileInfo = (String, i64, String, String, usize, file_data::CacheType);
|
||||
|
||||
mod processing_files {
|
||||
use hashbrown::HashSet;
|
||||
use parking_lot::RwLock;
|
||||
|
||||
use super::*;
|
||||
|
||||
static PROCESSING_FILES: Lazy<RwLock<HashSet<String>>> =
|
||||
Lazy::new(|| RwLock::new(HashSet::new()));
|
||||
|
||||
pub fn is_processing(file_name: &str) -> bool {
|
||||
PROCESSING_FILES.read().contains(file_name)
|
||||
}
|
||||
|
||||
pub fn add(file_name: &str) {
|
||||
PROCESSING_FILES.write().insert(file_name.to_string());
|
||||
}
|
||||
|
||||
pub fn remove(file_name: &str) {
|
||||
PROCESSING_FILES.write().remove(file_name);
|
||||
}
|
||||
}
|
||||
|
||||
mod queued_files {
|
||||
use hashbrown::HashSet;
|
||||
use parking_lot::RwLock;
|
||||
|
||||
use super::*;
|
||||
|
||||
static QUEUED_FILES: Lazy<RwLock<HashSet<String>>> = Lazy::new(|| RwLock::new(HashSet::new()));
|
||||
|
||||
/// Returns true if newly inserted (caller should enqueue), false if already present (skip).
|
||||
pub fn try_add(file_name: &str) -> bool {
|
||||
QUEUED_FILES.write().insert(file_name.to_string())
|
||||
}
|
||||
|
||||
pub fn remove(file_name: &str) {
|
||||
QUEUED_FILES.write().remove(file_name);
|
||||
}
|
||||
}
|
||||
|
||||
struct DownloadQueue {
|
||||
sender: Sender<FileInfo>,
|
||||
receiver: Arc<Mutex<Receiver<FileInfo>>>,
|
||||
}
|
||||
|
||||
impl DownloadQueue {
|
||||
fn new(sender: Sender<FileInfo>, receiver: Arc<Mutex<Receiver<FileInfo>>>) -> Self {
|
||||
Self { sender, receiver }
|
||||
}
|
||||
}
|
||||
|
||||
struct PriorityDownloadQueue {
|
||||
stack: Mutex<VecDeque<FileInfo>>,
|
||||
notify: tokio::sync::Notify,
|
||||
max_size: usize,
|
||||
}
|
||||
|
||||
impl PriorityDownloadQueue {
|
||||
fn new(max_size: usize) -> Self {
|
||||
Self {
|
||||
stack: Mutex::new(VecDeque::with_capacity(max_size)),
|
||||
notify: tokio::sync::Notify::new(),
|
||||
max_size,
|
||||
}
|
||||
}
|
||||
|
||||
// Returns false if at cap — caller must remove from queued_files and skip metric inc.
|
||||
async fn push(&self, file_info: FileInfo) -> bool {
|
||||
let mut stack = self.stack.lock().await;
|
||||
if stack.len() >= self.max_size {
|
||||
return false;
|
||||
}
|
||||
stack.push_back(file_info);
|
||||
drop(stack);
|
||||
self.notify.notify_one();
|
||||
true
|
||||
}
|
||||
|
||||
// Blocks until an item is available. LIFO via pop_back.
|
||||
// notified() is created before the lock check to avoid missed-wakeup race.
|
||||
// Chains notify_one() so all workers drain the queue under burst load.
|
||||
async fn pop(&self) -> FileInfo {
|
||||
loop {
|
||||
let notified = self.notify.notified();
|
||||
{
|
||||
let mut stack = self.stack.lock().await;
|
||||
if let Some(item) = stack.pop_back() {
|
||||
if !stack.is_empty() {
|
||||
self.notify.notify_one();
|
||||
}
|
||||
return item;
|
||||
}
|
||||
}
|
||||
notified.await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
const FILE_DOWNLOAD_QUEUE_SIZE: usize = 10000;
|
||||
static FILE_DOWNLOAD_CHANNEL: Lazy<DownloadQueue> = Lazy::new(|| {
|
||||
let (tx, rx) = tokio::sync::mpsc::channel::<FileInfo>(FILE_DOWNLOAD_QUEUE_SIZE);
|
||||
DownloadQueue::new(tx, Arc::new(Mutex::new(rx)))
|
||||
});
|
||||
|
||||
static PRIORITY_FILE_DOWNLOAD_CHANNEL: Lazy<PriorityDownloadQueue> =
|
||||
Lazy::new(|| PriorityDownloadQueue::new(FILE_DOWNLOAD_QUEUE_SIZE));
|
||||
|
||||
pub async fn run() -> Result<(), anyhow::Error> {
|
||||
let cfg = get_config();
|
||||
// worker tasks: handle normal queue download (FIFO via mpsc channel)
|
||||
for thread in 0..cfg.limit.file_download_thread_num {
|
||||
let rx = FILE_DOWNLOAD_CHANNEL.receiver.clone();
|
||||
tokio::spawn(async move {
|
||||
loop {
|
||||
let ret = rx.lock().await.recv().await;
|
||||
match ret {
|
||||
None => {
|
||||
log::debug!("[FILE_CACHE_DOWNLOAD:JOB:NORMAL] Receiving channel is closed");
|
||||
break;
|
||||
}
|
||||
Some((trace_id, id, account, file, file_size, cache)) => {
|
||||
// transition: queued → not-queued (must run before any skip path)
|
||||
queued_files::remove(&file);
|
||||
|
||||
// check if the file is already being downloaded
|
||||
if processing_files::is_processing(&file) {
|
||||
log::warn!(
|
||||
"[trace_id {trace_id}] [thread {thread}] search->storage: file {file} is already being downloaded, will skip it"
|
||||
);
|
||||
// update metrics
|
||||
metrics::FILE_DOWNLOADER_NORMAL_QUEUE_SIZE
|
||||
.with_label_values::<&str>(&[])
|
||||
.dec();
|
||||
continue;
|
||||
}
|
||||
|
||||
// add the file to processing set
|
||||
processing_files::add(&file);
|
||||
|
||||
// download the file
|
||||
match download_file(
|
||||
thread, &trace_id, id, &account, &file, file_size, cache,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(data_len) => {
|
||||
if data_len > 0 && data_len != file_size {
|
||||
log::warn!(
|
||||
"[FILE_CACHE_DOWNLOAD:JOB:NORMAL] download file {file} found size mismatch, expected: {file_size}, actual: {data_len}, will skip it",
|
||||
);
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
log::error!(
|
||||
"[FILE_CACHE_DOWNLOAD:JOB:NORMAL] download file {file} to cache {cache:?} err: {e}",
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
// remove the file from processing set
|
||||
processing_files::remove(&file);
|
||||
|
||||
// update metrics
|
||||
metrics::FILE_DOWNLOADER_NORMAL_QUEUE_SIZE
|
||||
.with_label_values::<&str>(&[])
|
||||
.dec();
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
// worker tasks: handle priority queue download (LIFO via VecDeque::pop_back)
|
||||
for thread in 0..cfg.limit.file_download_priority_queue_thread_num {
|
||||
tokio::spawn(async move {
|
||||
loop {
|
||||
let (trace_id, id, account, file, file_size, cache) =
|
||||
PRIORITY_FILE_DOWNLOAD_CHANNEL.pop().await;
|
||||
|
||||
// transition: queued → not-queued (must run before any skip path)
|
||||
queued_files::remove(&file);
|
||||
|
||||
// check if the file is already being downloaded
|
||||
if processing_files::is_processing(&file) {
|
||||
log::warn!(
|
||||
"[trace_id {trace_id}] [thread {thread}] search->storage: file {file} is already being downloaded, will skip it"
|
||||
);
|
||||
metrics::FILE_DOWNLOADER_PRIORITY_QUEUE_SIZE
|
||||
.with_label_values::<&str>(&[])
|
||||
.dec();
|
||||
continue;
|
||||
}
|
||||
|
||||
// add the file to processing set
|
||||
processing_files::add(&file);
|
||||
|
||||
// download the file
|
||||
match download_file(thread, &trace_id, id, &account, &file, file_size, cache).await
|
||||
{
|
||||
Ok(data_len) => {
|
||||
if data_len > 0 && data_len != file_size {
|
||||
log::warn!(
|
||||
"[FILE_CACHE_DOWNLOAD:JOB:PRIORITY] download file {file} found size mismatch, expected: {file_size}, actual: {data_len}, will skip it",
|
||||
);
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
log::error!(
|
||||
"[FILE_CACHE_DOWNLOAD:JOB:PRIORITY] download file {file} to cache {cache:?} err: {e}",
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
// remove the file from processing set
|
||||
processing_files::remove(&file);
|
||||
|
||||
// update metrics
|
||||
metrics::FILE_DOWNLOADER_PRIORITY_QUEUE_SIZE
|
||||
.with_label_values::<&str>(&[])
|
||||
.dec();
|
||||
}
|
||||
});
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn download_file(
|
||||
thread: usize,
|
||||
trace_id: &str,
|
||||
file_id: i64,
|
||||
account: &str,
|
||||
file_name: &str,
|
||||
file_size: usize,
|
||||
cache_type: file_data::CacheType,
|
||||
) -> Result<usize, anyhow::Error> {
|
||||
let cfg = get_config();
|
||||
|
||||
// download file from node
|
||||
if cfg.cache_latest_files.download_from_node
|
||||
&& let Ok(ok) =
|
||||
download_file_with_consistent_hash(file_id, account, file_name, file_size).await
|
||||
&& ok
|
||||
{
|
||||
return Ok(file_size);
|
||||
}
|
||||
|
||||
// download from object store
|
||||
let size = if file_size > 0 { Some(file_size) } else { None };
|
||||
let start = std::time::Instant::now();
|
||||
let ret = match cache_type {
|
||||
file_data::CacheType::Memory => {
|
||||
let mut disk_exists = false;
|
||||
let mem_exists = file_data::memory::exist(file_name).await;
|
||||
if !mem_exists && !cfg.memory_cache.skip_disk_check {
|
||||
disk_exists = file_data::disk::exist(file_name).await;
|
||||
}
|
||||
if !mem_exists && (cfg.memory_cache.skip_disk_check || !disk_exists) {
|
||||
file_data::memory::download(account, file_name, size).await
|
||||
} else {
|
||||
Ok(0)
|
||||
}
|
||||
}
|
||||
file_data::CacheType::Disk => {
|
||||
if !file_data::disk::exist(file_name).await {
|
||||
file_data::disk::download(account, file_name, size).await
|
||||
} else {
|
||||
Ok(0)
|
||||
}
|
||||
}
|
||||
_ => Ok(0),
|
||||
};
|
||||
log::debug!(
|
||||
"[FILE_CACHE_DOWNLOAD:JOB:{thread}] [trace_id {trace_id}] download file: {file_name}, ret: {:?}, took: {} ms",
|
||||
ret,
|
||||
start.elapsed().as_millis()
|
||||
);
|
||||
ret
|
||||
}
|
||||
|
||||
async fn download_file_with_consistent_hash(
|
||||
file_id: i64,
|
||||
account: &str,
|
||||
file_name: &str,
|
||||
file_size: usize,
|
||||
) -> Result<bool, anyhow::Error> {
|
||||
let role_group = if LOCAL_NODE.is_interactive_querier() {
|
||||
RoleGroup::Interactive
|
||||
} else {
|
||||
RoleGroup::Background
|
||||
};
|
||||
let Some(node_name) = cluster::get_node_from_consistent_hash(
|
||||
&file_id.to_string(),
|
||||
&Role::Querier,
|
||||
Some(role_group),
|
||||
)
|
||||
.await
|
||||
else {
|
||||
return Ok(false);
|
||||
};
|
||||
// get node by file_id
|
||||
let Some(node) = cluster::get_cached_node_by_name(&node_name).await else {
|
||||
return Ok(false);
|
||||
};
|
||||
// download file from node
|
||||
let Ok(failed) = download_from_node(
|
||||
&node.grpc_addr,
|
||||
&[(
|
||||
file_id,
|
||||
account.to_string(),
|
||||
file_name.to_string(),
|
||||
file_size as i64,
|
||||
0,
|
||||
)],
|
||||
)
|
||||
.await
|
||||
else {
|
||||
return Ok(false);
|
||||
};
|
||||
// failed is empty means download success
|
||||
Ok(failed.is_empty())
|
||||
}
|
||||
|
||||
// download files from node and return download failed files
|
||||
// file: (account, file, size, ts)
|
||||
pub async fn download_from_node(
|
||||
addr: &str,
|
||||
files: &[(i64, String, String, i64, i64)],
|
||||
) -> Result<Vec<(i64, String, String, i64, i64)>, anyhow::Error> {
|
||||
let start = std::time::Instant::now();
|
||||
let cfg = get_config();
|
||||
log::debug!(
|
||||
"[FILE_CACHE_DOWNLOAD:gRPC] Download files from node start, files: {:?}",
|
||||
files.len()
|
||||
);
|
||||
|
||||
let token: MetadataValue<_> = get_internal_grpc_token()
|
||||
.parse()
|
||||
.map_err(|_| anyhow::anyhow!("Invalid token"))?;
|
||||
|
||||
let channel = infra::client::grpc::get_cached_channel(addr).await?;
|
||||
let client = EventClient::with_interceptor(channel, move |mut req: tonic::Request<()>| {
|
||||
req.metadata_mut().insert("authorization", token.clone());
|
||||
Ok(req)
|
||||
});
|
||||
|
||||
let file_size_map = files
|
||||
.iter()
|
||||
.filter_map(|(_, _, f, s, _)| {
|
||||
if *s > cfg.cache_latest_files.download_node_size * 1024 * 1024 {
|
||||
None
|
||||
} else {
|
||||
Some((f, *s as usize))
|
||||
}
|
||||
})
|
||||
.collect::<HashMap<_, _>>();
|
||||
if file_size_map.is_empty() {
|
||||
return Ok(files.to_vec());
|
||||
}
|
||||
let request = tonic::Request::new(SimpleFileList {
|
||||
files: files.iter().map(|(_, _, f, ..)| f.to_string()).collect(),
|
||||
});
|
||||
|
||||
let resp = client
|
||||
.send_compressed(CompressionEncoding::Gzip)
|
||||
.accept_compressed(CompressionEncoding::Gzip)
|
||||
.max_decoding_message_size(cfg.grpc.max_message_size * 1024 * 1024)
|
||||
.max_encoding_message_size(cfg.grpc.max_message_size * 1024 * 1024)
|
||||
.get_files(request)
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("Failed to get files from {addr}, {e}"))?;
|
||||
|
||||
let mut file_contents = HashMap::new();
|
||||
let mut downloaded_files = HashSet::new();
|
||||
let mut resp_stream = resp.into_inner();
|
||||
while let Some(resp) = resp_stream.next().await {
|
||||
let resp = match resp {
|
||||
Ok(resp) => resp,
|
||||
Err(err) => {
|
||||
if err.code() == tonic::Code::NotFound {
|
||||
log::debug!(
|
||||
"[FILE_CACHE_DOWNLOAD:gRPC] Failed to download file {} from {}: file not found",
|
||||
err.message(),
|
||||
addr
|
||||
);
|
||||
continue;
|
||||
}
|
||||
return Err(anyhow::anyhow!(
|
||||
"Failed to download file from {addr}, {err}"
|
||||
));
|
||||
}
|
||||
};
|
||||
for content in resp.entries {
|
||||
let entry = file_contents
|
||||
.entry(content.filename.clone())
|
||||
.or_insert(bytes::BytesMut::new());
|
||||
entry.extend_from_slice(&content.content);
|
||||
downloaded_files.insert(content.filename);
|
||||
}
|
||||
}
|
||||
|
||||
log::debug!(
|
||||
"[FILE_CACHE_DOWNLOAD:gRPC] Successfully retrieved {} files from {} in {} ms",
|
||||
downloaded_files.len(),
|
||||
addr,
|
||||
start.elapsed().as_millis()
|
||||
);
|
||||
|
||||
// Cache the file contents
|
||||
for (file, content) in file_contents {
|
||||
let data = content.freeze();
|
||||
if let Some(size) = file_size_map.get(&file)
|
||||
&& *size != data.len()
|
||||
{
|
||||
log::warn!(
|
||||
"[FILE_CACHE_DOWNLOAD:gRPC] Failed to download file {} from {}: size mismatch, expected {} but got {}",
|
||||
file,
|
||||
addr,
|
||||
size,
|
||||
data.len()
|
||||
);
|
||||
downloaded_files.remove(&file);
|
||||
continue;
|
||||
}
|
||||
if let Err(e) = infra::cache::file_data::set(&file, data).await {
|
||||
log::error!("[FILE_CACHE_DOWNLOAD:gRPC] Failed to cache file {file}: {e}");
|
||||
downloaded_files.remove(&file);
|
||||
}
|
||||
}
|
||||
|
||||
// Return list of failed files
|
||||
let failed_files: Vec<_> = files
|
||||
.iter()
|
||||
.filter(|(_, f, ..)| !downloaded_files.contains(f))
|
||||
.cloned()
|
||||
.collect();
|
||||
|
||||
log::debug!(
|
||||
"[FILE_CACHE_DOWNLOAD:gRPC] Failed to retrieve {} files from {} in {} ms",
|
||||
failed_files.len(),
|
||||
addr,
|
||||
start.elapsed().as_millis()
|
||||
);
|
||||
|
||||
Ok(failed_files)
|
||||
}
|
||||
|
||||
pub async fn queue_download(
|
||||
trace_id: String,
|
||||
id: i64,
|
||||
account: String,
|
||||
file: String,
|
||||
size: i64,
|
||||
ts: i64,
|
||||
cache_type: file_data::CacheType,
|
||||
) -> Result<(), anyhow::Error> {
|
||||
log::debug!(
|
||||
"[FILE_CACHE_DOWNLOAD:JOB] [trace_id {trace_id}] enqueue file: {file}, size: {size}, ts: {ts}"
|
||||
);
|
||||
|
||||
// files with data older than the cache max age are not worth caching: with
|
||||
// time_lru they become the oldest bucket and are evicted first, so the
|
||||
// download would be wasted work. Skip the queue and let the query read
|
||||
// them directly from object storage.
|
||||
if exceeds_cache_max_age(ts, cache_type) {
|
||||
log::debug!(
|
||||
"[FILE_CACHE_DOWNLOAD:JOB] [trace_id {trace_id}] skip file: {file}, data is older than the cache max age"
|
||||
);
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let cfg = get_config();
|
||||
|
||||
// skip if already queued or already being downloaded
|
||||
if !queued_files::try_add(&file) {
|
||||
return Ok(());
|
||||
}
|
||||
if processing_files::is_processing(&file) {
|
||||
queued_files::remove(&file);
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
if cfg.limit.file_download_enable_priority_queue
|
||||
&& should_prioritize_file(ts, cfg.limit.file_download_priority_queue_window_secs)
|
||||
{
|
||||
if PRIORITY_FILE_DOWNLOAD_CHANNEL
|
||||
.push((
|
||||
trace_id,
|
||||
id,
|
||||
account,
|
||||
file.clone(),
|
||||
size as usize,
|
||||
cache_type,
|
||||
))
|
||||
.await
|
||||
{
|
||||
metrics::FILE_DOWNLOADER_PRIORITY_QUEUE_SIZE
|
||||
.with_label_values::<&str>(&[])
|
||||
.inc();
|
||||
} else {
|
||||
queued_files::remove(&file);
|
||||
log::warn!(
|
||||
"[FILE_CACHE_DOWNLOAD:JOB] priority queue full ({FILE_DOWNLOAD_QUEUE_SIZE}), dropping file: {file}"
|
||||
);
|
||||
}
|
||||
} else {
|
||||
if let Err(e) = FILE_DOWNLOAD_CHANNEL
|
||||
.sender
|
||||
.send((
|
||||
trace_id,
|
||||
id,
|
||||
account,
|
||||
file.clone(),
|
||||
size as usize,
|
||||
cache_type,
|
||||
))
|
||||
.await
|
||||
{
|
||||
queued_files::remove(&file);
|
||||
return Err(e.into());
|
||||
}
|
||||
metrics::FILE_DOWNLOADER_NORMAL_QUEUE_SIZE
|
||||
.with_label_values::<&str>(&[])
|
||||
.inc();
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
// if the file timestamp is in the past window, it should be prioritized
|
||||
fn should_prioritize_file(ts: i64, window_secs: i64) -> bool {
|
||||
let window_micros = window_secs * 1_000_000;
|
||||
let now = now_micros();
|
||||
ts > now - window_micros
|
||||
}
|
||||
|
||||
/// Returns true if the file's data is older than the cache max age and should
|
||||
/// not be downloaded into the cache.
|
||||
pub fn exceeds_cache_max_age(ts: i64, cache_type: file_data::CacheType) -> bool {
|
||||
let cfg = get_config();
|
||||
let max_age_days = match cache_type {
|
||||
file_data::CacheType::Memory => cfg.memory_cache.max_age_days,
|
||||
file_data::CacheType::Disk => cfg.disk_cache.max_age_days,
|
||||
file_data::CacheType::None => 0,
|
||||
};
|
||||
exceeds_max_age(ts, max_age_days)
|
||||
}
|
||||
|
||||
// if the file data is older than max_age_days, it should not be cached.
|
||||
// max_age_days == 0 means no limit, ts <= 0 means the timestamp is unknown.
|
||||
fn exceeds_max_age(ts: i64, max_age_days: i64) -> bool {
|
||||
if max_age_days <= 0 || ts <= 0 {
|
||||
return false;
|
||||
}
|
||||
ts < now_micros() - day_micros(max_age_days)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use config::utils::time::{day_micros, hour_micros, now_micros};
|
||||
|
||||
use super::{
|
||||
FileInfo, PriorityDownloadQueue, exceeds_max_age, file_data, processing_files, queued_files,
|
||||
};
|
||||
|
||||
#[test]
|
||||
fn test_exceeds_max_age_disabled() {
|
||||
// max_age_days == 0 means no limit, nothing is too old
|
||||
let one_year_ago = now_micros() - day_micros(365);
|
||||
assert!(!exceeds_max_age(one_year_ago, 0));
|
||||
assert!(!exceeds_max_age(one_year_ago, -1));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_exceeds_max_age_unknown_ts() {
|
||||
// unknown timestamp should not be skipped
|
||||
assert!(!exceeds_max_age(0, 3));
|
||||
assert!(!exceeds_max_age(-1, 3));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_exceeds_max_age_recent_file() {
|
||||
let one_hour_ago = now_micros() - hour_micros(1);
|
||||
assert!(!exceeds_max_age(one_hour_ago, 3));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_exceeds_max_age_old_file() {
|
||||
let thirty_days_ago = now_micros() - day_micros(30);
|
||||
assert!(exceeds_max_age(thirty_days_ago, 3));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_is_processing_false_initially() {
|
||||
assert!(!processing_files::is_processing(
|
||||
"unique_not_added_file_abc123.parquet"
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_add_and_is_processing() {
|
||||
let name = "test_add_file_xyz999.parquet";
|
||||
processing_files::add(name);
|
||||
assert!(processing_files::is_processing(name));
|
||||
processing_files::remove(name);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_remove_clears_processing() {
|
||||
let name = "test_remove_file_qqq888.parquet";
|
||||
processing_files::add(name);
|
||||
processing_files::remove(name);
|
||||
assert!(!processing_files::is_processing(name));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_normal_queue_cap_honored() {
|
||||
let cap = 3;
|
||||
let (tx, _rx) = tokio::sync::mpsc::channel::<FileInfo>(cap);
|
||||
|
||||
for i in 0..cap {
|
||||
let result = tx.try_send((
|
||||
format!("trace-{i}"),
|
||||
i as i64,
|
||||
"org".into(),
|
||||
format!("file-{i}.parquet"),
|
||||
1024,
|
||||
file_data::CacheType::Disk,
|
||||
));
|
||||
assert!(result.is_ok(), "send {i} should succeed under cap");
|
||||
}
|
||||
|
||||
let overflow = tx.try_send((
|
||||
"trace-overflow".into(),
|
||||
99,
|
||||
"org".into(),
|
||||
"overflow.parquet".into(),
|
||||
1024,
|
||||
file_data::CacheType::Disk,
|
||||
));
|
||||
assert!(overflow.is_err(), "send beyond cap must fail");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_queued_files_no_duplicates() {
|
||||
let file = "dedup_test_unique_zz1234.parquet";
|
||||
|
||||
// first add succeeds
|
||||
assert!(queued_files::try_add(file));
|
||||
// same file rejected while already queued
|
||||
assert!(!queued_files::try_add(file));
|
||||
|
||||
// after remove, can be queued again
|
||||
queued_files::remove(file);
|
||||
assert!(queued_files::try_add(file));
|
||||
|
||||
// cleanup
|
||||
queued_files::remove(file);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_priority_queue_cap_honored() {
|
||||
let cap = 3;
|
||||
let queue = PriorityDownloadQueue::new(cap);
|
||||
|
||||
// fill to cap — all succeed
|
||||
for i in 0..cap {
|
||||
let ok = queue
|
||||
.push((
|
||||
format!("trace-{i}"),
|
||||
i as i64,
|
||||
"org".into(),
|
||||
format!("file-{i}.parquet"),
|
||||
1024,
|
||||
file_data::CacheType::Disk,
|
||||
))
|
||||
.await;
|
||||
assert!(ok, "push {i} should succeed under cap");
|
||||
}
|
||||
|
||||
// one more must be rejected
|
||||
let overflow = queue
|
||||
.push((
|
||||
"trace-overflow".into(),
|
||||
99,
|
||||
"org".into(),
|
||||
"overflow.parquet".into(),
|
||||
1024,
|
||||
file_data::CacheType::Disk,
|
||||
))
|
||||
.await;
|
||||
assert!(!overflow, "push beyond cap must return false");
|
||||
}
|
||||
}
|
||||
pub use infra::cache::file_downloader::*;
|
||||
|
||||
@@ -29,7 +29,7 @@ use config::{
|
||||
ider::SnowflakeIdGenerator,
|
||||
meta::{
|
||||
alerts::alert::Alert,
|
||||
function::{VRLResultResolver, VRLRuntimeConfig},
|
||||
function::VRLResultResolver,
|
||||
self_reporting::usage::{RequestStats, TriggerData, TriggerDataStatus, TriggerDataType},
|
||||
stream::{PartitionTimeLevel, StreamParams, StreamPartition, StreamType},
|
||||
},
|
||||
@@ -41,10 +41,7 @@ use infra::{
|
||||
schema::STREAM_RECORD_ID_GENERATOR,
|
||||
};
|
||||
use proto::cluster_rpc::IngestionType;
|
||||
use vrl::{
|
||||
compiler::{CompilationResult, TargetValueRef, runtime::Runtime},
|
||||
prelude::state,
|
||||
};
|
||||
use vrl::compiler::{TargetValueRef, runtime::Runtime};
|
||||
|
||||
use super::{
|
||||
db::{alerts::alert, pipeline},
|
||||
@@ -55,11 +52,8 @@ use crate::{
|
||||
common::{
|
||||
infra::config::{REALTIME_ALERT_TRIGGERS, STREAM_ALERTS},
|
||||
meta::{ingestion::IngestionRequest, stream::SchemaRecords},
|
||||
utils::{
|
||||
functions::get_vrl_compiler_config,
|
||||
js::{
|
||||
JSRuntimeConfig, apply_js_fn as apply_js, compile_js_function as compile_js_func,
|
||||
},
|
||||
utils::js::{
|
||||
JSRuntimeConfig, apply_js_fn as apply_js, compile_js_function as compile_js_func,
|
||||
},
|
||||
},
|
||||
service::{
|
||||
@@ -72,6 +66,8 @@ use crate::{
|
||||
pub mod grpc;
|
||||
pub mod ingestion_service;
|
||||
|
||||
pub use openobserve_vrl::compile_vrl_function;
|
||||
|
||||
pub type TriggerAlertData = Vec<(Alert, Vec<Map<String, Value>>)>;
|
||||
|
||||
/// Global atomic counter for round-robin distribution of requests across memory table buckets.
|
||||
@@ -90,34 +86,6 @@ pub fn get_thread_id() -> usize {
|
||||
REQUEST_COUNTER.fetch_add(1, Ordering::Relaxed) % cfg.limit.http_worker_num
|
||||
}
|
||||
|
||||
pub fn compile_vrl_function(func: &str, org_id: &str) -> Result<VRLRuntimeConfig, std::io::Error> {
|
||||
if func.contains("get_env_var") {
|
||||
return Err(std::io::Error::other("get_env_var is not supported"));
|
||||
}
|
||||
|
||||
let external = state::ExternalEnv::default();
|
||||
let vrl_config = get_vrl_compiler_config(org_id);
|
||||
match vrl::compiler::compile_with_external(
|
||||
func,
|
||||
&vrl_config.functions,
|
||||
&external,
|
||||
vrl_config.config,
|
||||
) {
|
||||
Ok(CompilationResult {
|
||||
program,
|
||||
warnings: _,
|
||||
config,
|
||||
}) => Ok(VRLRuntimeConfig {
|
||||
program,
|
||||
config,
|
||||
fields: vec![],
|
||||
}),
|
||||
Err(e) => Err(std::io::Error::other(
|
||||
vrl::diagnostic::Formatter::new(func, e).to_string(),
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
/// Compile a JS function and return configuration
|
||||
/// This wraps the common::utils::js::compile_js_function
|
||||
pub fn compile_js_function(func: &str, org_id: &str) -> Result<JSRuntimeConfig, std::io::Error> {
|
||||
|
||||
+1
-51
@@ -35,7 +35,7 @@ use crate::{
|
||||
MultiCachedQueryResponse,
|
||||
result_utils::{get_ts_value, has_non_timestamp_ordering, is_timestamp_field},
|
||||
},
|
||||
sql::{RE_HISTOGRAM, Sql, visitor::histogram_interval::generate_histogram_interval},
|
||||
sql::Sql,
|
||||
},
|
||||
};
|
||||
|
||||
@@ -907,46 +907,6 @@ pub async fn delete_cache(
|
||||
Ok(true)
|
||||
}
|
||||
|
||||
pub fn handle_histogram(
|
||||
origin_sql: &mut String,
|
||||
q_time_range: (i64, i64),
|
||||
histogram_interval: i64,
|
||||
) {
|
||||
let caps = if let Some(caps) = RE_HISTOGRAM.captures(origin_sql.as_str()) {
|
||||
caps
|
||||
} else {
|
||||
return;
|
||||
};
|
||||
|
||||
// 0th capture is the whole histogram(...) ,
|
||||
// 1st capture is the comma-delimited list of args
|
||||
// ideally there should be at least one arg, otherwise df with anyways complain,
|
||||
// so we we return from here if capture[1] is None
|
||||
let args = match caps.get(1) {
|
||||
Some(v) => v
|
||||
.as_str()
|
||||
.split(',')
|
||||
.map(|v| v.trim().trim_matches(|v| v == '\'' || v == '"'))
|
||||
.collect::<Vec<&str>>(),
|
||||
None => return,
|
||||
};
|
||||
|
||||
let interval = if histogram_interval > 0 {
|
||||
format!("{histogram_interval} second")
|
||||
} else {
|
||||
args.get(1)
|
||||
.map_or_else(|| generate_histogram_interval(q_time_range), |v| *v)
|
||||
.to_string()
|
||||
};
|
||||
|
||||
let field = args.first().unwrap_or(&"_timestamp");
|
||||
|
||||
*origin_sql = origin_sql.replace(
|
||||
caps.get(0).unwrap().as_str(),
|
||||
&format!("histogram({field},'{interval}')"),
|
||||
);
|
||||
}
|
||||
|
||||
fn calculate_deltas_multi(
|
||||
results: &[CachedQueryResponse],
|
||||
start_time: i64,
|
||||
@@ -1164,16 +1124,6 @@ mod tests {
|
||||
assert_eq!(deltas, expected_deltas);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_handle_histogram() {
|
||||
// Test case 1: Basic histogram with numeric interval
|
||||
let mut sql = "SELECT histogram(_timestamp, '10 seconds') FROM logs".to_string();
|
||||
let time_range = (1640995200000000, 1641081600000000); // 2022-01-01 to 2022-01-02
|
||||
handle_histogram(&mut sql, time_range, 10);
|
||||
assert!(sql.contains("histogram(_timestamp,"));
|
||||
assert!(sql.contains("second"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_get_ts_col_order_by() {
|
||||
let sql = Sql {
|
||||
|
||||
+5
-1
@@ -568,7 +568,11 @@ pub async fn prepare_cache_response(
|
||||
} else {
|
||||
sql.time_range
|
||||
};
|
||||
cacher::handle_histogram(&mut origin_sql, q_time_range, req.query.histogram_interval);
|
||||
crate::service::search::sql::histogram::handle_histogram(
|
||||
&mut origin_sql,
|
||||
q_time_range,
|
||||
req.query.histogram_interval,
|
||||
);
|
||||
}
|
||||
|
||||
// calculate hash for the query with version (after normalizing histogram interval)
|
||||
|
||||
@@ -30,9 +30,7 @@ use config::{
|
||||
metrics,
|
||||
utils::{json, time::now_micros, took_watcher::TookWatcher},
|
||||
};
|
||||
use datafusion::{
|
||||
common::TableReference, physical_plan::visit_execution_plan, prelude::SessionContext,
|
||||
};
|
||||
use datafusion::{common::TableReference, physical_plan::visit_execution_plan};
|
||||
use hashbrown::{HashMap, HashSet};
|
||||
use infra::{
|
||||
cluster,
|
||||
@@ -56,16 +54,14 @@ use crate::service::{
|
||||
search::{
|
||||
SearchResult,
|
||||
datafusion::{
|
||||
exec::{DataFusionContextBuilder, register_udf},
|
||||
context::{SearchContextBuilder, register_table},
|
||||
optimizer::{
|
||||
context::{
|
||||
PhysicalOptimizerContext, RemoteScanContext, StreamingAggregationContext,
|
||||
},
|
||||
create_physical_plan, generate_analyzer_rules, generate_optimizer_rules,
|
||||
generate_physical_optimizer_rules,
|
||||
create_physical_plan,
|
||||
},
|
||||
plan_metrics::get_peak_memory_from_ctx,
|
||||
table_provider::{catalog::StreamTypeProvider, empty_table::NewEmptyTable},
|
||||
},
|
||||
inspector::{SearchInspectorFieldsBuilder, search_inspector_fields},
|
||||
sql::Sql,
|
||||
@@ -664,93 +660,6 @@ pub async fn partition_file_by_hash(
|
||||
partitions
|
||||
}
|
||||
|
||||
pub struct SearchContextBuilder {
|
||||
pub target_partitions: usize,
|
||||
pub contexts: Vec<PhysicalOptimizerContext>,
|
||||
}
|
||||
|
||||
impl Default for SearchContextBuilder {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
impl SearchContextBuilder {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
target_partitions: 0,
|
||||
contexts: vec![],
|
||||
}
|
||||
}
|
||||
|
||||
pub fn target_partitions(mut self, target_partitions: usize) -> Self {
|
||||
self.target_partitions = target_partitions;
|
||||
self
|
||||
}
|
||||
|
||||
pub fn add_context(mut self, context: PhysicalOptimizerContext) -> Self {
|
||||
self.contexts.push(context);
|
||||
self
|
||||
}
|
||||
|
||||
pub async fn build(self, req: &Request, sql: &Arc<Sql>) -> Result<SessionContext> {
|
||||
let analyzer_rules = generate_analyzer_rules(sql);
|
||||
let optimizer_rules = generate_optimizer_rules(sql);
|
||||
let physical_optimizer_rules = generate_physical_optimizer_rules(req, sql, self.contexts);
|
||||
let mut ctx = DataFusionContextBuilder::new()
|
||||
.trace_id(&req.trace_id)
|
||||
.work_group(req.work_group.clone())
|
||||
.analyzer_rules(analyzer_rules)
|
||||
.optimizer_rules(optimizer_rules)
|
||||
.physical_optimizer_rules(physical_optimizer_rules)
|
||||
.sorted_by_time(sql.sorted_by_time)
|
||||
.build(self.target_partitions)
|
||||
.await?;
|
||||
|
||||
// register udf
|
||||
register_udf(&ctx, &req.org_id)?;
|
||||
datafusion_functions_json::register_all(&mut ctx)?;
|
||||
|
||||
Ok(ctx)
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn register_table(ctx: &SessionContext, sql: &Sql) -> Result<()> {
|
||||
// register schema provider
|
||||
let mut registed_schema = HashSet::new();
|
||||
for (stream, _) in &sql.schemas {
|
||||
let stream_type = stream.stream_type();
|
||||
if !stream.has_stream_type() || registed_schema.contains(&stream_type) {
|
||||
continue;
|
||||
}
|
||||
registed_schema.insert(stream_type.clone());
|
||||
let schema_provider = StreamTypeProvider::create(&stream_type).await?;
|
||||
let _ = ctx
|
||||
.catalog("datafusion")
|
||||
.unwrap()
|
||||
.as_ref()
|
||||
.register_schema(&stream_type, schema_provider);
|
||||
}
|
||||
|
||||
// register table
|
||||
for (stream, schema) in &sql.schemas {
|
||||
let schema = schema
|
||||
.schema()
|
||||
.as_ref()
|
||||
.clone()
|
||||
.with_metadata(Default::default());
|
||||
let stream_name = stream.to_quoted_string();
|
||||
let table = Arc::new(
|
||||
NewEmptyTable::new(&stream_name, Arc::new(schema))
|
||||
.with_partitions(ctx.state().config().target_partitions())
|
||||
.with_sorted_by_time(sql.sorted_by_time),
|
||||
);
|
||||
ctx.register_table(&stream_name, table)?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tracing::instrument(name = "service:search:cluster:flight:get_file_id_lists", skip_all)]
|
||||
pub async fn get_file_id_lists(
|
||||
trace_id: &str,
|
||||
|
||||
@@ -0,0 +1,114 @@
|
||||
// Copyright 2026 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 <http://www.gnu.org/licenses/>.
|
||||
|
||||
use std::{collections::HashSet, sync::Arc};
|
||||
|
||||
use config::{datafusion::request::Request, meta::sql::TableReferenceExt};
|
||||
use datafusion::prelude::SessionContext;
|
||||
use infra::errors::Result;
|
||||
|
||||
use super::{
|
||||
exec::{DataFusionContextBuilder, register_udf},
|
||||
optimizer::{
|
||||
context::PhysicalOptimizerContext, generate_analyzer_rules, generate_optimizer_rules,
|
||||
generate_physical_optimizer_rules,
|
||||
},
|
||||
table_provider::{catalog::StreamTypeProvider, empty_table::NewEmptyTable},
|
||||
};
|
||||
use crate::service::search::sql::Sql;
|
||||
|
||||
pub struct SearchContextBuilder {
|
||||
pub target_partitions: usize,
|
||||
pub contexts: Vec<PhysicalOptimizerContext>,
|
||||
}
|
||||
|
||||
impl Default for SearchContextBuilder {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
impl SearchContextBuilder {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
target_partitions: 0,
|
||||
contexts: vec![],
|
||||
}
|
||||
}
|
||||
|
||||
pub fn target_partitions(mut self, target_partitions: usize) -> Self {
|
||||
self.target_partitions = target_partitions;
|
||||
self
|
||||
}
|
||||
|
||||
pub fn add_context(mut self, context: PhysicalOptimizerContext) -> Self {
|
||||
self.contexts.push(context);
|
||||
self
|
||||
}
|
||||
|
||||
pub async fn build(self, req: &Request, sql: &Arc<Sql>) -> Result<SessionContext> {
|
||||
let analyzer_rules = generate_analyzer_rules(sql);
|
||||
let optimizer_rules = generate_optimizer_rules(sql);
|
||||
let physical_optimizer_rules = generate_physical_optimizer_rules(req, sql, self.contexts);
|
||||
let mut ctx = DataFusionContextBuilder::new()
|
||||
.trace_id(&req.trace_id)
|
||||
.work_group(req.work_group.clone())
|
||||
.analyzer_rules(analyzer_rules)
|
||||
.optimizer_rules(optimizer_rules)
|
||||
.physical_optimizer_rules(physical_optimizer_rules)
|
||||
.sorted_by_time(sql.sorted_by_time)
|
||||
.build(self.target_partitions)
|
||||
.await?;
|
||||
|
||||
register_udf(&ctx, &sql.org_id)?;
|
||||
datafusion_functions_json::register_all(&mut ctx)?;
|
||||
|
||||
Ok(ctx)
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn register_table(ctx: &SessionContext, sql: &Sql) -> Result<()> {
|
||||
let mut registered_schemas = HashSet::new();
|
||||
for (stream, _) in &sql.schemas {
|
||||
let stream_type = stream.stream_type();
|
||||
if !stream.has_stream_type() || registered_schemas.contains(&stream_type) {
|
||||
continue;
|
||||
}
|
||||
registered_schemas.insert(stream_type.clone());
|
||||
let schema_provider = StreamTypeProvider::create(&stream_type).await?;
|
||||
let _ = ctx
|
||||
.catalog("datafusion")
|
||||
.unwrap()
|
||||
.as_ref()
|
||||
.register_schema(&stream_type, schema_provider);
|
||||
}
|
||||
|
||||
for (stream, schema) in &sql.schemas {
|
||||
let schema = schema
|
||||
.schema()
|
||||
.as_ref()
|
||||
.clone()
|
||||
.with_metadata(Default::default());
|
||||
let stream_name = stream.to_quoted_string();
|
||||
let table = Arc::new(
|
||||
NewEmptyTable::new(&stream_name, Arc::new(schema))
|
||||
.with_partitions(ctx.state().config().target_partitions())
|
||||
.with_sorted_by_time(sql.sorted_by_time),
|
||||
);
|
||||
ctx.register_table(&stream_name, table)?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
@@ -15,6 +15,7 @@
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use ::search::types::QueryParams;
|
||||
use config::meta::{inverted_index::IndexOptimizeMode, stream::FileKey};
|
||||
use datafusion::{
|
||||
arrow::datatypes::SchemaRef,
|
||||
@@ -32,8 +33,7 @@ use datafusion::{
|
||||
use o2_enterprise::enterprise::search::datafusion::distributed_plan::metadata_count_exec::MetadataCountExec;
|
||||
|
||||
use crate::service::search::{
|
||||
datafusion::plan::tantivy_optimize_exec::TantivyOptimizeExec, grpc::QueryParams,
|
||||
index::IndexCondition,
|
||||
datafusion::plan::tantivy_optimize_exec::TantivyOptimizeExec, index::IndexCondition,
|
||||
};
|
||||
|
||||
pub fn aggregate_optimize_rewrite(
|
||||
|
||||
@@ -15,6 +15,7 @@
|
||||
|
||||
use std::str::FromStr;
|
||||
|
||||
pub mod context;
|
||||
pub mod distributed_plan;
|
||||
pub mod exec;
|
||||
pub mod merge;
|
||||
|
||||
@@ -27,9 +27,8 @@ use datafusion::{
|
||||
use hashbrown::HashSet;
|
||||
|
||||
use crate::service::search::{
|
||||
cache::cacher::handle_histogram,
|
||||
cluster::flight::{SearchContextBuilder, register_table},
|
||||
sql::Sql,
|
||||
datafusion::context::{SearchContextBuilder, register_table},
|
||||
sql::{Sql, histogram::handle_histogram},
|
||||
};
|
||||
|
||||
/// Structure to store the result schema info
|
||||
|
||||
@@ -19,7 +19,7 @@ use config::{datafusion::request::Request, meta::projections::ProjectionColumnMa
|
||||
use datafusion::common::tree_node::TreeNode;
|
||||
|
||||
use crate::service::search::{
|
||||
cluster::flight::{SearchContextBuilder, register_table},
|
||||
datafusion::context::{SearchContextBuilder, register_table},
|
||||
sql::Sql,
|
||||
};
|
||||
|
||||
|
||||
@@ -15,6 +15,7 @@
|
||||
|
||||
use std::{collections::HashSet, sync::Arc};
|
||||
|
||||
use ::search::types::QueryParams;
|
||||
use arrow::array::{
|
||||
Array, Int64Array, TimestampMicrosecondArray, TimestampNanosecondArray, UInt64Array,
|
||||
};
|
||||
@@ -41,10 +42,7 @@ use datafusion::{
|
||||
};
|
||||
use futures::TryStreamExt;
|
||||
|
||||
use crate::service::search::{
|
||||
grpc::{QueryParams, storage::tantivy_search},
|
||||
index::IndexCondition,
|
||||
};
|
||||
use crate::service::search::{index::IndexCondition, tantivy::tantivy_search};
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct TantivyOptimizeExec {
|
||||
|
||||
@@ -18,6 +18,7 @@ const BASE64_CHARS: &str = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ
|
||||
|
||||
use std::sync::{Arc, LazyLock as Lazy};
|
||||
|
||||
use ::search::cipher::registry::REGISTRY;
|
||||
use arrow::array::{ArrayRef, StringArray};
|
||||
use config::utils::str;
|
||||
use datafusion::{
|
||||
@@ -32,8 +33,6 @@ use datafusion::{
|
||||
use o2_enterprise::enterprise::cipher::Cipher;
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::cipher::registry::REGISTRY;
|
||||
|
||||
/// The name of the decrypt UDF given to DataFusion.
|
||||
pub const DECRYPT_UDF_NAME: &str = "decrypt_path";
|
||||
/// The name of the decrypt UDF given to DataFusion.
|
||||
|
||||
@@ -26,11 +26,10 @@ use datafusion::{
|
||||
logical_expr::{ColumnarValue, ScalarUDF, Volatility},
|
||||
prelude::create_udf,
|
||||
};
|
||||
use openobserve_vrl::{QUERY_FUNCTIONS, compile_vrl_function};
|
||||
use vector_enrichment::TableRegistry;
|
||||
use vrl::compiler::{TargetValueRef, VrlRuntime, runtime::Runtime};
|
||||
|
||||
use crate::{common::infra::config::QUERY_FUNCTIONS, service::ingestion::compile_vrl_function};
|
||||
|
||||
type FnType = Arc<dyn Fn(&[ColumnarValue]) -> Result<ColumnarValue> + Sync + Send>;
|
||||
|
||||
fn create_user_df(fn_name: &str, num_args: u8, pow_scalar: FnType) -> ScalarUDF {
|
||||
@@ -49,13 +48,11 @@ fn create_user_df(fn_name: &str, num_args: u8, pow_scalar: FnType) -> ScalarUDF
|
||||
|
||||
pub fn get_all_transform(org_id: &str) -> Result<Vec<ScalarUDF>> {
|
||||
let mut udf_list = Vec::new();
|
||||
for transform in QUERY_FUNCTIONS.clone().iter() {
|
||||
let key = transform.key();
|
||||
// do not register ingest_time transforms
|
||||
if key.contains(org_id) {
|
||||
for transform in QUERY_FUNCTIONS.iter() {
|
||||
if transform.key().contains(org_id) {
|
||||
udf_list.push(get_udf_vrl(
|
||||
transform.name.to_owned(),
|
||||
transform.function.to_owned().as_str(),
|
||||
transform.name.clone(),
|
||||
transform.function.as_str(),
|
||||
&transform.params,
|
||||
transform.num_args,
|
||||
org_id,
|
||||
@@ -67,12 +64,12 @@ pub fn get_all_transform(org_id: &str) -> Result<Vec<ScalarUDF>> {
|
||||
|
||||
fn get_udf_vrl(
|
||||
fn_name: String,
|
||||
func: &str,
|
||||
source: &str,
|
||||
params: &str,
|
||||
num_args: u8,
|
||||
org_id: &str,
|
||||
) -> Result<ScalarUDF> {
|
||||
let local_func = func.trim().to_owned();
|
||||
let local_func = source.trim().to_owned();
|
||||
let local_fn_params = params.to_owned();
|
||||
let local_org_id = org_id.to_owned();
|
||||
|
||||
@@ -95,10 +92,10 @@ fn get_udf_vrl(
|
||||
}
|
||||
obj_str.push_str(&format!(" \n {}", local_func));
|
||||
match compile_vrl_function(&obj_str, &local_org_id) {
|
||||
Ok(res) => {
|
||||
let registry = res.config.get_custom::<TableRegistry>().unwrap();
|
||||
Ok(result) => {
|
||||
let registry = result.config.get_custom::<TableRegistry>().unwrap();
|
||||
registry.finish_load();
|
||||
let result = apply_vrl_fn(&mut runtime, res.program);
|
||||
let result = apply_vrl_fn(&mut runtime, result.program);
|
||||
if result != json::Value::Null {
|
||||
res_data_vec.insert(i, json::get_string_value(&result));
|
||||
} else {
|
||||
|
||||
@@ -16,14 +16,8 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use arrow_schema::Schema;
|
||||
use config::meta::{
|
||||
search::ScanStats,
|
||||
stream::{FileKey, StreamType},
|
||||
};
|
||||
use datafusion::{
|
||||
datasource::TableProvider, execution::cache::cache_manager::FileStatisticsCache,
|
||||
sql::TableReference,
|
||||
};
|
||||
use config::meta::{search::ScanStats, stream::FileKey};
|
||||
use datafusion::{datasource::TableProvider, execution::cache::cache_manager::FileStatisticsCache};
|
||||
use hashbrown::HashSet;
|
||||
use infra::errors::Result;
|
||||
|
||||
@@ -32,29 +26,12 @@ use super::{datafusion::exec::TableBuilder, index::IndexCondition};
|
||||
pub mod flight;
|
||||
pub mod storage;
|
||||
pub mod wal;
|
||||
pub use ::search::{file_cache::calc_target_partitions, types::QueryParams};
|
||||
|
||||
pub use super::tantivy::cache as tantivy_result_cache;
|
||||
|
||||
pub type SearchTable = Result<(Vec<Arc<dyn TableProvider>>, ScanStats, HashSet<u64>)>;
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct QueryParams {
|
||||
pub trace_id: String,
|
||||
pub org_id: String,
|
||||
pub stream: TableReference,
|
||||
pub stream_type: StreamType,
|
||||
pub stream_name: String,
|
||||
pub time_range: (i64, i64),
|
||||
pub work_group: Option<String>,
|
||||
pub use_inverted_index: bool,
|
||||
}
|
||||
|
||||
/// Linear interpolation: cached_ratio=0 -> query_thread_num, cached_ratio=1 -> cpu_num.
|
||||
pub fn calc_target_partitions(cpu_num: usize, query_thread_num: usize, cached_ratio: f64) -> usize {
|
||||
(cpu_num as i64
|
||||
+ ((query_thread_num as i64 - cpu_num as i64) as f64 * (1.0 - cached_ratio)) as i64)
|
||||
as usize
|
||||
}
|
||||
|
||||
/// Create tables from files, automatically splitting them based on time range overlap:
|
||||
/// - Files completely within the query time range: no timestamp filter applied
|
||||
/// - Files partially overlapping with the query time range: timestamp filter applied
|
||||
|
||||
@@ -15,10 +15,11 @@
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
pub use ::search::file_cache::{cache_files, calc_target_partitions};
|
||||
use arrow_schema::Schema;
|
||||
use config::{
|
||||
cluster::LOCAL_NODE,
|
||||
get_config, is_local_disk_storage,
|
||||
get_config,
|
||||
meta::{
|
||||
inverted_index::IndexOptimizeMode,
|
||||
search::{ScanStats, StorageType},
|
||||
@@ -36,7 +37,6 @@ use infra::{
|
||||
use itertools::Itertools;
|
||||
use tracing::Instrument;
|
||||
|
||||
pub use super::calc_target_partitions;
|
||||
pub use crate::service::search::tantivy::tantivy_search;
|
||||
use crate::service::{
|
||||
file_list,
|
||||
@@ -297,134 +297,6 @@ pub async fn search(
|
||||
Ok((tables, scan_stats, HashSet::new()))
|
||||
}
|
||||
|
||||
#[tracing::instrument(name = "service:search:grpc:storage:cache_files", skip_all)]
|
||||
pub async fn cache_files(
|
||||
trace_id: &str,
|
||||
files: &[(i64, &String, &String, i64, i64)],
|
||||
scan_stats: &mut ScanStats,
|
||||
file_type: &str,
|
||||
) -> (file_data::CacheType, u64, u64) {
|
||||
// check how many files already cached
|
||||
let mut cached_files = HashSet::with_capacity(files.len());
|
||||
let (mut cache_hits, mut cache_misses) = (0, 0);
|
||||
|
||||
let start = std::time::Instant::now();
|
||||
for (_id, _account, file, _size, max_ts) in files.iter() {
|
||||
if file_data::memory::exist(file).await {
|
||||
scan_stats.querier_memory_cached_files += 1;
|
||||
cached_files.insert(file);
|
||||
cache_hits += 1;
|
||||
} else if file_data::disk::exist(file).await {
|
||||
scan_stats.querier_disk_cached_files += 1;
|
||||
cached_files.insert(file);
|
||||
cache_hits += 1;
|
||||
} else {
|
||||
cache_misses += 1;
|
||||
};
|
||||
|
||||
// Record file access metrics
|
||||
let stream_type = if file_type == "index" {
|
||||
config::meta::stream::StreamType::Index
|
||||
} else {
|
||||
// Determine stream type from the file path
|
||||
if file.contains("/logs/") {
|
||||
config::meta::stream::StreamType::Logs
|
||||
} else if file.contains("/metrics/") {
|
||||
config::meta::stream::StreamType::Metrics
|
||||
} else if file.contains("/traces/") {
|
||||
config::meta::stream::StreamType::Traces
|
||||
} else {
|
||||
config::meta::stream::StreamType::Logs // Default
|
||||
}
|
||||
};
|
||||
|
||||
let current_time = chrono::Utc::now().timestamp_micros();
|
||||
let file_age_seconds = (current_time - max_ts) / 1_000_000;
|
||||
let file_age_hours = file_age_seconds as f64 / 3600.0;
|
||||
|
||||
if file_age_hours > 0.0 {
|
||||
config::metrics::FILE_ACCESS_TIME
|
||||
.with_label_values(&[&stream_type.to_string()])
|
||||
.observe(file_age_hours);
|
||||
}
|
||||
}
|
||||
|
||||
let check_cache_took = start.elapsed().as_millis() as usize;
|
||||
if check_cache_took > 1000 {
|
||||
log::warn!(
|
||||
"[trace_id {trace_id}] search->storage: check file cache took: {check_cache_took} ms",
|
||||
);
|
||||
}
|
||||
|
||||
let files_num = files.len() as i64;
|
||||
if files_num == scan_stats.querier_memory_cached_files + scan_stats.querier_disk_cached_files {
|
||||
// all files are cached
|
||||
return (file_data::CacheType::Disk, cache_hits, cache_misses);
|
||||
}
|
||||
|
||||
// check cache size
|
||||
let cfg = get_config();
|
||||
let cache_type = if cfg.memory_cache.enabled
|
||||
&& scan_stats.compressed_size < cfg.memory_cache.skip_size as i64
|
||||
{
|
||||
// if scan_compressed_size < ZO_MEMORY_CACHE_SKIP_SIZE, use memory cache
|
||||
file_data::CacheType::Memory
|
||||
} else if !is_local_disk_storage()
|
||||
&& cfg.disk_cache.enabled
|
||||
&& scan_stats.compressed_size < cfg.disk_cache.skip_size as i64
|
||||
{
|
||||
// if scan_compressed_size < ZO_DISK_CACHE_SKIP_SIZE, use disk cache
|
||||
file_data::CacheType::Disk
|
||||
} else {
|
||||
// no cache, the files are too big than cache size
|
||||
return (file_data::CacheType::None, cache_hits, cache_misses);
|
||||
};
|
||||
|
||||
let trace_id = trace_id.to_string();
|
||||
let files = files
|
||||
.iter()
|
||||
.filter_map(|(id, account, file, size, ts)| {
|
||||
if cached_files.contains(&file) {
|
||||
None
|
||||
} else {
|
||||
Some((*id, account.to_string(), file.to_string(), *size, *ts))
|
||||
}
|
||||
})
|
||||
.collect_vec();
|
||||
let file_type = file_type.to_string();
|
||||
tokio::spawn(async move {
|
||||
let files_num = files.len();
|
||||
for (id, account, file, size, ts) in files {
|
||||
if let Err(e) = crate::service::file_downloader::queue_download(
|
||||
trace_id.clone(),
|
||||
id,
|
||||
account,
|
||||
file.clone(),
|
||||
size,
|
||||
ts,
|
||||
cache_type,
|
||||
)
|
||||
.await
|
||||
{
|
||||
log::error!(
|
||||
"[trace_id {trace_id}] error in queuing file {file} for background download: {e}"
|
||||
);
|
||||
}
|
||||
}
|
||||
log::info!(
|
||||
"[trace_id {trace_id}] search->storage: successfully enqueued {files_num} files of {file_type} for background download into {cache_type:?}",
|
||||
);
|
||||
});
|
||||
|
||||
// if cached file less than 50% of the total files, return None
|
||||
if scan_stats.querier_memory_cached_files + scan_stats.querier_disk_cached_files < files_num / 2
|
||||
{
|
||||
(file_data::CacheType::None, cache_hits, cache_misses)
|
||||
} else {
|
||||
(cache_type, cache_hits, cache_misses)
|
||||
}
|
||||
}
|
||||
|
||||
/// Check bloom filter for the file list
|
||||
#[tracing::instrument(name = "service:search:grpc:storage:check_bloom_filter", skip_all)]
|
||||
pub async fn check_bloom_filter(
|
||||
|
||||
@@ -23,6 +23,45 @@ use sqlparser::{
|
||||
parser::Parser,
|
||||
};
|
||||
|
||||
use super::{RE_HISTOGRAM, visitor::histogram_interval::generate_histogram_interval};
|
||||
|
||||
pub fn handle_histogram(
|
||||
origin_sql: &mut String,
|
||||
q_time_range: (i64, i64),
|
||||
histogram_interval: i64,
|
||||
) {
|
||||
let Some(captures) = RE_HISTOGRAM.captures(origin_sql.as_str()) else {
|
||||
return;
|
||||
};
|
||||
|
||||
let args = match captures.get(1) {
|
||||
Some(args) => args
|
||||
.as_str()
|
||||
.split(',')
|
||||
.map(|value| {
|
||||
value
|
||||
.trim()
|
||||
.trim_matches(|value| value == '\'' || value == '"')
|
||||
})
|
||||
.collect::<Vec<_>>(),
|
||||
None => return,
|
||||
};
|
||||
|
||||
let interval = if histogram_interval > 0 {
|
||||
format!("{histogram_interval} second")
|
||||
} else {
|
||||
args.get(1)
|
||||
.map_or_else(|| generate_histogram_interval(q_time_range), |value| *value)
|
||||
.to_string()
|
||||
};
|
||||
let field = args.first().unwrap_or(&"_timestamp");
|
||||
|
||||
*origin_sql = origin_sql.replace(
|
||||
captures.get(0).unwrap().as_str(),
|
||||
&format!("histogram({field},'{interval}')"),
|
||||
);
|
||||
}
|
||||
|
||||
/// Converts an original query to a histogram query
|
||||
/// Extracts WHERE clause and builds histogram query with provided stream name
|
||||
pub fn convert_to_histogram_query(
|
||||
@@ -201,6 +240,16 @@ pub fn histogram_bucket_start(timestamp_us: i64, interval_us: i64) -> i64 {
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_handle_histogram() {
|
||||
// Test case 1: Basic histogram with numeric interval
|
||||
let mut sql = "SELECT histogram(_timestamp, '10 seconds') FROM logs".to_string();
|
||||
let time_range = (1640995200000000, 1641081600000000); // 2022-01-01 to 2022-01-02
|
||||
handle_histogram(&mut sql, time_range, 10);
|
||||
assert!(sql.contains("histogram(_timestamp,"));
|
||||
assert!(sql.contains("second"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_convert_simple_query() {
|
||||
let original_query = "SELECT * FROM \"logs\" WHERE status = 500";
|
||||
|
||||
@@ -28,7 +28,7 @@ use hashbrown::HashSet;
|
||||
use infra::errors::Error;
|
||||
|
||||
use crate::service::search::{
|
||||
cluster::flight::{SearchContextBuilder, register_table},
|
||||
datafusion::context::{SearchContextBuilder, register_table},
|
||||
sql::Sql,
|
||||
};
|
||||
|
||||
|
||||
@@ -44,10 +44,12 @@ use tracing_opentelemetry::OpenTelemetrySpanExt;
|
||||
|
||||
use crate::service::search::{
|
||||
SEARCH_SERVER, SearchResult,
|
||||
cluster::flight::{SearchContextBuilder, register_table},
|
||||
datafusion::optimizer::{
|
||||
context::{PhysicalOptimizerContext, RemoteScanContext, StreamingAggregationContext},
|
||||
create_physical_plan,
|
||||
datafusion::{
|
||||
context::{SearchContextBuilder, register_table},
|
||||
optimizer::{
|
||||
context::{PhysicalOptimizerContext, RemoteScanContext, StreamingAggregationContext},
|
||||
create_physical_plan,
|
||||
},
|
||||
},
|
||||
inspector::{SearchInspectorFieldsBuilder, search_inspector_fields},
|
||||
sql::Sql,
|
||||
|
||||
@@ -22,6 +22,10 @@ mod warm;
|
||||
|
||||
use std::{ops::Bound, sync::Arc};
|
||||
|
||||
use ::search::{
|
||||
file_cache::{cache_files, calc_target_partitions},
|
||||
types::QueryParams,
|
||||
};
|
||||
use arrow::{
|
||||
buffer::{BooleanBuffer, MutableBuffer},
|
||||
util::bit_util,
|
||||
@@ -62,7 +66,6 @@ use self::{
|
||||
warm::WarmPlan,
|
||||
};
|
||||
use crate::service::search::{
|
||||
grpc::{QueryParams, calc_target_partitions, storage::cache_files},
|
||||
index::IndexCondition,
|
||||
inspector::{SearchInspectorFieldsBuilder, search_inspector_fields},
|
||||
};
|
||||
|
||||
@@ -19,6 +19,7 @@ use std::{
|
||||
sync::{Arc, atomic::Ordering},
|
||||
};
|
||||
|
||||
use ::search::CAPPED_RESULTS_MSG;
|
||||
use config::meta::search::{PARTIAL_ERROR_RESPONSE_MESSAGE, ScanStats};
|
||||
use datafusion::physical_plan::{ExecutionPlan, ExecutionPlanVisitor};
|
||||
use infra::runtime::DATAFUSION_RUNTIME;
|
||||
@@ -26,7 +27,7 @@ use sqlparser::ast::{BinaryOperator, Expr};
|
||||
use tokio::sync::Mutex;
|
||||
|
||||
use super::datafusion::distributed_plan::remote_scan_exec::RemoteScanExec;
|
||||
use crate::{common::meta::search::CAPPED_RESULTS_MSG, service::search::sql::Sql};
|
||||
use crate::service::search::sql::Sql;
|
||||
|
||||
type Cleanup = Pin<Box<dyn Future<Output = ()> + Send>>;
|
||||
|
||||
|
||||
+733
@@ -0,0 +1,733 @@
|
||||
// Copyright 2026 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 <http://www.gnu.org/licenses/>.
|
||||
|
||||
use std::{
|
||||
collections::VecDeque,
|
||||
sync::{Arc, LazyLock as Lazy},
|
||||
};
|
||||
|
||||
use config::{
|
||||
cluster::LOCAL_NODE,
|
||||
get_config,
|
||||
meta::cluster::{Role, RoleGroup, get_internal_grpc_token},
|
||||
metrics,
|
||||
utils::time::{day_micros, now_micros},
|
||||
};
|
||||
use futures::StreamExt;
|
||||
use hashbrown::{HashMap, HashSet};
|
||||
use proto::cluster_rpc::{SimpleFileList, event_client::EventClient};
|
||||
use tokio::sync::{
|
||||
Mutex,
|
||||
mpsc::{Receiver, Sender},
|
||||
};
|
||||
use tonic::{codec::CompressionEncoding, metadata::MetadataValue};
|
||||
|
||||
use crate::{cache::file_data, cluster};
|
||||
|
||||
/// (trace_id, file_id, account, file, size, cache_type)
|
||||
type FileInfo = (String, i64, String, String, usize, file_data::CacheType);
|
||||
|
||||
mod processing_files {
|
||||
use hashbrown::HashSet;
|
||||
use parking_lot::RwLock;
|
||||
|
||||
use super::*;
|
||||
|
||||
static PROCESSING_FILES: Lazy<RwLock<HashSet<String>>> =
|
||||
Lazy::new(|| RwLock::new(HashSet::new()));
|
||||
|
||||
pub fn is_processing(file_name: &str) -> bool {
|
||||
PROCESSING_FILES.read().contains(file_name)
|
||||
}
|
||||
|
||||
pub fn add(file_name: &str) {
|
||||
PROCESSING_FILES.write().insert(file_name.to_string());
|
||||
}
|
||||
|
||||
pub fn remove(file_name: &str) {
|
||||
PROCESSING_FILES.write().remove(file_name);
|
||||
}
|
||||
}
|
||||
|
||||
mod queued_files {
|
||||
use hashbrown::HashSet;
|
||||
use parking_lot::RwLock;
|
||||
|
||||
use super::*;
|
||||
|
||||
static QUEUED_FILES: Lazy<RwLock<HashSet<String>>> = Lazy::new(|| RwLock::new(HashSet::new()));
|
||||
|
||||
/// Returns true if newly inserted (caller should enqueue), false if already present (skip).
|
||||
pub fn try_add(file_name: &str) -> bool {
|
||||
QUEUED_FILES.write().insert(file_name.to_string())
|
||||
}
|
||||
|
||||
pub fn remove(file_name: &str) {
|
||||
QUEUED_FILES.write().remove(file_name);
|
||||
}
|
||||
}
|
||||
|
||||
struct DownloadQueue {
|
||||
sender: Sender<FileInfo>,
|
||||
receiver: Arc<Mutex<Receiver<FileInfo>>>,
|
||||
}
|
||||
|
||||
impl DownloadQueue {
|
||||
fn new(sender: Sender<FileInfo>, receiver: Arc<Mutex<Receiver<FileInfo>>>) -> Self {
|
||||
Self { sender, receiver }
|
||||
}
|
||||
}
|
||||
|
||||
struct PriorityDownloadQueue {
|
||||
stack: Mutex<VecDeque<FileInfo>>,
|
||||
notify: tokio::sync::Notify,
|
||||
max_size: usize,
|
||||
}
|
||||
|
||||
impl PriorityDownloadQueue {
|
||||
fn new(max_size: usize) -> Self {
|
||||
Self {
|
||||
stack: Mutex::new(VecDeque::with_capacity(max_size)),
|
||||
notify: tokio::sync::Notify::new(),
|
||||
max_size,
|
||||
}
|
||||
}
|
||||
|
||||
// Returns false if at cap — caller must remove from queued_files and skip metric inc.
|
||||
async fn push(&self, file_info: FileInfo) -> bool {
|
||||
let mut stack = self.stack.lock().await;
|
||||
if stack.len() >= self.max_size {
|
||||
return false;
|
||||
}
|
||||
stack.push_back(file_info);
|
||||
drop(stack);
|
||||
self.notify.notify_one();
|
||||
true
|
||||
}
|
||||
|
||||
// Blocks until an item is available. LIFO via pop_back.
|
||||
// notified() is created before the lock check to avoid missed-wakeup race.
|
||||
// Chains notify_one() so all workers drain the queue under burst load.
|
||||
async fn pop(&self) -> FileInfo {
|
||||
loop {
|
||||
let notified = self.notify.notified();
|
||||
{
|
||||
let mut stack = self.stack.lock().await;
|
||||
if let Some(item) = stack.pop_back() {
|
||||
if !stack.is_empty() {
|
||||
self.notify.notify_one();
|
||||
}
|
||||
return item;
|
||||
}
|
||||
}
|
||||
notified.await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
const FILE_DOWNLOAD_QUEUE_SIZE: usize = 10000;
|
||||
static FILE_DOWNLOAD_CHANNEL: Lazy<DownloadQueue> = Lazy::new(|| {
|
||||
let (tx, rx) = tokio::sync::mpsc::channel::<FileInfo>(FILE_DOWNLOAD_QUEUE_SIZE);
|
||||
DownloadQueue::new(tx, Arc::new(Mutex::new(rx)))
|
||||
});
|
||||
|
||||
static PRIORITY_FILE_DOWNLOAD_CHANNEL: Lazy<PriorityDownloadQueue> =
|
||||
Lazy::new(|| PriorityDownloadQueue::new(FILE_DOWNLOAD_QUEUE_SIZE));
|
||||
|
||||
pub async fn run() -> Result<(), anyhow::Error> {
|
||||
let cfg = get_config();
|
||||
// worker tasks: handle normal queue download (FIFO via mpsc channel)
|
||||
for thread in 0..cfg.limit.file_download_thread_num {
|
||||
let rx = FILE_DOWNLOAD_CHANNEL.receiver.clone();
|
||||
tokio::spawn(async move {
|
||||
loop {
|
||||
let ret = rx.lock().await.recv().await;
|
||||
match ret {
|
||||
None => {
|
||||
log::debug!("[FILE_CACHE_DOWNLOAD:JOB:NORMAL] Receiving channel is closed");
|
||||
break;
|
||||
}
|
||||
Some((trace_id, id, account, file, file_size, cache)) => {
|
||||
// transition: queued → not-queued (must run before any skip path)
|
||||
queued_files::remove(&file);
|
||||
|
||||
// check if the file is already being downloaded
|
||||
if processing_files::is_processing(&file) {
|
||||
log::warn!(
|
||||
"[trace_id {trace_id}] [thread {thread}] search->storage: file {file} is already being downloaded, will skip it"
|
||||
);
|
||||
// update metrics
|
||||
metrics::FILE_DOWNLOADER_NORMAL_QUEUE_SIZE
|
||||
.with_label_values::<&str>(&[])
|
||||
.dec();
|
||||
continue;
|
||||
}
|
||||
|
||||
// add the file to processing set
|
||||
processing_files::add(&file);
|
||||
|
||||
// download the file
|
||||
match download_file(
|
||||
thread, &trace_id, id, &account, &file, file_size, cache,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(data_len) => {
|
||||
if data_len > 0 && data_len != file_size {
|
||||
log::warn!(
|
||||
"[FILE_CACHE_DOWNLOAD:JOB:NORMAL] download file {file} found size mismatch, expected: {file_size}, actual: {data_len}, will skip it",
|
||||
);
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
log::error!(
|
||||
"[FILE_CACHE_DOWNLOAD:JOB:NORMAL] download file {file} to cache {cache:?} err: {e}",
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
// remove the file from processing set
|
||||
processing_files::remove(&file);
|
||||
|
||||
// update metrics
|
||||
metrics::FILE_DOWNLOADER_NORMAL_QUEUE_SIZE
|
||||
.with_label_values::<&str>(&[])
|
||||
.dec();
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
// worker tasks: handle priority queue download (LIFO via VecDeque::pop_back)
|
||||
for thread in 0..cfg.limit.file_download_priority_queue_thread_num {
|
||||
tokio::spawn(async move {
|
||||
loop {
|
||||
let (trace_id, id, account, file, file_size, cache) =
|
||||
PRIORITY_FILE_DOWNLOAD_CHANNEL.pop().await;
|
||||
|
||||
// transition: queued → not-queued (must run before any skip path)
|
||||
queued_files::remove(&file);
|
||||
|
||||
// check if the file is already being downloaded
|
||||
if processing_files::is_processing(&file) {
|
||||
log::warn!(
|
||||
"[trace_id {trace_id}] [thread {thread}] search->storage: file {file} is already being downloaded, will skip it"
|
||||
);
|
||||
metrics::FILE_DOWNLOADER_PRIORITY_QUEUE_SIZE
|
||||
.with_label_values::<&str>(&[])
|
||||
.dec();
|
||||
continue;
|
||||
}
|
||||
|
||||
// add the file to processing set
|
||||
processing_files::add(&file);
|
||||
|
||||
// download the file
|
||||
match download_file(thread, &trace_id, id, &account, &file, file_size, cache).await
|
||||
{
|
||||
Ok(data_len) => {
|
||||
if data_len > 0 && data_len != file_size {
|
||||
log::warn!(
|
||||
"[FILE_CACHE_DOWNLOAD:JOB:PRIORITY] download file {file} found size mismatch, expected: {file_size}, actual: {data_len}, will skip it",
|
||||
);
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
log::error!(
|
||||
"[FILE_CACHE_DOWNLOAD:JOB:PRIORITY] download file {file} to cache {cache:?} err: {e}",
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
// remove the file from processing set
|
||||
processing_files::remove(&file);
|
||||
|
||||
// update metrics
|
||||
metrics::FILE_DOWNLOADER_PRIORITY_QUEUE_SIZE
|
||||
.with_label_values::<&str>(&[])
|
||||
.dec();
|
||||
}
|
||||
});
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn download_file(
|
||||
thread: usize,
|
||||
trace_id: &str,
|
||||
file_id: i64,
|
||||
account: &str,
|
||||
file_name: &str,
|
||||
file_size: usize,
|
||||
cache_type: file_data::CacheType,
|
||||
) -> Result<usize, anyhow::Error> {
|
||||
let cfg = get_config();
|
||||
|
||||
// download file from node
|
||||
if cfg.cache_latest_files.download_from_node
|
||||
&& let Ok(ok) =
|
||||
download_file_with_consistent_hash(file_id, account, file_name, file_size).await
|
||||
&& ok
|
||||
{
|
||||
return Ok(file_size);
|
||||
}
|
||||
|
||||
// download from object store
|
||||
let size = if file_size > 0 { Some(file_size) } else { None };
|
||||
let start = std::time::Instant::now();
|
||||
let ret = match cache_type {
|
||||
file_data::CacheType::Memory => {
|
||||
let mut disk_exists = false;
|
||||
let mem_exists = file_data::memory::exist(file_name).await;
|
||||
if !mem_exists && !cfg.memory_cache.skip_disk_check {
|
||||
disk_exists = file_data::disk::exist(file_name).await;
|
||||
}
|
||||
if !mem_exists && (cfg.memory_cache.skip_disk_check || !disk_exists) {
|
||||
file_data::memory::download(account, file_name, size).await
|
||||
} else {
|
||||
Ok(0)
|
||||
}
|
||||
}
|
||||
file_data::CacheType::Disk => {
|
||||
if !file_data::disk::exist(file_name).await {
|
||||
file_data::disk::download(account, file_name, size).await
|
||||
} else {
|
||||
Ok(0)
|
||||
}
|
||||
}
|
||||
_ => Ok(0),
|
||||
};
|
||||
log::debug!(
|
||||
"[FILE_CACHE_DOWNLOAD:JOB:{thread}] [trace_id {trace_id}] download file: {file_name}, ret: {:?}, took: {} ms",
|
||||
ret,
|
||||
start.elapsed().as_millis()
|
||||
);
|
||||
ret
|
||||
}
|
||||
|
||||
async fn download_file_with_consistent_hash(
|
||||
file_id: i64,
|
||||
account: &str,
|
||||
file_name: &str,
|
||||
file_size: usize,
|
||||
) -> Result<bool, anyhow::Error> {
|
||||
let role_group = if LOCAL_NODE.is_interactive_querier() {
|
||||
RoleGroup::Interactive
|
||||
} else {
|
||||
RoleGroup::Background
|
||||
};
|
||||
let Some(node_name) = cluster::get_node_from_consistent_hash(
|
||||
&file_id.to_string(),
|
||||
&Role::Querier,
|
||||
Some(role_group),
|
||||
)
|
||||
.await
|
||||
else {
|
||||
return Ok(false);
|
||||
};
|
||||
// get node by file_id
|
||||
let Some(node) = cluster::get_cached_node_by_name(&node_name).await else {
|
||||
return Ok(false);
|
||||
};
|
||||
// download file from node
|
||||
let Ok(failed) = download_from_node(
|
||||
&node.grpc_addr,
|
||||
&[(
|
||||
file_id,
|
||||
account.to_string(),
|
||||
file_name.to_string(),
|
||||
file_size as i64,
|
||||
0,
|
||||
)],
|
||||
)
|
||||
.await
|
||||
else {
|
||||
return Ok(false);
|
||||
};
|
||||
// failed is empty means download success
|
||||
Ok(failed.is_empty())
|
||||
}
|
||||
|
||||
// download files from node and return download failed files
|
||||
// file: (account, file, size, ts)
|
||||
pub async fn download_from_node(
|
||||
addr: &str,
|
||||
files: &[(i64, String, String, i64, i64)],
|
||||
) -> Result<Vec<(i64, String, String, i64, i64)>, anyhow::Error> {
|
||||
let start = std::time::Instant::now();
|
||||
let cfg = get_config();
|
||||
log::debug!(
|
||||
"[FILE_CACHE_DOWNLOAD:gRPC] Download files from node start, files: {:?}",
|
||||
files.len()
|
||||
);
|
||||
|
||||
let token: MetadataValue<_> = get_internal_grpc_token()
|
||||
.parse()
|
||||
.map_err(|_| anyhow::anyhow!("Invalid token"))?;
|
||||
|
||||
let channel = crate::client::grpc::get_cached_channel(addr).await?;
|
||||
let client = EventClient::with_interceptor(channel, move |mut req: tonic::Request<()>| {
|
||||
req.metadata_mut().insert("authorization", token.clone());
|
||||
Ok(req)
|
||||
});
|
||||
|
||||
let file_size_map = files
|
||||
.iter()
|
||||
.filter_map(|(_, _, f, s, _)| {
|
||||
if *s > cfg.cache_latest_files.download_node_size * 1024 * 1024 {
|
||||
None
|
||||
} else {
|
||||
Some((f, *s as usize))
|
||||
}
|
||||
})
|
||||
.collect::<HashMap<_, _>>();
|
||||
if file_size_map.is_empty() {
|
||||
return Ok(files.to_vec());
|
||||
}
|
||||
let request = tonic::Request::new(SimpleFileList {
|
||||
files: files.iter().map(|(_, _, f, ..)| f.to_string()).collect(),
|
||||
});
|
||||
|
||||
let resp = client
|
||||
.send_compressed(CompressionEncoding::Gzip)
|
||||
.accept_compressed(CompressionEncoding::Gzip)
|
||||
.max_decoding_message_size(cfg.grpc.max_message_size * 1024 * 1024)
|
||||
.max_encoding_message_size(cfg.grpc.max_message_size * 1024 * 1024)
|
||||
.get_files(request)
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("Failed to get files from {addr}, {e}"))?;
|
||||
|
||||
let mut file_contents = HashMap::new();
|
||||
let mut downloaded_files = HashSet::new();
|
||||
let mut resp_stream = resp.into_inner();
|
||||
while let Some(resp) = resp_stream.next().await {
|
||||
let resp = match resp {
|
||||
Ok(resp) => resp,
|
||||
Err(err) => {
|
||||
if err.code() == tonic::Code::NotFound {
|
||||
log::debug!(
|
||||
"[FILE_CACHE_DOWNLOAD:gRPC] Failed to download file {} from {}: file not found",
|
||||
err.message(),
|
||||
addr
|
||||
);
|
||||
continue;
|
||||
}
|
||||
return Err(anyhow::anyhow!(
|
||||
"Failed to download file from {addr}, {err}"
|
||||
));
|
||||
}
|
||||
};
|
||||
for content in resp.entries {
|
||||
let entry = file_contents
|
||||
.entry(content.filename.clone())
|
||||
.or_insert(bytes::BytesMut::new());
|
||||
entry.extend_from_slice(&content.content);
|
||||
downloaded_files.insert(content.filename);
|
||||
}
|
||||
}
|
||||
|
||||
log::debug!(
|
||||
"[FILE_CACHE_DOWNLOAD:gRPC] Successfully retrieved {} files from {} in {} ms",
|
||||
downloaded_files.len(),
|
||||
addr,
|
||||
start.elapsed().as_millis()
|
||||
);
|
||||
|
||||
// Cache the file contents
|
||||
for (file, content) in file_contents {
|
||||
let data = content.freeze();
|
||||
if let Some(size) = file_size_map.get(&file)
|
||||
&& *size != data.len()
|
||||
{
|
||||
log::warn!(
|
||||
"[FILE_CACHE_DOWNLOAD:gRPC] Failed to download file {} from {}: size mismatch, expected {} but got {}",
|
||||
file,
|
||||
addr,
|
||||
size,
|
||||
data.len()
|
||||
);
|
||||
downloaded_files.remove(&file);
|
||||
continue;
|
||||
}
|
||||
if let Err(e) = file_data::set(&file, data).await {
|
||||
log::error!("[FILE_CACHE_DOWNLOAD:gRPC] Failed to cache file {file}: {e}");
|
||||
downloaded_files.remove(&file);
|
||||
}
|
||||
}
|
||||
|
||||
// Return list of failed files
|
||||
let failed_files: Vec<_> = files
|
||||
.iter()
|
||||
.filter(|(_, f, ..)| !downloaded_files.contains(f))
|
||||
.cloned()
|
||||
.collect();
|
||||
|
||||
log::debug!(
|
||||
"[FILE_CACHE_DOWNLOAD:gRPC] Failed to retrieve {} files from {} in {} ms",
|
||||
failed_files.len(),
|
||||
addr,
|
||||
start.elapsed().as_millis()
|
||||
);
|
||||
|
||||
Ok(failed_files)
|
||||
}
|
||||
|
||||
pub async fn queue_download(
|
||||
trace_id: String,
|
||||
id: i64,
|
||||
account: String,
|
||||
file: String,
|
||||
size: i64,
|
||||
ts: i64,
|
||||
cache_type: file_data::CacheType,
|
||||
) -> Result<(), anyhow::Error> {
|
||||
log::debug!(
|
||||
"[FILE_CACHE_DOWNLOAD:JOB] [trace_id {trace_id}] enqueue file: {file}, size: {size}, ts: {ts}"
|
||||
);
|
||||
|
||||
// files with data older than the cache max age are not worth caching: with
|
||||
// time_lru they become the oldest bucket and are evicted first, so the
|
||||
// download would be wasted work. Skip the queue and let the query read
|
||||
// them directly from object storage.
|
||||
if exceeds_cache_max_age(ts, cache_type) {
|
||||
log::debug!(
|
||||
"[FILE_CACHE_DOWNLOAD:JOB] [trace_id {trace_id}] skip file: {file}, data is older than the cache max age"
|
||||
);
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let cfg = get_config();
|
||||
|
||||
// skip if already queued or already being downloaded
|
||||
if !queued_files::try_add(&file) {
|
||||
return Ok(());
|
||||
}
|
||||
if processing_files::is_processing(&file) {
|
||||
queued_files::remove(&file);
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
if cfg.limit.file_download_enable_priority_queue
|
||||
&& should_prioritize_file(ts, cfg.limit.file_download_priority_queue_window_secs)
|
||||
{
|
||||
if PRIORITY_FILE_DOWNLOAD_CHANNEL
|
||||
.push((
|
||||
trace_id,
|
||||
id,
|
||||
account,
|
||||
file.clone(),
|
||||
size as usize,
|
||||
cache_type,
|
||||
))
|
||||
.await
|
||||
{
|
||||
metrics::FILE_DOWNLOADER_PRIORITY_QUEUE_SIZE
|
||||
.with_label_values::<&str>(&[])
|
||||
.inc();
|
||||
} else {
|
||||
queued_files::remove(&file);
|
||||
log::warn!(
|
||||
"[FILE_CACHE_DOWNLOAD:JOB] priority queue full ({FILE_DOWNLOAD_QUEUE_SIZE}), dropping file: {file}"
|
||||
);
|
||||
}
|
||||
} else {
|
||||
if let Err(e) = FILE_DOWNLOAD_CHANNEL
|
||||
.sender
|
||||
.send((
|
||||
trace_id,
|
||||
id,
|
||||
account,
|
||||
file.clone(),
|
||||
size as usize,
|
||||
cache_type,
|
||||
))
|
||||
.await
|
||||
{
|
||||
queued_files::remove(&file);
|
||||
return Err(e.into());
|
||||
}
|
||||
metrics::FILE_DOWNLOADER_NORMAL_QUEUE_SIZE
|
||||
.with_label_values::<&str>(&[])
|
||||
.inc();
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
// if the file timestamp is in the past window, it should be prioritized
|
||||
fn should_prioritize_file(ts: i64, window_secs: i64) -> bool {
|
||||
let window_micros = window_secs * 1_000_000;
|
||||
let now = now_micros();
|
||||
ts > now - window_micros
|
||||
}
|
||||
|
||||
/// Returns true if the file's data is older than the cache max age and should
|
||||
/// not be downloaded into the cache.
|
||||
pub fn exceeds_cache_max_age(ts: i64, cache_type: file_data::CacheType) -> bool {
|
||||
let cfg = get_config();
|
||||
let max_age_days = match cache_type {
|
||||
file_data::CacheType::Memory => cfg.memory_cache.max_age_days,
|
||||
file_data::CacheType::Disk => cfg.disk_cache.max_age_days,
|
||||
file_data::CacheType::None => 0,
|
||||
};
|
||||
exceeds_max_age(ts, max_age_days)
|
||||
}
|
||||
|
||||
// if the file data is older than max_age_days, it should not be cached.
|
||||
// max_age_days == 0 means no limit, ts <= 0 means the timestamp is unknown.
|
||||
fn exceeds_max_age(ts: i64, max_age_days: i64) -> bool {
|
||||
if max_age_days <= 0 || ts <= 0 {
|
||||
return false;
|
||||
}
|
||||
ts < now_micros() - day_micros(max_age_days)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use config::utils::time::{day_micros, hour_micros, now_micros};
|
||||
|
||||
use super::{
|
||||
FileInfo, PriorityDownloadQueue, exceeds_max_age, file_data, processing_files, queued_files,
|
||||
};
|
||||
|
||||
#[test]
|
||||
fn test_exceeds_max_age_disabled() {
|
||||
// max_age_days == 0 means no limit, nothing is too old
|
||||
let one_year_ago = now_micros() - day_micros(365);
|
||||
assert!(!exceeds_max_age(one_year_ago, 0));
|
||||
assert!(!exceeds_max_age(one_year_ago, -1));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_exceeds_max_age_unknown_ts() {
|
||||
// unknown timestamp should not be skipped
|
||||
assert!(!exceeds_max_age(0, 3));
|
||||
assert!(!exceeds_max_age(-1, 3));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_exceeds_max_age_recent_file() {
|
||||
let one_hour_ago = now_micros() - hour_micros(1);
|
||||
assert!(!exceeds_max_age(one_hour_ago, 3));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_exceeds_max_age_old_file() {
|
||||
let thirty_days_ago = now_micros() - day_micros(30);
|
||||
assert!(exceeds_max_age(thirty_days_ago, 3));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_is_processing_false_initially() {
|
||||
assert!(!processing_files::is_processing(
|
||||
"unique_not_added_file_abc123.parquet"
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_add_and_is_processing() {
|
||||
let name = "test_add_file_xyz999.parquet";
|
||||
processing_files::add(name);
|
||||
assert!(processing_files::is_processing(name));
|
||||
processing_files::remove(name);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_remove_clears_processing() {
|
||||
let name = "test_remove_file_qqq888.parquet";
|
||||
processing_files::add(name);
|
||||
processing_files::remove(name);
|
||||
assert!(!processing_files::is_processing(name));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_normal_queue_cap_honored() {
|
||||
let cap = 3;
|
||||
let (tx, _rx) = tokio::sync::mpsc::channel::<FileInfo>(cap);
|
||||
|
||||
for i in 0..cap {
|
||||
let result = tx.try_send((
|
||||
format!("trace-{i}"),
|
||||
i as i64,
|
||||
"org".into(),
|
||||
format!("file-{i}.parquet"),
|
||||
1024,
|
||||
file_data::CacheType::Disk,
|
||||
));
|
||||
assert!(result.is_ok(), "send {i} should succeed under cap");
|
||||
}
|
||||
|
||||
let overflow = tx.try_send((
|
||||
"trace-overflow".into(),
|
||||
99,
|
||||
"org".into(),
|
||||
"overflow.parquet".into(),
|
||||
1024,
|
||||
file_data::CacheType::Disk,
|
||||
));
|
||||
assert!(overflow.is_err(), "send beyond cap must fail");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_queued_files_no_duplicates() {
|
||||
let file = "dedup_test_unique_zz1234.parquet";
|
||||
|
||||
// first add succeeds
|
||||
assert!(queued_files::try_add(file));
|
||||
// same file rejected while already queued
|
||||
assert!(!queued_files::try_add(file));
|
||||
|
||||
// after remove, can be queued again
|
||||
queued_files::remove(file);
|
||||
assert!(queued_files::try_add(file));
|
||||
|
||||
// cleanup
|
||||
queued_files::remove(file);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_priority_queue_cap_honored() {
|
||||
let cap = 3;
|
||||
let queue = PriorityDownloadQueue::new(cap);
|
||||
|
||||
// fill to cap — all succeed
|
||||
for i in 0..cap {
|
||||
let ok = queue
|
||||
.push((
|
||||
format!("trace-{i}"),
|
||||
i as i64,
|
||||
"org".into(),
|
||||
format!("file-{i}.parquet"),
|
||||
1024,
|
||||
file_data::CacheType::Disk,
|
||||
))
|
||||
.await;
|
||||
assert!(ok, "push {i} should succeed under cap");
|
||||
}
|
||||
|
||||
// one more must be rejected
|
||||
let overflow = queue
|
||||
.push((
|
||||
"trace-overflow".into(),
|
||||
99,
|
||||
"org".into(),
|
||||
"overflow.parquet".into(),
|
||||
1024,
|
||||
file_data::CacheType::Disk,
|
||||
))
|
||||
.await;
|
||||
assert!(!overflow, "push beyond cap must return false");
|
||||
}
|
||||
}
|
||||
Vendored
+1
@@ -15,6 +15,7 @@
|
||||
|
||||
pub mod bytes_cache;
|
||||
pub mod file_data;
|
||||
pub mod file_downloader;
|
||||
pub mod meta;
|
||||
pub mod stats;
|
||||
pub mod storage;
|
||||
|
||||
@@ -33,6 +33,7 @@ hashbrown.workspace = true
|
||||
infra.workspace = true
|
||||
ingester.workspace = true
|
||||
log.workspace = true
|
||||
openobserve-vrl.workspace = true
|
||||
o2_dex = { workspace = true, optional = true }
|
||||
o2_enterprise = { workspace = true, optional = true, default-features = false }
|
||||
o2_openfga = { workspace = true, optional = true }
|
||||
|
||||
@@ -16,7 +16,8 @@
|
||||
use std::sync::LazyLock;
|
||||
|
||||
use config::{
|
||||
MMDB_ASN_FILE_NAME, MMDB_CITY_FILE_NAME,
|
||||
GEO_IP_ASN_ENRICHMENT_TABLE, GEO_IP_CITY_ENRICHMENT_TABLE, MMDB_ASN_FILE_NAME,
|
||||
MMDB_CITY_FILE_NAME,
|
||||
utils::download_utils::{download_file, is_digest_different},
|
||||
};
|
||||
#[cfg(feature = "enterprise")]
|
||||
@@ -24,10 +25,7 @@ use o2_enterprise::enterprise::common::config::get_config as get_o2_config;
|
||||
use tokio::time;
|
||||
|
||||
use crate::{
|
||||
common::{
|
||||
infra::config::{GEOIP_ASN_TABLE, GEOIP_CITY_TABLE, MAXMIND_DB_CLIENT},
|
||||
meta::maxmind::MaxmindClient,
|
||||
},
|
||||
common::{infra::config::MAXMIND_DB_CLIENT, meta::maxmind::MaxmindClient},
|
||||
service::enrichment_table::geoip::{Geoip, GeoipConfig, MMDB_INIT_NOTIFIER},
|
||||
};
|
||||
|
||||
@@ -156,8 +154,8 @@ async fn update_maxmind_table(fname: &str) {
|
||||
Ok(_) => {
|
||||
#[cfg(feature = "enterprise")]
|
||||
if get_o2_config().common.enable_enterprise_mmdb {
|
||||
let mut geoip = crate::common::infra::config::GEOIP_ENT_TABLE.write();
|
||||
*geoip = Some(
|
||||
openobserve_vrl::register_global_enrichment_table(
|
||||
o2_enterprise::enterprise::common::config::GEO_IP_ENTERPRISE_ENRICHMENT_TABLE,
|
||||
Geoip::new(GeoipConfig::new(
|
||||
&get_o2_config().common.mmdb_enterprise_file_name,
|
||||
))
|
||||
@@ -166,11 +164,15 @@ async fn update_maxmind_table(fname: &str) {
|
||||
}
|
||||
|
||||
if fname.ends_with(MMDB_CITY_FILE_NAME) {
|
||||
let mut geoip_city = GEOIP_CITY_TABLE.write();
|
||||
*geoip_city = Some(Geoip::new(GeoipConfig::new(MMDB_CITY_FILE_NAME)).unwrap());
|
||||
openobserve_vrl::register_global_enrichment_table(
|
||||
GEO_IP_CITY_ENRICHMENT_TABLE,
|
||||
Geoip::new(GeoipConfig::new(MMDB_CITY_FILE_NAME)).unwrap(),
|
||||
);
|
||||
} else if fname.ends_with(MMDB_ASN_FILE_NAME) {
|
||||
let mut geoip_asn = GEOIP_ASN_TABLE.write();
|
||||
*geoip_asn = Some(Geoip::new(GeoipConfig::new(MMDB_ASN_FILE_NAME)).unwrap());
|
||||
openobserve_vrl::register_global_enrichment_table(
|
||||
GEO_IP_ASN_ENRICHMENT_TABLE,
|
||||
Geoip::new(GeoipConfig::new(MMDB_ASN_FILE_NAME)).unwrap(),
|
||||
);
|
||||
};
|
||||
}
|
||||
Err(e) => log::warn!("Failed to update maxmind table with path: {fname}, {e}"),
|
||||
|
||||
+13
-5
@@ -7,18 +7,26 @@ publish = false
|
||||
|
||||
[features]
|
||||
default = []
|
||||
enterprise = []
|
||||
enterprise = ["dep:hashbrown", "dep:o2_enterprise", "dep:parking_lot"]
|
||||
vectorscan = []
|
||||
vortex = []
|
||||
|
||||
[dependencies]
|
||||
async-trait.workspace = true
|
||||
chrono.workspace = true
|
||||
config.workspace = true
|
||||
datafusion.workspace = true
|
||||
hashbrown = { workspace = true, optional = true }
|
||||
infra.workspace = true
|
||||
log.workspace = true
|
||||
openobserve-vrl.workspace = true
|
||||
o2_enterprise = { workspace = true, optional = true, default-features = false }
|
||||
parking_lot = { workspace = true, optional = true }
|
||||
serde.workspace = true
|
||||
tokio.workspace = true
|
||||
tracing.workspace = true
|
||||
utoipa.workspace = true
|
||||
vector-enrichment.workspace = true
|
||||
vrl.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
arrow.workspace = true
|
||||
arrow-schema.workspace = true
|
||||
o2_enterprise.workspace = true
|
||||
serde_json.workspace = true
|
||||
|
||||
@@ -0,0 +1,16 @@
|
||||
// Copyright 2026 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 <http://www.gnu.org/licenses/>.
|
||||
|
||||
pub mod registry;
|
||||
@@ -5,7 +5,7 @@
|
||||
// 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
|
||||
// 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.
|
||||
@@ -13,209 +13,4 @@
|
||||
// You should have received a copy of the GNU Affero General Public License
|
||||
// along with this program. If not, see <http://www.gnu.org/licenses/>.
|
||||
|
||||
use std::{
|
||||
collections::BTreeMap,
|
||||
sync::{Arc, LazyLock},
|
||||
};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use config::{RwHashMap, stats::MemorySize, utils::time::parse_str_to_time};
|
||||
use vector_enrichment::{Case, IndexHandle, Table};
|
||||
use vrl::value::{KeyString, ObjectMap, Value};
|
||||
|
||||
pub static ENRICHMENT_TABLES: LazyLock<RwHashMap<String, StreamTable>> =
|
||||
LazyLock::new(Default::default);
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct StreamTable {
|
||||
pub org_id: String,
|
||||
pub stream_name: String,
|
||||
pub data: Arc<Vec<Value>>,
|
||||
}
|
||||
|
||||
impl MemorySize for StreamTable {
|
||||
fn mem_size(&self) -> usize {
|
||||
std::mem::size_of::<StreamTable>()
|
||||
+ self.org_id.mem_size()
|
||||
+ self.stream_name.mem_size()
|
||||
+ self.data.iter().map(|v| v.to_string().len()).sum::<usize>()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Table for StreamTable {
|
||||
fn find_table_row(
|
||||
&self,
|
||||
case: Case,
|
||||
conditions: &[vector_enrichment::Condition],
|
||||
select: Option<&[String]>,
|
||||
_wildcard: Option<&Value>,
|
||||
_index: Option<IndexHandle>,
|
||||
) -> Result<ObjectMap, vector_enrichment::Error> {
|
||||
let resp = get_data(self, conditions, select, case);
|
||||
Ok(resp.first().cloned().unwrap_or_default())
|
||||
}
|
||||
|
||||
fn find_table_rows(
|
||||
&self,
|
||||
case: Case,
|
||||
conditions: &[vector_enrichment::Condition],
|
||||
select: Option<&[String]>,
|
||||
_wildcard: Option<&Value>,
|
||||
_index: Option<IndexHandle>,
|
||||
) -> Result<Vec<ObjectMap>, vector_enrichment::Error> {
|
||||
Ok(get_data(self, conditions, select, case))
|
||||
}
|
||||
|
||||
fn add_index(
|
||||
&mut self,
|
||||
_case: Case,
|
||||
_fields: &[&str],
|
||||
) -> Result<IndexHandle, vector_enrichment::Error> {
|
||||
Ok(IndexHandle(1))
|
||||
}
|
||||
|
||||
fn index_fields(&self) -> Vec<(Case, Vec<String>)> {
|
||||
Vec::new()
|
||||
}
|
||||
|
||||
fn needs_reload(&self) -> bool {
|
||||
false
|
||||
}
|
||||
}
|
||||
|
||||
fn get_data(
|
||||
table: &StreamTable,
|
||||
conditions: &[vector_enrichment::Condition],
|
||||
select: Option<&[String]>,
|
||||
case: Case,
|
||||
) -> Vec<ObjectMap> {
|
||||
if conditions.is_empty() {
|
||||
return vec![];
|
||||
}
|
||||
|
||||
let filtered = table.data.iter().filter(|value| {
|
||||
let Value::Object(map) = value else {
|
||||
return false;
|
||||
};
|
||||
|
||||
conditions.iter().all(|condition| match condition {
|
||||
vector_enrichment::Condition::Equals { field, value } => match case {
|
||||
Case::Insensitive => match (map.get(*field), value) {
|
||||
(Some(Value::Bytes(left)), Value::Bytes(right)) => {
|
||||
match (std::str::from_utf8(left), std::str::from_utf8(right)) {
|
||||
(Ok(left), Ok(right)) => left.eq_ignore_ascii_case(right),
|
||||
(Err(_), Err(_)) => left == right,
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
_ => false,
|
||||
},
|
||||
Case::Sensitive => map.get(*field).is_some_and(|item| item == value),
|
||||
},
|
||||
vector_enrichment::Condition::FromDate { field, from } => map
|
||||
.get(*field)
|
||||
.and_then(|value| value.as_str())
|
||||
.and_then(|value| parse_str_to_time(value.as_ref()).ok())
|
||||
.is_some_and(|date| date >= *from),
|
||||
vector_enrichment::Condition::ToDate { field, to } => map
|
||||
.get(*field)
|
||||
.and_then(|value| value.as_str())
|
||||
.and_then(|value| parse_str_to_time(value.as_ref()).ok())
|
||||
.is_some_and(|date| date <= *to),
|
||||
vector_enrichment::Condition::BetweenDates { field, from, to } => map
|
||||
.get(*field)
|
||||
.and_then(|value| value.as_str())
|
||||
.and_then(|value| parse_str_to_time(value.as_ref()).ok())
|
||||
.is_some_and(|date| date >= *from && date <= *to),
|
||||
})
|
||||
});
|
||||
|
||||
filtered
|
||||
.filter_map(Value::as_object)
|
||||
.map(|map| {
|
||||
if let Some(fields) = select {
|
||||
fields
|
||||
.iter()
|
||||
.filter_map(|field| map.get_key_value(field.as_str()))
|
||||
.map(|(key, value)| (key.clone(), value.clone()))
|
||||
.collect::<BTreeMap<KeyString, Value>>()
|
||||
} else {
|
||||
map.clone()
|
||||
}
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use serde_json::json;
|
||||
use vector_enrichment::{Condition, Table};
|
||||
|
||||
use super::*;
|
||||
|
||||
fn table() -> StreamTable {
|
||||
StreamTable {
|
||||
org_id: "default".to_string(),
|
||||
stream_name: "users".to_string(),
|
||||
data: Arc::new(vec![
|
||||
json!({"name": "Alice", "country": "US", "role": "admin"}).into(),
|
||||
json!({"name": "Bob", "country": "GB", "role": "user"}).into(),
|
||||
Value::from("not an object"),
|
||||
]),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn matches_rows_using_requested_case_sensitivity() {
|
||||
let table = table();
|
||||
let conditions = [Condition::Equals {
|
||||
field: "name",
|
||||
value: Value::from("ALICE"),
|
||||
}];
|
||||
|
||||
let sensitive = table
|
||||
.find_table_rows(Case::Sensitive, &conditions, None, None, None)
|
||||
.unwrap();
|
||||
let insensitive = table
|
||||
.find_table_rows(Case::Insensitive, &conditions, None, None, None)
|
||||
.unwrap();
|
||||
|
||||
assert!(sensitive.is_empty());
|
||||
assert_eq!(insensitive.len(), 1);
|
||||
assert_eq!(insensitive[0].get("name"), Some(&Value::from("Alice")));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn combines_conditions_and_projects_selected_fields() {
|
||||
let table = table();
|
||||
let conditions = [
|
||||
Condition::Equals {
|
||||
field: "country",
|
||||
value: Value::from("US"),
|
||||
},
|
||||
Condition::Equals {
|
||||
field: "role",
|
||||
value: Value::from("admin"),
|
||||
},
|
||||
];
|
||||
let select = ["name".to_string()];
|
||||
|
||||
let rows = table
|
||||
.find_table_rows(Case::Sensitive, &conditions, Some(&select), None, None)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(rows.len(), 1);
|
||||
assert_eq!(rows[0].len(), 1);
|
||||
assert_eq!(rows[0].get("name"), Some(&Value::from("Alice")));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn returns_no_rows_for_empty_conditions() {
|
||||
let rows = table()
|
||||
.find_table_rows(Case::Sensitive, &[], None, None, None)
|
||||
.unwrap();
|
||||
|
||||
assert!(rows.is_empty());
|
||||
}
|
||||
}
|
||||
pub use openobserve_vrl::enrichment::*;
|
||||
|
||||
@@ -0,0 +1,155 @@
|
||||
// Copyright 2026 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 <http://www.gnu.org/licenses/>.
|
||||
|
||||
use std::collections::HashSet;
|
||||
|
||||
use config::{get_config, is_local_disk_storage, meta::search::ScanStats};
|
||||
use infra::cache::{file_data, file_downloader};
|
||||
|
||||
/// Linear interpolation: cached_ratio=0 -> query_thread_num, cached_ratio=1 -> cpu_num.
|
||||
pub fn calc_target_partitions(cpu_num: usize, query_thread_num: usize, cached_ratio: f64) -> usize {
|
||||
(cpu_num as i64
|
||||
+ ((query_thread_num as i64 - cpu_num as i64) as f64 * (1.0 - cached_ratio)) as i64)
|
||||
as usize
|
||||
}
|
||||
|
||||
#[tracing::instrument(name = "service:search:grpc:storage:cache_files", skip_all)]
|
||||
pub async fn cache_files(
|
||||
trace_id: &str,
|
||||
files: &[(i64, &String, &String, i64, i64)],
|
||||
scan_stats: &mut ScanStats,
|
||||
file_type: &str,
|
||||
) -> (file_data::CacheType, u64, u64) {
|
||||
let mut cached_files = HashSet::with_capacity(files.len());
|
||||
let (mut cache_hits, mut cache_misses) = (0, 0);
|
||||
|
||||
let start = std::time::Instant::now();
|
||||
for (_id, _account, file, _size, max_ts) in files.iter() {
|
||||
if file_data::memory::exist(file).await {
|
||||
scan_stats.querier_memory_cached_files += 1;
|
||||
cached_files.insert(file);
|
||||
cache_hits += 1;
|
||||
} else if file_data::disk::exist(file).await {
|
||||
scan_stats.querier_disk_cached_files += 1;
|
||||
cached_files.insert(file);
|
||||
cache_hits += 1;
|
||||
} else {
|
||||
cache_misses += 1;
|
||||
}
|
||||
|
||||
let stream_type = if file_type == "index" {
|
||||
config::meta::stream::StreamType::Index
|
||||
} else if file.contains("/logs/") {
|
||||
config::meta::stream::StreamType::Logs
|
||||
} else if file.contains("/metrics/") {
|
||||
config::meta::stream::StreamType::Metrics
|
||||
} else if file.contains("/traces/") {
|
||||
config::meta::stream::StreamType::Traces
|
||||
} else {
|
||||
config::meta::stream::StreamType::Logs
|
||||
};
|
||||
|
||||
let current_time = chrono::Utc::now().timestamp_micros();
|
||||
let file_age_seconds = (current_time - max_ts) / 1_000_000;
|
||||
let file_age_hours = file_age_seconds as f64 / 3600.0;
|
||||
|
||||
if file_age_hours > 0.0 {
|
||||
config::metrics::FILE_ACCESS_TIME
|
||||
.with_label_values(&[&stream_type.to_string()])
|
||||
.observe(file_age_hours);
|
||||
}
|
||||
}
|
||||
|
||||
let check_cache_took = start.elapsed().as_millis() as usize;
|
||||
if check_cache_took > 1000 {
|
||||
log::warn!(
|
||||
"[trace_id {trace_id}] search->storage: check file cache took: {check_cache_took} ms",
|
||||
);
|
||||
}
|
||||
|
||||
let files_num = files.len() as i64;
|
||||
if files_num == scan_stats.querier_memory_cached_files + scan_stats.querier_disk_cached_files {
|
||||
return (file_data::CacheType::Disk, cache_hits, cache_misses);
|
||||
}
|
||||
|
||||
let cfg = get_config();
|
||||
let cache_type = if cfg.memory_cache.enabled
|
||||
&& scan_stats.compressed_size < cfg.memory_cache.skip_size as i64
|
||||
{
|
||||
file_data::CacheType::Memory
|
||||
} else if !is_local_disk_storage()
|
||||
&& cfg.disk_cache.enabled
|
||||
&& scan_stats.compressed_size < cfg.disk_cache.skip_size as i64
|
||||
{
|
||||
file_data::CacheType::Disk
|
||||
} else {
|
||||
return (file_data::CacheType::None, cache_hits, cache_misses);
|
||||
};
|
||||
|
||||
let trace_id = trace_id.to_string();
|
||||
let files = files
|
||||
.iter()
|
||||
.filter_map(|(id, account, file, size, ts)| {
|
||||
if cached_files.contains(file) {
|
||||
None
|
||||
} else {
|
||||
Some((*id, account.to_string(), file.to_string(), *size, *ts))
|
||||
}
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
let file_type = file_type.to_string();
|
||||
tokio::spawn(async move {
|
||||
let files_num = files.len();
|
||||
for (id, account, file, size, ts) in files {
|
||||
if let Err(e) = file_downloader::queue_download(
|
||||
trace_id.clone(),
|
||||
id,
|
||||
account,
|
||||
file.clone(),
|
||||
size,
|
||||
ts,
|
||||
cache_type,
|
||||
)
|
||||
.await
|
||||
{
|
||||
log::error!(
|
||||
"[trace_id {trace_id}] error in queuing file {file} for background download: {e}"
|
||||
);
|
||||
}
|
||||
}
|
||||
log::info!(
|
||||
"[trace_id {trace_id}] search->storage: successfully enqueued {files_num} files of {file_type} for background download into {cache_type:?}",
|
||||
);
|
||||
});
|
||||
|
||||
if scan_stats.querier_memory_cached_files + scan_stats.querier_disk_cached_files < files_num / 2
|
||||
{
|
||||
(file_data::CacheType::None, cache_hits, cache_misses)
|
||||
} else {
|
||||
(cache_type, cache_hits, cache_misses)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::calc_target_partitions;
|
||||
|
||||
#[test]
|
||||
fn target_partitions_interpolates_cache_ratio() {
|
||||
assert_eq!(calc_target_partitions(8, 32, 0.0), 32);
|
||||
assert_eq!(calc_target_partitions(8, 32, 0.5), 20);
|
||||
assert_eq!(calc_target_partitions(8, 32, 1.0), 8);
|
||||
}
|
||||
}
|
||||
@@ -13,7 +13,11 @@
|
||||
// You should have received a copy of the GNU Affero General Public License
|
||||
// along with this program. If not, see <http://www.gnu.org/licenses/>.
|
||||
|
||||
#[cfg(feature = "enterprise")]
|
||||
pub mod cipher;
|
||||
pub mod enrichment;
|
||||
pub mod file_cache;
|
||||
pub mod types;
|
||||
|
||||
use config::meta::{
|
||||
search::{Response, ScanStats},
|
||||
|
||||
@@ -0,0 +1,29 @@
|
||||
// Copyright 2026 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 <http://www.gnu.org/licenses/>.
|
||||
|
||||
use config::meta::stream::StreamType;
|
||||
use datafusion::sql::TableReference;
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct QueryParams {
|
||||
pub trace_id: String,
|
||||
pub org_id: String,
|
||||
pub stream: TableReference,
|
||||
pub stream_type: StreamType,
|
||||
pub stream_name: String,
|
||||
pub time_range: (i64, i64),
|
||||
pub work_group: Option<String>,
|
||||
pub use_inverted_index: bool,
|
||||
}
|
||||
@@ -0,0 +1,16 @@
|
||||
[package]
|
||||
name = "openobserve-vrl"
|
||||
version.workspace = true
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
publish = false
|
||||
|
||||
[dependencies]
|
||||
async-trait.workspace = true
|
||||
config.workspace = true
|
||||
vector-enrichment.workspace = true
|
||||
vrl.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
serde_json.workspace = true
|
||||
tokio.workspace = true
|
||||
@@ -0,0 +1,219 @@
|
||||
// Copyright 2026 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 <http://www.gnu.org/licenses/>.
|
||||
|
||||
use std::{
|
||||
collections::BTreeMap,
|
||||
sync::{Arc, LazyLock},
|
||||
};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use config::{RwHashMap, stats::MemorySize, utils::time::parse_str_to_time};
|
||||
use vector_enrichment::{Case, IndexHandle, Table};
|
||||
use vrl::value::{KeyString, ObjectMap, Value};
|
||||
|
||||
pub static ENRICHMENT_TABLES: LazyLock<RwHashMap<String, StreamTable>> =
|
||||
LazyLock::new(Default::default);
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct StreamTable {
|
||||
pub org_id: String,
|
||||
pub stream_name: String,
|
||||
pub data: Arc<Vec<Value>>,
|
||||
}
|
||||
|
||||
impl MemorySize for StreamTable {
|
||||
fn mem_size(&self) -> usize {
|
||||
std::mem::size_of::<StreamTable>()
|
||||
+ self.org_id.mem_size()
|
||||
+ self.stream_name.mem_size()
|
||||
+ self.data.iter().map(|v| v.to_string().len()).sum::<usize>()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Table for StreamTable {
|
||||
fn find_table_row(
|
||||
&self,
|
||||
case: Case,
|
||||
conditions: &[vector_enrichment::Condition],
|
||||
select: Option<&[String]>,
|
||||
_wildcard: Option<&Value>,
|
||||
_index: Option<IndexHandle>,
|
||||
) -> Result<ObjectMap, vector_enrichment::Error> {
|
||||
let resp = get_data(self, conditions, select, case);
|
||||
Ok(resp.first().cloned().unwrap_or_default())
|
||||
}
|
||||
|
||||
fn find_table_rows(
|
||||
&self,
|
||||
case: Case,
|
||||
conditions: &[vector_enrichment::Condition],
|
||||
select: Option<&[String]>,
|
||||
_wildcard: Option<&Value>,
|
||||
_index: Option<IndexHandle>,
|
||||
) -> Result<Vec<ObjectMap>, vector_enrichment::Error> {
|
||||
Ok(get_data(self, conditions, select, case))
|
||||
}
|
||||
|
||||
fn add_index(
|
||||
&mut self,
|
||||
_case: Case,
|
||||
_fields: &[&str],
|
||||
) -> Result<IndexHandle, vector_enrichment::Error> {
|
||||
Ok(IndexHandle(1))
|
||||
}
|
||||
|
||||
fn index_fields(&self) -> Vec<(Case, Vec<String>)> {
|
||||
Vec::new()
|
||||
}
|
||||
|
||||
fn needs_reload(&self) -> bool {
|
||||
false
|
||||
}
|
||||
}
|
||||
|
||||
fn get_data(
|
||||
table: &StreamTable,
|
||||
conditions: &[vector_enrichment::Condition],
|
||||
select: Option<&[String]>,
|
||||
case: Case,
|
||||
) -> Vec<ObjectMap> {
|
||||
if conditions.is_empty() {
|
||||
return vec![];
|
||||
}
|
||||
|
||||
let filtered = table.data.iter().filter(|value| {
|
||||
let Value::Object(map) = value else {
|
||||
return false;
|
||||
};
|
||||
|
||||
conditions.iter().all(|condition| match condition {
|
||||
vector_enrichment::Condition::Equals { field, value } => match case {
|
||||
Case::Insensitive => match (map.get(*field), value) {
|
||||
(Some(Value::Bytes(left)), Value::Bytes(right)) => {
|
||||
match (std::str::from_utf8(left), std::str::from_utf8(right)) {
|
||||
(Ok(left), Ok(right)) => left.eq_ignore_ascii_case(right),
|
||||
(Err(_), Err(_)) => left == right,
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
_ => false,
|
||||
},
|
||||
Case::Sensitive => map.get(*field).is_some_and(|item| item == value),
|
||||
},
|
||||
vector_enrichment::Condition::FromDate { field, from } => map
|
||||
.get(*field)
|
||||
.and_then(|value| value.as_str())
|
||||
.and_then(|value| parse_str_to_time(value.as_ref()).ok())
|
||||
.is_some_and(|date| date >= *from),
|
||||
vector_enrichment::Condition::ToDate { field, to } => map
|
||||
.get(*field)
|
||||
.and_then(|value| value.as_str())
|
||||
.and_then(|value| parse_str_to_time(value.as_ref()).ok())
|
||||
.is_some_and(|date| date <= *to),
|
||||
vector_enrichment::Condition::BetweenDates { field, from, to } => map
|
||||
.get(*field)
|
||||
.and_then(|value| value.as_str())
|
||||
.and_then(|value| parse_str_to_time(value.as_ref()).ok())
|
||||
.is_some_and(|date| date >= *from && date <= *to),
|
||||
})
|
||||
});
|
||||
|
||||
filtered
|
||||
.filter_map(Value::as_object)
|
||||
.map(|map| {
|
||||
if let Some(fields) = select {
|
||||
fields
|
||||
.iter()
|
||||
.filter_map(|field| map.get_key_value(field.as_str()))
|
||||
.map(|(key, value)| (key.clone(), value.clone()))
|
||||
.collect::<BTreeMap<KeyString, Value>>()
|
||||
} else {
|
||||
map.clone()
|
||||
}
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use serde_json::json;
|
||||
use vector_enrichment::{Condition, Table};
|
||||
|
||||
use super::*;
|
||||
|
||||
fn table() -> StreamTable {
|
||||
StreamTable {
|
||||
org_id: "default".to_string(),
|
||||
stream_name: "users".to_string(),
|
||||
data: Arc::new(vec![
|
||||
json!({"name": "Alice", "country": "US", "role": "admin"}).into(),
|
||||
json!({"name": "Bob", "country": "GB", "role": "user"}).into(),
|
||||
Value::from("not an object"),
|
||||
]),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn matches_rows_using_requested_case_sensitivity() {
|
||||
let conditions = [Condition::Equals {
|
||||
field: "name",
|
||||
value: Value::from("ALICE"),
|
||||
}];
|
||||
|
||||
let sensitive = table()
|
||||
.find_table_rows(Case::Sensitive, &conditions, None, None, None)
|
||||
.unwrap();
|
||||
let insensitive = table()
|
||||
.find_table_rows(Case::Insensitive, &conditions, None, None, None)
|
||||
.unwrap();
|
||||
|
||||
assert!(sensitive.is_empty());
|
||||
assert_eq!(insensitive.len(), 1);
|
||||
assert_eq!(insensitive[0].get("name"), Some(&Value::from("Alice")));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn combines_conditions_and_projects_selected_fields() {
|
||||
let conditions = [
|
||||
Condition::Equals {
|
||||
field: "country",
|
||||
value: Value::from("US"),
|
||||
},
|
||||
Condition::Equals {
|
||||
field: "role",
|
||||
value: Value::from("admin"),
|
||||
},
|
||||
];
|
||||
let select = ["name".to_string()];
|
||||
|
||||
let rows = table()
|
||||
.find_table_rows(Case::Sensitive, &conditions, Some(&select), None, None)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(rows.len(), 1);
|
||||
assert_eq!(rows[0].len(), 1);
|
||||
assert_eq!(rows[0].get("name"), Some(&Value::from("Alice")));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn returns_no_rows_for_empty_conditions() {
|
||||
let rows = table()
|
||||
.find_table_rows(Case::Sensitive, &[], None, None, None)
|
||||
.unwrap();
|
||||
|
||||
assert!(rows.is_empty());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,209 @@
|
||||
// Copyright 2026 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 <http://www.gnu.org/licenses/>.
|
||||
|
||||
//! Shared ownership for OpenObserve transform definitions, enrichment tables,
|
||||
//! and VRL compilation. This crate deliberately has no dependency on core or
|
||||
//! the search engine so both can consume the same runtime state.
|
||||
|
||||
use std::{collections::HashMap, sync::LazyLock};
|
||||
|
||||
use config::{
|
||||
DEFAULT_ORG, RwHashMap,
|
||||
meta::function::{Transform, VRLCompilerConfig, VRLRuntimeConfig},
|
||||
};
|
||||
use vector_enrichment::{Table, TableRegistry};
|
||||
use vrl::compiler::{CompilationResult, runtime::Runtime};
|
||||
|
||||
pub mod enrichment;
|
||||
|
||||
use enrichment::ENRICHMENT_TABLES;
|
||||
|
||||
/// Query and ingestion transform definitions, keyed by `org_id/function_name`.
|
||||
pub static QUERY_FUNCTIONS: LazyLock<RwHashMap<String, Transform>> =
|
||||
LazyLock::new(Default::default);
|
||||
|
||||
/// Process-wide enrichment tables such as MaxMind databases.
|
||||
///
|
||||
/// Organization-owned stream enrichment tables live in [`ENRICHMENT_TABLES`].
|
||||
static GLOBAL_ENRICHMENT_TABLES: LazyLock<RwHashMap<String, Box<dyn Table + Send + Sync>>> =
|
||||
LazyLock::new(Default::default);
|
||||
|
||||
pub fn register_global_enrichment_table<T>(name: impl Into<String>, table: T)
|
||||
where
|
||||
T: Table + Send + Sync + 'static,
|
||||
{
|
||||
GLOBAL_ENRICHMENT_TABLES.insert(name.into(), Box::new(table));
|
||||
}
|
||||
|
||||
pub fn remove_global_enrichment_table(name: &str) {
|
||||
GLOBAL_ENRICHMENT_TABLES.remove(name);
|
||||
}
|
||||
|
||||
pub async fn get_all_transform_keys(org_id: &str) -> Vec<String> {
|
||||
let org_key = format!("{org_id}/");
|
||||
|
||||
QUERY_FUNCTIONS
|
||||
.iter()
|
||||
.filter_map(|transform| {
|
||||
let key = transform.key();
|
||||
key.strip_prefix(&org_key).map(ToOwned::to_owned)
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub fn init_vrl_runtime() -> Runtime {
|
||||
Runtime::new(vrl::prelude::state::RuntimeState::default())
|
||||
}
|
||||
|
||||
pub fn get_enrichment_tables(org_id: &str) -> HashMap<String, Box<dyn Table + Send + Sync>> {
|
||||
let mut tables = HashMap::new();
|
||||
|
||||
for table in ENRICHMENT_TABLES.iter() {
|
||||
if table.org_id == org_id || table.org_id == DEFAULT_ORG {
|
||||
tables.insert(
|
||||
table.stream_name.clone(),
|
||||
Box::new(table.value().clone()) as Box<dyn Table + Send + Sync>,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
for table in GLOBAL_ENRICHMENT_TABLES.iter() {
|
||||
tables.insert(table.key().clone(), table.value().clone());
|
||||
}
|
||||
|
||||
tables
|
||||
}
|
||||
|
||||
pub fn get_vrl_compiler_config(org_id: &str) -> VRLCompilerConfig {
|
||||
let mut functions = vrl::stdlib::all();
|
||||
functions.append(&mut vector_enrichment::vrl_functions());
|
||||
let registry = TableRegistry::default();
|
||||
registry.load(get_enrichment_tables(org_id));
|
||||
let mut config = vrl::compiler::CompileConfig::default();
|
||||
config.set_custom(registry);
|
||||
VRLCompilerConfig { config, functions }
|
||||
}
|
||||
|
||||
pub fn compile_vrl_function(
|
||||
source: &str,
|
||||
org_id: &str,
|
||||
) -> Result<VRLRuntimeConfig, std::io::Error> {
|
||||
if source.contains("get_env_var") {
|
||||
return Err(std::io::Error::other("get_env_var is not supported"));
|
||||
}
|
||||
|
||||
let external = vrl::prelude::state::ExternalEnv::default();
|
||||
let vrl_config = get_vrl_compiler_config(org_id);
|
||||
match vrl::compiler::compile_with_external(
|
||||
source,
|
||||
&vrl_config.functions,
|
||||
&external,
|
||||
vrl_config.config,
|
||||
) {
|
||||
Ok(CompilationResult {
|
||||
program,
|
||||
warnings: _,
|
||||
config,
|
||||
}) => Ok(VRLRuntimeConfig {
|
||||
program,
|
||||
config,
|
||||
fields: vec![],
|
||||
}),
|
||||
Err(error) => Err(std::io::Error::other(
|
||||
vrl::diagnostic::Formatter::new(source, error).to_string(),
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::Arc;
|
||||
|
||||
use config::meta::function::Transform;
|
||||
use vector_enrichment::TableRegistry;
|
||||
|
||||
use super::{
|
||||
enrichment::{ENRICHMENT_TABLES, StreamTable},
|
||||
*,
|
||||
};
|
||||
|
||||
#[test]
|
||||
fn creates_vrl_compiler_config_with_org_enrichment_tables() {
|
||||
let key = "vrl_crate_test/default/users";
|
||||
ENRICHMENT_TABLES.insert(
|
||||
key.to_string(),
|
||||
StreamTable {
|
||||
org_id: "vrl_crate_test".to_string(),
|
||||
stream_name: "users".to_string(),
|
||||
data: Arc::new(vec![]),
|
||||
},
|
||||
);
|
||||
|
||||
let config = get_vrl_compiler_config("vrl_crate_test");
|
||||
let registry = config.config.get_custom::<TableRegistry>().unwrap();
|
||||
assert!(registry.table_ids().contains(&"users".to_string()));
|
||||
|
||||
ENRICHMENT_TABLES.remove(key);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn includes_registered_global_enrichment_tables() {
|
||||
let name = "vrl_crate_test_global";
|
||||
register_global_enrichment_table(
|
||||
name,
|
||||
StreamTable {
|
||||
org_id: "default".to_string(),
|
||||
stream_name: name.to_string(),
|
||||
data: Arc::new(vec![]),
|
||||
},
|
||||
);
|
||||
|
||||
let config = get_vrl_compiler_config("vrl_crate_test");
|
||||
let registry = config.config.get_custom::<TableRegistry>().unwrap();
|
||||
assert!(registry.table_ids().contains(&name.to_string()));
|
||||
|
||||
remove_global_enrichment_table(name);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn returns_transform_keys_for_requested_org() {
|
||||
let key = "vrl_crate_test/test_transform";
|
||||
QUERY_FUNCTIONS.insert(
|
||||
key.to_string(),
|
||||
Transform {
|
||||
name: "test_transform".to_string(),
|
||||
function: ".".to_string(),
|
||||
params: "row".to_string(),
|
||||
num_args: 1,
|
||||
trans_type: Some(0),
|
||||
streams: None,
|
||||
},
|
||||
);
|
||||
|
||||
let keys = get_all_transform_keys("vrl_crate_test").await;
|
||||
assert_eq!(keys, vec!["test_transform"]);
|
||||
|
||||
QUERY_FUNCTIONS.remove(key);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_environment_access() {
|
||||
let error = match compile_vrl_function("get_env_var!(\"HOME\")", "default") {
|
||||
Ok(_) => panic!("get_env_var must be rejected"),
|
||||
Err(error) => error,
|
||||
};
|
||||
assert_eq!(error.to_string(), "get_env_var is not supported");
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user