From a03526eaef1155ec191b35fb1cbef2316e0280f2 Mon Sep 17 00:00:00 2001 From: Patrick Sullivan Date: Tue, 14 Jan 2025 18:04:10 -0600 Subject: [PATCH] feat: Decouple `super_cluster` from `table` via dependency injection (#5698) This is the follow-up PR to [this other PR](https://github.com/openobserve/o2-enterprise/pull/281#issue-2773724675). Please see the other PR description for a detailed explanation of the changes. This PR copies the supercluster queue message handling code from `o2-enterprise` which updates the database when different messages are received. In this PR these message-handling functions are injected into `o2-enterprise`'s various `SuperClusterQueue` structures so that the `SuperClusterQueue`s can execute the logic to update the `openobserve` database without the `o2-enterprise` crate needing a dependency on the `table` module. ## Summary by CodeRabbit - **New Features** - Added comprehensive message processing capabilities for various enterprise components including alerts, dashboards, distinct values, meta, pipelines, schemas, search jobs, and short URLs - Introduced a new super cluster queue system with modular message handling and initialization - **Chores** - Updated library and main application to support enterprise feature initialization - Added conditional compilation for enterprise-specific modules - **Infrastructure** - Implemented robust error handling and logging for message processing across different system components --- src/lib.rs | 3 + src/main.rs | 18 ++- src/super_cluster_queue/alerts.rs | 108 +++++++++++++ src/super_cluster_queue/dashboards.rs | 77 ++++++++++ src/super_cluster_queue/distinct_values.rs | 63 ++++++++ src/super_cluster_queue/meta.rs | 40 +++++ src/super_cluster_queue/mod.rs | 90 +++++++++++ src/super_cluster_queue/pipelines.rs | 57 +++++++ src/super_cluster_queue/scheduler.rs | 117 +++++++++++++++ src/super_cluster_queue/schemas.rs | 142 ++++++++++++++++++ src/super_cluster_queue/search_jobs/mod.rs | 54 +++++++ .../search_jobs/search_job_partitions.rs | 46 ++++++ .../search_jobs/search_job_results.rs | 33 ++++ .../search_jobs/search_jobs.rs | 77 ++++++++++ src/super_cluster_queue/short_urls.rs | 54 +++++++ 15 files changed, 978 insertions(+), 1 deletion(-) create mode 100644 src/super_cluster_queue/alerts.rs create mode 100644 src/super_cluster_queue/dashboards.rs create mode 100644 src/super_cluster_queue/distinct_values.rs create mode 100644 src/super_cluster_queue/meta.rs create mode 100644 src/super_cluster_queue/mod.rs create mode 100644 src/super_cluster_queue/pipelines.rs create mode 100644 src/super_cluster_queue/scheduler.rs create mode 100644 src/super_cluster_queue/schemas.rs create mode 100644 src/super_cluster_queue/search_jobs/mod.rs create mode 100644 src/super_cluster_queue/search_jobs/search_job_partitions.rs create mode 100644 src/super_cluster_queue/search_jobs/search_job_results.rs create mode 100644 src/super_cluster_queue/search_jobs/search_jobs.rs create mode 100644 src/super_cluster_queue/short_urls.rs diff --git a/src/lib.rs b/src/lib.rs index 3afaae87ec..3b14982d40 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -22,6 +22,9 @@ pub mod job; pub mod router; pub mod service; +#[cfg(feature = "enterprise")] +pub mod super_cluster_queue; + pub(crate) static USER_AGENT_REGEX_FILE: &[u8] = include_bytes!(concat!( env!("CARGO_MANIFEST_DIR"), "/ua_regex/regexes.yaml" diff --git a/src/main.rs b/src/main.rs index 7590c95423..6d56a2945e 100644 --- a/src/main.rs +++ b/src/main.rs @@ -211,7 +211,7 @@ async fn main() -> Result<(), anyhow::Error> { // init enterprise #[cfg(feature = "enterprise")] - if let Err(e) = o2_enterprise::enterprise::init().await { + if let Err(e) = init_enterprise().await { job_init_tx.send(false).ok(); panic!("enerprise init failed: {}", e); } @@ -890,3 +890,19 @@ fn enable_tracing() -> Result<(), anyhow::Error> { .init(); Ok(()) } + +/// Initializes enterprise features. +#[cfg(feature = "enterprise")] +async fn init_enterprise() -> Result<(), anyhow::Error> { + o2_enterprise::enterprise::search::init().await?; + + if o2_enterprise::enterprise::common::infra::config::get_config() + .super_cluster + .enabled + { + log::info!("init super cluster"); + o2_enterprise::enterprise::super_cluster::kv::init().await?; + openobserve::super_cluster_queue::init().await?; + } + Ok(()) +} diff --git a/src/super_cluster_queue/alerts.rs b/src/super_cluster_queue/alerts.rs new file mode 100644 index 0000000000..106a647d77 --- /dev/null +++ b/src/super_cluster_queue/alerts.rs @@ -0,0 +1,108 @@ +// Copyright 2024 Zinc Labs Inc. +// +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU Affero General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// This program is distributed in the hope that it will be useful +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU Affero General Public License for more details. +// +// You should have received a copy of the GNU Affero General Public License +// along with this program. If not, see . + +use config::{meta::alerts::alert::Alert, utils::json}; +use infra::{ + db::{connect_to_orm, ORM_CLIENT}, + errors::{Error, Result}, + table, +}; +use o2_enterprise::enterprise::super_cluster::queue::{Message, MessageType}; +use serde::{Deserialize, Serialize}; +use svix_ksuid::Ksuid; + +#[derive(Debug, Clone, Deserialize, Serialize)] +pub enum AlertMessage { + Create { + org_id: String, + folder_id: String, + alert: Alert, + }, + Update { + org_id: String, + folder_id: Option, + alert: Alert, + }, + Delete { + org_id: String, + alert_id: Ksuid, + }, +} + +impl TryFrom for AlertMessage { + type Error = Error; + + fn try_from(msg: Message) -> std::result::Result { + let bytes = msg + .value + .ok_or(Error::Message("Message missing value".to_string()))?; + let alert_msg: AlertMessage = json::from_slice(&bytes).inspect_err(|_| { + log::error!("[SUPER_CLUSTER:DB] Failed to parse alert message"); + })?; + Ok(alert_msg) + } +} + +pub(crate) async fn process(msg: Message) -> Result<()> { + match msg.message_type { + MessageType::AlertsTable => { + let msg = msg.try_into()?; + process_msg(msg).await?; + } + _ => { + // Try to process the message as an old event for the meta table for backward + // compatability. This logic can be removed after all logic reading and writing alerts + // to the meta table is removed from the application. + super::meta::process(msg).await?; + } + } + + Ok(()) +} + +pub(crate) async fn process_msg(msg: AlertMessage) -> Result<()> { + let conn = ORM_CLIENT.get_or_init(connect_to_orm).await; + match msg { + AlertMessage::Create { + org_id, + folder_id, + alert, + } => { + let alert = table::alerts::create(conn, &org_id, &folder_id, alert).await?; + infra::cluster_coordinator::alerts::emit_put_event(&org_id, &alert).await?; + } + AlertMessage::Update { + org_id, + folder_id, + alert, + } => { + let alert = table::alerts::update(conn, &org_id, folder_id.as_deref(), alert).await?; + infra::cluster_coordinator::alerts::emit_put_event(&org_id, &alert).await?; + } + AlertMessage::Delete { org_id, alert_id } => { + if let Some((_, alert)) = table::alerts::get_by_id(conn, &org_id, alert_id).await? { + table::alerts::delete_by_id(conn, &org_id, alert_id).await?; + infra::cluster_coordinator::alerts::emit_delete_event( + &org_id, + alert.stream_type, + &alert.stream_name, + &alert.name, + ) + .await?; + } + } + }; + Ok(()) +} diff --git a/src/super_cluster_queue/dashboards.rs b/src/super_cluster_queue/dashboards.rs new file mode 100644 index 0000000000..f6dd93c725 --- /dev/null +++ b/src/super_cluster_queue/dashboards.rs @@ -0,0 +1,77 @@ +// Copyright 2024 Zinc Labs Inc. +// +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU Affero General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// This program is distributed in the hope that it will be useful +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU Affero General Public License for more details. +// +// You should have received a copy of the GNU Affero General Public License +// along with this program. If not, see . + +use config::{meta::dashboards::Dashboard, utils::json}; +use infra::{ + errors::{Error, Result}, + table, +}; +use o2_enterprise::enterprise::super_cluster::queue::Message; +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Clone, Deserialize, Serialize)] +#[allow(clippy::large_enum_variant)] +pub enum DashboardMessage { + Put { + org_id: String, + folder_id: String, + dashboard: Dashboard, + }, + Delete { + org_id: String, + folder_id: String, + dashboard_id: String, + }, +} + +impl TryFrom for DashboardMessage { + type Error = Error; + + fn try_from(msg: Message) -> std::result::Result { + let bytes = msg + .value + .ok_or(Error::Message("Message missing value".to_string()))?; + let alert_msg: DashboardMessage = json::from_slice(&bytes).inspect_err(|_| { + log::error!("[SUPER_CLUSTER:DB] Failed to parse dashboard message"); + })?; + Ok(alert_msg) + } +} + +pub(crate) async fn process(msg: Message) -> Result<()> { + let msg = msg.try_into()?; + process_msg(msg).await?; + Ok(()) +} + +pub(crate) async fn process_msg(msg: DashboardMessage) -> Result<()> { + match msg { + DashboardMessage::Put { + org_id, + folder_id, + dashboard, + } => { + table::dashboards::put(&org_id, &folder_id, dashboard).await?; + } + DashboardMessage::Delete { + org_id, + folder_id, + dashboard_id, + } => { + table::dashboards::delete_from_folder(&org_id, &folder_id, &dashboard_id).await?; + } + }; + Ok(()) +} diff --git a/src/super_cluster_queue/distinct_values.rs b/src/super_cluster_queue/distinct_values.rs new file mode 100644 index 0000000000..3050d7d04e --- /dev/null +++ b/src/super_cluster_queue/distinct_values.rs @@ -0,0 +1,63 @@ +// Copyright 2024 Zinc Labs Inc. +// +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU Affero General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// This program is distributed in the hope that it will be useful +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU Affero General Public License for more details. +// +// You should have received a copy of the GNU Affero General Public License +// along with this program. If not, see . + +use infra::{ + errors::{Error, Result}, + table::distinct_values::{BatchDeleteMessage, DistinctFieldRecord}, +}; +use o2_enterprise::enterprise::super_cluster::queue::{Message, MessageType}; + +pub(crate) async fn process(msg: Message) -> Result<()> { + match msg.message_type { + MessageType::DistinctValuesPut => { + let record: DistinctFieldRecord = serde_json::from_slice(&msg.value.unwrap())?; + log::info!( + "[SUPER_CLUSTER:DB] adding distinct value {}/{}/{}", + record.org_name, + record.stream_name, + record.field_name + ); + infra::table::distinct_values::add(record).await?; + } + MessageType::DistinctValuesDelete(true) => { + let val: BatchDeleteMessage = serde_json::from_slice(&msg.value.unwrap())?; + log::info!( + "[SUPER_CLUSTER:DB] batch removing {:?}:{}", + val.origin_type, + val.id + ); + infra::table::distinct_values::batch_remove(val.origin_type, &val.id).await?; + } + MessageType::DistinctValuesDelete(false) => { + let record: DistinctFieldRecord = serde_json::from_slice(&msg.value.unwrap())?; + log::info!( + "[SUPER_CLUSTER:DB] removing distinct value {}/{}/{}", + record.org_name, + record.stream_name, + record.field_name + ); + infra::table::distinct_values::remove(record).await?; + } + _ => { + log::error!( + "[SUPER_CLUSTER:DB] Invalid message: type: {:?}, key: {}", + msg.message_type, + msg.key + ); + return Err(Error::Message("Invalid message type".to_string())); + } + } + Ok(()) +} diff --git a/src/super_cluster_queue/meta.rs b/src/super_cluster_queue/meta.rs new file mode 100644 index 0000000000..0588672041 --- /dev/null +++ b/src/super_cluster_queue/meta.rs @@ -0,0 +1,40 @@ +// Copyright 2024 Zinc Labs Inc. +// +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU Affero General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// This program is distributed in the hope that it will be useful +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU Affero General Public License for more details. +// +// You should have received a copy of the GNU Affero General Public License +// along with this program. If not, see . + +use infra::errors::{Error, Result}; +use o2_enterprise::enterprise::super_cluster::queue::{Message, MessageType}; + +pub(crate) async fn process(msg: Message) -> Result<()> { + let db = infra::db::get_db().await; + match msg.message_type { + MessageType::Put => { + db.put(&msg.key, msg.value.unwrap(), msg.need_watch, None) + .await?; + } + MessageType::Delete(with_prefix) => { + db.delete(&msg.key, with_prefix, msg.need_watch, None) + .await?; + } + _ => { + log::error!( + "[SUPER_CLUSTER:DB] Invalid message: type: {:?}, key: {}", + msg.message_type, + msg.key + ); + return Err(Error::Message("Invalid message type".to_string())); + } + } + Ok(()) +} diff --git a/src/super_cluster_queue/mod.rs b/src/super_cluster_queue/mod.rs new file mode 100644 index 0000000000..6574ab7f3f --- /dev/null +++ b/src/super_cluster_queue/mod.rs @@ -0,0 +1,90 @@ +// Copyright 2024 Zinc Labs 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 . + +mod alerts; +mod dashboards; +mod distinct_values; +mod meta; +mod pipelines; +mod scheduler; +mod schemas; +mod search_jobs; +mod short_urls; + +use config::cluster::{is_offline, LOCAL_NODE}; +use o2_enterprise::enterprise::super_cluster::queue::{ + AlertsQueue, DashboardsQueue, MetaQueue, PipelinesQueue, SchemasQueue, SearchJobsQueue, + SuperClusterQueueTrait, +}; + +/// Creates a super cluster queue for each super cluster topic and begins +/// polling messages from each queue in a separate thread. +pub async fn init() -> Result<(), anyhow::Error> { + let meta_queue = MetaQueue { + on_meta_msg: meta::process, + on_distinct_values_msg: distinct_values::process, + on_short_url_msg: short_urls::process, + on_schema_msg: schemas::process, + on_alert_msg: alerts::process, + on_scheduler_msg: scheduler::process, + on_search_job_msg: search_jobs::process, + on_dashboard_msg: dashboards::process, + on_pipeline_msg: pipelines::process, + }; + let schema_queue = SchemasQueue { + on_schema_msg: schemas::process, + }; + let alerts_queue = AlertsQueue { + on_alert_msg: alerts::process, + on_scheduler_msg: scheduler::process, + }; + let search_jobs_queue = SearchJobsQueue { + on_search_job_msg: search_jobs::process, + }; + let dashboards_queue = DashboardsQueue { + on_dashboard_msg: dashboards::process, + }; + let pipelines_queue = PipelinesQueue { + on_pipeline_msg: pipelines::process, + }; + + let queues: Vec> = vec![ + Box::new(meta_queue), + Box::new(schema_queue), + Box::new(alerts_queue), + Box::new(search_jobs_queue), + Box::new(dashboards_queue), + Box::new(pipelines_queue), + ]; + + for queue in queues { + queue.create_queue().await?; + + if LOCAL_NODE.is_compactor() { + tokio::task::spawn(async move { + loop { + if is_offline() { + break; + } + if let Err(e) = queue.subscribe().await { + log::error!("[SUPER_CLUSTER:sync] failed to subscribe: {}", e); + } + } + }); + } + } + + Ok(()) +} diff --git a/src/super_cluster_queue/pipelines.rs b/src/super_cluster_queue/pipelines.rs new file mode 100644 index 0000000000..b3455e29d0 --- /dev/null +++ b/src/super_cluster_queue/pipelines.rs @@ -0,0 +1,57 @@ +// Copyright 2024 Zinc Labs Inc. +// +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU Affero General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// This program is distributed in the hope that it will be useful +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU Affero General Public License for more details. +// +// You should have received a copy of the GNU Affero General Public License +// along with this program. If not, see . + +use config::{meta::pipeline::Pipeline, utils::json}; +use infra::errors::{Error, Result}; +use o2_enterprise::enterprise::super_cluster::queue::{Message, MessageType}; + +pub(crate) async fn process(msg: Message) -> Result<()> { + match msg.message_type { + MessageType::PipelinePut => { + let bytes = msg + .value + .ok_or(Error::Message("Message missing value".to_string()))?; + let pipeline: Pipeline = json::from_slice(&bytes).inspect_err(|e| { + log::error!( + "[SUPER_CLUSTER:PIPELINE] Failed to deserialize message value to pipeline: {}", + e + ); + })?; + infra::pipeline::put(&pipeline).await?; + infra::cluster_coordinator::pipelines::emit_put_event(&pipeline.id).await?; + } + MessageType::PipelineDelete => { + let pipeline_id = parse_key(&msg.key)?; + infra::cluster_coordinator::pipelines::emit_delete_event(&pipeline_id).await?; + } + _ => { + log::error!( + "[SUPER_CLUSTER:PIPELINE] Invalid message: type: {:?}, key: {}", + msg.message_type, + msg.key + ); + return Err(Error::Message("Invalid message type".to_string())); + } + } + Ok(()) +} + +fn parse_key(key: &str) -> Result { + let key_columns: Vec<&str> = key.split('/').collect(); + if key_columns.len() != 3 || key_columns[2].is_empty() { + return Err(Error::Message("Invalid key".to_string())); + } + Ok(key_columns[2].into()) +} diff --git a/src/super_cluster_queue/scheduler.rs b/src/super_cluster_queue/scheduler.rs new file mode 100644 index 0000000000..e27352df5b --- /dev/null +++ b/src/super_cluster_queue/scheduler.rs @@ -0,0 +1,117 @@ +// Copyright 2024 Zinc Labs Inc. +// +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU Affero General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// This program is distributed in the hope that it will be useful +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU Affero General Public License for more details. +// +// You should have received a copy of the GNU Affero General Public License +// along with this program. If not, see . + +use config::utils::json; +use infra::{ + errors::{Error, Result}, + scheduler::{self, Trigger}, +}; +use o2_enterprise::enterprise::super_cluster::queue::{Message, MessageType}; + +pub(crate) async fn process(msg: Message) -> Result<()> { + match msg.message_type { + MessageType::SchedulerPush => { + push(msg).await?; + } + MessageType::SchedulerUpdate => { + update(msg).await?; + } + MessageType::SchedulerUpdateStatus => { + update_status(msg).await?; + } + MessageType::SchedulerDelete => { + delete(msg).await?; + } + _ => { + log::error!( + "[SUPER_CLUSTER:DB] Invalid message: type: {:?}, key: {}", + msg.message_type, + msg.key + ); + return Err(Error::Message("Invalid message type".to_string())); + } + } + Ok(()) +} + +async fn push(msg: Message) -> Result<()> { + let trigger: Trigger = json::from_slice(&msg.value.unwrap())?; + if let Err(e) = scheduler::push(trigger.clone()).await { + log::error!( + "[SUPER_CLUSTER:sync] Failed to push scheduler: {}/{:?}/{}, error: {}", + trigger.org, + trigger.module, + trigger.module_key, + e + ); + return Err(e); + } + Ok(()) +} + +async fn update(msg: Message) -> Result<()> { + let trigger: Trigger = json::from_slice(&msg.value.unwrap())?; + if let Err(e) = scheduler::update_trigger(trigger.clone()).await { + log::error!( + "[SUPER_CLUSTER:sync] Failed to update scheduler: {}/{:?}/{}, error: {}", + trigger.org, + trigger.module, + trigger.module_key, + e + ); + return Err(e); + } + Ok(()) +} + +async fn update_status(msg: Message) -> Result<()> { + let trigger: Trigger = json::from_slice(&msg.value.unwrap())?; + if let Err(e) = scheduler::update_status( + &trigger.org, + trigger.module.clone(), + &trigger.module_key, + trigger.status, + trigger.retries, + ) + .await + { + log::error!( + "[SUPER_CLUSTER:sync] Failed to update status for scheduler: {}/{:?}/{}, error: {}", + trigger.org, + trigger.module, + trigger.module_key, + e + ); + return Err(e); + } + Ok(()) +} + +async fn delete(msg: Message) -> Result<()> { + let trigger: Trigger = json::from_slice(&msg.value.unwrap())?; + if let Err(e) = + scheduler::delete(&trigger.org, trigger.module.clone(), &trigger.module_key).await + { + log::error!( + "[SUPER_CLUSTER:sync] Failed to delete scheduler: {}/{:?}/{}, error: {}", + trigger.org, + trigger.module, + trigger.module_key, + e + ); + return Err(e); + } + Ok(()) +} diff --git a/src/super_cluster_queue/schemas.rs b/src/super_cluster_queue/schemas.rs new file mode 100644 index 0000000000..753def09bf --- /dev/null +++ b/src/super_cluster_queue/schemas.rs @@ -0,0 +1,142 @@ +// Copyright 2024 Zinc Labs Inc. +// +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU Affero General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// This program is distributed in the hope that it will be useful +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU Affero General Public License for more details. +// +// You should have received a copy of the GNU Affero General Public License +// along with this program. If not, see . + +use arrow::datatypes::Schema; +use async_recursion::async_recursion; +use config::{meta::stream::StreamType, utils::json}; +use infra::errors::{Error, Result}; +use o2_enterprise::enterprise::super_cluster::queue::{Message, MessageType}; + +pub(crate) async fn process(msg: Message) -> Result<()> { + let db = infra::db::get_db().await; + match msg.message_type { + MessageType::Delete(with_prefix) => { + log::debug!("[SUPER_CLUSTER:sync] Schema Delete: {}", msg.key); + db.delete(&msg.key, with_prefix, msg.need_watch, msg.start_dt) + .await?; + } + MessageType::Put => { + log::debug!("[SUPER_CLUSTER:sync] Schema Put: {}", msg.key); + merge(msg).await?; + } + MessageType::SchemaMerge => { + log::debug!("[SUPER_CLUSTER:sync] Schema Merge: {}", msg.key); + merge(msg).await?; + } + MessageType::SchemaSetting => { + log::debug!("[SUPER_CLUSTER:sync] Schema Setting: {}", msg.key); + setting(msg).await?; + } + MessageType::SchemaDeleteFields => { + log::debug!("[SUPER_CLUSTER:sync] Schema DeleteFields: {}", msg.key); + delete_fields(msg).await?; + } + _ => { + log::error!( + "[SUPER_CLUSTER:DB] Invalid message: type: {:?}, key: {}", + msg.message_type, + msg.key + ); + return Err(Error::Message("Invalid message type".to_string())); + } + } + Ok(()) +} + +#[async_recursion] +async fn merge(msg: Message) -> Result<()> { + let (org_id, stream_type, stream_name) = parse_key(&msg.key)?; + let schema: Schema = json::from_slice(msg.value.as_ref().unwrap())?; + if let Err(e) = + infra::schema::merge(&org_id, &stream_name, stream_type, &schema, msg.start_dt).await + { + if e.to_string().to_lowercase().contains("duplicate key") + && e.to_string() + .to_lowercase() + .contains("meta_module_start_dt_idx") + { + log::warn!("[SUPER_CLUSTER:sync] Duplicate start_dt, retrying with new start_dt"); + let msg = Message { + key: msg.key, + value: msg.value, + start_dt: Some(chrono::Utc::now().timestamp_micros()), + need_watch: msg.need_watch, + message_type: msg.message_type, + }; + return merge(msg).await; + } + log::error!( + "[SUPER_CLUSTER:sync] Failed to merge schema for stream: {}/{}/{}, start_dt: {:?}, error: {}", + org_id, + stream_type, + stream_name, + msg.start_dt, + e + ); + return Err(e); + } + Ok(()) +} + +async fn setting(msg: Message) -> Result<()> { + let (org_id, stream_type, stream_name) = parse_key(&msg.key)?; + let metadata: std::collections::HashMap = + json::from_slice(&msg.value.unwrap())?; + if let Err(e) = + infra::schema::update_setting(&org_id, &stream_name, stream_type, metadata.clone()).await + { + log::error!( + "[SUPER_CLUSTER:sync] Failed to update setting for stream: {}/{}/{}, metadata: {:?}, error: {}", + org_id, + stream_type, + stream_name, + metadata, + e + ); + return Err(e); + } + Ok(()) +} + +async fn delete_fields(msg: Message) -> Result<()> { + let (org_id, stream_type, stream_name) = parse_key(&msg.key)?; + let fields: Vec = json::from_slice(&msg.value.unwrap())?; + if let Err(e) = + infra::schema::delete_fields(&org_id, &stream_name, stream_type, fields.clone()).await + { + log::error!( + "[SUPER_CLUSTER:sync] Failed to delete fields for stream: {}/{}/{}, fields: {:?}, error: {}", + org_id, + stream_type, + stream_name, + fields, + e + ); + return Err(e); + } + Ok(()) +} + +fn parse_key(key: &str) -> Result<(String, StreamType, String)> { + let key_columns = key.split('/').collect::>(); + if key_columns.len() < 5 { + log::error!("[SUPER_CLUSTER:sync] Invalid key format: {}", key); + return Err(Error::Message("Invalid key format".to_string())); + } + let org_id = key_columns[2].into(); + let stream_type = key_columns[3].into(); + let stream_name = key_columns[4].into(); + Ok((org_id, stream_type, stream_name)) +} diff --git a/src/super_cluster_queue/search_jobs/mod.rs b/src/super_cluster_queue/search_jobs/mod.rs new file mode 100644 index 0000000000..57709a40b1 --- /dev/null +++ b/src/super_cluster_queue/search_jobs/mod.rs @@ -0,0 +1,54 @@ +// Copyright 2024 Zinc Labs 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 . + +pub mod search_job_partitions; +pub mod search_job_results; +pub mod search_jobs; + +use config::utils::json; +use infra::{ + errors::{Error, Result}, + table::search_job::{ + search_job_partitions::PartitionJobOperator, search_job_results::JobResultOperator, + search_jobs::JobOperator, + }, +}; +use o2_enterprise::enterprise::super_cluster::queue::{Message, MessageType}; + +pub(crate) async fn process(msg: Message) -> Result<()> { + match msg.message_type { + MessageType::SearchJob => { + let operator: JobOperator = json::from_slice(&msg.value.unwrap())?; + search_jobs::process(operator).await?; + } + MessageType::SearchJobPartition => { + let operator: PartitionJobOperator = json::from_slice(&msg.value.unwrap())?; + search_job_partitions::process(operator).await?; + } + MessageType::SearchJobResult => { + let operator: JobResultOperator = json::from_slice(&msg.value.unwrap())?; + search_job_results::process(operator).await?; + } + _ => { + log::error!( + "[SUPER_CLUSTER:DB] Invalid message for search job: type: {:?}, key: {}", + msg.message_type, + msg.key + ); + return Err(Error::Message("Invalid message type".to_string())); + } + } + Ok(()) +} diff --git a/src/super_cluster_queue/search_jobs/search_job_partitions.rs b/src/super_cluster_queue/search_jobs/search_job_partitions.rs new file mode 100644 index 0000000000..9d23e14950 --- /dev/null +++ b/src/super_cluster_queue/search_jobs/search_job_partitions.rs @@ -0,0 +1,46 @@ +// Copyright 2024 Zinc Labs Inc. +// +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU Affero General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// This program is distributed in the hope that it will be useful +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU Affero General Public License for more details. +// +// You should have received a copy of the GNU Affero General Public License +// along with this program. If not, see . + +use infra::{errors::Result, table::search_job::search_job_partitions::*}; + +pub(crate) async fn process(operator: PartitionJobOperator) -> Result<()> { + match operator { + PartitionJobOperator::Submit { job_id, jobs } => { + if let Err(e) = submit_partitions(job_id.as_str(), *jobs).await { + log::error!( + "[SUPER_CLUSTER:DB] Failed to submit partition job: {job_id}, error: {e}", + ); + return Err(e); + } + } + PartitionJobOperator::Set(operator) => { + if let Err(e) = set(operator.clone()).await { + log::error!( + "[SUPER_CLUSTER:DB] Failed to set partition job: {operator:?}, error: {e}", + ); + return Err(e); + } + } + PartitionJobOperator::Delete { job_id } => { + if let Err(e) = clean_deleted_partition_job(job_id.as_str()).await { + log::error!( + "[SUPER_CLUSTER:DB] Failed to clean deleted partition job: {job_id}, error: {e}", + ); + return Err(e); + } + } + } + Ok(()) +} diff --git a/src/super_cluster_queue/search_jobs/search_job_results.rs b/src/super_cluster_queue/search_jobs/search_job_results.rs new file mode 100644 index 0000000000..011dc2099c --- /dev/null +++ b/src/super_cluster_queue/search_jobs/search_job_results.rs @@ -0,0 +1,33 @@ +// Copyright 2024 Zinc Labs Inc. +// +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU Affero General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// This program is distributed in the hope that it will be useful +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU Affero General Public License for more details. +// +// You should have received a copy of the GNU Affero General Public License +// along with this program. If not, see . + +use infra::{ + errors::Result, + table::search_job::search_job_results::{JobResultOperator, *}, +}; + +pub(crate) async fn process(operator: JobResultOperator) -> Result<()> { + match operator { + JobResultOperator::Delete { job_id } => { + if let Err(e) = clean_deleted_job_result(job_id.as_str()).await { + log::error!( + "[SUPER_CLUSTER:DB] Failed to clean deleted job result: {job_id}, error: {e}", + ); + return Err(e); + } + } + } + Ok(()) +} diff --git a/src/super_cluster_queue/search_jobs/search_jobs.rs b/src/super_cluster_queue/search_jobs/search_jobs.rs new file mode 100644 index 0000000000..fe5c1852a6 --- /dev/null +++ b/src/super_cluster_queue/search_jobs/search_jobs.rs @@ -0,0 +1,77 @@ +// Copyright 2024 Zinc Labs Inc. +// +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU Affero General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// This program is distributed in the hope that it will be useful +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU Affero General Public License for more details. +// +// You should have received a copy of the GNU Affero General Public License +// along with this program. If not, see . + +use infra::{ + errors::Result, + table::search_job::search_jobs::{JobOperator, *}, +}; + +pub(crate) async fn process(operator: JobOperator) -> Result<()> { + match operator { + JobOperator::Submit(job) => { + if let Err(e) = submit((*job).clone().into()).await { + log::error!( + "[SUPER_CLUSTER:DB] Failed to submit job: {}, error: {e}", + job.id, + ); + return Err(e); + } + } + JobOperator::Set(operator) => { + if let Err(e) = set(operator.clone()).await { + log::error!("[SUPER_CLUSTER:DB] Failed to set job: {operator:?}, error: {e}",); + return Err(e); + } + } + JobOperator::Cancel { job_id, updated_at } => { + if let Err(e) = cancel_job(job_id.as_str(), updated_at).await { + log::error!("[SUPER_CLUSTER:DB] Failed to cancel job: {job_id}, error: {e}",); + return Err(e); + } + } + JobOperator::GetJob { + job_id, + cluster, + node, + updated_at, + } => { + if let Err(e) = + set_job_start(job_id.as_str(), cluster.as_str(), node.as_str(), updated_at).await + { + log::error!("[SUPER_CLUSTER:DB] Failed to set job start: {job_id}, error: {e}",); + return Err(e); + } + } + JobOperator::Delete { job_id } => { + if let Err(e) = clean_deleted_job(job_id.as_str()).await { + log::error!("[SUPER_CLUSTER:DB] Failed to clean deleted job: {job_id}, error: {e}",); + return Err(e); + } + } + JobOperator::Retry { + job_id, + new_trace_id, + updated_at, + } => { + if let Err(e) = + retry_search_job(job_id.as_str(), new_trace_id.as_str(), updated_at).await + { + log::error!("[SUPER_CLUSTER:DB] Failed to retry job: {job_id}, error: {e}",); + return Err(e); + } + } + } + Ok(()) +} diff --git a/src/super_cluster_queue/short_urls.rs b/src/super_cluster_queue/short_urls.rs new file mode 100644 index 0000000000..1ed3d3ffad --- /dev/null +++ b/src/super_cluster_queue/short_urls.rs @@ -0,0 +1,54 @@ +// Copyright 2024 Zinc Labs Inc. +// +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU Affero General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// This program is distributed in the hope that it will be useful +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU Affero General Public License for more details. +// +// You should have received a copy of the GNU Affero General Public License +// along with this program. If not, see . + +use infra::errors::{Error, Result}; +use o2_enterprise::enterprise::super_cluster::queue::{Message, MessageType}; + +pub(crate) async fn process(msg: Message) -> Result<()> { + match msg.message_type { + MessageType::ShortUrlPut => { + let short_id = parse_key(&msg.key)?; + let original_url: String = match msg.value { + Some(ref value) => String::from_utf8_lossy(value).to_string(), + None => String::new(), + }; + if original_url.is_empty() { + return Err(Error::Message("Invalid message value".to_string())); + } + infra::table::short_urls::add(&short_id, &original_url).await?; + } + MessageType::ShortUrlDelete => { + let short_id = parse_key(&msg.key)?; + infra::table::short_urls::remove(&short_id).await?; + } + _ => { + log::error!( + "[SUPER_CLUSTER:DB] Invalid message: type: {:?}, key: {}", + msg.message_type, + msg.key + ); + return Err(Error::Message("Invalid message type".to_string())); + } + } + Ok(()) +} + +fn parse_key(key: &str) -> Result { + let key_columns: Vec<&str> = key.split('/').collect(); + if key_columns.len() < 3 || key_columns[2].is_empty() { + return Err(Error::Message("Invalid key".to_string())); + } + Ok(key_columns[2].into()) +}