diff --git a/src/config/src/config.rs b/src/config/src/config.rs index 623510c057..5efc931288 100644 --- a/src/config/src/config.rs +++ b/src/config/src/config.rs @@ -53,7 +53,7 @@ pub type RwAHashSet = tokio::sync::RwLock>; pub type RwBTreeMap = tokio::sync::RwLock>; // for DDL commands and migrations -pub const DB_SCHEMA_VERSION: u64 = 13; +pub const DB_SCHEMA_VERSION: u64 = 14; pub const DB_SCHEMA_KEY: &str = "/db_schema_version/"; // global version variables diff --git a/src/infra/src/table/entity/mod.rs b/src/infra/src/table/entity/mod.rs index b4026fcdb5..adb531baf1 100644 --- a/src/infra/src/table/entity/mod.rs +++ b/src/infra/src/table/entity/mod.rs @@ -24,6 +24,7 @@ pub mod search_job_partitions; pub mod search_job_results; pub mod search_jobs; pub mod search_queue; +pub mod sessions; pub mod system_prompts; pub mod templates; pub mod timed_annotation_panels; diff --git a/src/infra/src/table/entity/prelude.rs b/src/infra/src/table/entity/prelude.rs index 9d126ebaa8..13d8066c1c 100644 --- a/src/infra/src/table/entity/prelude.rs +++ b/src/infra/src/table/entity/prelude.rs @@ -12,7 +12,7 @@ pub use super::{ report_dashboards::Entity as ReportDashboards, reports::Entity as Reports, search_job_partitions::Entity as SearchJobPartitions, search_job_results::Entity as SearchJobResults, search_jobs::Entity as SearchJobs, - search_queue::Entity as SearchQueue, templates::Entity as Templates, - timed_annotation_panels::Entity as TimedAnnotationPanels, + search_queue::Entity as SearchQueue, sessions::Entity as Sessions, + templates::Entity as Templates, timed_annotation_panels::Entity as TimedAnnotationPanels, timed_annotations::Entity as TimedAnnotations, users::Entity as Users, }; diff --git a/src/infra/src/table/entity/sessions.rs b/src/infra/src/table/entity/sessions.rs new file mode 100644 index 0000000000..65220719ec --- /dev/null +++ b/src/infra/src/table/entity/sessions.rs @@ -0,0 +1,34 @@ +// Copyright 2025 OpenObserve Inc. +// +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU Affero General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// This program is distributed in the hope that it will be useful +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU Affero General Public License for more details. +// +// You should have received a copy of the GNU Affero General Public License +// along with this program. If not, see . + +//! `SeaORM` Entity for sessions table + +use sea_orm::entity::prelude::*; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Eq, Serialize, Deserialize)] +#[sea_orm(table_name = "sessions")] +pub struct Model { + #[sea_orm(primary_key, auto_increment = false)] + pub session_id: String, + pub access_token: String, + pub created_at: i64, + pub updated_at: i64, +} + +#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)] +pub enum Relation {} + +impl ActiveModelBehavior for ActiveModel {} diff --git a/src/infra/src/table/migration/m20251118_000002_create_sessions_table.rs b/src/infra/src/table/migration/m20251118_000002_create_sessions_table.rs new file mode 100644 index 0000000000..cfecf59fcf --- /dev/null +++ b/src/infra/src/table/migration/m20251118_000002_create_sessions_table.rs @@ -0,0 +1,71 @@ +// Copyright 2025 OpenObserve Inc. +// +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU Affero General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// This program is distributed in the hope that it will be useful +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU Affero General Public License for more details. +// +// You should have received a copy of the GNU Affero General Public License +// along with this program. If not, see . + +use sea_orm_migration::prelude::*; + +use super::get_text_type; + +#[derive(DeriveMigrationName)] +pub struct Migration; + +#[async_trait::async_trait] +impl MigrationTrait for Migration { + async fn up(&self, manager: &SchemaManager) -> Result<(), DbErr> { + manager + .create_table(create_sessions_table_statement()) + .await?; + Ok(()) + } + + async fn down(&self, manager: &SchemaManager) -> Result<(), DbErr> { + manager + .drop_table(Table::drop().table(Sessions::Table).to_owned()) + .await?; + Ok(()) + } +} + +/// Statement to create the sessions table. +fn create_sessions_table_statement() -> TableCreateStatement { + let text_type = get_text_type(); + Table::create() + .table(Sessions::Table) + .if_not_exists() + .col( + // session_id is the primary key - enforces uniqueness without extra index + ColumnDef::new(Sessions::SessionId) + .string_len(36) + .not_null() + .primary_key(), + ) + .col( + ColumnDef::new(Sessions::AccessToken) + .custom(Alias::new(&text_type)) + .not_null(), + ) + .col(ColumnDef::new(Sessions::CreatedAt).big_integer().not_null()) + .col(ColumnDef::new(Sessions::UpdatedAt).big_integer().not_null()) + .to_owned() +} + +/// Identifiers used in queries on the sessions table. +#[derive(DeriveIden)] +pub(super) enum Sessions { + Table, + SessionId, + AccessToken, + CreatedAt, + UpdatedAt, +} diff --git a/src/infra/src/table/migration/m20251118_000003_delete_meta_sessions.rs b/src/infra/src/table/migration/m20251118_000003_delete_meta_sessions.rs new file mode 100644 index 0000000000..764b29de4e --- /dev/null +++ b/src/infra/src/table/migration/m20251118_000003_delete_meta_sessions.rs @@ -0,0 +1,71 @@ +// Copyright 2025 OpenObserve Inc. +// +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU Affero General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// This program is distributed in the hope that it will be useful +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU Affero General Public License for more details. +// +// You should have received a copy of the GNU Affero General Public License +// along with this program. If not, see . + +//! Deletes user session records from the meta table. +//! Sessions are ephemeral and users can simply login again. + +use sea_orm::{ColumnTrait, EntityTrait, QueryFilter, TransactionTrait}; +use sea_orm_migration::prelude::*; + +#[derive(DeriveMigrationName)] +pub struct Migration; + +#[async_trait::async_trait] +impl MigrationTrait for Migration { + async fn up(&self, manager: &SchemaManager) -> Result<(), DbErr> { + let txn = manager.get_connection().begin().await?; + match meta::Entity::delete_many() + .filter(meta::Column::Module.eq("user_sessions")) + .exec(&txn) + .await + { + Ok(_) => { + txn.commit().await?; + Ok(()) + } + Err(e) => { + txn.rollback().await?; + Err(e) + } + } + } + + async fn down(&self, _manager: &SchemaManager) -> Result<(), DbErr> { + // The deletion of records from the meta table is not reversable. + Ok(()) + } +} + +/// Representation of the meta table at the time this migration executes. +mod meta { + use sea_orm::entity::prelude::*; + + #[derive(Clone, Debug, PartialEq, DeriveEntityModel, Eq)] + #[sea_orm(table_name = "meta")] + pub struct Model { + #[sea_orm(primary_key)] + pub id: i64, + pub module: String, + pub key1: String, + pub key2: String, + pub start_dt: i64, + pub value: String, + } + + #[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)] + pub enum Relation {} + + impl ActiveModelBehavior for ActiveModel {} +} diff --git a/src/infra/src/table/migration/mod.rs b/src/infra/src/table/migration/mod.rs index a24078f488..c976232461 100644 --- a/src/infra/src/table/migration/mod.rs +++ b/src/infra/src/table/migration/mod.rs @@ -72,6 +72,8 @@ mod m20250930_000001_create_pipeline_last_errors_table; mod m20251024_000001_add_alert_deduplication; mod m20251105_000001_update_enrichment_table_created_at_mysql; mod m20251118_000001_add_alert_row_template_type; +mod m20251118_000002_create_sessions_table; +mod m20251118_000003_delete_meta_sessions; pub struct Migrator; @@ -133,6 +135,8 @@ impl MigratorTrait for Migrator { Box::new(m20251024_000001_add_alert_deduplication::Migration), Box::new(m20251105_000001_update_enrichment_table_created_at_mysql::Migration), Box::new(m20251118_000001_add_alert_row_template_type::Migration), + Box::new(m20251118_000002_create_sessions_table::Migration), + Box::new(m20251118_000003_delete_meta_sessions::Migration), ] } } diff --git a/src/infra/src/table/mod.rs b/src/infra/src/table/mod.rs index 0ca27d717b..772f2daaea 100644 --- a/src/infra/src/table/mod.rs +++ b/src/infra/src/table/mod.rs @@ -41,6 +41,7 @@ pub mod re_pattern_stream_map; pub mod reports; pub mod search_job; pub mod search_queue; +pub mod sessions; pub mod short_urls; pub mod system_prompts; pub mod templates; diff --git a/src/infra/src/table/sessions.rs b/src/infra/src/table/sessions.rs new file mode 100644 index 0000000000..cab321fc06 --- /dev/null +++ b/src/infra/src/table/sessions.rs @@ -0,0 +1,74 @@ +// Copyright 2025 OpenObserve Inc. +// +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU Affero General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// This program is distributed in the hope that it will be useful +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU Affero General Public License for more details. +// +// You should have received a copy of the GNU Affero General Public License +// along with this program. If not, see . + +use sea_orm::{ColumnTrait, EntityTrait, QueryFilter, Set, sea_query::OnConflict}; + +use super::entity::sessions::{ActiveModel, Column, Entity, Model}; +use crate::{ + db::{ORM_CLIENT, connect_to_orm}, + errors, +}; + +/// Gets a session by session_id +pub async fn get(session_id: &str) -> Result, errors::Error> { + let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let session = Entity::find() + .filter(Column::SessionId.eq(session_id)) + .one(client) + .await?; + Ok(session) +} + +/// Creates or updates a session atomically using upsert +pub async fn set(session_id: &str, access_token: &str) -> Result<(), errors::Error> { + let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let now = chrono::Utc::now().timestamp_micros(); + + // Use atomic upsert to avoid race conditions + let active_model = ActiveModel { + session_id: Set(session_id.to_string()), + access_token: Set(access_token.to_string()), + created_at: Set(now), + updated_at: Set(now), + }; + + Entity::insert(active_model) + .on_conflict( + OnConflict::column(Column::SessionId) + .update_columns([Column::AccessToken, Column::UpdatedAt]) + .to_owned(), + ) + .exec(client) + .await?; + + Ok(()) +} + +/// Deletes a session by session_id +pub async fn delete(session_id: &str) -> Result<(), errors::Error> { + let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + Entity::delete_many() + .filter(Column::SessionId.eq(session_id)) + .exec(client) + .await?; + Ok(()) +} + +/// Lists all sessions +pub async fn list() -> Result, errors::Error> { + let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let sessions = Entity::find().all(client).await?; + Ok(sessions) +} diff --git a/src/service/db/session.rs b/src/service/db/session.rs index 5f0717cf9a..204bcf1683 100644 --- a/src/service/db/session.rs +++ b/src/service/db/session.rs @@ -15,14 +15,14 @@ use std::sync::Arc; -use config::utils::json; +use bytes::Bytes; +use infra::db::{delete_from_db_coordinator, put_into_db_coordinator}; +#[cfg(feature = "enterprise")] +use o2_enterprise::enterprise::common::config::get_config as get_o2_config; -use crate::{ - common::infra::config::USER_SESSIONS, - service::db::{self}, -}; +use crate::common::infra::config::USER_SESSIONS; -// DBKey to set settings for an org +// Key prefix for session events in coordinator pub const USER_SESSION_KEY: &str = "/user_sessions/"; pub async fn get(session_id: &str) -> Result { @@ -32,43 +32,68 @@ pub async fn get(session_id: &str) -> Result { // get from db log::warn!("Cache miss for user session, read from db: {}", session_id); - let val = db::get(&format!("{USER_SESSION_KEY}{session_id}")).await?; - // json format: convert bytes to string and trim any quotes - let val = String::from_utf8(val.to_vec()) - .unwrap() - .trim_matches('"') - .to_string(); - // cache it in memory - if !val.is_empty() { - USER_SESSIONS.insert(session_id.to_string(), val.to_string()); + + let session = infra::table::sessions::get(session_id).await?; + + match session { + Some(session) => { + let access_token = session.access_token.clone(); + // Cache it in memory + if !access_token.is_empty() { + USER_SESSIONS.insert(session_id.to_string(), access_token.clone()); + } + Ok(access_token) + } + None => Err(anyhow::anyhow!("Session not found: {}", session_id)), } - Ok(val) } pub async fn set(session_id: &str, val: &str) -> Result<(), anyhow::Error> { - db::put( - &format!("{USER_SESSION_KEY}{session_id}"), - json::to_vec(&val).unwrap().into(), - db::NEED_WATCH, - None, - ) - .await?; + infra::table::sessions::set(session_id, val).await?; + let key = format!("{USER_SESSION_KEY}{session_id}"); + if let Err(e) = put_into_db_coordinator(&key, Bytes::new(), true, None).await { + log::error!("[SESSION] Failed to sync session to coordinator: {key} - {e}"); + } + + #[cfg(feature = "enterprise")] + if get_o2_config().super_cluster.enabled { + let _ = + o2_enterprise::enterprise::super_cluster::queue::put(&key, Bytes::new(), true, None) + .await + .inspect_err(|e| { + log::error!("[SESSION] put to super cluster failed: {key} - {e}"); + }); + } + + USER_SESSIONS.insert(session_id.to_string(), val.to_string()); + Ok(()) } pub async fn delete(session_id: &str) -> Result<(), anyhow::Error> { - Ok(db::delete( - &format!("{USER_SESSION_KEY}{session_id}"), - false, - db::NEED_WATCH, - None, - ) - .await?) + infra::table::sessions::delete(session_id).await?; + let key = format!("{USER_SESSION_KEY}{session_id}"); + if let Err(e) = delete_from_db_coordinator(&key, false, true, None).await { + log::error!("[SESSION] Failed to delete session from coordinator: {key} - {e}"); + } + + #[cfg(feature = "enterprise")] + if get_o2_config().super_cluster.enabled { + let _ = o2_enterprise::enterprise::super_cluster::queue::delete(&key, false, true, None) + .await + .inspect_err(|e| { + log::error!("[SESSION] delete to super cluster failed: {key} - {e}"); + }); + } + + USER_SESSIONS.remove(session_id); + + Ok(()) } pub async fn watch() -> Result<(), anyhow::Error> { let key = USER_SESSION_KEY; - let cluster_coordinator = db::get_coordinator().await; + let cluster_coordinator = infra::db::get_coordinator().await; let mut events = cluster_coordinator.watch(key).await?; let events = Arc::get_mut(&mut events).unwrap(); log::info!("Start watching user sessions"); @@ -76,56 +101,57 @@ pub async fn watch() -> Result<(), anyhow::Error> { let ev = match events.recv().await { Some(ev) => ev, None => { - log::error!("watch_user: event channel closed"); + log::error!("watch_sessions: event channel closed"); return Ok(()); } }; match ev { - db::Event::Put(ev) => { - let item_key = ev.key.strip_prefix(key).unwrap(); - let item_value: String = match db::get(&ev.key).await { - Ok(val) => match json::from_slice(&val) { - Ok(val) => val, - Err(e) => { - log::error!("Error getting value: {e}"); - continue; + infra::db::Event::Put(ev) => { + let session_id = ev.key.strip_prefix(key).unwrap(); + match infra::table::sessions::get(session_id).await { + Ok(Some(session)) => { + if !session.access_token.is_empty() { + USER_SESSIONS.insert(session_id.to_string(), session.access_token); + log::debug!("Session added to cache: {}", session_id); } - }, - Err(e) => { - log::error!("Error getting value: {e}"); - continue; } - }; - if item_value.is_empty() { - continue; + Ok(None) => { + log::warn!( + "Got coordinator message, but session id not found in DB: {}", + session_id + ); + } + Err(e) => { + log::error!( + "Error fetching session after coordinator message {}: {}", + session_id, + e + ); + } } - USER_SESSIONS.insert(item_key.to_string(), item_value); } - db::Event::Delete(ev) => { - let item_key = ev.key.strip_prefix(key).unwrap(); - USER_SESSIONS.remove(item_key); + infra::db::Event::Delete(ev) => { + let session_id = ev.key.strip_prefix(key).unwrap(); + USER_SESSIONS.remove(session_id); + log::debug!("Session removed from cache: {}", session_id); } - db::Event::Empty => {} + infra::db::Event::Empty => {} } } } pub async fn cache() -> Result<(), anyhow::Error> { - let key = USER_SESSION_KEY; - let ret = db::list(key).await?; - for (item_key, item_value) in ret { - let session_id = item_key.strip_prefix(key).unwrap(); - let json_val: String = match json::from_slice(&item_value) { - Ok(val) => val, - Err(_) => { - continue; - } - }; - if json_val.is_empty() { - continue; + let sessions_list = infra::table::sessions::list().await?; + + for session in sessions_list { + if !session.access_token.is_empty() { + USER_SESSIONS.insert(session.session_id, session.access_token); } - USER_SESSIONS.insert(session_id.to_owned(), json_val); } - log::info!("User Sessions Cached"); + + log::info!( + "User Sessions Cached: {} sessions loaded", + USER_SESSIONS.len() + ); Ok(()) }