diff --git a/src/api/common/src/auth/validator.rs b/src/api/common/src/auth/validator.rs index f0298d6b8c..cc81ba43c0 100644 --- a/src/api/common/src/auth/validator.rs +++ b/src/api/common/src/auth/validator.rs @@ -1273,7 +1273,7 @@ fn _extract_full_url(req: &Request) -> String { mod tests { use infra::{ db as infra_db, - db::{ORM_CLIENT, connect_to_orm}, + db::{get_orm_client_ro, get_orm_client_rw}, table as infra_table, }; use openobserve_core::{organization, users}; @@ -1575,7 +1575,7 @@ mod tests { let sa_email = "sa-nostatic@example.com"; let token = "sa_static_token_nostatic_test"; - let _ = ORM_CLIENT.get_or_init(connect_to_orm).await; + let _ = get_orm_client_rw().await; let _ = infra_db::create_table().await; let _ = infra_table::create_user_tables().await; let _ = organization::check_and_create_org_without_ofga(org_id).await; @@ -1631,7 +1631,7 @@ mod tests { let pwd = "Complexpass#123"; // Initialize ORM client and clear database tables for test isolation - let _ = ORM_CLIENT.get_or_init(connect_to_orm).await; + let _ = get_orm_client_ro().await; let _ = infra::table::org_users::clear().await; let _ = infra::table::users::clear().await; let _ = infra::table::organizations::clear().await; diff --git a/src/api/management/src/request/alerts/dedup_stats.rs b/src/api/management/src/request/alerts/dedup_stats.rs index 4b300a729b..b3e65d8833 100644 --- a/src/api/management/src/request/alerts/dedup_stats.rs +++ b/src/api/management/src/request/alerts/dedup_stats.rs @@ -15,6 +15,8 @@ use axum::{extract::Path, response::Response}; use common::meta::http::HttpResponse as MetaHttpResponse; +#[cfg(feature = "enterprise")] +use infra::db::get_orm_client_ro; use serde::{Deserialize, Serialize}; use utoipa::ToSchema; @@ -157,9 +159,7 @@ async fn get_dedup_counts_from_db(org_id: &str) -> Result<(i64, i64), anyhow::Er use infra::table::entity::alert_dedup_state; use sea_orm::{ColumnTrait, EntityTrait, QueryFilter}; - let db = infra::db::ORM_CLIENT - .get() - .ok_or_else(|| anyhow::anyhow!("Database not initialized"))?; + let db = get_orm_client_ro().await; // Get all dedup states for this org let states = alert_dedup_state::Entity::find() diff --git a/src/api/management/src/request/alerts/history.rs b/src/api/management/src/request/alerts/history.rs index 56d88ce839..ff0a5d011a 100644 --- a/src/api/management/src/request/alerts/history.rs +++ b/src/api/management/src/request/alerts/history.rs @@ -424,9 +424,7 @@ pub async fn get_alert_history( // If alert_id filter is provided, validate it exists let _folder_id = if let Some(alert_id) = query.alert_id { // Verify the alert exists in the organization - let conn = infra::db::ORM_CLIENT - .get_or_init(infra::db::connect_to_orm) - .await; + let conn = infra::db::get_orm_client_ro().await; match get_by_id(conn, &org_id, alert_id).await { Ok((f, _)) => Some(f.folder_id), Err(_) => { diff --git a/src/api/management/src/request/alerts/mod.rs b/src/api/management/src/request/alerts/mod.rs index 77efa62e59..4afe841d0c 100644 --- a/src/api/management/src/request/alerts/mod.rs +++ b/src/api/management/src/request/alerts/mod.rs @@ -28,7 +28,7 @@ use config::meta::{ }; use db::scheduler; use hashbrown::HashMap; -use infra::db::{ORM_CLIENT, connect_to_orm}; +use infra::db::{get_orm_client_ro, get_orm_client_rw}; use openobserve_api_common::extractors::Headers; use openobserve_core::{ alerts::{ @@ -132,7 +132,7 @@ pub async fn create_alert( } alert.last_edited_by = Some(user_email.user_id); - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; match alert::create(client, &org_id, &folder_id, alert, overwrite).await { Ok(v) => MetaHttpResponse::json( MetaHttpResponse::message(StatusCode::OK, "Alert saved") @@ -335,7 +335,7 @@ async fn composite_detail_response( ) .await .is_ok(); - let db = ORM_CLIENT.get_or_init(connect_to_orm).await; + let db = get_orm_client_ro().await; // Evaluate exactly as the scheduler would (system context, no child query). // Inaccessible children are masked below; a failed evaluation degrades to @@ -693,7 +693,7 @@ pub async fn validate_composite_alert( inaccessible.push(id.clone()); } } - let db = ORM_CLIENT.get_or_init(connect_to_orm).await; + let db = get_orm_client_ro().await; // Resolve every reference in one batched query instead of per-ID round trips. let resolutions = infra::table::alert_composites::resolve_many(db, &org_id, &references) .await @@ -835,7 +835,7 @@ pub async fn get_composite_references( Path((org_id, alert_id)): Path<(String, String)>, #[cfg(feature = "enterprise")] Headers(user_email): Headers, ) -> Response { - let db = ORM_CLIENT.get_or_init(connect_to_orm).await; + let db = get_orm_client_ro().await; let subject = openobserve_core::alerts::composite::get_composite(&org_id, &alert_id) .await .ok() @@ -950,7 +950,7 @@ pub async fn get_composite_timeline( Query(query): Query>, #[cfg(feature = "enterprise")] Headers(user_email): Headers, ) -> Response { - let db = ORM_CLIENT.get_or_init(connect_to_orm).await; + let db = get_orm_client_ro().await; let Some(composite) = openobserve_core::alerts::composite::get_composite(&org_id, &alert_id) .await .ok() @@ -1223,7 +1223,7 @@ pub async fn list_alert_groups(Path((org_id, alert_id)): Path<(String, String)>) let Ok(ksuid) = Ksuid::from_str(&alert_id) else { return MetaHttpResponse::not_found(format!("invalid alert id {alert_id}")); }; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; // Resolve through the alert itself so this endpoint inherits the same // org scoping and not-found behaviour as every other alert read — a raw // state-table query would happily serve another org's group labels, which @@ -1329,7 +1329,7 @@ pub async fn list_alert_group_transitions( let Ok(ksuid) = Ksuid::from_str(&alert_id) else { return MetaHttpResponse::not_found(format!("invalid alert id {alert_id}")); }; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; if let Err(e) = alert::get_by_id(client, &org_id, ksuid).await { return e.into(); } @@ -1407,7 +1407,7 @@ pub async fn get_alert( return MetaHttpResponse::not_found(format!("invalid alert id {alert_id}")); } }; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; match alert::get_by_id(client, &org_id, alert_id).await { Ok((_, alert)) => { let key = alert.get_unique_key(); @@ -1488,7 +1488,7 @@ pub async fn export_alert(Path((org_id, alert_id)): Path<(String, String)>) -> R return MetaHttpResponse::not_found(format!("invalid alert id {alert_id}")); } }; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; match alert::get_by_id(client, &org_id, alert_id).await { Ok((_, alert)) => { let key = alert.get_unique_key(); @@ -1586,7 +1586,7 @@ pub async fn clone_alert( return MetaHttpResponse::not_found(format!("invalid alert id {alert_id}")); } }; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; // Check if this is a regular alert first match alert::get_by_id(client, &org_id, alert_id).await { @@ -1739,7 +1739,7 @@ pub async fn update_alert( alert.last_edited_by = Some(user_email.user_id.clone()); alert.id = Some(alert_id); - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; match alert::update(client, &org_id, None, alert).await { Ok(_) => MetaHttpResponse::ok("Alert Updated"), Err(AlertError::AlertNotFound) => { @@ -1865,7 +1865,7 @@ pub async fn delete_alert( return MetaHttpResponse::not_found(format!("invalid alert id {alert_id}")); } }; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; // Check whether this ID belongs to a regular alert before attempting delete. // delete_by_id silently returns Ok(()) when the record is not found (required @@ -2062,7 +2062,7 @@ pub async fn delete_alert_bulk( let mut err = None; let mut conflicts = Vec::new(); - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; for id in req.ids { // already checked this is valid, so ok to unwrap let alert_id = Ksuid::from_str(&id).unwrap(); @@ -2221,7 +2221,7 @@ pub async fn list_alert_tags( if let Some(folder) = query.folder.clone() { params = params.in_folder(&folder); } - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let visible_ids: Vec = match alert::list_v2(client, user_id, params).await { Ok(list) => list .into_iter() @@ -2386,7 +2386,7 @@ pub async fn list_alerts( // Composite definitions are fetched once here — used both to decide whether // the `All` filter must merge (and so drop SQL pagination) and to build the // merged rows below. - let db = ORM_CLIENT.get_or_init(connect_to_orm).await; + let db = get_orm_client_ro().await; let composite_definitions = if matches!( alert_type, AlertTypeFilter::All | AlertTypeFilter::Composite @@ -2436,7 +2436,7 @@ pub async fn list_alerts( alert_type, AlertTypeFilter::AnomalyDetection | AlertTypeFilter::Composite ) { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let mut scheduled_jobs: HashMap = scheduler::list_by_org(&org_id, Some(TriggerModule::Alert)) .await @@ -2653,7 +2653,7 @@ async fn enrich_with_composite_metadata( visibility: &Option<(bool, hashbrown::HashSet)>, list: &mut [ListAlertsResponseBodyItem], ) { - let db = ORM_CLIENT.get_or_init(connect_to_orm).await; + let db = get_orm_client_ro().await; // Split the page by kind and resolve everything in bulk: one child-count // query for composites and two reverse-reference queries, instead of a @@ -2854,7 +2854,7 @@ pub async fn enable_alert( } }; let should_enable = query.value; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; match alert::enable_by_id(client, &org_id, alert_id, should_enable).await { Ok(_) => { let resp_body = EnableAlertResponseBody { @@ -2992,7 +2992,7 @@ pub async fn enable_alert_bulk( let mut unsuccessful = Vec::with_capacity(req.ids.len()); let mut err = None; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; for id in req.ids { match alert::enable_by_id(client, &org_id, id, should_enable).await { Ok(_) => { @@ -3110,7 +3110,7 @@ pub async fn trigger_alert( return MetaHttpResponse::not_found(format!("invalid alert id {alert_id}")); } }; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; match alert::trigger_by_id(client, &org_id, alert_id).await { Ok(_) => MetaHttpResponse::ok("Alert triggered"), Err(AlertError::AlertNotFound) => { @@ -3204,7 +3204,7 @@ pub async fn retrain_alert(Path((org_id, alert_id)): Path<(String, String)>) -> return MetaHttpResponse::not_found(format!("invalid alert id {alert_id}")); } }; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; // Check if this is a regular alert — if so, return 400 match alert::get_by_id(client, &org_id, alert_id).await { Ok(_) => { @@ -3267,7 +3267,7 @@ pub async fn move_alerts( Headers(user_email): Headers, Json(req_body): Json, ) -> Response { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let total_ids = req_body.alert_ids.len() + req_body.anomaly_config_ids.len(); // anomaly_config_ids is now a required Vec (defaults to empty), so no diff --git a/src/cli/basic/cli.rs b/src/cli/basic/cli.rs index 7541f1410c..93283cb8a2 100644 --- a/src/cli/basic/cli.rs +++ b/src/cli/basic/cli.rs @@ -20,10 +20,7 @@ use clap::{Arg, ArgAction, Command}; use common::{infra::config::USERS, meta}; use config::{DEFAULT_ORG, utils::file::set_permission}; use db; -use infra::{ - db::{ORM_CLIENT, connect_to_orm}, - file_list as infra_file_list, table, -}; +use infra::{db::get_orm_client_rw, file_list as infra_file_list, table}; use openobserve_core::users; use crate::{ @@ -277,7 +274,7 @@ pub async fn cli() -> Result { table::dashboards::delete_all().await?; } "report" => { - let conn = ORM_CLIENT.get_or_init(connect_to_orm).await; + let conn = get_orm_client_rw().await; db::dashboards::reports::reset(conn).await?; } "function" => { diff --git a/src/core/src/alerts/alert.rs b/src/core/src/alerts/alert.rs index 9d0c7aee8f..43f57dde8c 100644 --- a/src/core/src/alerts/alert.rs +++ b/src/core/src/alerts/alert.rs @@ -55,11 +55,7 @@ use db::{ }; #[cfg(feature = "enterprise")] use infra::table::workflows::WorkflowTriggerEntity; -use infra::{ - db::{ORM_CLIENT, connect_to_orm}, - schema::unwrap_stream_settings, - table, -}; +use infra::{db::get_orm_client_ro, schema::unwrap_stream_settings, table}; use itertools::Itertools; use lettre::{ AsyncTransport, Message, @@ -1068,7 +1064,7 @@ pub async fn get_by_id( } pub async fn get_by_id_db(org_id: &str, alert_id: Ksuid) -> Result { - let conn = ORM_CLIENT.get_or_init(connect_to_orm).await; + let conn = get_orm_client_ro().await; get_by_id(conn, org_id, alert_id).await.map(|f_a| f_a.1) } @@ -1120,11 +1116,11 @@ pub async fn list( } } -/// Gets a list of alerts from the database `ORM_CLIENT`. +/// Gets a list of alerts from the database. pub async fn list_with_folders_db( params: ListAlertsParams, ) -> Result, AlertError> { - let conn = ORM_CLIENT.get_or_init(connect_to_orm).await; + let conn = get_orm_client_ro().await; db::alerts::alert::list_with_folders(conn, params) .await .map_err(|e| e.into()) @@ -1801,9 +1797,7 @@ pub async fn delete_by_name( .map_err(|_| AlertError::AlertNotFound)? .ok_or(AlertError::AlertNotFound)?; let alert_id = alert.id.ok_or(AlertError::AlertNotFound)?; - let client = infra::db::ORM_CLIENT - .get_or_init(infra::db::connect_to_orm) - .await; + let client = infra::db::get_orm_client_rw().await; delete_by_id_user(client, org_id, alert_id).await } @@ -3807,8 +3801,8 @@ mod threshold_validation_tests { // every one of them green. // // Closing that needs `prepare_alert` itself, which reaches for folders, - // `get_by_id_db` and `infra::schema::get` — all through the global - // `ORM_CLIENT` rather than an injectable connection. The infra SQLite + // `get_by_id_db` and `infra::schema::get` — all through the global ORM + // clients rather than an injectable connection. The infra SQLite // harness cannot reach it without first making that global overridable in // tests; until then the wiring is verified by review, not by test. @@ -7019,11 +7013,7 @@ async fn validate_slo_alert_wiring( .map_err(AlertError::InvalidSloAlert); }; - let db = infra::db::ORM_CLIENT.get().ok_or_else(|| { - AlertError::InfraError(infra::errors::Error::Message( - "database not initialized".into(), - )) - })?; + let db = get_orm_client_ro().await; let slo = infra::table::slos::get(db, org_id, &cond.slo_id) .await diff --git a/src/core/src/alerts/composite/service.rs b/src/core/src/alerts/composite/service.rs index 9a40daf008..027a32d75d 100644 --- a/src/core/src/alerts/composite/service.rs +++ b/src/core/src/alerts/composite/service.rs @@ -16,12 +16,11 @@ use config::meta::{ }; use db::authz::{remove_ownership, set_ownership}; use infra::{ - db::{ORM_CLIENT, connect_to_orm}, + db::{get_orm_client_ro, get_orm_client_rw}, scheduler, table::{ alert_composites, entity::{alert_composite_children, alert_composites as composite_entity}, - get_lock, }, }; #[cfg(feature = "enterprise")] @@ -151,9 +150,7 @@ pub async fn set_composite_enabled( ensure_mutation_allowed()?; } use sea_orm::{ColumnTrait, EntityTrait, QueryFilter, sea_query::Expr}; - let db = ORM_CLIENT.get_or_init(connect_to_orm).await; - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; + let db = get_orm_client_rw().await; let result = composite_entity::Entity::update_many() .col_expr(composite_entity::Column::Enabled, Expr::value(enabled)) .col_expr( @@ -164,8 +161,6 @@ pub async fn set_composite_enabled( .filter(composite_entity::Column::Id.eq(id)) .exec(db) .await?; - // released before the scheduler follow-ups below, which take this lock themselves - drop(_lock); if result.rows_affected != 1 { return Err(CompositeServiceError::NotFound); } @@ -193,7 +188,7 @@ async fn increment_and_advance( next_run_at: i64, ) -> Result<(), CompositeServiceError> { ensure_mutation_allowed()?; - let db = ORM_CLIENT.get_or_init(connect_to_orm).await; + let db = get_orm_client_rw().await; alert_composites::increment_evaluation_generation(db, org, id).await?; match scheduler::get(org, TriggerModule::CompositeAlert, id).await { Ok(mut trigger) => { @@ -221,7 +216,7 @@ pub async fn move_composite( ) -> Result<(), CompositeServiceError> { ensure_mutation_allowed()?; use sea_orm::{ColumnTrait, EntityTrait, QueryFilter, sea_query::Expr}; - let db = ORM_CLIENT.get_or_init(connect_to_orm).await; + let db = get_orm_client_rw().await; // Resolve the current folder for the authorization check and OpenFGA // relation rewrite below, mirroring alert::move_to_folder. #[cfg(feature = "enterprise")] @@ -247,8 +242,6 @@ pub async fn move_composite( { return Err(CompositeServiceError::PermissionDenied); } - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; let result = composite_entity::Entity::update_many() .col_expr(composite_entity::Column::FolderId, Expr::value(folder_id)) .col_expr(composite_entity::Column::LastEditedBy, Expr::value(editor)) @@ -265,7 +258,6 @@ pub async fn move_composite( .exec(db) .await?; // released before the OpenFGA relation rewrites below - drop(_lock); #[cfg(feature = "enterprise")] if get_openfga_config().enabled { set_parent_relation( @@ -342,7 +334,7 @@ async fn persist( let graph_guard = composite_graph_lock::lock(&request.org) .await .map_err(|error| CompositeServiceError::Lock(error.to_string()))?; - let db = ORM_CLIENT.get_or_init(connect_to_orm).await; + let db = get_orm_client_rw().await; let mutation = persist_under_lock(db, &id, request, expression, references, update).await; let unlock = graph_guard .release() @@ -550,7 +542,7 @@ pub async fn get_composite( org: &str, id: &str, ) -> Result, CompositeServiceError> { - let db = ORM_CLIENT.get_or_init(connect_to_orm).await; + let db = get_orm_client_ro().await; Ok(alert_composites::get_by_id(db, org, id).await?) } @@ -724,7 +716,7 @@ pub async fn validate_composite_graph( let parsed = parse_expr(expression).map_err(map_expression_error)?; let references = collect_references(&parsed).map_err(map_expression_error)?; validate_children(&references).map_err(map_expression_error)?; - let db = ORM_CLIENT.get_or_init(connect_to_orm).await; + let db = get_orm_client_rw().await; let resolved = resolve_children(db, org, &references).await?; let mut graph = alert_composites::load_graph(db, org).await?; let candidate_id = composite_id.unwrap_or("__composite_preview__"); @@ -747,7 +739,7 @@ pub async fn delete_composite(org: &str, id: &str) -> Result<(), CompositeServic let graph_guard = composite_graph_lock::lock(org) .await .map_err(|error| CompositeServiceError::Lock(error.to_string()))?; - let db = ORM_CLIENT.get_or_init(connect_to_orm).await; + let db = get_orm_client_rw().await; let result = async { let parents = alert_composites::list_parents(db, org, alert_composites::ChildKind::Composite, id) @@ -837,7 +829,7 @@ pub async fn startup_preflight() -> anyhow::Result<()> { return Ok(()); } - let db = ORM_CLIENT.get_or_init(connect_to_orm).await; + let db = get_orm_client_ro().await; let definition_count = alert_composites::count_all(db).await? as usize; // `list` (not `len_module`) so a scheduler read failure fails the preflight // closed rather than silently reporting zero jobs. diff --git a/src/core/src/alerts/deduplication.rs b/src/core/src/alerts/deduplication.rs index f8c4a30083..5a8a341dcf 100644 --- a/src/core/src/alerts/deduplication.rs +++ b/src/core/src/alerts/deduplication.rs @@ -25,7 +25,7 @@ use config::{ }, utils::json::{Map, Value}, }; -use infra::table::{entity::alert_dedup_state, get_lock}; +use infra::table::entity::alert_dedup_state; use sea_orm::{ActiveModelTrait, ColumnTrait, DatabaseConnection, EntityTrait, QueryFilter, Set}; /// Append the evaluated level to a fingerprint for MULTI-LEVEL alerts. @@ -161,8 +161,6 @@ pub async fn confirm_notification_sent( if fingerprints.is_empty() { return Ok(()); } - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; alert_dedup_state::Entity::update_many() .col_expr( alert_dedup_state::Column::NotificationSent, @@ -199,8 +197,6 @@ pub async fn save_dedup_state( db: &DatabaseConnection, params: DedupStateParams<'_>, ) -> Result { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; // Try to find existing record if let Some(existing) = get_dedup_state(db, params.fingerprint).await? { // Update existing @@ -241,8 +237,6 @@ pub async fn cleanup_expired_state( let cutoff_time = o2_enterprise::enterprise::alerts::dedup::current_timestamp_micros() - (older_than_minutes * 60 * 1_000_000); - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; let query = alert_dedup_state::Entity::delete_many() .filter(alert_dedup_state::Column::LastSeenAt.lt(cutoff_time)); diff --git a/src/core/src/alerts/scheduler/handlers.rs b/src/core/src/alerts/scheduler/handlers.rs index 69280333e2..ef8ace876e 100644 --- a/src/core/src/alerts/scheduler/handlers.rs +++ b/src/core/src/alerts/scheduler/handlers.rs @@ -38,9 +38,8 @@ use config::{ }; use cron::Schedule; use infra::{ - db::{ORM_CLIENT, connect_to_orm}, + db::{get_orm_client_ro, get_orm_client_rw}, scheduler::get_scheduler_max_retries, - table::get_lock, }; #[cfg(feature = "enterprise")] use o2_enterprise::enterprise::recommendations::service::QueryRecommendationService; @@ -136,7 +135,7 @@ async fn persist_alert_run_state( .unwrap_or(false) }); if transition_changed || stale_to_fresh { - let db = ORM_CLIENT.get_or_init(connect_to_orm).await; + let db = get_orm_client_ro().await; nudge_composite_parents( db, &alert.org_id, @@ -224,7 +223,7 @@ async fn persist_alert_run_state( .unwrap_or(false) }); if transition_changed || stale_to_fresh { - let db = ORM_CLIENT.get_or_init(connect_to_orm).await; + let db = get_orm_client_ro().await; nudge_composite_parents( db, &alert.org_id, @@ -537,9 +536,7 @@ async fn confirm_dedup_reservations(fingerprints: &[String], delivered: bool) { } #[cfg(feature = "enterprise")] { - let db = infra::db::ORM_CLIENT - .get_or_init(infra::db::connect_to_orm) - .await; + let db = infra::db::get_orm_client_rw().await; if let Err(e) = crate::alerts::deduplication::confirm_notification_sent(db, fingerprints).await { @@ -681,7 +678,7 @@ async fn handle_composite_alert_trigger( // transaction used for state fencing, a missing/stale definition is // completed safely through the epoch-fenced scheduler seam. This path is // also the required lifecycle behavior after deletion. - let db = ORM_CLIENT.get_or_init(connect_to_orm).await; + let db = get_orm_client_rw().await; let definition = infra::table::alert_composites::get_by_id(db, &trigger.org, &trigger.module_key).await?; let Some(definition) = definition else { @@ -816,8 +813,6 @@ async fn handle_composite_alert_trigger( use sea_orm::{ ColumnTrait, ConnectionTrait, EntityTrait, QueryFilter, QuerySelect, TransactionTrait, }; - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; let transaction = db.begin().await?; let lease_deadline = now + config::get_config().limit.alert_schedule_timeout * 1_000_000; if !infra::scheduler::renew_claim_in_transaction( @@ -843,8 +838,6 @@ async fn handle_composite_alert_trigger( || current.evaluation_generation != definition.definition.evaluation_generation }) { transaction.rollback().await?; - // released before complete_claim, which takes this lock itself on sqlite - drop(_lock); trigger.status = db::scheduler::TriggerStatus::Waiting; trigger.next_run_at = now + composite_debounce_secs() * 1_000_000; let _ = infra::scheduler::complete_claim(trigger).await?; @@ -852,8 +845,6 @@ async fn handle_composite_alert_trigger( } infra::table::alert_states::persist_update_in_transaction(&transaction, &update).await?; transaction.commit().await?; - // released before delivery and complete_claim below, which take this lock themselves - drop(_lock); let mut delivery_retry_at = None; if evaluated.result { @@ -1144,9 +1135,7 @@ async fn handle_anomaly_detection_triggers( return Ok(()); } - let db = infra::db::ORM_CLIENT - .get() - .ok_or_else(|| anyhow::anyhow!("Database not initialized"))?; + let db = get_orm_client_rw().await; let config = anomaly_config_table::get_by_id(db, &trigger.org, &anomaly_id) .await @@ -1284,8 +1273,6 @@ async fn handle_anomaly_detection_triggers( let mut active = config.into_active_model(); active.status = Set(AnomalyStatus::Active.to_i32()); active.updated_at = Set(run_end_us); - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; if let Err(e) = active.update(db).await { log::warn!( "[anomaly_detection] failed to reset status to Active for {anomaly_id}: {e}" @@ -1444,7 +1431,7 @@ async fn handle_alert_triggers( // here it can be alert id or alert name let alert = if let Ok(alert_id) = svix_ksuid::Ksuid::from_str(&trigger.module_key) { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; match db::alerts::alert::get_by_id(client, &trigger.org, alert_id).await { Ok(Some((_, alert))) => alert, Ok(None) => { @@ -2264,7 +2251,8 @@ async fn handle_alert_triggers( // Apply deduplication if enabled (enterprise-only feature) #[cfg(feature = "enterprise")] - let data = if let Some(db) = ORM_CLIENT.get() { + let data = { + let db = get_orm_client_rw().await; match crate::alerts::deduplication::apply_deduplication( db, &alert, @@ -2331,11 +2319,6 @@ async fn handle_alert_triggers( data } } - } else { - log::warn!( - "[SCHEDULER trace_id {scheduler_trace_id}] Could not connect to ORM for deduplication, continuing without it" - ); - data }; // [ENTERPRISE] Collect alert events for batched incident creation @@ -2924,7 +2907,7 @@ async fn handle_report_triggers( trace_id: &str, trigger: db::scheduler::Trigger, ) -> Result<(), anyhow::Error> { - let conn = ORM_CLIENT.get_or_init(connect_to_orm).await; + let conn = get_orm_client_rw().await; let query_trace_id = ider::generate_trace_id(); let scheduler_trace_id = format!("{trace_id}/{query_trace_id}"); let (_, max_retries) = get_scheduler_max_retries(); @@ -5304,9 +5287,7 @@ async fn handle_slo_triggers(mut trigger: db::scheduler::Trigger) -> Result<(), let slo_id = trigger.module_key.clone(); - let db = infra::db::ORM_CLIENT - .get() - .ok_or_else(|| anyhow::anyhow!("Database not initialized"))?; + let db = get_orm_client_ro().await; let slo = infra::table::slos::get(db, &trigger.org, &slo_id) .await .map_err(|e| anyhow::anyhow!(e.to_string()))?; @@ -5366,9 +5347,7 @@ async fn handle_slo_backfill_triggers( let slo_id = trigger.module_key.clone(); - let db = infra::db::ORM_CLIENT - .get() - .ok_or_else(|| anyhow::anyhow!("Database not initialized"))?; + let db = get_orm_client_ro().await; let Some(slo) = infra::table::slos::get(db, &trigger.org, &slo_id) .await .map_err(|e| anyhow::anyhow!(e.to_string()))? diff --git a/src/core/src/anomaly_detection.rs b/src/core/src/anomaly_detection.rs index d496ce2d83..82ece7754c 100644 --- a/src/core/src/anomaly_detection.rs +++ b/src/core/src/anomaly_detection.rs @@ -26,8 +26,8 @@ use config::{ }; use db::authz::{remove_ownership, set_ownership}; use infra::{ - db::ORM_CLIENT, - table::{anomaly_detection::config as anomaly_config_table, get_lock}, + db::{get_orm_client_ro, get_orm_client_rw}, + table::anomaly_detection::config as anomaly_config_table, }; use sea_orm::{ActiveModelTrait, IntoActiveModel, Set}; use search_service as search; @@ -220,9 +220,7 @@ pub async fn list_configs( folder_name: Option<&str>, name_substring: Option<&str>, ) -> Result> { - let db = ORM_CLIENT - .get() - .ok_or_else(|| anyhow::anyhow!("Database not initialized"))?; + let db = get_orm_client_rw().await; let configs = anomaly_config_table::list_by_org(db, org_id) .await @@ -336,9 +334,7 @@ pub async fn list_configs( /// Get a specific anomaly detection configuration pub async fn get_config(org_id: &str, anomaly_id: &str) -> Result> { - let db = ORM_CLIENT - .get() - .ok_or_else(|| anyhow::anyhow!("Database not initialized"))?; + let db = get_orm_client_rw().await; let config = anomaly_config_table::get_by_id(db, org_id, anomaly_id) .await @@ -371,9 +367,7 @@ pub async fn create_config( let normalized_tags = config::meta::alerts::tags::normalize_tags(&req.tags).map_err(anyhow::Error::new)?; - let db = ORM_CLIENT - .get() - .ok_or_else(|| anyhow::anyhow!("Database not initialized"))?; + let db = get_orm_client_rw().await; let anomaly_id = svix_ksuid::Ksuid::new(None, None).to_string(); let now_us = Utc::now().timestamp_micros(); @@ -552,9 +546,7 @@ pub async fn update_config( anomaly_id: &str, req: UpdateAnomalyConfigRequest, ) -> Result { - let db = ORM_CLIENT - .get() - .ok_or_else(|| anyhow::anyhow!("Database not initialized"))?; + let db = get_orm_client_rw().await; // Fetch existing config — into_active_model() on a DB-fetched model sets the PK as // Unchanged, which is required for SeaORM to generate UPDATE … WHERE anomaly_id = ? @@ -696,11 +688,7 @@ pub async fn update_config( active_model.updated_at = Set(Utc::now().timestamp_micros()); - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; let updated = active_model.update(db).await?; - // released before the retrain/scheduler follow-ups below, which take this lock themselves - drop(_lock); // Evict any cached model for this config — training params may have changed, // so the next detection run should load a fresh model from S3. @@ -851,9 +839,7 @@ pub async fn update_config( /// Delete an anomaly detection configuration pub async fn delete_config(org_id: &str, anomaly_id: &str) -> Result<()> { - let db = ORM_CLIENT - .get() - .ok_or_else(|| anyhow::anyhow!("Database not initialized"))?; + let db = get_orm_client_rw().await; // Verify the config exists before deleting (returns 404 if missing). let existing = anomaly_config_table::get_by_id(db, org_id, anomaly_id) @@ -927,9 +913,7 @@ pub async fn clone_config( new_name: Option, folder_id: Option, ) -> Result { - let db = ORM_CLIENT - .get() - .ok_or_else(|| anyhow::anyhow!("Database not initialized"))?; + let db = get_orm_client_rw().await; let src = anomaly_config_table::get_by_id(db, org_id, anomaly_id) .await @@ -1060,9 +1044,7 @@ pub async fn clone_config( /// threshold changes are recomputed in place with no retrain. #[cfg(feature = "enterprise")] async fn force_retrain_for_threshold(org_id: &str, anomaly_id: &str) -> Result<()> { - let db = ORM_CLIENT - .get() - .ok_or_else(|| anyhow::anyhow!("Database not initialized"))?; + let db = get_orm_client_rw().await; let config = anomaly_config_table::get_by_id(db, org_id, anomaly_id) .await @@ -1086,10 +1068,7 @@ async fn force_retrain_for_threshold(org_id: &str, anomaly_id: &str) -> Result<( active.status = Set(0i32); active.is_trained = Set(false); active.updated_at = Set(Utc::now().timestamp_micros()); - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; active.update(db).await?; - drop(_lock); // Nudge the training scheduler so the retrain fires promptly rather than on its next // periodic sweep. @@ -1107,9 +1086,7 @@ async fn force_retrain_for_threshold(org_id: &str, anomaly_id: &str) -> Result<( /// which the user can then re-cancel if needed. In practice the task will notice the /// status change and bail early on the next DB write. pub async fn cancel_training(org_id: &str, anomaly_id: &str) -> Result<()> { - let db = ORM_CLIENT - .get() - .ok_or_else(|| anyhow::anyhow!("Database not initialized"))?; + let db = get_orm_client_rw().await; let config = anomaly_config_table::get_by_id(db, org_id, anomaly_id) .await @@ -1122,8 +1099,6 @@ pub async fn cancel_training(org_id: &str, anomaly_id: &str) -> Result<()> { active.training_started_at = Set(None); active.last_error = Set(Some("Training cancelled by user.".to_string())); active.updated_at = Set(Utc::now().timestamp_micros()); - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; active.update(db).await?; log::info!("[anomaly_detection {anomaly_id}] training cancelled by user"); @@ -1133,9 +1108,7 @@ pub async fn cancel_training(org_id: &str, anomaly_id: &str) -> Result<()> { /// Train a model for a configuration pub async fn train_model(org_id: &str, anomaly_id: &str) -> Result { // Verify the config exists and belongs to this org before delegating. - let db = ORM_CLIENT - .get() - .ok_or_else(|| anyhow::anyhow!("Database not initialized"))?; + let db = get_orm_client_ro().await; anomaly_config_table::get_by_id(db, org_id, anomaly_id) .await .map_err(|e| anyhow::anyhow!(e.to_string()))? @@ -1158,9 +1131,7 @@ pub async fn train_model(org_id: &str, anomaly_id: &str) -> Result Result { - let db = ORM_CLIENT - .get() - .ok_or_else(|| anyhow::anyhow!("Database not initialized"))?; + let db = get_orm_client_ro().await; // Fetch config let config = anomaly_config_table::get_by_id(db, org_id, anomaly_id) @@ -1365,13 +1336,7 @@ pub struct DetectionHistoryItem { /// scheduler's `watch_timeout()` job already resets any `Processing` rows whose /// `end_time` has passed, exactly as it does for alert triggers. pub async fn recover_detection_triggers_on_startup() { - let db = match ORM_CLIENT.get() { - Some(db) => db, - None => { - log::warn!("[anomaly_detection] DB not available for startup trigger recovery"); - return; - } - }; + let db = get_orm_client_ro().await; let configs = match anomaly_config_table::list_all_enabled(db).await { Ok(c) => c, diff --git a/src/core/src/dashboards/reports.rs b/src/core/src/dashboards/reports.rs index 57ca778093..aa90ac6f48 100644 --- a/src/core/src/dashboards/reports.rs +++ b/src/core/src/dashboards/reports.rs @@ -40,7 +40,7 @@ use db::{ }; use futures::{StreamExt, future::try_join_all}; use infra::{ - db::{ORM_CLIENT, connect_to_orm}, + db::{get_orm_client_ro, get_orm_client_rw}, table, }; use itertools::Itertools; @@ -117,7 +117,7 @@ pub async fn save( mut report: Report, create: bool, ) -> Result<(), ReportError> { - let conn = ORM_CLIENT.get_or_init(connect_to_orm).await; + let conn = get_orm_client_rw().await; let cfg = get_config(); if cfg.common.report_server_url.is_empty() { // Check if SMTP is enabled, otherwise don't save the report @@ -250,14 +250,14 @@ pub async fn save( } pub async fn get(org_id: &str, folder_id: &str, name: &str) -> Result { - let conn = ORM_CLIENT.get_or_init(connect_to_orm).await; + let conn = get_orm_client_rw().await; db::dashboards::reports::get(conn, org_id, folder_id, name) .await .map_err(|_| ReportError::ReportNotFound) } pub async fn get_by_id(org_id: &str, report_id: &str) -> Result<(Folder, Report), ReportError> { - let conn = ORM_CLIENT.get_or_init(connect_to_orm).await; + let conn = get_orm_client_ro().await; match table::reports::get_by_id(conn, report_id).await { Ok(Some((folder, report))) if report.org_id == org_id => Ok((folder, report)), Ok(_) => Err(ReportError::ReportNotFound), @@ -270,7 +270,7 @@ pub async fn list( filters: ReportListFilters, permitted: Option>, ) -> Result, ReportError> { - let conn = ORM_CLIENT.get_or_init(connect_to_orm).await; + let conn = get_orm_client_rw().await; let params = filters.into_params(org_id); let reports = db::dashboards::reports::list(conn, ¶ms) .await @@ -293,7 +293,7 @@ pub async fn list( } pub async fn delete(org_id: &str, folder_id: &str, name: &str) -> Result<(), ReportError> { - let conn = ORM_CLIENT.get_or_init(connect_to_orm).await; + let conn = get_orm_client_rw().await; // Existence check — db::delete returns empty string if not found. if db::dashboards::reports::get(conn, org_id, folder_id, name) @@ -324,7 +324,7 @@ pub async fn delete(org_id: &str, folder_id: &str, name: &str) -> Result<(), Rep } pub async fn delete_by_id(org_id: &str, report_id: &str) -> Result<(), ReportError> { - let conn = ORM_CLIENT.get_or_init(connect_to_orm).await; + let conn = get_orm_client_rw().await; let (folder, _) = get_by_id(org_id, report_id).await?; db::dashboards::reports::delete_by_id(conn, org_id, report_id) @@ -345,7 +345,7 @@ pub async fn delete_by_id(org_id: &str, report_id: &str) -> Result<(), ReportErr } pub async fn trigger(org_id: &str, folder_id: &str, name: &str) -> Result<(), ReportError> { - let conn = ORM_CLIENT.get_or_init(connect_to_orm).await; + let conn = get_orm_client_rw().await; let report = match db::dashboards::reports::get(conn, org_id, folder_id, name).await { Ok(report) => report, _ => { @@ -368,7 +368,7 @@ pub async fn enable( name: &str, value: bool, ) -> Result<(), ReportError> { - let conn = ORM_CLIENT.get_or_init(connect_to_orm).await; + let conn = get_orm_client_rw().await; // TODO: The "get" and "update" operations should be in a transaction. let mut report = match db::dashboards::reports::get(conn, org_id, folder_id, name).await { @@ -384,7 +384,7 @@ pub async fn enable( } pub async fn enable_by_id(org_id: &str, report_id: &str, value: bool) -> Result<(), ReportError> { - let conn = ORM_CLIENT.get_or_init(connect_to_orm).await; + let conn = get_orm_client_rw().await; let (_, mut report) = get_by_id(org_id, report_id).await?; report.enabled = value; db::dashboards::reports::update_by_id(conn, report_id, None, report) @@ -398,7 +398,7 @@ pub async fn update_by_id( new_folder_id: Option<&str>, mut report: Report, ) -> Result<(), ReportError> { - let conn = ORM_CLIENT.get_or_init(connect_to_orm).await; + let conn = get_orm_client_rw().await; let cfg = get_config(); if cfg.common.report_server_url.is_empty() { @@ -483,7 +483,7 @@ pub async fn move_to_folder( report_ids: &[String], dst_folder_id: &str, ) -> Result<(), ReportError> { - let conn = ORM_CLIENT.get_or_init(connect_to_orm).await; + let conn = get_orm_client_rw().await; for report_id in report_ids { let (curr_folder, report) = get_by_id(org_id, report_id).await?; diff --git a/src/core/src/ofga/migrations/mod.rs b/src/core/src/ofga/migrations/mod.rs index 5bedcd21c8..2fb7f1afb0 100644 --- a/src/core/src/ofga/migrations/mod.rs +++ b/src/core/src/ofga/migrations/mod.rs @@ -1,4 +1,4 @@ -use infra::db::{ORM_CLIENT, connect_to_orm}; +use infra::db::get_orm_client_rw; pub(super) mod alert_folders; pub(super) mod anomaly_detection; @@ -6,21 +6,21 @@ pub(super) mod report_folders; pub(super) mod stream_name_migration; pub async fn migrate_alert_folders() -> Result<(), anyhow::Error> { - let db = ORM_CLIENT.get_or_init(connect_to_orm).await; + let db = get_orm_client_rw().await; alert_folders::migrate_alert_folders(db).await } pub async fn migrate_anomaly_detection() -> Result<(), anyhow::Error> { - let db = ORM_CLIENT.get_or_init(connect_to_orm).await; + let db = get_orm_client_rw().await; anomaly_detection::migrate_anomaly_detection(db).await } pub async fn migrate_report_folders() -> Result<(), anyhow::Error> { - let db = ORM_CLIENT.get_or_init(connect_to_orm).await; + let db = get_orm_client_rw().await; report_folders::migrate_report_folders(db).await } pub async fn migrate_stream_names() -> Result<(), anyhow::Error> { - let db = ORM_CLIENT.get_or_init(connect_to_orm).await; + let db = get_orm_client_rw().await; stream_name_migration::migrate_stream_names(db).await } diff --git a/src/core/src/org_cleanup.rs b/src/core/src/org_cleanup.rs index d7e9018d2d..22d9cf9731 100644 --- a/src/core/src/org_cleanup.rs +++ b/src/core/src/org_cleanup.rs @@ -17,7 +17,7 @@ use std::sync::Arc; use async_trait::async_trait; use config::{meta::stream::StreamType, spawn_pausable_job}; -use infra::{dist_lock, table::org_cleanup_tasks}; +use infra::{db::get_orm_client_rw, dist_lock, table::org_cleanup_tasks}; const LOCK_KEY: &str = "/org_cleanup/worker_lock"; const MAX_ATTEMPTS: i32 = 10; @@ -378,8 +378,7 @@ async fn step_delete_file_list(org_id: &str) -> Result<(), anyhow::Error> { /// each alert's scheduler trigger + OFGA ownership is removed — none of which a /// raw table wipe would do. async fn delete_org_alerts(org_id: &str) -> Result<(), anyhow::Error> { - use infra::db::{ORM_CLIENT, connect_to_orm}; - let conn = ORM_CLIENT.get_or_init(connect_to_orm).await; + let conn = get_orm_client_rw().await; // Organization teardown is the sole graph-guard bypass. Remove composite // child-index rows and definitions before ordinary alerts so no dangling @@ -436,15 +435,12 @@ async fn step_delete_db_resources(org_id: &str) -> Result<(), anyhow::Error> { // build through the storage layer instead, and imports the table module only // under #[cfg(not(feature = "enterprise"))] above. Importing it twice would // collide on the OSS build. - use infra::{ - db::{ORM_CLIENT, connect_to_orm}, - table::{ - action_scripts, alert_incidents, backfill_jobs, compactor_manual_jobs, dashboards, - destinations, distinct_values, enrichment_table_urls, enrichment_tables, folders, - incident_events, kv_store, org_ingestion_tokens, org_storage_providers, re_pattern, - re_pattern_stream_map, reports, search_queue, short_urls, slo, slo_backfill_jobs, - slo_budget, slos, system_settings, templates, timed_annotations, - }, + use infra::table::{ + action_scripts, alert_incidents, backfill_jobs, compactor_manual_jobs, dashboards, + destinations, distinct_values, enrichment_table_urls, enrichment_tables, folders, + incident_events, kv_store, org_ingestion_tokens, org_storage_providers, re_pattern, + re_pattern_stream_map, reports, search_queue, short_urls, slo, slo_backfill_jobs, + slo_budget, slos, system_settings, templates, timed_annotations, }; // FK-constrained children must be deleted before their parents. @@ -463,7 +459,7 @@ async fn step_delete_db_resources(org_id: &str) -> Result<(), anyhow::Error> { // Within the SLO tables the order is also fixed: status rows and backfill jobs // carry no org column and resolve their org through `slos`, so they are swept // before the definitions that identify them. - let slo_conn = ORM_CLIENT.get_or_init(connect_to_orm).await; + let slo_conn = get_orm_client_rw().await; slo::delete_by_org(slo_conn, org_id) .await .map_err(|e| anyhow::anyhow!("step_delete_db_resources/slo_status: {e}"))?; diff --git a/src/core/src/slo/backfill.rs b/src/core/src/slo/backfill.rs index 000d880a82..5fd53f0150 100644 --- a/src/core/src/slo/backfill.rs +++ b/src/core/src/slo/backfill.rs @@ -32,7 +32,10 @@ use config::{ window::{align_down, align_up}, }, }; -use infra::table::{slo as slo_table, slo_backfill_jobs as jobs}; +use infra::{ + db::{get_orm_client_ro, get_orm_client_rw}, + table::{slo as slo_table, slo_backfill_jobs as jobs}, +}; /// Whether the job has more work. #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -72,9 +75,7 @@ pub fn next_chunk( /// Fill one chunk for `slo`. pub async fn run_chunk(slo: &Slo) -> Result { let cfg = get_config(); - let db = infra::db::ORM_CLIENT - .get() - .ok_or_else(|| anyhow::anyhow!("database not initialized"))?; + let db = get_orm_client_rw().await; let Some(job) = jobs::get(db, &slo.id, slo.definition_generation).await? else { return Ok(ChunkOutcome::Done); @@ -168,9 +169,7 @@ pub fn clamp_backfill_start( /// Whether a status row still needs its history filled. pub async fn is_needed(slo: &Slo) -> Result { - let db = infra::db::ORM_CLIENT - .get() - .ok_or_else(|| anyhow::anyhow!("database not initialized"))?; + let db = get_orm_client_ro().await; let status = slo_table::load_status(db, &slo.id, "").await?; Ok(status.is_some_and(|s| s.definition_generation == slo.definition_generation)) } diff --git a/src/core/src/slo/evaluate.rs b/src/core/src/slo/evaluate.rs index 7b0cb000c6..c70c424f41 100644 --- a/src/core/src/slo/evaluate.rs +++ b/src/core/src/slo/evaluate.rs @@ -43,7 +43,10 @@ use config::{ window::{expected_slices, watermark_is_stale_or_absent}, }, }; -use infra::table::{slo as slo_table, slos as slos_table}; +use infra::{ + db::get_orm_client_ro, + table::{slo as slo_table, slos as slos_table}, +}; /// Everything one SLO-alert evaluation produced. #[derive(Debug, Clone)] @@ -85,9 +88,7 @@ pub async fn evaluate( ); } - let db = infra::db::ORM_CLIENT - .get() - .ok_or_else(|| anyhow::anyhow!("database not initialized"))?; + let db = get_orm_client_ro().await; let Some(slo) = slos_table::get(db, org, &cond.slo_id).await? else { anyhow::bail!("SLO {} not found", cond.slo_id); diff --git a/src/core/src/slo/job.rs b/src/core/src/slo/job.rs index 1a1690cfcd..8ccce888f2 100644 --- a/src/core/src/slo/job.rs +++ b/src/core/src/slo/job.rs @@ -41,7 +41,10 @@ use config::{ }, utils::{json, time::now_micros}, }; -use infra::table::{slo as slo_table, slos as slos_table}; +use infra::{ + db::{get_orm_client_ro, get_orm_client_rw}, + table::{slo as slo_table, slos as slos_table}, +}; use super::{ ingest::{ @@ -70,9 +73,7 @@ pub enum PassOutcome { /// Run one SLI ingest pass for `slo`. pub async fn run_pass(slo: &Slo, now_secs: i64) -> Result { let cfg = get_config(); - let db = infra::db::ORM_CLIENT - .get() - .ok_or_else(|| anyhow::anyhow!("database not initialized"))?; + let db = get_orm_client_ro().await; let status = slo_table::load_status(db, &slo.id, "").await?; // A status row whose generation has already moved on means this pass was @@ -896,9 +897,7 @@ pub async fn run_range( writer: config::meta::slo::slice::Writer, ) -> Result { let cfg = get_config(); - let db = infra::db::ORM_CLIENT - .get() - .ok_or_else(|| anyhow::anyhow!("database not initialized"))?; + let db = get_orm_client_rw().await; let group_by = slo.definition.group_by.clone().unwrap_or_default(); let params = PassParams { diff --git a/src/core/src/slo/reconcile.rs b/src/core/src/slo/reconcile.rs index abe7277046..9d54e6d600 100644 --- a/src/core/src/slo/reconcile.rs +++ b/src/core/src/slo/reconcile.rs @@ -36,7 +36,7 @@ use config::{ }, utils::json, }; -use infra::table::slo as slo_table; +use infra::{db::get_orm_client_rw, table::slo as slo_table}; /// The per-group figures a rebuild produced. #[derive(Debug, Clone, PartialEq)] @@ -49,9 +49,7 @@ pub struct Rebuilt { /// Rebuild every group's aggregate for one SLO and write the results. pub async fn reconcile(slo: &Slo) -> Result, anyhow::Error> { - let db = infra::db::ORM_CLIENT - .get() - .ok_or_else(|| anyhow::anyhow!("database not initialized"))?; + let db = get_orm_client_rw().await; let Some(status) = slo_table::load_status(db, &slo.id, "").await? else { return Ok(Vec::new()); diff --git a/src/core/src/slo/service.rs b/src/core/src/slo/service.rs index ca8966e8d9..2fa4a53c5c 100644 --- a/src/core/src/slo/service.rs +++ b/src/core/src/slo/service.rs @@ -44,9 +44,12 @@ use config::{ }, utils::time::now_micros, }; -use infra::table::{ - folders, slo_backfill_jobs as backfill_jobs, slo_budget, slos as slos_table, - slos::GenerationEffect, +use infra::{ + db::{get_orm_client_ro, get_orm_client_rw}, + table::{ + folders, slo_backfill_jobs as backfill_jobs, slo_budget, slos as slos_table, + slos::GenerationEffect, + }, }; use serde::Serialize; use svix_ksuid::Ksuid; @@ -355,9 +358,7 @@ pub fn source_alert_condition_changed(before: &Alert, after: &Alert) -> bool { /// SLO, and refusing a save because the ORM was not warm yet would be a new way /// for alerts to fail. pub async fn slos_sourced_from_alert(org: &str, alert_id: &str) -> Result, SloError> { - let db = infra::db::ORM_CLIENT - .get_or_init(infra::db::connect_to_orm) - .await; + let db = infra::db::get_orm_client_ro().await; slos_table::list_by_source_alert(db, org, alert_id) .await .map_err(Into::into) @@ -374,9 +375,7 @@ pub async fn redefine_for_source_alert(org: &str, slo_ids: &[String]) { // runs on the alert write path, and silently skipping the bump would leave // one window mixing two definitions — the one corruption D59 exists to // prevent. - let db = infra::db::ORM_CLIENT - .get_or_init(infra::db::connect_to_orm) - .await; + let db = infra::db::get_orm_client_rw().await; let now = now_micros() / 1_000_000; for id in slo_ids { let mut slo = match slos_table::get(db, org, id).await { @@ -469,9 +468,7 @@ pub async fn list_slo_eligible_alerts( org: &str, user_id: Option<&str>, ) -> Result, SloError> { - let db = infra::db::ORM_CLIENT - .get() - .ok_or_else(|| SloError::Db("database not initialized".into()))?; + let db = get_orm_client_ro().await; let params = config::meta::alerts::alert::ListAlertsParams::new(org); let alerts = crate::alerts::alert::list_v2(db, user_id, params) .await @@ -714,9 +711,7 @@ pub fn reservation(slo: &Slo) -> (i64, i64) { pub async fn create(slo: &mut Slo) -> Result<(), SloError> { let cfg = get_config(); - let db = infra::db::ORM_CLIENT - .get() - .ok_or_else(|| SloError::Db("database not initialized".into()))?; + let db = get_orm_client_rw().await; validate(slo).await?; // Before the budget charge: a bad folder should cost nothing, and the @@ -757,9 +752,7 @@ pub async fn create(slo: &mut Slo) -> Result<(), SloError> { } pub async fn update(slo: &mut Slo) -> Result<(), SloError> { - let db = infra::db::ORM_CLIENT - .get() - .ok_or_else(|| SloError::Db("database not initialized".into()))?; + let db = get_orm_client_rw().await; validate(slo).await?; // An update carries `folder_id` too, so it is also a move. @@ -845,9 +838,7 @@ pub async fn move_to_folder( dst_folder_id: &str, editor: Option<&str>, ) -> Result { - let db = infra::db::ORM_CLIENT - .get() - .ok_or_else(|| SloError::Db("database not initialized".into()))?; + let db = get_orm_client_rw().await; ensure_folder(org, dst_folder_id).await?; @@ -863,9 +854,7 @@ pub async fn move_to_folder( } pub async fn delete(org: &str, id: &str) -> Result { - let db = infra::db::ORM_CLIENT - .get() - .ok_or_else(|| SloError::Db("database not initialized".into()))?; + let db = get_orm_client_rw().await; let Some(slo) = slos_table::get(db, org, id).await? else { return Ok(false); @@ -934,9 +923,7 @@ fn plan_alert_cascade(dependents: &[(String, String)]) -> Vec<(Ksuid, String)> { } pub async fn set_enabled(org: &str, id: &str, enabled: bool) -> Result { - let db = infra::db::ORM_CLIENT - .get() - .ok_or_else(|| SloError::Db("database not initialized".into()))?; + let db = get_orm_client_rw().await; let now = now_micros() / 1_000_000; let changed = slos_table::set_enabled(db, org, id, enabled, now).await?; if changed && let Some(slo) = slos_table::get(db, org, id).await? { @@ -954,9 +941,7 @@ pub async fn get_with_status( org: &str, id: &str, ) -> Result)>, anyhow::Error> { - let db = infra::db::ORM_CLIENT - .get() - .ok_or_else(|| anyhow::anyhow!("database not initialized"))?; + let db = get_orm_client_ro().await; let Some(slo) = slos_table::get(db, org, id).await? else { return Ok(None); }; @@ -969,9 +954,7 @@ pub async fn list_with_status( org: &str, folder: Option<&str>, ) -> Result)>, anyhow::Error> { - let db = infra::db::ORM_CLIENT - .get() - .ok_or_else(|| anyhow::anyhow!("database not initialized"))?; + let db = get_orm_client_ro().await; let mut out = Vec::new(); for slo in slos_table::list(db, org, folder).await? { let status = rollup_view(db, &slo).await?; @@ -989,9 +972,7 @@ pub async fn group_status( org: &str, id: &str, ) -> Result, anyhow::Error> { - let db = infra::db::ORM_CLIENT - .get() - .ok_or_else(|| anyhow::anyhow!("database not initialized"))?; + let db = get_orm_client_ro().await; let Some(slo) = slos_table::get(db, org, id).await? else { return Ok(Vec::new()); }; @@ -1105,10 +1086,7 @@ fn view_of( async fn schedule(slo: &Slo, now: i64) { sync_ingest_trigger(slo).await; - let db = match infra::db::ORM_CLIENT.get() { - Some(db) => db, - None => return, - }; + let db = get_orm_client_rw().await; // A new generation has no history, so its window is meaningless until // backfill fills it. The range ends where the incremental writer begins, // which is what keeps the two writers off each other's slices (§6b.9). diff --git a/src/core/src/stream.rs b/src/core/src/stream.rs index 36ea22e8d8..a5d3f0d2b0 100644 --- a/src/core/src/stream.rs +++ b/src/core/src/stream.rs @@ -63,9 +63,7 @@ async fn cleanup_related_resources( .iter() .filter_map(|alert| alert.id) .collect::>(); - let client = infra::db::ORM_CLIENT - .get_or_init(infra::db::connect_to_orm) - .await; + let client = infra::db::get_orm_client_rw().await; crate::alerts::alert::delete_many_for_cascade(client, &org_id, &alert_ids) .await .map_err(alert_cleanup_error)?; diff --git a/src/core/src/synthetics.rs b/src/core/src/synthetics.rs index 2034562bdb..08c89ddad4 100644 --- a/src/core/src/synthetics.rs +++ b/src/core/src/synthetics.rs @@ -486,9 +486,7 @@ pub async fn location_staleness_watcher() { continue; } - let conn = infra::db::ORM_CLIENT - .get_or_init(infra::db::connect_to_orm) - .await; + let conn = infra::db::get_orm_client_ro().await; let checks = infra::table::synthetics_checks::list_referencing_location(conn, &org_id, &loc.id) .await diff --git a/src/core/src/users.rs b/src/core/src/users.rs index 04cfb39b12..d32a4c466d 100644 --- a/src/core/src/users.rs +++ b/src/core/src/users.rs @@ -1338,7 +1338,7 @@ mod tests { use common::{infra::config::USERS, meta::user::get_default_user_role}; use config::meta::user::{UserRole, UserType}; use infra::{ - db::{self as infra_db, ORM_CLIENT, connect_to_orm}, + db::{self as infra_db, get_orm_client_rw}, table as infra_table, }; use tokio::sync::Mutex; @@ -1408,14 +1408,16 @@ mod tests { assert_eq!(org.name, "org2"); } - async fn set_up() { - // Acquire lock to serialize database setup across concurrent tests + /// Returned guard must be held for the whole test: `set_up` truncates the + /// shared user tables, so a concurrent test's setup would clear rows out + /// from under a test body that only serialized the setup itself. + async fn set_up() -> tokio::sync::MutexGuard<'static, ()> { let lock = TEST_SETUP_LOCK .get_or_init(|| async { Mutex::new(()) }) .await; - let _guard = lock.lock().await; + let guard = lock.lock().await; - let _ = ORM_CLIENT.get_or_init(connect_to_orm).await; + let _ = get_orm_client_rw().await; // clear the table here as previous tests could have written to it let _ = infra::table::org_users::clear().await; let _ = infra::table::users::clear().await; @@ -1474,29 +1476,31 @@ mod tests { password_ext: Some("root_password_ext_hash".to_string()), }, ); + + guard } #[tokio::test] async fn test_list_users() { - set_up().await; + let _guard = set_up().await; assert!(list_users("", "dummy", None, None, false).await.is_ok()) } #[tokio::test] async fn test_root_user_exists() { - set_up().await; + let _guard = set_up().await; assert!(!root_user_exists().await); } #[tokio::test] async fn test_get_user() { - set_up().await; + let _guard = set_up().await; assert!(get_user(Some("dummy"), "admin@zo.dev").await.is_some()) } #[tokio::test] async fn test_post_user() { - set_up().await; + let _guard = set_up().await; let resp = post_user( "dummy", @@ -1520,7 +1524,7 @@ mod tests { #[tokio::test] async fn test_user() { - set_up().await; + let _guard = set_up().await; let resp = update_user( "dummy", @@ -1586,7 +1590,7 @@ mod tests { #[tokio::test] async fn test_create_new_user() { - set_up().await; + let _guard = set_up().await; let db_user = DBUser { email: "newuser@example.com".to_string(), @@ -1635,7 +1639,7 @@ mod tests { #[tokio::test] async fn test_delete_user() { - set_up().await; + let _guard = set_up().await; let resp = delete_user("nonexistent@example.com").await; assert!(resp.is_ok()); @@ -1646,7 +1650,7 @@ mod tests { #[tokio::test] async fn test_root_user_exists_edge_cases() { - set_up().await; + let _guard = set_up().await; let exists = root_user_exists().await; assert!(!exists); @@ -1674,7 +1678,7 @@ mod tests { #[tokio::test] async fn test_create_root_user_if_not_exists() { - set_up().await; + let _guard = set_up().await; let user_req = UserRequest { email: "root@example.com".to_string(), @@ -1718,7 +1722,7 @@ mod tests { #[tokio::test] async fn test_post_user_validation() { - set_up().await; + let _guard = set_up().await; let invalid_email_req = UserRequest { email: "invalid-email".to_string(), @@ -1777,7 +1781,7 @@ mod tests { #[tokio::test] async fn test_update_user_validation() { - set_up().await; + let _guard = set_up().await; let resp = update_user( "dummy", @@ -1822,7 +1826,7 @@ mod tests { #[tokio::test] async fn test_add_user_to_org_validation() { - set_up().await; + let _guard = set_up().await; let resp = add_user_to_org( "dummy", diff --git a/src/db/src/alerts/alert.rs b/src/db/src/alerts/alert.rs index 3e6e6a91a1..6c9716f6ee 100644 --- a/src/db/src/alerts/alert.rs +++ b/src/db/src/alerts/alert.rs @@ -30,7 +30,7 @@ use config::{ #[cfg(feature = "enterprise")] use infra::table::workflows::WorkflowTriggerEntity; use infra::{ - db::{ORM_CLIENT, connect_to_orm}, + db::{get_orm_client_ro, get_orm_client_rw}, table::alerts as table, }; use sea_orm::{ConnectionTrait, TransactionTrait}; @@ -63,7 +63,7 @@ pub async fn get_by_name( stream_name: &str, name: &str, ) -> Result, infra::errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let alert = table::get_by_name(client, org_id, "default", stream_type, stream_name, name).await?; let value = alert.map(|(_f, a)| a); @@ -71,7 +71,7 @@ pub async fn get_by_name( } pub async fn set(org_id: &str, alert: Alert, create: bool) -> Result { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; match table::put(client, org_id, "default", alert).await { Ok(alert) => { infra::coordinator::alerts::emit_put_event(org_id, &alert, None).await?; @@ -154,7 +154,7 @@ pub async fn set_without_updating_trigger(org_id: &str, alert: Alert) -> Result< let Some(alert_id) = alert.id else { return Err(anyhow::anyhow!("Alert ID is required")); }; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let Some((_f, _)) = table::get_by_id(client, org_id, alert_id).await? else { return Err(anyhow::anyhow!("Alert not found")); }; @@ -329,7 +329,7 @@ pub async fn delete_by_name( stream_name: &str, name: &str, ) -> Result<(), infra::errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let Some(alert_id) = table::get_by_name(client, org_id, "default", stream_type, stream_name, name) @@ -373,7 +373,7 @@ pub async fn list( params }; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let alerts = table::list(client, params) .await? .into_iter() @@ -394,7 +394,7 @@ pub async fn watch() -> Result<(), anyhow::Error> { } pub async fn cache() -> Result<(), anyhow::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let mut alerts: Vec<(Folder, Alert)> = Vec::new(); let r = infra::schema::STREAM_SCHEMAS.read().await; let mut orgs = HashSet::new(); @@ -439,7 +439,7 @@ async fn put_into_cache( alert_id: String, _folder_id: Option, ) -> Result<(), anyhow::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let Ok(alert_id_ksuid) = svix_ksuid::Ksuid::from_str(&alert_id) else { log::error!("Error parsing alert id into Ksuid while putting the alert into cache"); return Ok(()); diff --git a/src/db/src/folders.rs b/src/db/src/folders.rs index 80d26ddb34..2df27184a5 100644 --- a/src/db/src/folders.rs +++ b/src/db/src/folders.rs @@ -23,10 +23,7 @@ use config::{ synthetics::ListSyntheticsParams, }, }; -use infra::{ - db::{ORM_CLIENT, connect_to_orm}, - table, -}; +use infra::{db::get_orm_client_rw, table}; #[cfg(feature = "enterprise")] use o2_openfga::meta::mapping::OFGA_MODELS; @@ -249,7 +246,7 @@ pub async fn delete_folder( folder_id: &str, folder_type: FolderType, ) -> Result<(), FolderError> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; match folder_type { FolderType::Dashboards => { let params = ListDashboardsParams::new(org_id).with_folder_id(folder_id); diff --git a/src/db/src/pipeline_errors.rs b/src/db/src/pipeline_errors.rs index c05b82c400..71bb49849d 100644 --- a/src/db/src/pipeline_errors.rs +++ b/src/db/src/pipeline_errors.rs @@ -15,7 +15,7 @@ use config::{meta::self_reporting::error::PipelineError, utils::time::now_micros}; use infra::{ - db::{ORM_CLIENT, connect_to_orm}, + db::{get_orm_client_ro, get_orm_client_rw}, table::entity::{pipeline_last_errors, prelude::PipelineLastErrors}, }; use sea_orm::{ @@ -33,7 +33,7 @@ pub async fn upsert( timestamp: i64, error_data: &PipelineError, ) -> Result<(), infra::errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; // Serialize node_errors to JSON before taking the write lock let node_errors_json = if !error_data.node_errors.is_empty() { @@ -42,10 +42,6 @@ pub async fn upsert( None }; - // make sure only one client is writing to the database(only for sqlite); - // taken before the existence check so the read-modify-write is serialized - let _lock = infra::table::get_lock().await; - // Check if record exists let existing = PipelineLastErrors::find_by_id(pipeline_id) .one(client) @@ -151,9 +147,7 @@ pub async fn batch_upsert( )); } - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; - // make sure only one client is writing to the database(only for sqlite) - let _lock = infra::table::get_lock().await; + let client = get_orm_client_rw().await; let txn = client.begin().await?; // Collect all pipeline_ids to check which exist @@ -218,7 +212,7 @@ pub async fn batch_upsert( pub async fn list_by_org( org_id: &str, ) -> Result, infra::errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let errors = PipelineLastErrors::find() .filter(pipeline_last_errors::Column::OrgId.eq(org_id)) @@ -233,7 +227,7 @@ pub async fn list_by_org( pub async fn get_by_pipeline_id( pipeline_id: &str, ) -> Result, infra::errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let error = PipelineLastErrors::find_by_id(pipeline_id) .one(client) @@ -244,9 +238,7 @@ pub async fn get_by_pipeline_id( /// Deletes a pipeline error record by pipeline_id. pub async fn delete(pipeline_id: &str) -> Result<(), infra::errors::Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = infra::table::get_lock().await; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; PipelineLastErrors::delete_many() .filter(pipeline_last_errors::Column::PipelineId.eq(pipeline_id)) @@ -258,9 +250,7 @@ pub async fn delete(pipeline_id: &str) -> Result<(), infra::errors::Error> { /// Deletes all pipeline error records for an organization. pub async fn delete_by_org(org_id: &str) -> Result<(), infra::errors::Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = infra::table::get_lock().await; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; PipelineLastErrors::delete_many() .filter(pipeline_last_errors::Column::OrgId.eq(org_id)) @@ -274,8 +264,7 @@ pub async fn delete_by_org(org_id: &str) -> Result<(), infra::errors::Error> { /// /// This is used for periodic cleanup of stale errors. pub async fn delete_older_than(cutoff_timestamp: i64) -> Result { - let _lock = infra::table::get_lock().await; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let result = PipelineLastErrors::delete_many() .filter(pipeline_last_errors::Column::LastErrorTimestamp.lt(cutoff_timestamp)) diff --git a/src/infra/src/db/mod.rs b/src/infra/src/db/mod.rs index 0b3b75923e..21524b0583 100644 --- a/src/infra/src/db/mod.rs +++ b/src/infra/src/db/mod.rs @@ -40,25 +40,44 @@ static SUPER_CLUSTER: OnceCell> = OnceCell::const_new(); pub const SQLITE_STORE: &str = "sqlite"; -pub static ORM_CLIENT: OnceCell = OnceCell::const_new(); -pub static ORM_CLIENT_DDL: OnceCell = OnceCell::const_new(); +static ORM_CLIENT_RO: OnceCell = OnceCell::const_new(); +pub(crate) static ORM_CLIENT_RW: OnceCell = OnceCell::const_new(); +static ORM_CLIENT_DDL: OnceCell = OnceCell::const_new(); -pub async fn connect_to_orm() -> DatabaseConnection { +/// ORM client for reads. +pub async fn get_orm_client_ro() -> &'static DatabaseConnection { + ORM_CLIENT_RO.get_or_init(connect_to_orm_ro).await +} + +/// ORM client for writes. On SQLite they serialize on a single connection. +pub async fn get_orm_client_rw() -> &'static DatabaseConnection { + ORM_CLIENT_RW.get_or_init(connect_to_orm_rw).await +} + +/// ORM client for schema migrations. +pub async fn get_orm_client_ddl() -> &'static DatabaseConnection { + ORM_CLIENT_DDL.get_or_init(connect_to_orm_ddl).await +} + +async fn connect_to_orm_ro() -> DatabaseConnection { match get_config().common.meta_store.as_str().into() { MetaStore::PostgreSQL => { - let pool = postgres::CLIENT.clone(); - SqlxPostgresConnector::from_sqlx_postgres_pool(pool) - } - _ => { - // clone the pool handle directly, NOT through the CLIENT_RW mutex: - // callers may lazily initialize ORM_CLIENT while already holding - // `table::get_lock()`, and locking here again would self-deadlock - SqlxSqliteConnector::from_sqlx_sqlite_pool(sqlite::CLIENT_RW_POOL.clone()) + SqlxPostgresConnector::from_sqlx_postgres_pool(postgres::CLIENT_RO.clone()) } + _ => SqlxSqliteConnector::from_sqlx_sqlite_pool(sqlite::CLIENT_RO.clone()), } } -pub async fn connect_to_orm_ddl() -> DatabaseConnection { +async fn connect_to_orm_rw() -> DatabaseConnection { + match get_config().common.meta_store.as_str().into() { + MetaStore::PostgreSQL => { + SqlxPostgresConnector::from_sqlx_postgres_pool(postgres::CLIENT_RW.clone()) + } + _ => SqlxSqliteConnector::from_sqlx_sqlite_pool(sqlite::CLIENT_RW.clone()), + } +} + +async fn connect_to_orm_ddl() -> DatabaseConnection { match get_config().common.meta_store.as_str().into() { MetaStore::PostgreSQL => { let pool = postgres::CLIENT_DDL.clone(); @@ -66,7 +85,7 @@ pub async fn connect_to_orm_ddl() -> DatabaseConnection { } _ => { // for sqlite, there is no separate ddl client, use the common one - SqlxSqliteConnector::from_sqlx_sqlite_pool(sqlite::CLIENT_RW_POOL.clone()) + SqlxSqliteConnector::from_sqlx_sqlite_pool(sqlite::CLIENT_RW.clone()) } } } @@ -455,14 +474,8 @@ mod tests { } #[tokio::test] - async fn connect_to_orm_never_takes_the_sqlite_write_mutex() { - // Regression: a task may lazily initialize ORM_CLIENT while already - // holding table::get_lock() (the CLIENT_RW mutex). connect_to_orm must - // clone the pool handle without locking CLIENT_RW itself, or the first - // cold-start write from any entry point self-deadlocks forever. - let _guard = sqlite::CLIENT_RW.lock().await; - tokio::time::timeout(std::time::Duration::from_secs(5), connect_to_orm()) - .await - .expect("connect_to_orm deadlocked on the held CLIENT_RW mutex"); + async fn write_pool_is_single_connection_on_sqlite() { + // a second write connection reintroduces SQLITE_BUSY_SNAPSHOT (517) + assert_eq!(sqlite::CLIENT_RW.options().get_max_connections(), 1); } } diff --git a/src/infra/src/db/postgres.rs b/src/infra/src/db/postgres.rs index 6e8c9ed611..6d0bf042fe 100644 --- a/src/infra/src/db/postgres.rs +++ b/src/infra/src/db/postgres.rs @@ -36,9 +36,9 @@ use tokio::sync::{OnceCell, mpsc}; use super::{DBIndex, IndexStatement}; use crate::errors::*; -pub static CLIENT: Lazy> = Lazy::new(|| connect(false, false)); -pub static CLIENT_RO: Lazy> = Lazy::new(|| connect(true, false)); -pub static CLIENT_DDL: Lazy> = Lazy::new(|| connect(false, true)); +pub(crate) static CLIENT_RO: Lazy> = Lazy::new(|| connect(true, false)); +pub(crate) static CLIENT_RW: Lazy> = Lazy::new(|| connect(false, false)); +pub(crate) static CLIENT_DDL: Lazy> = Lazy::new(|| connect(false, true)); static INDICES: OnceCell> = OnceCell::const_new(); /// The three pools come from separately configured DSNs whose roles can have @@ -47,7 +47,7 @@ static INDICES: OnceCell> = OnceCell::const_new(); pub async fn verify_schema_consistency() -> Result<()> { let mut schemas = Vec::new(); for (label, pool) in [ - ("read-write", &*CLIENT), + ("read-write", &*CLIENT_RW), ("read-only", &*CLIENT_RO), ("ddl", &*CLIENT_DDL), ] { @@ -99,7 +99,7 @@ fn connect(readonly: bool, ddl: bool) -> Pool { } async fn cache_indices() -> HashSet { - let client = CLIENT.clone(); + let client = CLIENT_RW.clone(); DB_QUERY_NUMS .with_label_values(&["select", "pg_indexes"]) .inc(); @@ -188,7 +188,7 @@ impl super::Db for PostgresDb { start_dt: Option, ) -> Result<()> { let (module, key1, key2) = super::parse_key(key); - let pool = CLIENT.clone(); + let pool = CLIENT_RW.clone(); let local_start_dt = start_dt.unwrap_or_default(); let mut tx = pool.begin().await?; DB_QUERY_NUMS.with_label_values(&["insert", "meta"]).inc(); @@ -260,7 +260,7 @@ impl super::Db for PostgresDb { update_fn: Box, ) -> Result<()> { let (module, key1, key2) = super::parse_key(key); - let pool = CLIENT.clone(); + let pool = CLIENT_RW.clone(); let mut tx = pool.begin().await?; let lock_key = format!("get_for_update_{key}"); let lock_id = config::utils::hash::gxhash::new().sum64(&lock_key); @@ -451,7 +451,7 @@ impl super::Db for PostgresDb { } let (module, key1, key2) = super::parse_key(key); - let pool = CLIENT.clone(); + let pool = CLIENT_RW.clone(); DB_QUERY_NUMS.with_label_values(&["delete", "meta"]).inc(); if with_prefix { diff --git a/src/infra/src/db/sqlite.rs b/src/infra/src/db/sqlite.rs index e42cdff038..e1dbf60a6f 100644 --- a/src/infra/src/db/sqlite.rs +++ b/src/infra/src/db/sqlite.rs @@ -31,7 +31,7 @@ use sqlx::{ SqliteSynchronous, }, }; -use tokio::sync::{Mutex, OnceCell, RwLock, mpsc}; +use tokio::sync::{OnceCell, RwLock, mpsc}; use super::{DBIndex, IndexStatement}; use crate::{ @@ -39,12 +39,9 @@ use crate::{ errors::*, }; -pub static CLIENT_RO: Lazy> = Lazy::new(connect_ro); -/// Same pool as [`CLIENT_RW`], cloneable without taking the mutex — so the ORM -/// can be lazily initialized while `table::get_lock()` is held. -pub static CLIENT_RW_POOL: Lazy> = Lazy::new(connect_rw); -pub static CLIENT_RW: Lazy>>> = - Lazy::new(|| Arc::new(Mutex::new(CLIENT_RW_POOL.clone()))); +pub(crate) static CLIENT_RO: Lazy> = Lazy::new(connect_ro); +/// Holds a single connection, so every writer serializes on it. +pub(crate) static CLIENT_RW: Lazy> = Lazy::new(connect_rw); static INDICES: OnceCell> = OnceCell::const_new(); pub static CHANNEL: Lazy = Lazy::new(SqliteDbChannel::new); @@ -75,9 +72,12 @@ fn connect_rw() -> Pool { .busy_timeout(Duration::from_secs(acquire_timeout)) .create_if_missing(true); + // SQLite has one writer anyway. A second write connection only lets a + // deferred transaction upgrade onto a moved snapshot, which fails as + // SQLITE_BUSY_SNAPSHOT (517) without ever consulting busy_timeout. SqlitePoolOptions::new() - .min_connections(cfg.limit.sql_db_connections_min) - .max_connections(cfg.limit.sql_db_connections_max) + .min_connections(1) + .max_connections(1) .acquire_timeout(Duration::from_secs(acquire_timeout)) .idle_timeout(Some(Duration::from_secs(idle_timeout))) .max_lifetime(Some(Duration::from_secs(max_lifetime))) @@ -264,7 +264,6 @@ impl super::Db for SqliteDb { let (module, key1, key2) = super::parse_key(key); let local_start_dt = start_dt.unwrap_or_default(); let client = CLIENT_RW.clone(); - let client = client.lock().await; let mut tx = client.begin().await?; if let Err(e) = sqlx::query( r#"INSERT OR IGNORE INTO meta (module, key1, key2, start_dt, value) VALUES ($1, $2, $3, $4, '');"# @@ -339,7 +338,6 @@ impl super::Db for SqliteDb { ) -> Result<()> { let (module, key1, key2) = super::parse_key(key); let client = CLIENT_RW.clone(); - let client = client.lock().await; let mut tx = client.begin().await?; let mut need_watch_dt = 0; let row = if let Some(start_dt) = start_dt { @@ -529,7 +527,6 @@ impl super::Db for SqliteDb { let (module, key1, key2) = super::parse_key(key); let client = CLIENT_RW.clone(); - let client = client.lock().await; if with_prefix { if key1.is_empty() { @@ -538,12 +535,12 @@ impl super::Db for SqliteDb { sqlx::query("DELETE FROM meta WHERE module = $1 AND start_dt = $2") .bind(&module) .bind(dt) - .execute(&*client) + .execute(&client) .await?; } else { sqlx::query("DELETE FROM meta WHERE module = $1") .bind(&module) - .execute(&*client) + .execute(&client) .await?; } } else if key2.is_empty() { @@ -555,13 +552,13 @@ impl super::Db for SqliteDb { .bind(&module) .bind(&key1) .bind(dt) - .execute(&*client) + .execute(&client) .await?; } else { sqlx::query("DELETE FROM meta WHERE module = $1 AND key1 = $2") .bind(&module) .bind(&key1) - .execute(&*client) + .execute(&client) .await?; } } else { @@ -573,7 +570,7 @@ impl super::Db for SqliteDb { .bind(&key2) .bind(format!("{}/%", key2)) .bind(dt) - .execute(&*client) + .execute(&client) .await?; } else { sqlx::query("DELETE FROM meta WHERE module = $1 AND key1 = $2 AND (key2 = $3 OR key2 LIKE $4)") @@ -581,7 +578,7 @@ impl super::Db for SqliteDb { .bind(&key1) .bind(&key2) .bind(format!("{}/%", key2)) - .execute(&*client) + .execute(&client) .await?; } } @@ -593,14 +590,14 @@ impl super::Db for SqliteDb { .bind(&key1) .bind(&key2) .bind(dt) - .execute(&*client) + .execute(&client) .await?; } else { sqlx::query("DELETE FROM meta WHERE module = $1 AND key1 = $2 AND key2 = $3") .bind(&module) .bind(&key1) .bind(&key2) - .execute(&*client) + .execute(&client) .await?; } } @@ -781,7 +778,6 @@ impl super::Db for SqliteDb { async fn create_table() -> Result<()> { let client = CLIENT_RW.clone(); - let client = client.lock().await; // create table sqlx::query( r#" @@ -796,7 +792,7 @@ CREATE TABLE IF NOT EXISTS meta ); "#, ) - .execute(&*client) + .execute(&client) .await?; drop(client); @@ -831,7 +827,6 @@ CREATE TABLE IF NOT EXISTS meta async fn add_start_dt_column() -> Result<()> { let client = CLIENT_RW.clone(); - let client = client.lock().await; add_column(&client, "meta", "start_dt", "INTEGER NOT NULL DEFAULT 0").await?; drop(client); @@ -850,7 +845,6 @@ async fn add_start_dt_column() -> Result<()> { async fn create_meta_backup() -> Result<()> { let client = CLIENT_RW.clone(); - let client = client.lock().await; let mut tx = client.begin().await?; if let Err(e) = sqlx::query(r#"CREATE TABLE IF NOT EXISTS meta_backup_20240330 AS SELECT * FROM meta;"#) @@ -871,7 +865,6 @@ async fn create_meta_backup() -> Result<()> { pub async fn create_index(index: IndexStatement<'_>) -> Result<()> { let client = CLIENT_RW.clone(); - let client = client.lock().await; let indices = INDICES.get_or_init(cache_indices).await; if indices.contains(&DBIndex { name: index.idx_name.into(), @@ -892,14 +885,13 @@ pub async fn create_index(index: IndexStatement<'_>) -> Result<()> { index.table, index.fields.join(",") ); - sqlx::query(&sql).execute(&*client).await?; + sqlx::query(&sql).execute(&client).await?; log::info!("[SQLITE] index {} created successfully", index.idx_name); Ok(()) } pub async fn delete_index(idx_name: &str, table: &str) -> Result<()> { let client = CLIENT_RW.clone(); - let client = client.lock().await; let indices = INDICES.get_or_init(cache_indices).await; if !indices.contains(&DBIndex { name: idx_name.into(), @@ -909,7 +901,7 @@ pub async fn delete_index(idx_name: &str, table: &str) -> Result<()> { } log::info!("[SQLITE] deleting index {idx_name} on table {table}"); let sql = format!("DROP INDEX IF EXISTS {idx_name};"); - sqlx::query(&sql).execute(&*client).await?; + sqlx::query(&sql).execute(&client).await?; log::info!("[SQLITE] index {idx_name} deleted successfully"); Ok(()) } diff --git a/src/infra/src/file_list/postgres.rs b/src/infra/src/file_list/postgres.rs index 4685124922..379d34de4c 100644 --- a/src/infra/src/file_list/postgres.rs +++ b/src/infra/src/file_list/postgres.rs @@ -36,7 +36,7 @@ use sqlx::{Executor, PgConnection, Postgres, QueryBuilder, Row}; use crate::{ db::{ IndexStatement, - postgres::{CLIENT, CLIENT_DDL, CLIENT_RO, add_column, create_index, delete_index}, + postgres::{CLIENT_DDL, CLIENT_RO, CLIENT_RW, add_column, create_index, delete_index}, }, errors::{Error, Result}, file_list::FileRecord, @@ -69,7 +69,7 @@ impl Default for PostgresFileList { #[async_trait] impl super::FileList for PostgresFileList { async fn health_check(&self) -> Result<()> { - let pool = CLIENT.clone(); + let pool = CLIENT_RW.clone(); let is_writable: bool = sqlx::query_scalar( r#"SELECT NOT pg_is_in_recovery() AND current_setting('transaction_read_only')::bool = false @@ -105,7 +105,7 @@ impl super::FileList for PostgresFileList { } async fn remove(&self, file: &str) -> Result<()> { - let pool = CLIENT.clone(); + let pool = CLIENT_RW.clone(); let (stream_key, date_key, file_name) = parse_file_key_columns(file).map_err(|e| Error::Message(e.to_string()))?; @@ -138,7 +138,7 @@ impl super::FileList for PostgresFileList { } async fn update_dump_records(&self, dump_file: &FileKey, dumped_ids: &[i64]) -> Result<()> { - let pool = CLIENT.clone(); + let pool = CLIENT_RW.clone(); // insert the dump file into file_list table let (stream_key, date_key, file_name) = @@ -230,7 +230,7 @@ impl super::FileList for PostgresFileList { if files.is_empty() { return Ok(()); } - let pool = CLIENT.clone(); + let pool = CLIENT_RW.clone(); let chunks = files.chunks(100); for files in chunks { // we don't care the id here, because the id is from file_list table not for this table @@ -274,7 +274,7 @@ impl super::FileList for PostgresFileList { let chunks = files.chunks(100); for files in chunks { // get ids of the files - let pool = CLIENT.clone(); + let pool = CLIENT_RW.clone(); let mut ids = Vec::with_capacity(files.len()); for file in files { if file.id > 0 { @@ -380,7 +380,7 @@ SELECT min_ts, max_ts, records, original_size, compressed_size, index_size, bloo } async fn update_flattened(&self, file: &str, flattened: bool) -> Result<()> { - let pool = CLIENT.clone(); + let pool = CLIENT_RW.clone(); let (stream_key, date_key, file_name) = parse_file_key_columns(file).map_err(|e| Error::Message(e.to_string()))?; DB_QUERY_NUMS @@ -399,7 +399,7 @@ SELECT min_ts, max_ts, records, original_size, compressed_size, index_size, bloo } async fn update_compressed_size(&self, file: &str, size: i64) -> Result<()> { - let pool = CLIENT.clone(); + let pool = CLIENT_RW.clone(); let (stream_key, date_key, file_name) = parse_file_key_columns(file).map_err(|e| Error::Message(e.to_string()))?; DB_QUERY_NUMS @@ -421,7 +421,7 @@ SELECT min_ts, max_ts, records, original_size, compressed_size, index_size, bloo if ids.is_empty() { return Ok(()); } - let pool = CLIENT.clone(); + let pool = CLIENT_RW.clone(); // Use ANY($2) with a bigint array — postgres handles arbitrary length // without needing chunking, and the partition pruning works on the // file_list_*_pkey index. @@ -876,7 +876,7 @@ SELECT date return Ok(Vec::new()); } - let pool = CLIENT.clone(); + let pool = CLIENT_RW.clone(); let mut tx = pool.begin().await?; let lock_key = "file_list_deleted:query_deleted"; let lock_id = config::utils::hash::gxhash::new().sum64(lock_key); @@ -1174,7 +1174,7 @@ GROUP BY stream; stream_name: &str, ) -> Result<()> { let stream_key = format!("{org_id}/{stream_type}/{stream_name}"); - let pool = CLIENT.clone(); + let pool = CLIENT_RW.clone(); DB_QUERY_NUMS .with_label_values(&["delete", "stream_stats"]) .inc(); @@ -1194,7 +1194,7 @@ GROUP BY stream; is_recent: bool, ) -> Result<()> { let stream_key = format!("{org_id}/{stream_type}/{stream_name}"); - let pool = CLIENT.clone(); + let pool = CLIENT_RW.clone(); let mut tx = pool.begin().await?; let start = std::time::Instant::now(); DB_QUERY_NUMS @@ -1249,7 +1249,7 @@ DO UPDATE SET } async fn reset_stream_stats(&self) -> Result<()> { - let pool = CLIENT.clone(); + let pool = CLIENT_RW.clone(); DB_QUERY_NUMS .with_label_values(&["update", "stream_stats"]) .inc(); @@ -1298,7 +1298,7 @@ DO UPDATE SET offset: i64, ) -> Result { let stream_key = format!("{org_id}/{stream_type}/{stream}"); - let pool = CLIENT.clone(); + let pool = CLIENT_RW.clone(); let mut tx = pool.begin().await?; DB_QUERY_NUMS .with_label_values(&["insert", "file_list_jobs"]) @@ -1381,7 +1381,7 @@ DO UPDATE SET return Ok(Vec::new()); } - let pool = CLIENT.clone(); + let pool = CLIENT_RW.clone(); let mut tx = pool.begin().await?; let lock_key = "file_list_jobs:get_pending_jobs"; let lock_id = config::utils::hash::gxhash::new().sum64(lock_key); @@ -1491,7 +1491,7 @@ DO UPDATE SET offsets: i64, stream: Option<&str>, ) -> Result { - let pool = CLIENT.clone(); + let pool = CLIENT_RW.clone(); let mut conditions: Vec = Vec::new(); if !ids.is_empty() { conditions.push(format!( @@ -1528,7 +1528,7 @@ DO UPDATE SET async fn set_job_done(&self, ids: &[i64]) -> Result<()> { let cfg = get_config(); - let pool = CLIENT.clone(); + let pool = CLIENT_RW.clone(); DB_QUERY_NUMS .with_label_values(&["update", "file_list_jobs"]) .inc(); @@ -1551,7 +1551,7 @@ DO UPDATE SET } async fn update_running_jobs(&self, ids: &[i64]) -> Result<()> { - let pool = CLIENT.clone(); + let pool = CLIENT_RW.clone(); DB_QUERY_NUMS .with_label_values(&["update", "file_list_jobs"]) .inc(); @@ -1570,7 +1570,7 @@ DO UPDATE SET } async fn check_running_jobs(&self, before_date: i64) -> Result<()> { - let pool = CLIENT.clone(); + let pool = CLIENT_RW.clone(); DB_QUERY_NUMS .with_label_values(&["update", "file_list_jobs"]) .inc(); @@ -1610,7 +1610,7 @@ DO UPDATE SET } async fn clean_done_jobs(&self, before_date: i64) -> Result<()> { - let pool = CLIENT.clone(); + let pool = CLIENT_RW.clone(); DB_QUERY_NUMS .with_label_values(&["delete", "file_list_jobs"]) .inc(); @@ -1629,7 +1629,7 @@ DO UPDATE SET } async fn get_pending_jobs_count(&self) -> Result>> { - let pool = CLIENT.clone(); + let pool = CLIENT_RW.clone(); DB_QUERY_NUMS .with_label_values(&["select", "file_list_jobs"]) @@ -1673,7 +1673,7 @@ DO UPDATE SET node: &str, limit: i64, ) -> Result> { - let pool = CLIENT.clone(); + let pool = CLIENT_RW.clone(); let mut tx = pool.begin().await?; let lock_key = "file_list_jobs:get_pending_dump_jobs"; let lock_id = config::utils::hash::gxhash::new().sum64(lock_key); @@ -1752,7 +1752,7 @@ DO UPDATE SET } async fn set_job_dumped_status(&self, ids: &[i64], dumped: bool) -> Result<()> { - let pool = CLIENT.clone(); + let pool = CLIENT_RW.clone(); DB_QUERY_NUMS .with_label_values(&["update", "file_list_jobs"]) .inc(); @@ -1771,7 +1771,7 @@ DO UPDATE SET let (stream_key, date_key, file_name) = parse_file_key_columns(file).expect("parse file key failed"); let org_id = stream_key[..stream_key.find('/').unwrap()].to_string(); - let pool = CLIENT.clone(); + let pool = CLIENT_RW.clone(); // Ensure partition exists for this date ensure_file_list_partition(&pool, "file_list_dump_stats", &date_key).await?; @@ -1814,7 +1814,7 @@ DO UPDATE SET async fn delete_dump_stats(&self, file: &str) -> Result<()> { let (stream_key, date_key, file_name) = parse_file_key_columns(file).expect("parse file key failed"); - let pool = CLIENT.clone(); + let pool = CLIENT_RW.clone(); DB_QUERY_NUMS .with_label_values(&["delete", "file_list_dump_stats"]) .inc(); @@ -1909,7 +1909,7 @@ WHERE org = $1 AND account = $2;"#; } async fn delete_by_org(&self, org_id: &str) -> Result<()> { - let pool = CLIENT.clone(); + let pool = CLIENT_RW.clone(); let created_at = now_micros(); let mut tx = pool.begin().await?; // Move remaining rows into file_list_deleted first so the file GC removes @@ -1949,7 +1949,7 @@ impl PostgresFileList { meta: &FileMeta, ) -> Result { let now_ts = now_micros(); - let pool = CLIENT.clone(); + let pool = CLIENT_RW.clone(); let (stream_key, date_key, file_name) = parse_file_key_columns(file).map_err(|e| Error::Message(e.to_string()))?; let org_id = stream_key[..stream_key.find('/').unwrap()].to_string(); @@ -2002,7 +2002,7 @@ INSERT INTO {table} (account, org, stream, date, file, deleted, min_ts, max_ts, return Ok(()); } - let pool = CLIENT.clone(); + let pool = CLIENT_RW.clone(); // Ensure partitions exist for all distinct date keys before batch INSERT let mut date_keys = HashSet::new(); @@ -3291,7 +3291,7 @@ pub async fn spawn_maintenance_task() -> std::result::Result<(), anyhow::Error> } async fn run_maintenance() -> Result<()> { - let pool = CLIENT.clone(); + let pool = CLIENT_RW.clone(); // Acquire the distributed lock; another node acquiring it first will run (or // has run) this cycle. wait_ttl=0 waits for the current holder to finish. @@ -3693,7 +3693,7 @@ mod tests { let _result = postgres_list.add("test_account", file_key, &meta).await; - // This test would pass with proper CLIENT mocking + // This test would pass with proper CLIENT_RW mocking // assert!(result.is_ok()); // assert!(result.unwrap() > 0); } diff --git a/src/infra/src/file_list/sqlite.rs b/src/infra/src/file_list/sqlite.rs index 2a91863274..9df2697222 100644 --- a/src/infra/src/file_list/sqlite.rs +++ b/src/infra/src/file_list/sqlite.rs @@ -77,7 +77,6 @@ impl super::FileList for SqliteFileList { async fn remove(&self, file: &str) -> Result<()> { let client = CLIENT_RW.clone(); - let client = client.lock().await; let pool = client.clone(); let (stream_key, date_key, file_name) = parse_file_key_columns(file).map_err(|e| Error::Message(e.to_string()))?; @@ -109,7 +108,6 @@ impl super::FileList for SqliteFileList { async fn update_dump_records(&self, file: &FileKey, dumped_ids: &[i64]) -> Result<()> { let client = CLIENT_RW.clone(); - let client = client.lock().await; let mut tx = client.begin().await?; // insert the dump file into file_list table @@ -183,7 +181,6 @@ impl super::FileList for SqliteFileList { for files in chunks { // we don't care the id here, because the id is from file_list table not for this table let client = CLIENT_RW.clone(); - let client = client.lock().await; let mut tx = client.begin().await?; let mut query_builder: QueryBuilder = QueryBuilder::new( "INSERT INTO file_list_deleted (account, org, stream, date, file, index_file, flattened, created_at)", @@ -222,7 +219,6 @@ impl super::FileList for SqliteFileList { for files in chunks { // get ids of the files let client = CLIENT_RW.clone(); - let client = client.lock().await; let pool = client.clone(); let mut ids = Vec::with_capacity(files.len()); for file in files { @@ -305,7 +301,6 @@ SELECT min_ts, max_ts, records, original_size, compressed_size, index_size, bloo async fn update_flattened(&self, file: &str, flattened: bool) -> Result<()> { let client = CLIENT_RW.clone(); - let client = client.lock().await; let (stream_key, date_key, file_name) = parse_file_key_columns(file).map_err(|e| Error::Message(e.to_string()))?; sqlx::query( @@ -315,14 +310,13 @@ SELECT min_ts, max_ts, records, original_size, compressed_size, index_size, bloo .bind(stream_key) .bind(date_key) .bind(file_name) - .execute(&*client) + .execute(&client) .await?; Ok(()) } async fn update_compressed_size(&self, file: &str, size: i64) -> Result<()> { let client = CLIENT_RW.clone(); - let client = client.lock().await; let (stream_key, date_key, file_name) = parse_file_key_columns(file).map_err(|e| Error::Message(e.to_string()))?; sqlx::query( @@ -332,7 +326,7 @@ SELECT min_ts, max_ts, records, original_size, compressed_size, index_size, bloo .bind(stream_key) .bind(date_key) .bind(file_name) - .execute(&*client) + .execute(&client) .await?; Ok(()) } @@ -342,7 +336,6 @@ SELECT min_ts, max_ts, records, original_size, compressed_size, index_size, bloo return Ok(()); } let client = CLIENT_RW.clone(); - let client = client.lock().await; // Chunk in case of very long lists; SQLite caps placeholders at 999 // and we'd rather not blow it up unexpectedly. for chunk in ids.chunks(900) { @@ -354,7 +347,7 @@ SELECT min_ts, max_ts, records, original_size, compressed_size, index_size, bloo for id in chunk { q = q.bind(*id); } - q.execute(&*client).await?; + q.execute(&client).await?; } Ok(()) } @@ -799,10 +792,9 @@ SELECT date async fn clean_by_min_update_at(&self, val: i64) -> Result<()> { let client = CLIENT_RW.clone(); - let client = client.lock().await; sqlx::query("DELETE FROM file_list WHERE updated_at < $1;") .bind(val) - .execute(&*client) + .execute(&client) .await?; Ok(()) } @@ -907,10 +899,9 @@ GROUP BY stream; ) -> Result<()> { let stream_key = format!("{org_id}/{stream_type}/{stream_name}"); let client = CLIENT_RW.clone(); - let client = client.lock().await; sqlx::query("DELETE FROM stream_stats WHERE stream = ?;") .bind(&stream_key) - .execute(&*client) + .execute(&client) .await?; Ok(()) } @@ -925,7 +916,6 @@ GROUP BY stream; ) -> Result<()> { let stream_key = format!("{org_id}/{stream_type}/{stream_name}"); let client = CLIENT_RW.clone(); - let client = client.lock().await; let mut tx = client.begin().await?; if let Err(e) = sqlx::query( r#" @@ -973,9 +963,8 @@ DO UPDATE SET async fn reset_stream_stats(&self) -> Result<()> { let client = CLIENT_RW.clone(); - let client = client.lock().await; sqlx::query(r#"UPDATE stream_stats SET file_num = 0, min_ts = 0, max_ts = 0, records = 0, original_size = 0, compressed_size = 0, index_size = 0;"#) - .execute(&*client) + .execute(&client) .await?; Ok(()) } @@ -1015,7 +1004,6 @@ DO UPDATE SET ) -> Result { let stream_key = format!("{org_id}/{stream_type}/{stream}"); let client = CLIENT_RW.clone(); - let client = client.lock().await; let mut tx = client.begin().await?; match sqlx::query( "INSERT INTO file_list_jobs (org, stream, offsets, status, node, started_at, updated_at) VALUES ($1, $2, $3, $4, '', 0, 0);", @@ -1085,7 +1073,6 @@ DO UPDATE SET fast_mode: bool, ) -> Result> { let client = CLIENT_RW.clone(); - let client = client.lock().await; let mut tx = client.begin().await?; // get pending jobs group by stream and order by num desc @@ -1171,7 +1158,6 @@ SELECT stream, max(id) as id, COUNT(*) AS num stream: Option<&str>, ) -> Result { let client = CLIENT_RW.clone(); - let client = client.lock().await; let mut conditions: Vec = Vec::new(); if !ids.is_empty() { conditions.push(format!( @@ -1198,7 +1184,7 @@ SELECT stream, max(id) as id, COUNT(*) AS num }; let ret = sqlx::query(&sql) .bind(super::FileListJobStatus::Pending) - .execute(&*client) + .execute(&client) .await?; Ok(ret.rows_affected()) } @@ -1206,7 +1192,6 @@ SELECT stream, max(id) as id, COUNT(*) AS num async fn set_job_done(&self, ids: &[i64]) -> Result<()> { let cfg = get_config(); let client = CLIENT_RW.clone(); - let client = client.lock().await; let sql = format!( "UPDATE file_list_jobs SET status = $1, updated_at = $2, dumped = $3, node = '' WHERE id IN ({});", ids.iter() @@ -1220,14 +1205,13 @@ SELECT stream, max(id) as id, COUNT(*) AS num .bind(super::FileListJobStatus::Done) .bind(config::utils::time::now_micros()) .bind(!cfg.compact.file_list_dump_enabled) - .execute(&*client) + .execute(&client) .await?; Ok(()) } async fn update_running_jobs(&self, ids: &[i64]) -> Result<()> { let client = CLIENT_RW.clone(); - let client = client.lock().await; let sql = format!( r#"UPDATE file_list_jobs SET updated_at = $1 WHERE id IN ({})"#, ids.iter() @@ -1237,14 +1221,13 @@ SELECT stream, max(id) as id, COUNT(*) AS num ); sqlx::query(&sql) .bind(config::utils::time::now_micros()) - .execute(&*client) + .execute(&client) .await?; Ok(()) } async fn check_running_jobs(&self, before_date: i64) -> Result<()> { let client = CLIENT_RW.clone(); - let client = client.lock().await; // reset running jobs status to pending let ret = sqlx::query( @@ -1253,7 +1236,7 @@ SELECT stream, max(id) as id, COUNT(*) AS num .bind(super::FileListJobStatus::Pending) .bind(super::FileListJobStatus::Running) .bind(before_date) - .execute(&*client) + .execute(&client) .await?; let rows_affected = ret.rows_affected(); if rows_affected > 0 { @@ -1269,7 +1252,7 @@ SELECT stream, max(id) as id, COUNT(*) AS num .bind(super::FileListJobStatus::Done) .bind(false) .bind(before_date) - .execute(&*client) + .execute(&client) .await?; let rows_affected = ret.rows_affected(); if rows_affected > 0 { @@ -1280,14 +1263,13 @@ SELECT stream, max(id) as id, COUNT(*) AS num async fn clean_done_jobs(&self, before_date: i64) -> Result<()> { let client = CLIENT_RW.clone(); - let client = client.lock().await; let ret = sqlx::query( r#"DELETE FROM file_list_jobs WHERE status = $1 AND dumped = $2 AND updated_at < $3;"#, ) .bind(super::FileListJobStatus::Done) .bind(true) .bind(before_date) - .execute(&*client) + .execute(&client) .await?; if ret.rows_affected() > 0 { log::warn!("[SQLITE] clean done jobs"); @@ -1335,7 +1317,6 @@ SELECT stream, max(id) as id, COUNT(*) AS num limit: i64, ) -> Result> { let client = CLIENT_RW.clone(); - let client = client.lock().await; let mut tx = client.begin().await?; // get pending dump jobs by updated_at asc let ret = match sqlx::query_as::<_, (i64, String, i64)>( @@ -1392,7 +1373,6 @@ SELECT stream, max(id) as id, COUNT(*) AS num async fn set_job_dumped_status(&self, ids: &[i64], dumped: bool) -> Result<()> { let client = CLIENT_RW.clone(); - let client = client.lock().await; let sql = format!( "UPDATE file_list_jobs SET dumped = $1, node = '' WHERE id IN ({});", ids.iter() @@ -1400,7 +1380,7 @@ SELECT stream, max(id) as id, COUNT(*) AS num .collect::>() .join(",") ); - sqlx::query(&sql).bind(dumped).execute(&*client).await?; + sqlx::query(&sql).bind(dumped).execute(&client).await?; Ok(()) } @@ -1409,7 +1389,6 @@ SELECT stream, max(id) as id, COUNT(*) AS num parse_file_key_columns(file).expect("parse file key failed"); let org_id = stream_key[..stream_key.find('/').unwrap()].to_string(); let client = CLIENT_RW.clone(); - let client = client.lock().await; sqlx::query( r#" INSERT INTO file_list_dump_stats @@ -1437,7 +1416,7 @@ DO UPDATE SET .bind(stats.storage_size as i64) .bind(stats.compressed_size as i64) .bind(stats.index_size as i64) - .execute(&*client) + .execute(&client) .await?; Ok(()) } @@ -1446,14 +1425,13 @@ DO UPDATE SET let (stream_key, date_key, file_name) = parse_file_key_columns(file).expect("parse file key failed"); let client = CLIENT_RW.clone(); - let client = client.lock().await; sqlx::query( r#"DELETE FROM file_list_dump_stats WHERE stream = $1 AND date = $2 AND file = $3;"#, ) .bind(stream_key) .bind(date_key) .bind(file_name) - .execute(&*client) + .execute(&client) .await?; Ok(()) } @@ -1519,7 +1497,6 @@ WHERE org = $1 AND account = $2;"#; async fn delete_by_org(&self, org_id: &str) -> Result<()> { let client = CLIENT_RW.clone(); - let client = client.lock().await; let created_at = now_micros(); let mut tx = client.begin().await?; // Move remaining rows into file_list_deleted first so the file GC removes @@ -1569,7 +1546,6 @@ impl SqliteFileList { parse_file_key_columns(file).map_err(|e| Error::Message(e.to_string()))?; let org_id = stream_key[..stream_key.find('/').unwrap()].to_string(); let client = CLIENT_RW.clone(); - let client = client.lock().await; if meta.min_ts == 0 || meta.max_ts == 0 { log::warn!("[SQLITE] min_ts or max_ts is 0 for file: {file}"); } @@ -1595,7 +1571,7 @@ INSERT INTO {table} (id, account, org, stream, date, file, deleted, min_ts, max_ .bind(meta.bloom_ver) .bind(meta.flattened) .bind(now_ts) - .execute(&*client) + .execute(&client) .await { Err(sqlx::Error::Database(e)) => if e.is_unique_violation() { Ok(0) @@ -1634,7 +1610,6 @@ INSERT INTO {table} (id, account, org, stream, date, file, deleted, min_ts, max_ .collect::>(); let client = CLIENT_RW.clone(); - let client = client.lock().await; let mut tx = client.begin().await?; if !add_items.is_empty() { @@ -1746,7 +1721,6 @@ INSERT INTO {table} (id, account, org, stream, date, file, deleted, min_ts, max_ pub async fn create_table() -> Result<()> { let client = CLIENT_RW.clone(); - let client = client.lock().await; sqlx::query( r#" CREATE TABLE IF NOT EXISTS file_list @@ -1770,7 +1744,7 @@ CREATE TABLE IF NOT EXISTS file_list ); "#, ) - .execute(&*client) + .execute(&client) .await?; sqlx::query( @@ -1796,7 +1770,7 @@ CREATE TABLE IF NOT EXISTS file_list_history ); "#, ) - .execute(&*client) + .execute(&client) .await?; sqlx::query( @@ -1815,7 +1789,7 @@ CREATE TABLE IF NOT EXISTS file_list_deleted ); "#, ) - .execute(&*client) + .execute(&client) .await?; sqlx::query( @@ -1834,7 +1808,7 @@ CREATE TABLE IF NOT EXISTS file_list_jobs ); "#, ) - .execute(&*client) + .execute(&client) .await?; sqlx::query( @@ -1855,7 +1829,7 @@ CREATE TABLE IF NOT EXISTS stream_stats ); "#, ) - .execute(&*client) + .execute(&client) .await?; sqlx::query( @@ -1877,7 +1851,7 @@ CREATE TABLE IF NOT EXISTS file_list_dump_stats ); "#, ) - .execute(&*client) + .execute(&client) .await?; // create column flattened for old version <= 0.10.5 @@ -2058,10 +2032,9 @@ pub async fn create_table_index() -> Result<()> { // delete duplicate records log::warn!("[SQLITE] starting delete duplicate records"); let client = CLIENT_RW.clone(); - let client = client.lock().await; let ret = sqlx::query( r#"SELECT stream, date, file, min(id) as id FROM file_list GROUP BY stream, date, file HAVING COUNT(*) > 1;"#, - ).fetch_all(&*client).await?; + ).fetch_all(&client).await?; log::warn!("[SQLITE] total: {} duplicate records", ret.len()); for (i, r) in ret.iter().enumerate() { let stream = r.get::("stream"); @@ -2070,7 +2043,7 @@ pub async fn create_table_index() -> Result<()> { let id = r.get::("id"); sqlx::query( r#"DELETE FROM file_list WHERE id != $1 AND stream = $2 AND date = $3 AND file = $4;"#, - ).bind(id).bind(stream).bind(date).bind(file).execute(&*client).await?; + ).bind(id).bind(stream).bind(date).bind(file).execute(&client).await?; if i.is_multiple_of(1000) { log::warn!("[SQLITE] delete duplicate records: {}/{}", i, ret.len()); } @@ -2098,9 +2071,8 @@ pub async fn create_table_index() -> Result<()> { // delete trigger for old version // compatible for old version <= 0.6.4 let client = CLIENT_RW.clone(); - let client = client.lock().await; sqlx::query(r#"DROP TRIGGER IF EXISTS update_stream_stats_delete;"#) - .execute(&*client) + .execute(&client) .await?; Ok(()) diff --git a/src/infra/src/pipeline/postgres.rs b/src/infra/src/pipeline/postgres.rs index 3ac39945e3..1bba6f7469 100644 --- a/src/infra/src/pipeline/postgres.rs +++ b/src/infra/src/pipeline/postgres.rs @@ -23,7 +23,7 @@ use config::{ }; use crate::{ - db::postgres::{CLIENT, CLIENT_DDL, CLIENT_RO, drop_column}, + db::postgres::{CLIENT_DDL, CLIENT_RO, CLIENT_RW, drop_column}, errors::{DbError, Error, Result}, }; @@ -101,7 +101,7 @@ CREATE TABLE IF NOT EXISTS pipeline } async fn put(&self, pipeline: &Pipeline) -> Result<()> { - let pool = CLIENT.clone(); + let pool = CLIENT_RW.clone(); let mut tx = pool.begin().await?; if let Err(e) = match &pipeline.source { @@ -187,7 +187,7 @@ INSERT INTO pipeline (id, version, enabled, name, description, org, source_type, } async fn update(&self, pipeline: &Pipeline) -> Result<()> { - let pool = CLIENT.clone(); + let pool = CLIENT_RW.clone(); let mut tx = pool.begin().await?; if let Err(e) = match &pipeline.source { @@ -375,7 +375,7 @@ SELECT * FROM pipeline WHERE org = $1 AND source_type = $2 ORDER BY id; } async fn delete(&self, pipeline_id: &str) -> Result { - let pool = CLIENT.clone(); + let pool = CLIENT_RW.clone(); let pipeline = sqlx::query_as::<_, Pipeline>("SELECT * FROM pipeline WHERE id = $1;") .bind(pipeline_id) diff --git a/src/infra/src/pipeline/sqlite.rs b/src/infra/src/pipeline/sqlite.rs index 8e0e4e3189..c063c701a4 100644 --- a/src/infra/src/pipeline/sqlite.rs +++ b/src/infra/src/pipeline/sqlite.rs @@ -45,7 +45,6 @@ impl Default for SqlitePipelineTable { impl super::PipelineTable for SqlitePipelineTable { async fn create_table(&self) -> Result<()> { let client = CLIENT_RW.clone(); - let client = client.lock().await; sqlx::query( r#" CREATE TABLE IF NOT EXISTS pipeline @@ -66,7 +65,7 @@ CREATE TABLE IF NOT EXISTS pipeline ); "#, ) - .execute(&*client) + .execute(&client) .await?; // drop created_at column for old version <= 0.40.0 @@ -77,30 +76,27 @@ CREATE TABLE IF NOT EXISTS pipeline async fn create_table_index(&self) -> Result<()> { let client = CLIENT_RW.clone(); - let client = client.lock().await; let queries = vec![ "CREATE INDEX IF NOT EXISTS pipeline_org_idx ON pipeline (org);", "CREATE INDEX IF NOT EXISTS pipeline_org_src_type_stream_params_idx ON pipeline (org, source_type, stream_org, stream_name, stream_type);", ]; for query in queries { - sqlx::query(query).execute(&*client).await?; + sqlx::query(query).execute(&client).await?; } Ok(()) } async fn drop_table(&self) -> Result<()> { let client = CLIENT_RW.clone(); - let client = client.lock().await; sqlx::query("DROP TABLE IF EXISTS pipeline;") - .execute(&*client) + .execute(&client) .await?; Ok(()) } async fn put(&self, pipeline: &Pipeline) -> Result<()> { let client = CLIENT_RW.clone(); - let client = client.lock().await; let mut tx = client.begin().await?; if let Err(e) = match &pipeline.source { @@ -188,7 +184,6 @@ INSERT INTO pipeline (id, version, enabled, name, description, org, source_type, async fn update(&self, pipeline: &Pipeline) -> Result<()> { let client = CLIENT_RW.clone(); - let client = client.lock().await; let mut tx = client.begin().await?; if let Err(e) = match &pipeline.source { @@ -379,7 +374,6 @@ SELECT * FROM pipeline WHERE org = $1 AND source_type = $2 ORDER BY id; async fn delete(&self, pipeline_id: &str) -> Result { let client = CLIENT_RW.clone(); - let client = client.lock().await; let mut tx = client.begin().await?; let pipeline = sqlx::query_as::<_, Pipeline>("SELECT * FROM pipeline WHERE id = $1;") diff --git a/src/infra/src/scheduler/postgres.rs b/src/infra/src/scheduler/postgres.rs index acaf57a446..3b252364b3 100644 --- a/src/infra/src/scheduler/postgres.rs +++ b/src/infra/src/scheduler/postgres.rs @@ -25,7 +25,7 @@ use super::{TRIGGERS_KEY, Trigger, TriggerModule, TriggerStatus, get_scheduler_m use crate::{ db::{ self, IndexStatement, - postgres::{CLIENT, CLIENT_DDL, CLIENT_RO, add_column, create_index, drop_column}, + postgres::{CLIENT_DDL, CLIENT_RO, CLIENT_RW, add_column, create_index, drop_column}, }, errors::{DbError, Error, Result}, }; @@ -244,7 +244,7 @@ SELECT COUNT(*)::BIGINT AS num FROM scheduled_jobs WHERE module = $1;"#, /// Pushes a Trigger job into the queue async fn push(&self, trigger: Trigger) -> Result<()> { // let db = db::get_db().await; - let pool = CLIENT.clone(); + let pool = CLIENT_RW.clone(); let mut tx = pool.begin().await?; DB_QUERY_NUMS .with_label_values(&["insert", "scheduled_jobs"]) @@ -287,7 +287,7 @@ INSERT INTO scheduled_jobs (org, module, module_key, is_realtime, is_silenced, s /// Deletes the Trigger job matching the given parameters async fn delete(&self, org: &str, module: TriggerModule, key: &str) -> Result<()> { - let pool = CLIENT.clone(); + let pool = CLIENT_RW.clone(); log::debug!("Deleting scheduled job for org: {org}, module: {module}, key: {key}",); DB_QUERY_NUMS .with_label_values(&["delete", "scheduled_jobs"]) @@ -323,7 +323,7 @@ INSERT INTO scheduled_jobs (org, module, module_key, is_realtime, is_silenced, s retries: i32, data: Option<&str>, ) -> Result<()> { - let pool = CLIENT.clone(); + let pool = CLIENT_RW.clone(); DB_QUERY_NUMS .with_label_values(&["update", "scheduled_jobs"]) .inc(); @@ -359,7 +359,7 @@ INSERT INTO scheduled_jobs (org, module, module_key, is_realtime, is_silenced, s } async fn update_trigger(&self, trigger: Trigger, clone: bool) -> Result<()> { - let pool = CLIENT.clone(); + let pool = CLIENT_RW.clone(); DB_QUERY_NUMS .with_label_values(&["update", "scheduled_jobs"]) .inc(); @@ -408,7 +408,7 @@ INSERT INTO scheduled_jobs (org, module, module_key, is_realtime, is_silenced, s return Ok(()); } - let pool = CLIENT.clone(); + let pool = CLIENT_RW.clone(); // Build bulk update query using UNNEST let query_builder = String::from( @@ -503,7 +503,7 @@ INSERT INTO scheduled_jobs (org, module, module_key, is_realtime, is_silenced, s return Ok(()); } - let pool = CLIENT.clone(); + let pool = CLIENT_RW.clone(); // Separate updates with and without data let mut updates_with_data = Vec::new(); @@ -640,7 +640,7 @@ INSERT INTO scheduled_jobs (org, module, module_key, is_realtime, is_silenced, s .collect::>() .join(",") ); - let pool = CLIENT.clone(); + let pool = CLIENT_RW.clone(); sqlx::query(&sql) .bind(TriggerModule::Report) .bind(report_max_time) @@ -675,7 +675,7 @@ INSERT INTO scheduled_jobs (org, module, module_key, is_realtime, is_silenced, s .bind(claim.id) .bind(claim.claim_epoch) .bind(TriggerStatus::Processing) - .execute(&CLIENT.clone()) + .execute(&CLIENT_RW.clone()) .await?; Ok(result.rows_affected() == 1) } @@ -691,7 +691,7 @@ INSERT INTO scheduled_jobs (org, module, module_key, is_realtime, is_silenced, s .bind(trigger.id) .bind(trigger.claim_epoch) .bind(TriggerStatus::Processing) - .execute(&CLIENT.clone()) + .execute(&CLIENT_RW.clone()) .await?; Ok(result.rows_affected() == 1) } @@ -723,7 +723,7 @@ INSERT INTO scheduled_jobs (org, module, module_key, is_realtime, is_silenced, s report_timeout: i64, module: Option, ) -> Result> { - let pool = CLIENT.clone(); + let pool = CLIENT_RW.clone(); let now = now_micros(); let report_max_time = now @@ -872,7 +872,7 @@ WHERE org = $1 AND module = $2 AND module_key = $3;"#; /// Background job that frequently (30 secs interval) cleans "Completed" jobs or jobs with /// retries >= threshold set through environment async fn clean_complete(&self) -> Result<()> { - let pool = CLIENT.clone(); + let pool = CLIENT_RW.clone(); log::debug!("[SCHEDULER] cleaning completed jobs"); let (include_max, mut max_retries) = get_scheduler_max_retries(); if include_max { @@ -899,7 +899,7 @@ WHERE org = $1 AND module = $2 AND module_key = $3;"#; /// - Lock candidates in id order, skipping rows being handled by another transaction /// - Update their status back to "Waiting" and increase their "retries" by 1 async fn watch_timeout(&self) -> Result<()> { - let pool = CLIENT.clone(); + let pool = CLIENT_RW.clone(); DB_QUERY_NUMS .with_label_values(&["update", "scheduled_jobs"]) .inc(); @@ -946,7 +946,7 @@ SELECT COUNT(*)::BIGINT AS num FROM scheduled_jobs;"#, } async fn clear(&self) -> Result<()> { - let pool = CLIENT.clone(); + let pool = CLIENT_RW.clone(); log::debug!("[SCHEDULER] clearing scheduled_jobs table"); DB_QUERY_NUMS .with_label_values(&["delete", "scheduled_jobs"]) diff --git a/src/infra/src/scheduler/sqlite.rs b/src/infra/src/scheduler/sqlite.rs index 645a6ff2c1..b139094b7f 100644 --- a/src/infra/src/scheduler/sqlite.rs +++ b/src/infra/src/scheduler/sqlite.rs @@ -77,7 +77,6 @@ impl super::Scheduler for SqliteScheduler { /// Creates the Scheduled Jobs table async fn create_table(&self) -> Result<()> { let client = CLIENT_RW.clone(); - let client = client.lock().await; sqlx::query( r#" CREATE TABLE IF NOT EXISTS scheduled_jobs @@ -98,7 +97,7 @@ CREATE TABLE IF NOT EXISTS scheduled_jobs ); "#, ) - .execute(&*client) + .execute(&client) .await?; // create data column for old version <= 0.10.9 @@ -175,7 +174,6 @@ SELECT COUNT(*) as num FROM scheduled_jobs WHERE module = $1;"#, /// Pushes a Trigger job into the queue async fn push(&self, trigger: Trigger) -> Result<()> { let client = CLIENT_RW.clone(); - let client = client.lock().await; let mut tx = client.begin().await?; if let Err(e) = sqlx::query( @@ -224,14 +222,13 @@ INSERT INTO scheduled_jobs (org, module, module_key, is_realtime, is_silenced, s /// Deletes the Trigger job matching the given parameters async fn delete(&self, org: &str, module: TriggerModule, key: &str) -> Result<()> { let client = CLIENT_RW.clone(); - let client = client.lock().await; sqlx::query( r#"DELETE FROM scheduled_jobs WHERE org = $1 AND module_key = $2 AND module = $3;"#, ) .bind(org) .bind(key) .bind(&module) - .execute(&*client) + .execute(&client) .await?; drop(client); @@ -259,7 +256,6 @@ INSERT INTO scheduled_jobs (org, module, module_key, is_realtime, is_silenced, s data: Option<&str>, ) -> Result<()> { let client = CLIENT_RW.clone(); - let client = client.lock().await; let query = match data { Some(data) => { sqlx::query( @@ -283,7 +279,7 @@ INSERT INTO scheduled_jobs (org, module, module_key, is_realtime, is_silenced, s .bind(&module) }, }; - query.execute(&*client).await?; + query.execute(&client).await?; drop(client); @@ -295,7 +291,6 @@ INSERT INTO scheduled_jobs (org, module, module_key, is_realtime, is_silenced, s async fn update_trigger(&self, trigger: Trigger, clone: bool) -> Result<()> { let client = CLIENT_RW.clone(); - let client = client.lock().await; let query = if clone { sqlx::query( r#"UPDATE scheduled_jobs @@ -330,7 +325,7 @@ INSERT INTO scheduled_jobs (org, module, module_key, is_realtime, is_silenced, s .bind(&trigger.module) }; - query.execute(&*client).await?; + query.execute(&client).await?; // release lock drop(client); @@ -381,12 +376,11 @@ INSERT INTO scheduled_jobs (org, module, module_key, is_realtime, is_silenced, s .join(",") ); let client = CLIENT_RW.clone(); - let client = client.lock().await; sqlx::query(&sql) .bind(TriggerModule::Report) .bind(report_max_time) .bind(alert_max_time) - .execute(&*client) + .execute(&client) .await?; Ok(()) @@ -410,7 +404,6 @@ INSERT INTO scheduled_jobs (org, module, module_key, is_realtime, is_silenced, s .num_microseconds() .ok_or_else(|| Error::Message("alert timeout overflow".into()))?; let client = CLIENT_RW.clone(); - let client = client.lock().await; let result = sqlx::query(KEEP_ALIVE_CLAIM_QUERY) .bind(TriggerModule::Report) .bind(report_max_time) @@ -418,14 +411,13 @@ INSERT INTO scheduled_jobs (org, module, module_key, is_realtime, is_silenced, s .bind(claim.id) .bind(claim.claim_epoch) .bind(TriggerStatus::Processing) - .execute(&*client) + .execute(&client) .await?; Ok(result.rows_affected() == 1) } async fn complete_claim(&self, trigger: Trigger) -> Result { let client = CLIENT_RW.clone(); - let client = client.lock().await; let result = sqlx::query(COMPLETE_CLAIM_QUERY) .bind(&trigger.status) .bind(trigger.retries) @@ -436,7 +428,7 @@ INSERT INTO scheduled_jobs (org, module, module_key, is_realtime, is_silenced, s .bind(trigger.id) .bind(trigger.claim_epoch) .bind(TriggerStatus::Processing) - .execute(&*client) + .execute(&client) .await?; Ok(result.rows_affected() == 1) } @@ -457,7 +449,6 @@ INSERT INTO scheduled_jobs (org, module, module_key, is_realtime, is_silenced, s module: Option, ) -> Result> { let client = CLIENT_RW.clone(); - let client = client.lock().await; let now = now_micros(); let report_max_time = now @@ -494,7 +485,7 @@ INSERT INTO scheduled_jobs (org, module, module_key, is_realtime, is_silenced, s q = q.bind(m); } q = q.bind(concurrency); - let jobs: Vec = q.fetch_all(&*client).await?; + let jobs: Vec = q.fetch_all(&client).await?; Ok(jobs) } @@ -562,7 +553,6 @@ WHERE org = $1 AND module = $2 AND module_key = $3;"#; /// retries >= threshold set through environment async fn clean_complete(&self) -> Result<()> { let client = CLIENT_RW.clone(); - let client = client.lock().await; let (include_max, mut max_retries) = get_scheduler_max_retries(); if include_max { max_retries += 1; @@ -574,7 +564,7 @@ WHERE org = $1 AND module = $2 AND module_key = $3;"#; .bind(TriggerStatus::Completed) .bind(max_retries) .bind(TriggerModule::Alert) - .execute(&*client) + .execute(&client) .await?; Ok(()) } @@ -587,7 +577,6 @@ WHERE org = $1 AND module = $2 AND module_key = $3;"#; /// - Update their status back to "Waiting" and increase their "retries" by 1 async fn watch_timeout(&self) -> Result<()> { let client = CLIENT_RW.clone(); - let client = client.lock().await; let now = now_micros(); sqlx::query( r#"UPDATE scheduled_jobs @@ -598,7 +587,7 @@ WHERE status = $2 AND end_time <= $3; .bind(TriggerStatus::Waiting) .bind(TriggerStatus::Processing) .bind(now) - .execute(&*client) + .execute(&client) .await?; Ok(()) } @@ -630,9 +619,8 @@ SELECT COUNT(*) as num FROM scheduled_jobs;"#, async fn clear(&self) -> Result<()> { let client = CLIENT_RW.clone(); - let client = client.lock().await; match sqlx::query(r#"DELETE FROM scheduled_jobs;"#) - .execute(&*client) + .execute(&client) .await { Ok(_) => log::info!("[SCHEDULER] scheduled_jobs table cleared"), diff --git a/src/infra/src/schema/history/postgres.rs b/src/infra/src/schema/history/postgres.rs index f85c0d506b..a33fe26af6 100644 --- a/src/infra/src/schema/history/postgres.rs +++ b/src/infra/src/schema/history/postgres.rs @@ -20,7 +20,7 @@ use datafusion::arrow::datatypes::Schema; use crate::{ db::{ IndexStatement, - postgres::{CLIENT, CLIENT_DDL, create_index}, + postgres::{CLIENT_DDL, CLIENT_RW, create_index}, }, errors::{Error, Result}, }; @@ -58,7 +58,7 @@ impl super::SchemaHistory for PostgresSchemaHistory { schema: Schema, ) -> Result<()> { let value = json::to_string(&schema)?; - let pool = CLIENT.clone(); + let pool = CLIENT_RW.clone(); DB_QUERY_NUMS .with_label_values(&["insert", "schema_history"]) .inc(); diff --git a/src/infra/src/schema/history/sqlite.rs b/src/infra/src/schema/history/sqlite.rs index 689cf6b96b..966529b945 100644 --- a/src/infra/src/schema/history/sqlite.rs +++ b/src/infra/src/schema/history/sqlite.rs @@ -59,7 +59,6 @@ impl super::SchemaHistory for SqliteSchemaHistory { ) -> Result<()> { let value = json::to_string(&schema)?; let client = CLIENT_RW.clone(); - let client = client.lock().await; match sqlx::query( r#" INSERT INTO schema_history (org, stream_type, stream_name, start_dt, value) @@ -71,7 +70,7 @@ INSERT INTO schema_history (org, stream_type, stream_name, start_dt, value) .bind(stream_name) .bind(start_dt) .bind(value) - .execute(&*client) + .execute(&client) .await { Err(sqlx::Error::Database(e)) => { @@ -89,7 +88,6 @@ INSERT INTO schema_history (org, stream_type, stream_name, start_dt, value) pub async fn create_table() -> Result<()> { let client = CLIENT_RW.clone(); - let client = client.lock().await; sqlx::query( r#" CREATE TABLE IF NOT EXISTS schema_history @@ -103,7 +101,7 @@ CREATE TABLE IF NOT EXISTS schema_history ); "#, ) - .execute(&*client) + .execute(&client) .await?; Ok(()) diff --git a/src/infra/src/table/action_scripts.rs b/src/infra/src/table/action_scripts.rs index 1c4a38da98..1b56058947 100644 --- a/src/infra/src/table/action_scripts.rs +++ b/src/infra/src/table/action_scripts.rs @@ -22,12 +22,9 @@ use sea_orm::{ QuerySelect, Schema, Set, Unchanged, }; -use super::{ - entity::action_scripts::{ActiveModel, Model}, - get_lock, -}; +use super::entity::action_scripts::{ActiveModel, Model}; use crate::{ - db::{ORM_CLIENT, connect_to_orm}, + db::{get_orm_client_ro, get_orm_client_rw}, errors::{self, DbError, Error}, table::entity::action_scripts::{Column, Entity}, }; @@ -67,7 +64,7 @@ impl TryFrom for Action { } } pub async fn create_table() -> Result<(), errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let builder = client.get_database_backend(); let schema = Schema::new(builder); @@ -105,10 +102,7 @@ pub async fn add(action: &Action) -> Result { service_account: Set(action.service_account.clone()), }; - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let uuid = Entity::insert(record).exec(client).await?.last_insert_id; Ok(uuid) @@ -139,19 +133,13 @@ pub async fn update(action: &Action) -> Result { service_account: Set(action.service_account.clone()), }; - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let action = Entity::update(record).exec(client).await?; Ok(action.id) } pub async fn remove(org_id: &str, id: &str) -> Result<(), errors::Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let _ = Entity::delete_many() .filter(Column::OrgId.eq(org_id)) .filter(Column::Id.eq(id)) @@ -162,7 +150,7 @@ pub async fn remove(org_id: &str, id: &str) -> Result<(), errors::Error> { } pub async fn get(id: &str, org_id: &str) -> Result { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let record = Entity::find_by_id(id) .filter(Column::OrgId.eq(org_id)) .one(client) @@ -176,7 +164,7 @@ pub async fn get(id: &str, org_id: &str) -> Result { pub async fn list(org_id: &str, limit: Option) -> Result, errors::Error> { let limit = limit.unwrap_or(100); - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let res = Entity::find() .filter(Column::OrgId.eq(org_id)) .limit(limit as u64) @@ -193,7 +181,7 @@ pub async fn list(org_id: &str, limit: Option) -> Result, error } pub async fn contains(id: &str, org_id: &str) -> Result { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let record = Entity::find() .filter(Column::Id.eq(id)) .filter(Column::OrgId.eq(org_id)) @@ -204,16 +192,13 @@ pub async fn contains(id: &str, org_id: &str) -> Result { } pub async fn len() -> Result { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let len = Entity::find().count(client).await?; Ok(len as usize) } pub async fn clear() -> Result<(), errors::Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::delete_many().exec(client).await?; Ok(()) @@ -225,8 +210,7 @@ pub async fn is_empty() -> Result { /// Deletes all action scripts belonging to the given org. pub async fn delete_by_org(org_id: &str) -> Result<(), errors::Error> { - let _lock = get_lock().await; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::delete_many() .filter(Column::OrgId.eq(org_id)) .exec(client) diff --git a/src/infra/src/table/alert_composites.rs b/src/infra/src/table/alert_composites.rs index 49e3d40a7b..24fa9610c2 100644 --- a/src/infra/src/table/alert_composites.rs +++ b/src/infra/src/table/alert_composites.rs @@ -11,10 +11,7 @@ use sea_orm::{ sea_query::{Expr, LockType}, }; -use super::{ - entity::{alert_composite_children, alert_composites, alerts}, - get_lock, -}; +use super::entity::{alert_composite_children, alert_composites, alerts}; #[derive(Clone, Copy, Debug, PartialEq, Eq)] #[repr(i16)] @@ -51,9 +48,6 @@ where let id = active_string(&definition.id, "composite id")?; validate_child_ownership(&children, &id)?; - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - let txn = conn.begin().await?; let definition = definition.insert(&txn).await?; if !children.is_empty() { @@ -81,9 +75,6 @@ where let org = active_string(&replacement.org, "composite org")?; validate_child_ownership(&children, &id)?; - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - let txn = conn.begin().await?; let mut current_query = alert_composites::Entity::find_by_id(&id).filter(alert_composites::Column::Org.eq(&org)); @@ -316,9 +307,6 @@ pub async fn increment_evaluation_generation( org: &str, id: &str, ) -> Result { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - let backend = conn.get_database_backend(); let sql = match backend { sea_orm::DatabaseBackend::Postgres => { @@ -366,9 +354,6 @@ pub async fn delete_by_id(conn: &C, org: &str, id: &str) -> Result<(), sea_or where C: ConnectionTrait + TransactionTrait, { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - let txn = conn.begin().await?; let exists = alert_composites::Entity::find_by_id(id) .filter(alert_composites::Column::Org.eq(org)) @@ -398,9 +383,6 @@ pub async fn delete_by_org(conn: &C, org: &str) -> Result<(), sea_orm::DbErr> where C: ConnectionTrait + TransactionTrait, { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - let txn = conn.begin().await?; let ids = alert_composites::Entity::find() .select_only() diff --git a/src/infra/src/table/alert_eval_intervals.rs b/src/infra/src/table/alert_eval_intervals.rs index 049d32d157..2b9f0d517b 100644 --- a/src/infra/src/table/alert_eval_intervals.rs +++ b/src/infra/src/table/alert_eval_intervals.rs @@ -28,9 +28,9 @@ use config::meta::alerts::{level::AlertLevel, state::EvalLedgerWrite}; use sea_orm::{ActiveModelTrait, ColumnTrait, EntityTrait, QueryFilter, QueryOrder, Set}; -use super::{entity::alert_eval_intervals, get_lock}; +use super::entity::alert_eval_intervals; use crate::{ - db::{ORM_CLIENT, connect_to_orm}, + db::{get_orm_client_ro, get_orm_client_rw}, errors, }; @@ -175,7 +175,7 @@ pub async fn list_overlapping( range_start_us: i64, range_end_us: i64, ) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; list_overlapping_with(client, alert_id, range_start_us, range_end_us).await } @@ -216,7 +216,7 @@ pub async fn list_overlapping_with( /// behaviour before this instant, so measuring earlier would fabricate coverage /// out of retention's blind spot. pub async fn earliest_from_us(alert_id: &str) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; earliest_from_us_with(client, alert_id).await } @@ -241,10 +241,7 @@ pub async fn earliest_from_us_with( /// cutoff but began long before it — survives. Deleting on `from_us` would /// delete the interval an alert has been sitting in for months. pub async fn delete_before(cutoff_us: i64) -> Result { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; delete_before_with(client, cutoff_us).await } @@ -263,10 +260,7 @@ pub async fn delete_before_with( /// Remove an alert's whole ledger. Called when the alert is deleted — these /// rows are owned by the alert's lifecycle, exactly like its state rows. pub async fn delete_by_alert(alert_id: &str) -> Result<(), errors::Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; delete_by_alert_with(client, alert_id).await } diff --git a/src/infra/src/table/alert_incidents.rs b/src/infra/src/table/alert_incidents.rs index 9bfd1c7567..05e1ae60b5 100644 --- a/src/infra/src/table/alert_incidents.rs +++ b/src/infra/src/table/alert_incidents.rs @@ -24,18 +24,15 @@ use sea_orm::{ }; use svix_ksuid::KsuidLike; -use super::{ - entity::{alert_incident_alerts, alert_incidents}, - get_lock, -}; +use super::entity::{alert_incident_alerts, alert_incidents}; use crate::{ - db::{ORM_CLIENT, connect_to_orm}, + db::{get_orm_client_ro, get_orm_client_rw}, errors::{self, DbError, Error}, }; /// Get incident by ID pub async fn get(org_id: &str, id: &str) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; alert_incidents::Entity::find_by_id(id) .filter(alert_incidents::Column::OrgId.eq(org_id)) @@ -53,7 +50,7 @@ pub async fn create( first_alert_at: i64, title: Option, ) -> Result { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let now = chrono::Utc::now().timestamp_micros(); let id = svix_ksuid::Ksuid::new(None, None).to_string(); @@ -75,9 +72,6 @@ pub async fn create( updated_at: Set(now), }; - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - model .insert(client) .await @@ -98,12 +92,9 @@ pub async fn add_alert_to_incident( alert_fired_at: i64, correlation_reason: &str, ) -> Result { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let now = chrono::Utc::now().timestamp_micros(); - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - // Use transaction for atomic update let txn = client .begin() @@ -172,12 +163,9 @@ pub async fn update_status( id: &str, status: &str, ) -> Result { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let now = chrono::Utc::now().timestamp_micros(); - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - let incident = get(org_id, id) .await? .ok_or_else(|| Error::DbError(DbError::SeaORMError("Incident not found".to_string())))?; @@ -202,12 +190,9 @@ pub async fn update_title( id: &str, title: &str, ) -> Result { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let now = chrono::Utc::now().timestamp_micros(); - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - let incident = get(org_id, id) .await? .ok_or_else(|| Error::DbError(DbError::SeaORMError("Incident not found".to_string())))?; @@ -228,12 +213,9 @@ pub async fn update_severity( id: &str, severity: &str, ) -> Result { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let now = chrono::Utc::now().timestamp_micros(); - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - let incident = get(org_id, id) .await? .ok_or_else(|| Error::DbError(DbError::SeaORMError("Incident not found".to_string())))?; @@ -255,7 +237,7 @@ pub async fn list( limit: u64, offset: u64, ) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let mut query = alert_incidents::Entity::find() .select_only() @@ -292,7 +274,7 @@ pub async fn list( pub async fn get_incident_alerts( incident_id: &str, ) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; alert_incident_alerts::Entity::find() .filter(alert_incident_alerts::Column::IncidentId.eq(incident_id)) @@ -310,7 +292,7 @@ pub async fn find_open_incident_by_alert_id( org_id: &str, alert_id: &str, ) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; // Step 1: find all incident IDs that reference this alert_id in the junction table let rows = alert_incident_alerts::Entity::find() @@ -345,7 +327,7 @@ pub async fn find_open_incident_containing_alert( org_id: &str, alert_id: &str, ) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let rows = alert_incident_alerts::Entity::find() .filter(alert_incident_alerts::Column::AlertId.eq(alert_id)) @@ -381,7 +363,7 @@ pub async fn get_alert_counts( return Ok(std::collections::HashMap::new()); } - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; #[derive(Debug, FromQueryResult)] struct CountResult { @@ -411,7 +393,7 @@ pub async fn get_alert_counts( /// Count open incidents for an org pub async fn count_open(org_id: &str) -> Result { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; alert_incidents::Entity::find() .filter(alert_incidents::Column::OrgId.eq(org_id)) @@ -423,7 +405,7 @@ pub async fn count_open(org_id: &str) -> Result { /// Count incidents with optional status filter pub async fn count(org_id: &str, status: Option<&str>) -> Result { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let mut query = alert_incidents::Entity::find().filter(alert_incidents::Column::OrgId.eq(org_id)); @@ -481,12 +463,9 @@ pub async fn update_topology( id: &str, topology: &config::meta::alerts::incidents::IncidentTopology, ) -> Result<(), errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let now = chrono::Utc::now().timestamp_micros(); - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - let incident = get(org_id, id) .await? .ok_or_else(|| Error::DbError(DbError::SeaORMError("Incident not found".to_string())))?; @@ -525,12 +504,9 @@ pub async fn update_incident_metadata( group_values: Option, key_type: Option<&str>, ) -> Result<(), errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let now = chrono::Utc::now().timestamp_micros(); - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - let incident = get(org_id, id) .await? .ok_or_else(|| Error::DbError(DbError::SeaORMError("Incident not found".to_string())))?; @@ -585,7 +561,7 @@ pub async fn find_open_incidents_filtered( created_after: Option, limit: Option, ) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let mut query = alert_incidents::Entity::find() .filter(alert_incidents::Column::OrgId.eq(org_id)) @@ -626,12 +602,9 @@ pub async fn upgrade_incident_group_values( alert_count: i32, last_alert_at: i64, ) -> Result<(), errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let now = chrono::Utc::now().timestamp_micros(); - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - let incident = get(org_id, id) .await? .ok_or_else(|| Error::DbError(DbError::SeaORMError("Incident not found".to_string())))?; @@ -667,17 +640,12 @@ pub async fn auto_resolve_stale( ) -> Result<(u64, Vec<(String, String)>), errors::Error> { const PAGE_SIZE: u64 = 500; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let now = chrono::Utc::now().timestamp_micros(); let cutoff = now - stale_threshold_micros; let mut resolved_ids = Vec::new(); loop { - // make sure only one client is writing to the database(only for sqlite). - // Re-acquired per page so a large backlog doesn't monopolize the global - // write lock; each pass shrinks the non-resolved candidate set. - let _lock = get_lock().await; - // Find open/acknowledged incidents with last_alert_at older than threshold let stale_incidents = alert_incidents::Entity::find() .filter(alert_incidents::Column::Status.ne("resolved")) @@ -718,10 +686,7 @@ pub async fn auto_resolve_stale( /// Deletes all alert incidents belonging to the given org. pub async fn delete_by_org(org_id: &str) -> Result<(), errors::Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; alert_incidents::Entity::delete_many() .filter(alert_incidents::Column::OrgId.eq(org_id)) .exec(client) diff --git a/src/infra/src/table/alert_states.rs b/src/infra/src/table/alert_states.rs index e48ad7ec0d..ecdb986124 100644 --- a/src/infra/src/table/alert_states.rs +++ b/src/infra/src/table/alert_states.rs @@ -33,12 +33,9 @@ use sea_orm::{ TransactionTrait, }; -use super::{ - entity::{alert_state_transitions, alert_states, alerts}, - get_lock, -}; +use super::entity::{alert_state_transitions, alert_states, alerts}; use crate::{ - db::{ORM_CLIENT, connect_to_orm}, + db::{get_orm_client_ro, get_orm_client_rw}, errors, }; @@ -72,7 +69,7 @@ impl From for AlertState { /// Fetch the state row for one `(alert_id, group_key)`. pub async fn get(alert_id: &str, group_key: &str) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; get_with(client, alert_id, group_key).await } @@ -101,7 +98,7 @@ pub async fn get_rollups(alert_ids: &[String]) -> Result, errors if alert_ids.is_empty() { return Ok(vec![]); } - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; Ok(alert_states::Entity::find() .filter(alert_states::Column::AlertId.is_in(alert_ids.to_vec())) .filter(alert_states::Column::GroupKey.eq(ROLLUP_GROUP_KEY)) @@ -114,7 +111,7 @@ pub async fn get_rollups(alert_ids: &[String]) -> Result, errors /// All per-group rows for one alert. pub async fn list_groups(alert_id: &str) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; list_groups_with(client, alert_id).await } @@ -151,7 +148,7 @@ pub async fn persist( if update.state.is_none() && ledger.is_none() { return Ok(()); } - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; persist_with(client, update, ledger).await } @@ -165,9 +162,6 @@ pub async fn persist_with( return Ok(()); } - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - let txn = conn.begin().await?; write_update(&txn, update).await?; if let Some(write) = ledger { @@ -192,7 +186,7 @@ pub async fn persist_with( /// region that has not yet seen the opt-out would materialise exactly the rows /// the toggle promised to remove. pub async fn persist_group_plan(plan: &GroupPlan, alert_id: &str) -> Result { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; persist_group_plan_with(client, plan, alert_id).await } @@ -202,9 +196,6 @@ pub async fn persist_group_plan_with Result { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - let txn = conn.begin().await?; // §5.3 opt-out race: this evaluation may have read the alert BEFORE a save @@ -258,7 +249,7 @@ pub async fn advance_delivery_state( episode: DeliveryEpisode, outcome: DeliveryOutcome, ) -> Result { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; advance_delivery_state_with(client, alert_id, group_key, episode, outcome).await } @@ -270,9 +261,6 @@ pub async fn advance_delivery_state_with Result { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - let txn = conn.begin().await?; let Some(current) = @@ -425,7 +413,7 @@ pub enum DeliveryOutcome { /// completeness gate, so a cluster full of frozen alerts costs one query /// instead of a full table read every tick. pub async fn list_alert_ids_with_groups() -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; Ok(alert_states::Entity::find() .select_only() .column(alert_states::Column::AlertId) @@ -445,7 +433,7 @@ pub async fn delete_groups(alert_id: &str, group_keys: &[String]) -> Result<(), if group_keys.is_empty() { return Ok(()); } - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; delete_groups_with(client, alert_id, group_keys).await } @@ -459,9 +447,6 @@ pub async fn delete_groups_with( return Ok(()); } - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - alert_states::Entity::delete_many() .filter(alert_states::Column::AlertId.eq(alert_id)) .filter(alert_states::Column::GroupKey.is_in(group_keys.to_vec())) @@ -520,7 +505,7 @@ async fn multi_alert_still_enabled( /// stale firing rows visible for K x interval plus the grace period, which is /// the opposite of the immediate rollback the toggle promises. pub async fn delete_all_groups(alert_id: &str) -> Result { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; delete_all_groups_with(client, alert_id).await } @@ -529,9 +514,6 @@ pub async fn delete_all_groups_with( conn: &C, alert_id: &str, ) -> Result { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - let res = alert_states::Entity::delete_many() .filter(alert_states::Column::AlertId.eq(alert_id)) .filter(alert_states::Column::GroupKey.ne(ROLLUP_GROUP_KEY)) @@ -693,7 +675,7 @@ pub async fn list_transitions_filtered( group_key: Option<&str>, limit: u64, ) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let mut query = alert_state_transitions::Entity::find() .filter(alert_state_transitions::Column::AlertId.eq(alert_id)); if let Some(key) = group_key { @@ -732,7 +714,7 @@ pub async fn list_transitions_between( to: i64, limit: u64, ) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let mut query = alert_state_transitions::Entity::find() .filter(alert_state_transitions::Column::AlertId.eq(alert_id)) .filter(alert_state_transitions::Column::At.gte(from)) @@ -778,7 +760,7 @@ pub async fn list_transitions_between_many( if alert_ids.is_empty() { return Ok(grouped); } - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let mut query = alert_state_transitions::Entity::find() .filter(alert_state_transitions::Column::AlertId.is_in(alert_ids.iter().cloned())) .filter(alert_state_transitions::Column::At.gte(from)) @@ -817,7 +799,7 @@ pub async fn list_transitions_between_many( /// Remove all state for an alert. Called when the alert itself is deleted — /// unlike `scheduled_jobs`, these rows are owned by the alert's lifecycle. pub async fn delete_by_alert(alert_id: &str) -> Result<(), errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; delete_by_alert_with(client, alert_id).await } @@ -826,9 +808,6 @@ pub async fn delete_by_alert_with Result<(), errors::Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - let txn = conn.begin().await?; alert_states::Entity::delete_many() .filter(alert_states::Column::AlertId.eq(alert_id)) @@ -845,10 +824,7 @@ pub async fn delete_by_alert_with Result<(), errors::Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; alert_state_transitions::Entity::delete_many() .filter(alert_state_transitions::Column::At.lt(cutoff)) .exec(client) diff --git a/src/infra/src/table/alerts/mod.rs b/src/infra/src/table/alerts/mod.rs index b7fe513712..dd4c66c24f 100644 --- a/src/infra/src/table/alerts/mod.rs +++ b/src/infra/src/table/alerts/mod.rs @@ -42,7 +42,7 @@ use super::{ folders::folder_type_into_i16, }; use crate::{ - db::{ORM_CLIENT, connect_to_orm}, + db::get_orm_client_rw, errors::{self, FromStrError, PutAlertError}, }; @@ -223,7 +223,6 @@ pub async fn get_by_id( org_id: &str, alert_id: Ksuid, ) -> Result, errors::Error> { - let _lock = super::get_lock().await; let models = get_model_by_id(conn, org_id, alert_id).await?; if let Some((folder_model, alert_model)) = models { @@ -244,7 +243,6 @@ pub async fn get_by_name( stream_name: &str, alert_name: &str, ) -> Result, errors::Error> { - let _lock = super::get_lock().await; let models = get_model_by_name( conn, org_id, @@ -273,7 +271,6 @@ pub async fn put( folder_id: &str, alert: MetaAlert, ) -> Result { - let _lock = super::get_lock().await; let txn = conn.begin().await?; let rslt: Result = match get_model_by_name( &txn, @@ -341,7 +338,6 @@ pub async fn create( alert: MetaAlert, use_given_id: bool, ) -> Result { - let _lock = super::get_lock().await; let txn = conn.begin().await?; // Get the destination folder. @@ -396,7 +392,6 @@ pub async fn update( return Err(errors::DbError::PutAlert(PutAlertError::UpdateAlertMissingID).into()); }; - let _lock = super::get_lock().await; let txn = conn.begin().await?; // Try to get the new parent folder if a folder ID is provided. @@ -413,7 +408,7 @@ pub async fn update( }; // Try to get the alert to update. - let Some((_, alert_m)) = get_model_by_id(conn, org_id, alert_id).await? else { + let Some((_, alert_m)) = get_model_by_id(&txn, org_id, alert_id).await? else { return Err(errors::DbError::PutAlert(PutAlertError::UpdateAlertNotFound).into()); }; @@ -438,7 +433,6 @@ pub async fn delete_by_id( org_id: &str, alert_id: Ksuid, ) -> Result<(), errors::Error> { - let _lock = super::get_lock().await; alerts::Entity::delete_many() .filter(alerts::Column::Org.eq(org_id)) .filter(alerts::Column::Id.eq(alert_id.to_string())) @@ -456,7 +450,6 @@ pub async fn delete_by_name( stream_name: &str, alert_name: &str, ) -> Result<(), errors::Error> { - let _lock = super::get_lock().await; let model = get_model_by_name( conn, org_id, @@ -480,7 +473,6 @@ pub async fn list( conn: &C, params: ListAlertsParams, ) -> Result, errors::Error> { - let _lock = super::get_lock().await; let alerts = list_models(conn, params) .await? .into_iter() @@ -520,7 +512,6 @@ pub async fn list_slo_burn_window_pairs( slo_id: &str, exclude_alert_id: Option<&str>, ) -> Result, errors::Error> { - let _lock = super::get_lock().await; let query = alerts::Entity::find() .filter(alerts::Column::Org.eq(org)) .filter(alerts::Column::SloId.eq(slo_id)) @@ -563,7 +554,6 @@ pub async fn list_alerts_by_slo( org: &str, slo_id: &str, ) -> Result, errors::Error> { - let _lock = super::get_lock().await; Ok(alerts::Entity::find() .filter(alerts::Column::Org.eq(org)) .filter(alerts::Column::SloId.eq(slo_id)) @@ -594,7 +584,6 @@ pub async fn last_written_us( org: &str, alert_id: &str, ) -> Result, errors::Error> { - let _lock = super::get_lock().await; Ok(alerts::Entity::find_by_id(alert_id.to_string()) .filter(alerts::Column::Org.eq(org)) .one(conn) @@ -604,7 +593,6 @@ pub async fn last_written_us( /// Lists all alerts. pub async fn list_all(conn: &C) -> Result, errors::Error> { - let _lock = super::get_lock().await; let alerts = list_all_models(conn) .await? .into_iter() @@ -1055,10 +1043,7 @@ fn update_mutable_fields( /// Deletes all alerts belonging to the given org. pub async fn delete_by_org(org_id: &str) -> Result<(), errors::Error> { - // Init the ORM client BEFORE taking the lock: connect_to_orm acquires the - // same lock internally, so locking first can deadlock on the initial connect. - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; - let _lock = super::get_lock().await; + let client = get_orm_client_rw().await; alerts::Entity::delete_many() .filter(alerts::Column::Org.eq(org_id)) .exec(client) diff --git a/src/infra/src/table/anomaly_detection/config.rs b/src/infra/src/table/anomaly_detection/config.rs index f9685ff076..d6ad705c7f 100644 --- a/src/infra/src/table/anomaly_detection/config.rs +++ b/src/infra/src/table/anomaly_detection/config.rs @@ -39,7 +39,6 @@ pub async fn get_by_id( org_id: &str, anomaly_id: &str, ) -> Result> { - let _lock = super::super::get_lock().await; let model = anomaly_detection_config::Entity::find_by_id(anomaly_id) .filter(anomaly_detection_config::Column::OrgId.eq(org_id)) .one(conn) @@ -50,7 +49,6 @@ pub async fn get_by_id( /// Lists all configs for an organisation, ordered by `created_at` descending. pub async fn list_by_org(conn: &C, org_id: &str) -> Result> { - let _lock = super::super::get_lock().await; let models = anomaly_detection_config::Entity::find() .filter(anomaly_detection_config::Column::OrgId.eq(org_id)) .order_by_desc(anomaly_detection_config::Column::CreatedAt) @@ -63,7 +61,6 @@ pub async fn list_by_org(conn: &C, org_id: &str) -> Result(conn: &C) -> Result> { - let _lock = super::super::get_lock().await; let models = anomaly_detection_config::Entity::find() .filter(anomaly_detection_config::Column::Enabled.eq(true)) .all(conn) @@ -83,7 +80,6 @@ pub async fn create( conn: &C, config: Model, ) -> Result { - let _lock = super::super::get_lock().await; let txn = conn.begin().await?; let folder_exists = folders::Entity::find_by_id(&config.folder_id) @@ -114,7 +110,6 @@ pub async fn create_if_not_exists( conn: &C, config: Model, ) -> Result<()> { - let _lock = super::super::get_lock().await; let txn = conn.begin().await?; let anomaly_id = config.anomaly_id.clone(); @@ -150,7 +145,6 @@ pub async fn update( conn: &C, incoming: Model, ) -> Result { - let _lock = super::super::get_lock().await; let txn = conn.begin().await?; let existing = anomaly_detection_config::Entity::find_by_id(&incoming.anomaly_id) @@ -198,7 +192,6 @@ pub async fn put( conn: &C, incoming: Model, ) -> Result { - let _lock = super::super::get_lock().await; let txn = conn.begin().await?; let anomaly_id = incoming.anomaly_id.clone(); @@ -233,7 +226,6 @@ pub async fn put( /// Deletes a config by `anomaly_id` + `org_id`. pub async fn delete(conn: &C, org_id: &str, anomaly_id: &str) -> Result<()> { - let _lock = super::super::get_lock().await; anomaly_detection_config::Entity::delete_many() .filter(anomaly_detection_config::Column::AnomalyId.eq(anomaly_id)) .filter(anomaly_detection_config::Column::OrgId.eq(org_id)) diff --git a/src/infra/src/table/anomaly_detection/models.rs b/src/infra/src/table/anomaly_detection/models.rs index 751dcd8f90..742049ea34 100644 --- a/src/infra/src/table/anomaly_detection/models.rs +++ b/src/infra/src/table/anomaly_detection/models.rs @@ -36,7 +36,6 @@ type Result = std::result::Result; /// object storage. Returns the inserted row so callers can inspect the stored /// values. pub async fn insert(conn: &C, model: Model) -> Result { - let _lock = super::super::get_lock().await; let anomaly_id = model.anomaly_id.clone(); let version = model.version; log::info!("[anomaly_detection_models] insert: anomaly_id={anomaly_id} version={version}"); @@ -55,7 +54,6 @@ pub async fn list_by_anomaly_id_desc( conn: &C, anomaly_id: &str, ) -> Result> { - let _lock = super::super::get_lock().await; let models = anomaly_detection_models::Entity::find() .filter(anomaly_detection_models::Column::AnomalyId.eq(anomaly_id)) .order_by_desc(anomaly_detection_models::Column::Version) @@ -71,7 +69,6 @@ pub async fn delete_by_key( anomaly_id: &str, version: i64, ) -> Result<()> { - let _lock = super::super::get_lock().await; anomaly_detection_models::Entity::delete_many() .filter(anomaly_detection_models::Column::AnomalyId.eq(anomaly_id)) .filter(anomaly_detection_models::Column::Version.eq(version)) diff --git a/src/infra/src/table/backfill_jobs.rs b/src/infra/src/table/backfill_jobs.rs index 17d0350dee..90f6fd1e85 100644 --- a/src/infra/src/table/backfill_jobs.rs +++ b/src/infra/src/table/backfill_jobs.rs @@ -16,10 +16,9 @@ use sea_orm::{ActiveModelTrait, ColumnTrait, EntityTrait, QueryFilter, Set}; use serde::{Deserialize, Serialize}; -use super::get_lock; pub use crate::table::entity::backfill_jobs::{ActiveModel, Column, Entity, Model, Relation}; use crate::{ - db::{ORM_CLIENT, connect_to_orm}, + db::{get_orm_client_ro, get_orm_client_rw}, errors, orm_err, }; @@ -55,7 +54,7 @@ impl From for BackfillJob { } pub async fn get(org: &str, job_id: &str) -> Result { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let res = Entity::find() .filter(Column::Org.eq(org)) .filter(Column::Id.eq(job_id)) @@ -69,7 +68,7 @@ pub async fn get(org: &str, job_id: &str) -> Result } pub async fn list_by_org(org: &str) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let res = Entity::find().filter(Column::Org.eq(org)).all(client).await; match res { Ok(models) => Ok(models.into_iter().map(|model| model.into()).collect()), @@ -81,7 +80,7 @@ pub async fn list_by_pipeline( org: &str, pipeline_id: &str, ) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let res = Entity::find() .filter(Column::Org.eq(org)) .filter(Column::PipelineId.eq(pipeline_id)) @@ -107,9 +106,7 @@ pub async fn add(job: BackfillJob) -> Result<(), errors::Error> { enabled: Set(job.enabled), }; - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let res = active.insert(client).await; match res { Ok(_) => Ok(()), @@ -118,9 +115,7 @@ pub async fn add(job: BackfillJob) -> Result<(), errors::Error> { } pub async fn delete(org: &str, job_id: &str) -> Result<(), errors::Error> { - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let res = Entity::delete_many() .filter(Column::Org.eq(org)) .filter(Column::Id.eq(job_id)) @@ -134,8 +129,7 @@ pub async fn delete(org: &str, job_id: &str) -> Result<(), errors::Error> { /// Deletes all backfill jobs belonging to the given org. pub async fn delete_by_org(org: &str) -> Result<(), errors::Error> { - let _lock = get_lock().await; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::delete_many() .filter(Column::Org.eq(org)) .exec(client) @@ -144,9 +138,7 @@ pub async fn delete_by_org(org: &str) -> Result<(), errors::Error> { } pub async fn update(job: &BackfillJob) -> Result<(), errors::Error> { - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; // Find existing model let existing = Entity::find() @@ -177,9 +169,7 @@ pub async fn update(job: &BackfillJob) -> Result<(), errors::Error> { } pub async fn update_enabled(org: &str, job_id: &str, enabled: bool) -> Result<(), errors::Error> { - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; // Find existing model let existing = Entity::find() diff --git a/src/infra/src/table/cipher.rs b/src/infra/src/table/cipher.rs index faa2994a46..fe6dec4640 100644 --- a/src/infra/src/table/cipher.rs +++ b/src/infra/src/table/cipher.rs @@ -26,9 +26,9 @@ use config::{ use sea_orm::{ColumnTrait, EntityTrait, Order, QueryFilter, QueryOrder, QuerySelect, Set, SqlErr}; use serde::{Deserialize, Serialize}; -use super::{entity::cipher_keys::*, get_lock}; +use super::entity::cipher_keys::*; use crate::{ - db::{ORM_CLIENT, connect_to_orm}, + db::{get_orm_client_ro, get_orm_client_rw}, errors, }; @@ -165,25 +165,18 @@ pub async fn add(entry: CipherEntry) -> Result<(), errors::Error> { is_system: Set(entry.is_system), }; - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; match Entity::insert(record).exec(client).await { Ok(_) => {} - Err(e) => { - drop(_lock); - match e.sql_err() { - Some(SqlErr::UniqueConstraintViolation(_)) => { - return Err(errors::Error::DbError(errors::DbError::UniqueViolation)); - } - _ => { - return Err(e.into()); - } + Err(e) => match e.sql_err() { + Some(SqlErr::UniqueConstraintViolation(_)) => { + return Err(errors::Error::DbError(errors::DbError::UniqueViolation)); } - } + _ => { + return Err(e.into()); + } + }, } - drop(_lock); Ok(()) } @@ -200,21 +193,14 @@ pub async fn update(entry: CipherEntry) -> Result<(), errors::Error> { is_system: Set(entry.is_system), }; - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::update(record).exec(client).await?; - drop(_lock); Ok(()) } pub async fn remove(org: &str, kind: EntryKind, name: &str) -> Result<(), errors::Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::delete_many() .filter(Column::Org.eq(org)) .filter(Column::Kind.eq(kind.to_string())) @@ -222,8 +208,6 @@ pub async fn remove(org: &str, kind: EntryKind, name: &str) -> Result<(), errors .exec(client) .await?; - drop(_lock); - Ok(()) } @@ -260,7 +244,7 @@ pub async fn get_data( } async fn get(org: &str, kind: EntryKind, name: &str) -> Result, errors::DbError> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; Entity::find() .filter(Column::Org.eq(org)) .filter(Column::Name.eq(name)) @@ -272,7 +256,7 @@ async fn get(org: &str, kind: EntryKind, name: &str) -> Result, er } pub async fn list_all(limit: Option) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let mut res = Entity::find().order_by(Column::CreatedAt, Order::Desc); if let Some(limit) = limit { res = res.limit(limit as u64); @@ -297,7 +281,7 @@ pub async fn list_filtered( filter: ListFilter, limit: Option, ) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let mut res = Entity::find().order_by(Column::CreatedAt, Order::Desc); if let Some(ref org) = filter.org { res = res.filter(Column::Org.eq(org)); @@ -329,10 +313,7 @@ pub async fn list_filtered( } pub async fn clear() -> Result<(), errors::Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::delete_many().exec(client).await?; Ok(()) @@ -403,8 +384,7 @@ pub async fn get_dek(org: &str) -> Result, errors::Error> { /// Deletes all cipher key entries belonging to the given org. pub async fn delete_by_org(org: &str) -> Result<(), errors::Error> { - let _lock = get_lock().await; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::delete_many() .filter(Column::Org.eq(org)) .exec(client) diff --git a/src/infra/src/table/compactor_manual_jobs.rs b/src/infra/src/table/compactor_manual_jobs.rs index 7da18d833e..126548beee 100644 --- a/src/infra/src/table/compactor_manual_jobs.rs +++ b/src/infra/src/table/compactor_manual_jobs.rs @@ -16,12 +16,12 @@ use sea_orm::{ColumnTrait, EntityTrait, QueryFilter, Set, TransactionTrait}; use serde::{Deserialize, Serialize}; -use super::{get_lock, migration::Expr}; +use super::migration::Expr; pub use crate::table::entity::compactor_manual_jobs::{ ActiveModel, Column, Entity, Model, Relation, }; use crate::{ - db::{ORM_CLIENT, connect_to_orm}, + db::{get_orm_client_ro, get_orm_client_rw}, errors, orm_err, }; @@ -91,10 +91,7 @@ pub struct CompactorManualJobStatusRes { } pub async fn get(ksuid: &str) -> Result { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let res = Entity::find() .filter(Column::Id.eq(ksuid)) .one(client) @@ -110,7 +107,7 @@ pub async fn get_by_key( key: &str, status: Option, ) -> Result { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let res; if let Some(status) = status { res = Entity::find() @@ -129,7 +126,7 @@ pub async fn get_by_key( } pub async fn list_by_key(key: &str) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let res = Entity::find().filter(Column::Key.eq(key)).all(client).await; match res { Ok(models) => Ok(models.into_iter().map(|model| model.into()).collect()), @@ -146,10 +143,7 @@ pub async fn add(job: CompactorManualJob) -> Result<(), errors::Error> { status: Set(job.status as i64), }; - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::insert(active).exec(client).await.map_err(|e| { log::error!("[COMPACTOR] add job error: {e}"); e @@ -159,10 +153,7 @@ pub async fn add(job: CompactorManualJob) -> Result<(), errors::Error> { } pub async fn update(ksuid: &str, ended_at: i64, status: Status) -> Result<(), errors::Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let res = Entity::update_many() .col_expr(Column::EndedAt, Expr::value(ended_at)) .col_expr(Column::Status, Expr::value(status as i64)) @@ -177,10 +168,7 @@ pub async fn update(ksuid: &str, ended_at: i64, status: Status) -> Result<(), er } pub async fn bulk_update(jobs: Vec) -> Result<(), errors::Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let tx = client.begin().await?; for job in jobs { @@ -199,10 +187,7 @@ pub async fn bulk_update(jobs: Vec) -> Result<(), errors::Er /// Delete all compactor manual jobs whose key starts with `{org_id}/`. /// The key format is `{org_id}/{stream_type}/{stream_name}/...`. pub async fn delete_by_org(org_id: &str) -> Result<(), errors::Error> { - // make sure only one client is writing to the database (only for SQLite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let prefix = format!("{org_id}/%"); let res = Entity::delete_many() .filter(Column::Key.like(&prefix)) diff --git a/src/infra/src/table/dashboards.rs b/src/infra/src/table/dashboards.rs index 0cf2608636..000b96d714 100644 --- a/src/infra/src/table/dashboards.rs +++ b/src/infra/src/table/dashboards.rs @@ -34,10 +34,9 @@ use super::{ distinct_values::{self, OriginType}, entity::{dashboards, folders, reports}, folders::folder_type_into_i16, - get_lock, }; use crate::{ - db::{ORM_CLIENT, connect_to_orm}, + db::{get_orm_client_ro, get_orm_client_rw}, errors::{self, GetDashboardError}, }; @@ -118,7 +117,7 @@ pub async fn get_from_folder( folder_id: &str, dashboard_id: &str, ) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let model = get_model_from_folder(client, org_id, folder_id, dashboard_id) .await? .and_then(|(_folder, maybe_dash)| maybe_dash); @@ -140,7 +139,7 @@ pub async fn get_by_id( org_id: &str, dashboard_id: &str, ) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let Some((folder_m, dash_m)) = get_model_by_id(client, org_id, dashboard_id).await? else { return Ok(None); }; @@ -151,7 +150,7 @@ pub async fn get_by_id( /// Lists dashboards. pub async fn list(params: ListDashboardsParams) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let dashboards = list_models(client, params) .await? .into_iter() @@ -166,7 +165,7 @@ pub async fn list(params: ListDashboardsParams) -> Result Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let dashboards = list_all_models(client) .await? .into_iter() @@ -188,7 +187,7 @@ pub async fn put( mut dashboard: Dashboard, clone: bool, ) -> Result { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; // Get the fields that will be inserted into or updated in the database. let dashboard_id = dashboard @@ -222,9 +221,6 @@ pub async fn put( let updated_at = dashboard.updated_at; let data = inner_data_as_json(dashboard)?; - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - // Try to get the new folder if the dashboard is being moved to a new folder. let maybe_new_folder_model = if let Some(new_folder_id) = new_folder_id { let new_folder_model = folders::Entity::find() @@ -309,10 +305,7 @@ pub async fn delete_from_folder( folder_id: &str, dashboard_id: &str, ) -> Result<(), errors::Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let txn = ORM_CLIENT.get_or_init(connect_to_orm).await.begin().await?; + let txn = get_orm_client_rw().await.begin().await?; let maybe_dashboard_model = get_model_from_folder(&txn, org_id, folder_id, dashboard_id) .await? .and_then(|(_folder, maybe_dash)| maybe_dash); @@ -337,7 +330,7 @@ pub async fn delete_from_folder( /// Deletes all dashboards. pub async fn delete_all() -> Result<(), errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let dashboards = list_all().await?; let ids: Vec<_> = dashboards @@ -349,10 +342,6 @@ pub async fn delete_all() -> Result<(), errors::Error> { distinct_values::batch_remove(OriginType::Dashboard, id).await?; } - // make sure only one client is writing to the database(only for sqlite) - // NOTE: acquired after batch_remove() calls above, which take this lock internally - let _lock = get_lock().await; - dashboards::Entity::delete_many().exec(client).await?; Ok(()) } @@ -543,10 +532,7 @@ fn inner_data_as_json(dashboard: Dashboard) -> Result /// Deletes all dashboards belonging to the given org (via folder join). pub async fn delete_by_org(org_id: &str) -> Result<(), errors::Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; // Collect folder PKs for this org let folder_ids: Vec = folders::Entity::find() .filter(folders::Column::Org.eq(org_id)) diff --git a/src/infra/src/table/destinations.rs b/src/infra/src/table/destinations.rs index f32d2cc10f..100d5bec31 100644 --- a/src/infra/src/table/destinations.rs +++ b/src/infra/src/table/destinations.rs @@ -22,14 +22,11 @@ use sea_orm::{ }; use crate::{ - db::{ORM_CLIENT, connect_to_orm}, + db::{get_orm_client_ro, get_orm_client_rw}, errors::{DestinationError, Error}, - table::{ - entity::{ - destinations::{ActiveModel, Column, Entity, Model}, - templates, - }, - get_lock, + table::entity::{ + destinations::{ActiveModel, Column, Entity, Model}, + templates, }, }; @@ -77,9 +74,7 @@ pub async fn put( None }; - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let (model, template): (Model, Option) = match get_model_and_template(client, &destination.org_id, &destination.name).await? { @@ -144,7 +139,7 @@ pub async fn put( } pub async fn get(org_id: &str, name: &str) -> Result, Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; match get_model_and_template(client, org_id, name).await? { Some((model, template)) => Ok(Some(model.try_into(template)?)), None => Ok(None), @@ -155,7 +150,7 @@ pub async fn list( org_id: &str, module: Option<&str>, ) -> Result, Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let destinations = list_models(client, Some(org_id), module) .await? .into_iter() @@ -165,7 +160,7 @@ pub async fn list( } pub async fn list_all() -> Result, Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let destinations = list_models(client, None, None) .await? .into_iter() @@ -175,10 +170,7 @@ pub async fn list_all() -> Result, Error> { } pub async fn delete(org_id: &str, name: &str) -> Result<(), Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let model = get_model_and_template(client, org_id, name).await?; if let Some((model, ..)) = model { @@ -226,8 +218,7 @@ async fn list_models( /// Deletes all destinations belonging to the given org. pub async fn delete_by_org(org_id: &str) -> Result<(), Error> { - let _lock = get_lock().await; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::delete_many() .filter(Column::Org.eq(org_id)) .exec(client) diff --git a/src/infra/src/table/distinct_values.rs b/src/infra/src/table/distinct_values.rs index 7ac793fc18..3814a3749f 100644 --- a/src/infra/src/table/distinct_values.rs +++ b/src/infra/src/table/distinct_values.rs @@ -19,9 +19,8 @@ use sea_orm::{ }; use serde::{Deserialize, Serialize}; -use super::get_lock; use crate::{ - db::{ORM_CLIENT, ORM_CLIENT_DDL, connect_to_orm, connect_to_orm_ddl}, + db::{get_orm_client_ddl, get_orm_client_ro, get_orm_client_rw}, errors::{self, DbError, Error}, }; @@ -109,7 +108,7 @@ pub async fn init() -> Result<(), errors::Error> { } pub async fn create_table() -> Result<(), errors::Error> { - let client = ORM_CLIENT_DDL.get_or_init(connect_to_orm_ddl).await; + let client = get_orm_client_ddl().await; let builder = client.get_database_backend(); let schema = Schema::new(builder); @@ -136,10 +135,7 @@ pub async fn add(record: DistinctFieldRecord) -> Result<(), errors::Error> { field_name: Set(record.field_name), }; - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let res = Entity::insert(record).exec(client).await; match res { @@ -158,9 +154,6 @@ pub async fn add(record: DistinctFieldRecord) -> Result<(), errors::Error> { } pub async fn remove(record: DistinctFieldRecord) -> Result<(), errors::Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - let record = ActiveModel { origin: Set(record.origin), origin_id: Set(record.origin_id), @@ -170,7 +163,7 @@ pub async fn remove(record: DistinctFieldRecord) -> Result<(), errors::Error> { field_name: Set(record.field_name), }; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::delete(record) .exec(client) .await @@ -185,7 +178,7 @@ pub async fn check_field_use( stream_type: &str, field_name: &str, ) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let records = Entity::find() .filter(Column::OrgName.eq(org_name)) .filter(Column::StreamName.eq(stream_name)) @@ -201,10 +194,7 @@ pub async fn check_field_use( /// This is specifically for the case when a dashboard is deleted, we can bulk remove /// the dependencies, without having to go through one by one pub async fn batch_remove(origin: OriginType, origin_id: &str) -> Result<(), errors::Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::delete_many() .filter(Column::Origin.eq(origin)) .filter(Column::OriginId.eq(origin_id)) @@ -217,10 +207,7 @@ pub async fn batch_remove(origin: OriginType, origin_id: &str) -> Result<(), err /// Delete all distinct value field records for the given org. pub async fn delete_by_org(org_id: &str) -> Result<(), errors::Error> { - // make sure only one client is writing to the database (only for SQLite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::delete_many() .filter(Column::OrgName.eq(org_id)) .exec(client) @@ -231,8 +218,7 @@ pub async fn delete_by_org(org_id: &str) -> Result<(), errors::Error> { } pub async fn len() -> Result { - let _lock = get_lock().await; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; Entity::find() .count(client) .await diff --git a/src/infra/src/table/enrichment_table_urls.rs b/src/infra/src/table/enrichment_table_urls.rs index 885cab3b26..6fece88c5e 100644 --- a/src/infra/src/table/enrichment_table_urls.rs +++ b/src/infra/src/table/enrichment_table_urls.rs @@ -16,8 +16,6 @@ use sea_orm::{ColumnTrait, EntityTrait, QueryFilter, QuerySelect, Set, entity::prelude::*}; use serde::{Deserialize, Serialize}; -use super::get_lock; - // Status constants for enrichment table URL jobs const STATUS_PROCESSING: i16 = 1; // Re-export the entity for convenience @@ -25,7 +23,7 @@ pub use crate::table::entity::enrichment_table_urls::{ ActiveModel, Column, Entity, Model, Relation, }; use crate::{ - db::{ORM_CLIENT, connect_to_orm}, + db::{get_orm_client_ro, get_orm_client_rw}, errors, }; @@ -78,9 +76,7 @@ impl From for EnrichmentTableUrlRecord { /// # Returns /// * `Result<(), errors::Error>` - Success or error pub async fn put(record: EnrichmentTableUrlRecord) -> Result<(), errors::Error> { - // make sure only one client is writing to the database(only for sqlite) - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; - let _lock = get_lock().await; + let client = get_orm_client_rw().await; // Try to find existing record by ID let existing = Entity::find() @@ -140,7 +136,7 @@ pub async fn get( org: &str, table_name: &str, ) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let record = Entity::find() .filter(Column::Org.eq(org)) .filter(Column::Name.eq(table_name)) @@ -159,9 +155,7 @@ pub async fn get( /// # Returns /// * `Result<(), errors::Error>` - Success or error pub async fn delete(org: &str, table_name: &str) -> Result<(), errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; + let client = get_orm_client_rw().await; Entity::delete_many() .filter(Column::Org.eq(org)) @@ -180,7 +174,7 @@ pub async fn delete(org: &str, table_name: &str) -> Result<(), errors::Error> { /// # Returns /// * `Result, errors::Error>` - List of records or error pub async fn list_by_org(org: &str) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let records = Entity::find() .filter(Column::Org.eq(org)) .all(client) @@ -203,7 +197,7 @@ pub async fn get_all_for_table( ) -> Result, errors::Error> { use sea_orm::QueryOrder; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let records = Entity::find() .filter(Column::Org.eq(org)) .filter(Column::Name.eq(table_name)) @@ -222,7 +216,7 @@ pub async fn get_all_for_table( /// # Returns /// * `Result, errors::Error>` - The record or None pub async fn get_by_id(job_id: &str) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let record = Entity::find() .filter(Column::Id.eq(job_id)) .one(client) @@ -239,8 +233,7 @@ pub async fn get_by_id(job_id: &str) -> Result, /// # Returns /// * `Result<(), errors::Error>` - Success or error pub async fn delete_by_id(job_id: &str) -> Result<(), errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; - let _lock = get_lock().await; + let client = get_orm_client_rw().await; Entity::delete_many() .filter(Column::Id.eq(job_id)) @@ -259,7 +252,7 @@ pub async fn delete_by_id(job_id: &str) -> Result<(), errors::Error> { /// # Returns /// * `Result` - true if any job is processing pub async fn has_processing_jobs(org: &str, table_name: &str) -> Result { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let count = Entity::find() .filter(Column::Org.eq(org)) @@ -280,8 +273,7 @@ pub async fn has_processing_jobs(org: &str, table_name: &str) -> Result` - Success or error pub async fn delete_all_for_table(org: &str, table_name: &str) -> Result<(), errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; - let _lock = get_lock().await; + let client = get_orm_client_rw().await; Entity::delete_many() .filter(Column::Org.eq(org)) @@ -300,10 +292,9 @@ pub async fn delete_all_for_table(org: &str, table_name: &str) -> Result<(), err /// # How it works /// /// ## SQLite (single-node) -/// 1. Uses `get_lock()` for process-level synchronization -/// 2. Finds up to `limit` stale jobs using Sea-ORM queries -/// 3. Updates each job to Pending status sequentially -/// 4. Returns the updated jobs +/// 1. Finds up to `limit` stale jobs using Sea-ORM queries +/// 2. Updates each job to Pending status sequentially +/// 3. Returns the updated jobs /// /// Since SQLite is single-node, only one ingester exists, so no distributed coordination needed. /// @@ -331,13 +322,12 @@ pub async fn claim_stale_jobs( ) -> Result, errors::Error> { use sea_orm::{ConnectionTrait, Statement}; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let backend = client.get_database_backend(); - let _lock = get_lock().await; let now = chrono::Utc::now().timestamp_micros(); - // For SQLite (single-node), use simple Sea-ORM API since get_lock() provides sufficient - // protection + // SQLite serializes writers on the single write connection, so the plain + // Sea-ORM path is enough here; postgres needs the locking query below. if backend == sea_orm::DatabaseBackend::Sqlite { // Find stale jobs (only local region jobs) let stale_jobs: Vec = Entity::find() @@ -405,8 +395,7 @@ pub async fn claim_stale_jobs( /// Deletes all enrichment table URL entries belonging to the given org. pub async fn delete_by_org(org: &str) -> Result<(), errors::Error> { - let _lock = get_lock().await; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::delete_many() .filter(Column::Org.eq(org)) .exec(client) diff --git a/src/infra/src/table/enrichment_tables.rs b/src/infra/src/table/enrichment_tables.rs index 465eafdabd..01181bf2d5 100644 --- a/src/infra/src/table/enrichment_tables.rs +++ b/src/infra/src/table/enrichment_tables.rs @@ -23,11 +23,10 @@ use sea_orm::{ }; use serde::{Deserialize, Serialize}; -use super::get_lock; // Re-export the entity for convenience pub use crate::table::entity::enrichment_tables::{ActiveModel, Column, Entity, Model, Relation}; use crate::{ - db::{ORM_CLIENT, connect_to_orm, get_db}, + db::{get_db, get_orm_client_ro, get_orm_client_rw}, errors, }; @@ -98,10 +97,7 @@ pub async fn add( ..Default::default() }; - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::insert(record).exec(client).await?; Ok(()) @@ -116,10 +112,7 @@ pub async fn add( /// # Returns /// * `Result<(), errors::Error>` - Success or error pub async fn delete(org: &str, table_name: &str) -> Result<(), errors::Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::delete_many() .filter(Column::Org.eq(org)) .filter(Column::Name.eq(table_name)) @@ -141,7 +134,7 @@ pub async fn get_by_org_and_name( org: &str, table_name: &str, ) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let records = Entity::find() .select_only() .column(Column::Org) @@ -178,7 +171,7 @@ pub async fn get_by_org_and_name_with_end_time( table_name: &str, end_time_exclusive: Option, ) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let mut query = Entity::find() .select_only() .column(Column::Org) @@ -252,7 +245,7 @@ pub async fn get_meta_stats(org: &str, table_name: &str) -> Option, errors::Error>` - List of records or error pub async fn get_by_org(org: &str) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let records = Entity::find() .select_only() .column(Column::Org) @@ -277,7 +270,7 @@ pub async fn get_by_org(org: &str) -> Result, errors: /// # Returns /// * `Result` - True if records exist, false otherwise pub async fn contains(org: &str, table_name: &str) -> Result { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let record = Entity::find() .filter(Column::Org.eq(org)) .filter(Column::Name.eq(table_name)) @@ -297,7 +290,7 @@ pub async fn contains(org: &str, table_name: &str) -> Result` - Count of records pub async fn count_by_org_and_name(org: &str, table_name: &str) -> Result { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let count = Entity::find() .filter(Column::Org.eq(org)) .filter(Column::Name.eq(table_name)) @@ -312,7 +305,7 @@ pub async fn count_by_org_and_name(org: &str, table_name: &str) -> Result` - Total count of records pub async fn count() -> Result { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let count = Entity::find().count(client).await?; Ok(count as usize) @@ -323,10 +316,7 @@ pub async fn count() -> Result { /// # Returns /// * `Result<(), errors::Error>` - Success or error pub async fn clear() -> Result<(), errors::Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::delete_many().exec(client).await?; Ok(()) @@ -346,7 +336,7 @@ pub async fn is_empty() -> Result { /// # Returns /// * `Result, errors::Error>` - List of (org_id, table_name) pairs pub async fn list() -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let records = Entity::find() .select_only() .column(Column::Org) @@ -363,8 +353,7 @@ pub async fn list() -> Result, errors::Error> { /// Deletes all enrichment table entries belonging to the given org. pub async fn delete_by_org(org: &str) -> Result<(), errors::Error> { - let _lock = get_lock().await; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::delete_many() .filter(Column::Org.eq(org)) .exec(client) diff --git a/src/infra/src/table/external_alerts.rs b/src/infra/src/table/external_alerts.rs index 45aa03804c..df0f433bc3 100644 --- a/src/infra/src/table/external_alerts.rs +++ b/src/infra/src/table/external_alerts.rs @@ -22,12 +22,9 @@ use config::meta::alerts::incidents::{ExternalAlertEvent, ExternalAlertStatus}; use sea_orm::{ColumnTrait, EntityTrait, QueryFilter, Set, SqlErr, sea_query::Expr}; use svix_ksuid::KsuidLike; -use super::{ - entity::external_alerts::{ActiveModel, Column, Entity, Model}, - get_lock, -}; +use super::entity::external_alerts::{ActiveModel, Column, Entity, Model}; use crate::{ - db::{ORM_CLIENT, connect_to_orm}, + db::{get_orm_client_ro, get_orm_client_rw}, errors::{self, DbError, Error}, }; @@ -144,8 +141,7 @@ pub async fn upsert_event( detected_source: &str, ev: &ExternalAlertEvent, ) -> Result<(ExternalAlertRecord, UpsertOutcome), errors::Error> { - let _lock = get_lock().await; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let existing = Entity::find() .filter(Column::OrgId.eq(org_id)) @@ -236,7 +232,7 @@ async fn insert_new( detected_source: &str, ev: &ExternalAlertEvent, ) -> Result { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let id = svix_ksuid::Ksuid::new(None, None).to_string(); let labels_json = serde_json::to_string(&ev.labels).unwrap_or_else(|_| "{}".into()); let last_payload_json = ev.raw.to_string(); @@ -376,7 +372,7 @@ pub async fn get_by_id( org_id: &str, id: &str, ) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let record = Entity::find() .filter(Column::OrgId.eq(org_id)) .filter(Column::Id.eq(id)) @@ -390,7 +386,7 @@ pub async fn get_by_ids( org_id: &str, ids: &[String], ) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let records = Entity::find() .filter(Column::OrgId.eq(org_id)) .filter(Column::Id.is_in(ids.to_vec())) diff --git a/src/infra/src/table/folders.rs b/src/infra/src/table/folders.rs index fa330d3be4..be332eabfd 100644 --- a/src/infra/src/table/folders.rs +++ b/src/infra/src/table/folders.rs @@ -20,12 +20,9 @@ use sea_orm::{ }; use svix_ksuid::{Ksuid, KsuidLike}; -use super::{ - entity::folders::{ActiveModel, Column, Entity, Model}, - get_lock, -}; +use super::entity::folders::{ActiveModel, Column, Entity, Model}; use crate::{ - db::{ORM_CLIENT, connect_to_orm}, + db::{get_orm_client_ro, get_orm_client_rw}, errors::{self, FromStrError}, }; @@ -56,7 +53,7 @@ pub async fn get( folder_id: &str, folder_type: FolderType, ) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let folder = get_model(client, org_id, folder_id, folder_type) .await .map(|f| f.map(Folder::from))?; @@ -69,7 +66,7 @@ pub async fn get_by_name( folder_name: &str, folder_type: FolderType, ) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let folder = get_model_by_name(client, org_id, folder_name, folder_type) .await .map(|f| f.map(Folder::from))?; @@ -91,7 +88,7 @@ pub async fn list_folders( org_id: &str, folder_type: FolderType, ) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let folders = list_models(client, org_id, folder_type) .await? .into_iter() @@ -108,9 +105,7 @@ pub async fn put( folder: Folder, folder_type: FolderType, ) -> Result<(Ksuid, Folder), errors::Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let model = match get_model(client, org_id, &folder.folder_id, folder_type).await? { // If a folder with the given folder_id already exists, get that folder @@ -158,9 +153,7 @@ pub async fn delete( folder_id: &str, folder_type: FolderType, ) -> Result<(), errors::Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let model = get_model(client, org_id, folder_id, folder_type).await?; if let Some(model) = model { @@ -180,7 +173,7 @@ pub async fn get_pk_by_name( name: &str, folder_type: FolderType, ) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; Ok(get_model(client, org_id, name, folder_type) .await? .map(|m| m.id)) @@ -191,7 +184,7 @@ pub async fn get_pk_by_name( /// Used to translate the stored PK back to the user-visible name when building /// API responses for anomaly detection configs. pub async fn get_name_by_pk(pk: &str) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; Ok(Entity::find_by_id(pk) .one(client) .await? @@ -205,7 +198,7 @@ pub async fn get_name_by_pk(pk: &str) -> Result, errors::Error> { pub async fn get_name_and_display_name_by_pk( pk: &str, ) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; Ok(Entity::find_by_id(pk) .one(client) .await? @@ -258,9 +251,7 @@ async fn list_models( /// Deletes all folders belonging to the given org. pub async fn delete_by_org(org_id: &str) -> Result<(), errors::Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::delete_many() .filter(Column::Org.eq(org_id)) .exec(client) diff --git a/src/infra/src/table/gen_ai_agents.rs b/src/infra/src/table/gen_ai_agents.rs index fe7b7ca497..bd021d1cb3 100644 --- a/src/infra/src/table/gen_ai_agents.rs +++ b/src/infra/src/table/gen_ai_agents.rs @@ -19,9 +19,8 @@ use sea_orm::{ ActiveModelTrait, ColumnTrait, EntityTrait, Order, PaginatorTrait, QueryFilter, QueryOrder, Set, }; -use super::get_lock; use crate::{ - db::{ORM_CLIENT, connect_to_orm}, + db::{get_orm_client_ro, get_orm_client_rw}, errors, table::entity::gen_ai_agents::{ActiveModel, Column, Entity, Model}, }; @@ -79,7 +78,7 @@ pub async fn existing_keys( return Ok(HashSet::new()); } - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let rows = Entity::find() .filter(Column::OrgId.eq(org_id)) .filter(Column::AgentKey.is_in(agent_keys.to_vec())) @@ -90,7 +89,7 @@ pub async fn existing_keys( } pub async fn count_by_org(org_id: &str) -> Result { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; Ok(Entity::find() .filter(Column::OrgId.eq(org_id)) .count(client) @@ -102,8 +101,7 @@ pub async fn upsert_many(records: Vec) -> Result) -> Result Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let mut query = Entity::find().filter(Column::OrgId.eq(org_id)); if let Some(end_time) = filter.end_time { @@ -170,8 +168,7 @@ pub async fn delete( source_stream: Option<&str>, source_stream_type: Option<&str>, ) -> Result { - let _lock = get_lock().await; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let mut query = Entity::delete_many().filter(Column::OrgId.eq(org_id)); if let Some(source_stream) = source_stream { diff --git a/src/infra/src/table/incident_events.rs b/src/infra/src/table/incident_events.rs index 3d1cfbf5d3..749cb55674 100644 --- a/src/infra/src/table/incident_events.rs +++ b/src/infra/src/table/incident_events.rs @@ -16,9 +16,8 @@ use config::meta::alerts::incidents::IncidentEvent; use sea_orm::{ActiveModelTrait, ColumnTrait, EntityTrait, QueryFilter, Set, TransactionTrait}; -use super::get_lock; use crate::{ - db::{ORM_CLIENT, connect_to_orm}, + db::{get_orm_client_ro, get_orm_client_rw}, table::entity::incident_events, }; @@ -112,13 +111,10 @@ fn decode_events(org_id: &str, incident_id: &str, raw: &serde_json::Value) -> Ve /// Initialize events row for a new incident with a Created event pub async fn init(org_id: &str, incident_id: &str) -> Result<(), sea_orm::DbErr> { - let db = ORM_CLIENT.get_or_init(connect_to_orm).await; + let db = get_orm_client_rw().await; let event = IncidentEvent::created(); let events_json = encode_events(&[event])?; - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - let model = incident_events::ActiveModel { org_id: Set(org_id.to_string()), incident_id: Set(incident_id.to_string()), @@ -130,7 +126,7 @@ pub async fn init(org_id: &str, incident_id: &str) -> Result<(), sea_orm::DbErr> /// Get all events for an incident pub async fn get(org_id: &str, incident_id: &str) -> Result, sea_orm::DbErr> { - let db = ORM_CLIENT.get_or_init(connect_to_orm).await; + let db = get_orm_client_ro().await; let row = incident_events::Entity::find() .filter(incident_events::Column::OrgId.eq(org_id)) .filter(incident_events::Column::IncidentId.eq(incident_id)) @@ -226,10 +222,7 @@ pub async fn reconstruct_if_empty( org_id: &str, incident_id: &str, ) -> Result, sea_orm::DbErr> { - let db = ORM_CLIENT.get_or_init(connect_to_orm).await; - - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; + let db = get_orm_client_rw().await; let txn = db.begin().await?; @@ -303,10 +296,7 @@ pub async fn append( incident_id: &str, event: IncidentEvent, ) -> Result<(), sea_orm::DbErr> { - let db = ORM_CLIENT.get_or_init(connect_to_orm).await; - - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; + let db = get_orm_client_rw().await; let txn = db.begin().await?; @@ -353,10 +343,7 @@ pub async fn record_alert( alert_name: &str, triggered_at: i64, ) -> Result<(), sea_orm::DbErr> { - let db = ORM_CLIENT.get_or_init(connect_to_orm).await; - - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; + let db = get_orm_client_rw().await; let txn = db.begin().await?; @@ -415,10 +402,7 @@ pub async fn record_alert( pub async fn delete_by_org(org_id: &str) -> Result<(), sea_orm::DbErr> { use sea_orm::{ColumnTrait, EntityTrait, QueryFilter}; - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let db = ORM_CLIENT.get_or_init(connect_to_orm).await; + let db = get_orm_client_rw().await; incident_events::Entity::delete_many() .filter(incident_events::Column::OrgId.eq(org_id)) .exec(db) diff --git a/src/infra/src/table/incident_integrations.rs b/src/infra/src/table/incident_integrations.rs index a9abec4309..de7fca5eaf 100644 --- a/src/infra/src/table/incident_integrations.rs +++ b/src/infra/src/table/incident_integrations.rs @@ -22,18 +22,15 @@ use sea_orm::{ColumnTrait, EntityTrait, QueryFilter, Set, SqlErr, sea_query::Expr}; -use super::{ - entity::{ - incident_integration_senders::{ - ActiveModel as SenderActiveModel, Column as SenderColumn, Entity as SenderEntity, - Model as SenderModel, - }, - incident_integrations::{ActiveModel, Column, Entity, Model}, +use super::entity::{ + incident_integration_senders::{ + ActiveModel as SenderActiveModel, Column as SenderColumn, Entity as SenderEntity, + Model as SenderModel, }, - get_lock, + incident_integrations::{ActiveModel, Column, Entity, Model}, }; use crate::{ - db::{ORM_CLIENT, connect_to_orm}, + db::{get_orm_client_ro, get_orm_client_rw}, errors::{self, DbError, Error}, }; @@ -136,7 +133,6 @@ pub fn generate_token() -> String { /// Insert a new integration row. Fails with a clear message if the token is /// already in use (the unique constraint). pub async fn add(record: &IncidentIntegrationRecord) -> Result<(), errors::Error> { - let _lock = get_lock().await; let now = chrono::Utc::now().timestamp_micros(); let model = ActiveModel { id: Set(record.id.clone()), @@ -155,7 +151,7 @@ pub async fn add(record: &IncidentIntegrationRecord) -> Result<(), errors::Error created_at: Set(now), updated_at: Set(now), }; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; match Entity::insert(model).exec(client).await { Ok(_) => Ok(()), Err(e) => match e.sql_err() { @@ -175,7 +171,7 @@ pub async fn add(record: &IncidentIntegrationRecord) -> Result<(), errors::Error pub async fn find_by_token( token: &str, ) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let record = Entity::find() .filter(Column::Token.eq(token)) .filter(Column::Enabled.eq(true)) @@ -187,7 +183,7 @@ pub async fn find_by_token( /// List all integrations for an org (enabled + disabled). pub async fn list_by_org(org_id: &str) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let records = Entity::find() .filter(Column::OrgId.eq(org_id)) .all(client) @@ -204,7 +200,7 @@ pub async fn ensure_default_for_org( org_id: &str, created_by: &str, ) -> Result { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let existing = Entity::find() .filter(Column::OrgId.eq(org_id)) .filter(Column::Name.eq(DEFAULT_INTEGRATION_NAME)) @@ -233,7 +229,7 @@ pub async fn ensure_default_for_org( Ok(_) => Ok(record), Err(_) => { // Concurrent create raced us — re-fetch and return the winner. - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let winner = Entity::find() .filter(Column::OrgId.eq(org_id)) .filter(Column::Name.eq(DEFAULT_INTEGRATION_NAME)) @@ -252,9 +248,8 @@ pub async fn ensure_default_for_org( /// Enable or disable an integration by `(org_id, id)`. pub async fn set_enabled(org_id: &str, id: &str, enabled: bool) -> Result<(), errors::Error> { - let _lock = get_lock().await; let now = chrono::Utc::now().timestamp_micros(); - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::update_many() .col_expr(Column::Enabled, Expr::value(enabled)) .col_expr(Column::UpdatedAt, Expr::value(now)) @@ -277,7 +272,6 @@ pub async fn update( name: Option<&str>, destinations: Option<&[String]>, ) -> Result<(), errors::Error> { - let _lock = get_lock().await; let now = chrono::Utc::now().timestamp_micros(); let mut q = Entity::update_many().col_expr(Column::UpdatedAt, Expr::value(now)); if let Some(name) = name { @@ -295,7 +289,7 @@ pub async fn update( }; q = q.col_expr(Column::Destinations, Expr::value(encoded)); } - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; q.filter(Column::OrgId.eq(org_id)) .filter(Column::Id.eq(id)) .exec(client) @@ -306,10 +300,9 @@ pub async fn update( /// Rotate an integration's token to a new random value, returning it. pub async fn rotate_token(org_id: &str, id: &str) -> Result { - let _lock = get_lock().await; let new_token = generate_token(); let now = chrono::Utc::now().timestamp_micros(); - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::update_many() .col_expr(Column::Token, Expr::value(new_token.clone())) .col_expr(Column::UpdatedAt, Expr::value(now)) @@ -347,8 +340,7 @@ pub async fn touch_sender(params: TouchSenderParams<'_>) -> Result<(), errors::E saw_resolved, } = params; - let _lock = get_lock().await; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let mut existing_query = SenderEntity::find() .filter(SenderColumn::IntegrationId.eq(integration_id)) .filter(SenderColumn::DetectedSource.eq(detected_source)); @@ -448,8 +440,7 @@ async fn update_sender( /// Returns `Ok(false)` if no matching row was found (already deleted / wrong /// org), `Ok(true)` on success. pub async fn delete(org_id: &str, id: &str) -> Result { - let _lock = get_lock().await; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let existing = Entity::find() .filter(Column::OrgId.eq(org_id)) @@ -479,7 +470,7 @@ pub async fn delete(org_id: &str, id: &str) -> Result { /// List all observed senders for an integration. pub async fn list_senders(integration_id: &str) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let records = SenderEntity::find() .filter(SenderColumn::IntegrationId.eq(integration_id)) .all(client) diff --git a/src/infra/src/table/kv_store.rs b/src/infra/src/table/kv_store.rs index 0109433c1f..7df3afbf8e 100644 --- a/src/infra/src/table/kv_store.rs +++ b/src/infra/src/table/kv_store.rs @@ -15,18 +15,15 @@ use sea_orm::{ColumnTrait, EntityTrait, QueryFilter, QuerySelect, Set}; -use super::{ - entity::kv_store::{ActiveModel, Column, Entity, Model}, - get_lock, -}; +use super::entity::kv_store::{ActiveModel, Column, Entity, Model}; use crate::{ - db::{ORM_CLIENT, connect_to_orm}, + db::{get_orm_client_ro, get_orm_client_rw}, errors, }; /// Gets a KV value by org_id and key pub async fn get(org_id: &str, key: &str) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let kv_entry = Entity::find() .filter(Column::OrgId.eq(org_id)) .filter(Column::Key.eq(key)) @@ -39,9 +36,7 @@ pub async fn get(org_id: &str, key: &str) -> Result, errors::Error pub async fn set(org_id: &str, key: &str, value: &[u8]) -> Result<(), errors::Error> { // Get client first, then acquire lock to prevent deadlock // (get_or_init may internally acquire locks during connection) - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; - // make sure only one client is writing to the database (only for sqlite) - let _lock = get_lock().await; + let client = get_orm_client_rw().await; let now = chrono::Utc::now().timestamp_micros(); let active_model = ActiveModel { @@ -69,9 +64,7 @@ pub async fn set(org_id: &str, key: &str, value: &[u8]) -> Result<(), errors::Er pub async fn delete(org_id: &str, key: &str) -> Result<(), errors::Error> { // Get client first, then acquire lock to prevent deadlock // (get_or_init may internally acquire locks during connection) - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; - // make sure only one client is writing to the database (only for sqlite) - let _lock = get_lock().await; + let client = get_orm_client_rw().await; Entity::delete_many() .filter(Column::OrgId.eq(org_id)) .filter(Column::Key.eq(key)) @@ -82,7 +75,7 @@ pub async fn delete(org_id: &str, key: &str) -> Result<(), errors::Error> { /// Lists all keys for an org_id with optional prefix filter pub async fn list(org_id: &str, prefix: &str) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let mut query = Entity::find() .filter(Column::OrgId.eq(org_id)) @@ -102,17 +95,14 @@ pub async fn list(org_id: &str, prefix: &str) -> Result, errors::Err /// Clears all KV entries from the table pub async fn clear() -> Result<(), errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; - // make sure only one client is writing to the database (only for sqlite) - let _lock = get_lock().await; + let client = get_orm_client_rw().await; Entity::delete_many().exec(client).await?; Ok(()) } /// Deletes all KV entries belonging to the given org. pub async fn delete_by_org(org_id: &str) -> Result<(), errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; - let _lock = get_lock().await; + let client = get_orm_client_rw().await; Entity::delete_many() .filter(Column::OrgId.eq(org_id)) .exec(client) diff --git a/src/infra/src/table/mod.rs b/src/infra/src/table/mod.rs index d5193b9d05..c2192fe894 100644 --- a/src/infra/src/table/mod.rs +++ b/src/infra/src/table/mod.rs @@ -13,14 +13,10 @@ // You should have received a copy of the GNU Affero General Public License // along with this program. If not, see . -use config::{get_config, meta::meta_store::MetaStore}; use migration::Migrator; use sea_orm_migration::MigratorTrait; -use crate::{ - db::{ORM_CLIENT_DDL, connect_to_orm_ddl, sqlite::CLIENT_RW}, - dist_lock, -}; +use crate::{db::get_orm_client_ddl, dist_lock}; pub mod action_scripts; pub mod alert_composites; @@ -97,7 +93,7 @@ pub async fn init() -> Result<(), anyhow::Error> { pub async fn migrate() -> Result<(), anyhow::Error> { let locker = dist_lock::lock("/database/migration", 0).await?; - let client = ORM_CLIENT_DDL.get_or_init(connect_to_orm_ddl).await; + let client = get_orm_client_ddl().await; // This is a hack to fix the failing alerts migration // For postgres, we need to run the migration that populates the alerts table first. // Otherwise, the `m20250109_092400_recreate_tables_with_ksuids` migration will fail. @@ -111,7 +107,7 @@ pub async fn migrate() -> Result<(), anyhow::Error> { /// Get the index of the migration that populates the alerts table. /// This index is used as the first stage of the migration process. async fn get_alerts_populate_migration_index() -> Result { - let client = ORM_CLIENT_DDL.get_or_init(connect_to_orm_ddl).await; + let client = get_orm_client_ddl().await; let migrations = Migrator::get_pending_migrations(client).await?; let mut index: u32 = 0; for (i, migration) in migrations.iter().enumerate() { @@ -128,7 +124,7 @@ async fn get_alerts_populate_migration_index() -> Result { } pub async fn down(steps: Option) -> Result<(), anyhow::Error> { - let client = ORM_CLIENT_DDL.get_or_init(connect_to_orm_ddl).await; + let client = get_orm_client_ddl().await; Migrator::down(client, steps).await?; Ok(()) } @@ -141,32 +137,6 @@ pub async fn create_user_tables() -> Result<(), anyhow::Error> { Ok(()) } -/// Acquires a lock on the SQLite client if the ORM runs on SQLite. -/// -/// `connect_to_orm` uses the shared SQLite pool for every meta store except -/// PostgreSQL, so the predicate here must match it — keying off -/// `ZO_META_STORE=sqlite` alone would silently disable every lock site for -/// other non-postgres values while the ORM still writes SQLite. -/// -/// The guard wraps one process-wide, non-reentrant mutex serializing every -/// SQLite write (raw meta-store puts, the coordinator, scheduler backends, and -/// all ORM table writers). Acquire it at most once per call chain: the function -/// that owns the write or transaction takes it; helpers that receive an open -/// transaction must not. Drop the guard before calling anything that acquires -/// it internally (coordinator emits, `db::scheduler` calls, other table write -/// functions) — a nested acquisition deadlocks the process. -/// -/// # Returns -/// - `Some(MutexGuard)` if the ORM runs on SQLite -/// - `None` if PostgreSQL is configured -pub async fn get_lock() -> Option>> { - if MetaStore::from(get_config().common.meta_store.as_str()) != MetaStore::PostgreSQL { - Some(CLIENT_RW.lock().await) - } else { - None - } -} - #[macro_export] macro_rules! orm_err { ($e:expr) => { diff --git a/src/infra/src/table/model_pricing.rs b/src/infra/src/table/model_pricing.rs index defaf341c3..96bb09491e 100644 --- a/src/infra/src/table/model_pricing.rs +++ b/src/infra/src/table/model_pricing.rs @@ -26,12 +26,9 @@ use sea_orm::{ }; use crate::{ - db::{ORM_CLIENT, connect_to_orm}, + db::{get_orm_client_ro, get_orm_client_rw}, errors::Error, - table::{ - entity::model_pricing::{ActiveModel, Column, Entity, Model}, - get_lock, - }, + table::entity::model_pricing::{ActiveModel, Column, Entity, Model}, }; impl TryFrom for ModelPricingDefinition { @@ -62,8 +59,7 @@ impl TryFrom for ModelPricingDefinition { } pub async fn put(item: ModelPricingDefinition) -> Result { - let _lock = get_lock().await; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let now = chrono::Utc::now().timestamp_micros(); let tiers_json = json::to_value(&item.tiers)?; @@ -140,7 +136,7 @@ pub async fn put(item: ModelPricingDefinition) -> Result Result, Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; match get_by_org_and_name(client, org_id, name).await? { Some(model) => Ok(Some(model.try_into()?)), None => Ok(None), @@ -148,7 +144,7 @@ pub async fn get(org_id: &str, name: &str) -> Result Result, Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; match Entity::find_by_id(id).one(client).await? { Some(model) => Ok(Some(model.try_into()?)), None => Ok(None), @@ -158,7 +154,7 @@ pub async fn get_by_id(id: &str) -> Result, Error /// Return the distinct org IDs that have at least one model pricing definition. /// Used for startup cache warming. pub async fn list_orgs() -> Result, Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let rows: Vec<(String,)> = Entity::find() .select_only() .column(Column::Org) @@ -170,7 +166,7 @@ pub async fn list_orgs() -> Result, Error> { } pub async fn list(org_id: &str) -> Result, Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; Entity::find() .filter(Column::Org.eq(org_id)) .order_by(Column::Name, sea_orm::Order::Asc) @@ -185,8 +181,7 @@ pub async fn list(org_id: &str) -> Result, Error> { /// Returns `Ok(true)` if a row was deleted, `Ok(false)` if not found. /// Built-in entries (`source = 'built_in'`) are never deleted — returns an error. pub async fn delete_by_id(org_id: &str, id: &str) -> Result { - let _lock = get_lock().await; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; if let Some(model) = Entity::find_by_id(id) .filter(Column::Org.eq(org_id)) .one(client) @@ -207,7 +202,7 @@ pub async fn delete_by_id(org_id: &str, id: &str) -> Result { /// Return all model names belonging to the built-in org. /// Used by the sync job to detect models removed from the upstream source. pub async fn list_built_in_names() -> Result, Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let rows: Vec<(String,)> = Entity::find() .select_only() .column(Column::Name) diff --git a/src/infra/src/table/online_eval_jobs/mod.rs b/src/infra/src/table/online_eval_jobs/mod.rs index 3d33d5b9b4..ab7ec3c826 100644 --- a/src/infra/src/table/online_eval_jobs/mod.rs +++ b/src/infra/src/table/online_eval_jobs/mod.rs @@ -33,9 +33,8 @@ pub mod span_selector; pub use completion_window::*; pub use span_selector::*; -use super::get_lock; use crate::{ - db::{ORM_CLIENT, connect_to_orm}, + db::{get_orm_client_ro, get_orm_client_rw}, errors, table::entity::online_eval_jobs::{ActiveModel, Column, Entity, Model}, }; @@ -407,7 +406,7 @@ impl From for OnlineEvalJob { } pub async fn create_table() -> Result<(), errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let builder = client.get_database_backend(); let schema = Schema::new(builder); @@ -422,8 +421,6 @@ pub async fn create_table() -> Result<(), errors::Error> { } pub async fn add(job: &OnlineEvalJob) -> Result<(), errors::Error> { - let _lock = get_lock().await; - let record = ActiveModel { id: Set(job.id.clone()), org_id: Set(job.org_id.clone()), @@ -457,15 +454,13 @@ pub async fn add(job: &OnlineEvalJob) -> Result<(), errors::Error> { updated_at: Set(job.updated_at), }; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::insert(record).exec(client).await?; Ok(()) } pub async fn update(job: &OnlineEvalJob) -> Result<(), errors::Error> { - let _lock = get_lock().await; - let record = ActiveModel { id: Set(job.id.clone()), org_id: Set(job.org_id.clone()), @@ -499,14 +494,14 @@ pub async fn update(job: &OnlineEvalJob) -> Result<(), errors::Error> { updated_at: Set(job.updated_at), }; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::update(record).exec(client).await?; Ok(()) } pub async fn get(id: &str) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let record = Entity::find() .filter(Column::Id.eq(id)) @@ -518,7 +513,7 @@ pub async fn get(id: &str) -> Result, errors::Error> { } pub async fn get_by_org(id: &str, org_id: &str) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let record = Entity::find() .filter(Column::Id.eq(id)) @@ -531,7 +526,7 @@ pub async fn get_by_org(id: &str, org_id: &str) -> Result, } pub async fn get_all_by_org(org_id: &str) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let records = Entity::find() .filter(Column::OrgId.eq(org_id)) @@ -549,7 +544,7 @@ pub async fn get_by_status( org_id: &str, status: &str, ) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let records = Entity::find() .filter(Column::OrgId.eq(org_id)) @@ -565,7 +560,7 @@ pub async fn get_by_status( } pub async fn get_active_trace_session() -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let records = Entity::find() .filter(Column::Status.eq("active")) @@ -585,7 +580,7 @@ pub async fn get_by_stream( stream: &str, stream_type: &str, ) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let records = Entity::find() .filter(Column::OrgId.eq(org_id)) @@ -609,7 +604,7 @@ pub async fn has_non_archived_by_scorer_ref( org_id: &str, scorer_entity_id: &str, ) -> Result { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let jobs = Entity::find() .filter(Column::OrgId.eq(org_id)) @@ -625,9 +620,7 @@ pub async fn has_non_archived_by_scorer_ref( } pub async fn delete(id: &str) -> Result<(), errors::Error> { - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::delete_many() .filter(Column::Id.eq(id)) .exec(client) @@ -637,9 +630,7 @@ pub async fn delete(id: &str) -> Result<(), errors::Error> { } pub async fn delete_all_by_org(org_id: &str) -> Result<(), errors::Error> { - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::delete_many() .filter(Column::OrgId.eq(org_id)) .exec(client) @@ -649,7 +640,7 @@ pub async fn delete_all_by_org(org_id: &str) -> Result<(), errors::Error> { } pub async fn exists(id: &str) -> Result { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let record = Entity::find().filter(Column::Id.eq(id)).one(client).await?; Ok(record.is_some()) @@ -662,9 +653,7 @@ pub async fn update_status( pipeline_id: Option<&str>, updated_at: i64, ) -> Result<(), errors::Error> { - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let model = Entity::find() .filter(Column::Id.eq(id)) .one(client) diff --git a/src/infra/src/table/org_ai_toolsets.rs b/src/infra/src/table/org_ai_toolsets.rs index e04a1fff7b..cc5d409a4a 100644 --- a/src/infra/src/table/org_ai_toolsets.rs +++ b/src/infra/src/table/org_ai_toolsets.rs @@ -25,9 +25,9 @@ use sea_orm::{ }; use svix_ksuid::{Ksuid, KsuidLike}; -use super::{entity::org_ai_toolsets::*, get_lock}; +use super::entity::org_ai_toolsets::*; use crate::{ - db::{ORM_CLIENT, connect_to_orm}, + db::{get_orm_client_ro, get_orm_client_rw}, errors, table::cipher, }; @@ -130,27 +130,22 @@ pub async fn add(mut entry: OrgToolset) -> Result<(), errors::Error> { updated_at: Set(entry.updated_at), }; - let _lock = get_lock().await; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; match Entity::insert(record).exec(client).await { Ok(_) => {} - Err(e) => { - drop(_lock); - match e.sql_err() { - Some(SqlErr::UniqueConstraintViolation(_)) => { - return Err(errors::Error::DbError(errors::DbError::UniqueViolation)); - } - _ => return Err(e.into()), + Err(e) => match e.sql_err() { + Some(SqlErr::UniqueConstraintViolation(_)) => { + return Err(errors::Error::DbError(errors::DbError::UniqueViolation)); } - } + _ => return Err(e.into()), + }, } - drop(_lock); Ok(()) } /// Fetch a toolset by its KSUID `id`. Returns the record with plaintext `data`. pub async fn get(org: &str, id: &str) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let row = Entity::find() .filter(Column::Org.eq(org)) .filter(Column::Id.eq(id)) @@ -170,7 +165,7 @@ pub async fn get(org: &str, id: &str) -> Result, errors::Erro /// Fetch a toolset by its unique `(org, name)` pair. Returns the record with /// plaintext `data`. pub async fn get_by_name(org: &str, name: &str) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let row = Entity::find() .filter(Column::Org.eq(org)) .filter(Column::Name.eq(name)) @@ -200,7 +195,7 @@ pub async fn list( filter: ListFilter, limit: Option, ) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let mut query = Entity::find() .filter(Column::Org.eq(org)) .order_by(Column::CreatedAt, Order::Desc); @@ -249,23 +244,19 @@ pub async fn update(entry: OrgToolset) -> Result<(), errors::Error> { updated_at: Set(entry.updated_at), }; - let _lock = get_lock().await; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; record.update(client).await?; - drop(_lock); Ok(()) } /// Delete a toolset by its KSUID `id`. pub async fn remove(org: &str, id: &str) -> Result<(), errors::Error> { - let _lock = get_lock().await; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::delete_many() .filter(Column::Org.eq(org)) .filter(Column::Id.eq(id)) .exec(client) .await?; - drop(_lock); Ok(()) } diff --git a/src/infra/src/table/org_cleanup_tasks.rs b/src/infra/src/table/org_cleanup_tasks.rs index ace3d71b6e..75b4f231d1 100644 --- a/src/infra/src/table/org_cleanup_tasks.rs +++ b/src/infra/src/table/org_cleanup_tasks.rs @@ -17,12 +17,9 @@ use sea_orm::{ ColumnTrait, Condition, EntityTrait, Order, QueryFilter, QueryOrder, Set, entity::prelude::*, }; -use super::{ - entity::org_cleanup_tasks::{ActiveModel, Column, Entity, Model}, - get_lock, -}; +use super::entity::org_cleanup_tasks::{ActiveModel, Column, Entity, Model}; use crate::{ - db::{ORM_CLIENT, connect_to_orm}, + db::{get_orm_client_ro, get_orm_client_rw}, errors, }; @@ -39,10 +36,7 @@ pub async fn add_batch(tasks: &[NewCleanupTask]) -> Result<(), errors::Error> { if tasks.is_empty() { return Ok(()); } - // Init the ORM client BEFORE taking the lock: connect_to_orm acquires the - // same lock internally, so locking first can deadlock on the initial connect. - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; - let _lock = get_lock().await; + let client = get_orm_client_rw().await; let now = config::utils::time::now_micros(); let models: Vec = tasks .iter() @@ -73,7 +67,7 @@ pub async fn add_batch(tasks: &[NewCleanupTask]) -> Result<(), errors::Error> { } pub async fn list_pending(max_attempts: i32) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let tasks = Entity::find() .filter( Condition::any() @@ -91,8 +85,7 @@ pub async fn list_pending(max_attempts: i32) -> Result, errors: /// Atomically transition a task from any non-running state to 'running'. /// Returns true if this node won the CAS; false if another node beat us. pub async fn mark_running(id: &str) -> Result { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; - let _lock = get_lock().await; + let client = get_orm_client_rw().await; let now = config::utils::time::now_micros(); let result = Entity::update_many() .col_expr(Column::Status, Expr::value("running")) @@ -107,8 +100,7 @@ pub async fn mark_running(id: &str) -> Result { } pub async fn mark_done(id: &str) -> Result<(), errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; - let _lock = get_lock().await; + let client = get_orm_client_rw().await; let now = config::utils::time::now_micros(); Entity::update_many() .col_expr(Column::Status, Expr::value("done")) @@ -121,8 +113,7 @@ pub async fn mark_done(id: &str) -> Result<(), errors::Error> { } pub async fn mark_failed(id: &str, error: &str) -> Result<(), errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; - let _lock = get_lock().await; + let client = get_orm_client_rw().await; let now = config::utils::time::now_micros(); Entity::update_many() .col_expr(Column::Status, Expr::value("failed")) @@ -141,8 +132,7 @@ pub async fn mark_failed(id: &str, error: &str) -> Result<(), errors::Error> { /// return it again — the deletion would stall forever. `attempts` is left as-is so /// the MAX_ATTEMPTS ceiling still applies across restarts. pub async fn requeue_running() -> Result { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; - let _lock = get_lock().await; + let client = get_orm_client_rw().await; let now = config::utils::time::now_micros(); let result = Entity::update_many() .col_expr(Column::Status, Expr::value("pending")) @@ -158,7 +148,7 @@ pub async fn list_by_org_status( org_id: &str, status: Option<&str>, ) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let mut q = Entity::find().filter(Column::OrgId.eq(org_id)); if let Some(s) = status { q = q.filter(Column::Status.eq(s)); @@ -173,8 +163,7 @@ pub async fn list_by_org_status( /// completes (org gone) or when an org is resurrected — the ledger for a /// nonexistent/restored org has no value. pub async fn delete_by_org(org_id: &str) -> Result { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; - let _lock = get_lock().await; + let client = get_orm_client_rw().await; let result = Entity::delete_many() .filter(Column::OrgId.eq(org_id)) .exec(client) diff --git a/src/infra/src/table/org_ingestion_tokens.rs b/src/infra/src/table/org_ingestion_tokens.rs index 1f9ecaf146..6704ff77cf 100644 --- a/src/infra/src/table/org_ingestion_tokens.rs +++ b/src/infra/src/table/org_ingestion_tokens.rs @@ -25,12 +25,9 @@ use sea_orm::{ }; use serde::{Deserialize, Serialize}; -use super::{ - entity::org_ingestion_tokens::{ActiveModel, Column, Entity, Model}, - get_lock, -}; +use super::entity::org_ingestion_tokens::{ActiveModel, Column, Entity, Model}; use crate::{ - db::{ORM_CLIENT, connect_to_orm}, + db::{get_orm_client_ro, get_orm_client_rw}, errors::{self, DbError, Error}, }; @@ -82,7 +79,7 @@ pub async fn find_default_enabled( return Ok(entry.0.clone()); } - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let record = Entity::find() .filter(Column::OrgId.eq(org_id)) .filter(Column::Enabled.eq(true)) @@ -150,7 +147,7 @@ pub fn generate_token() -> String { pub async fn find_enabled_token_global( token: &str, ) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let record = Entity::find() .filter(Column::Token.eq(token)) .filter(Column::Enabled.eq(true)) @@ -162,7 +159,6 @@ pub async fn find_enabled_token_global( /// Insert a new org ingestion token row. pub async fn add(record: &OrgIngestionTokenRecord) -> Result<(), errors::Error> { - let _lock = get_lock().await; let now = chrono::Utc::now().timestamp_micros(); let model = ActiveModel { id: Set(record.id.clone()), @@ -177,7 +173,7 @@ pub async fn add(record: &OrgIngestionTokenRecord) -> Result<(), errors::Error> updated_at: Set(now), }; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; match Entity::insert(model).exec(client).await { Ok(_) => { // A new token can be the org default, so the cached pick is stale. @@ -202,7 +198,6 @@ pub async fn add(record: &OrgIngestionTokenRecord) -> Result<(), errors::Error> /// on the originating cluster, preserving timestamps. On conflict on `id` the /// mutable columns are overwritten. pub async fn upsert(record: &OrgIngestionTokenRecord) -> Result<(), errors::Error> { - let _lock = get_lock().await; let model = ActiveModel { id: Set(record.id.clone()), org_id: Set(record.org_id.clone()), @@ -216,7 +211,7 @@ pub async fn upsert(record: &OrgIngestionTokenRecord) -> Result<(), errors::Erro updated_at: Set(record.updated_at), }; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::insert(model) .on_conflict( OnConflict::column(Column::Id) @@ -243,7 +238,7 @@ pub async fn find_enabled_token( org_id: &str, token: &str, ) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let record = Entity::find() .filter(Column::OrgId.eq(org_id)) .filter(Column::Token.eq(token)) @@ -257,7 +252,7 @@ pub async fn find_enabled_token( /// List all tokens for an org (token values masked). pub async fn list_by_org(org_id: &str) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let records = Entity::find() .filter(Column::OrgId.eq(org_id)) .order_by(Column::IsDefault, Order::Desc) @@ -283,7 +278,7 @@ pub async fn get_by_name( org_id: &str, name: &str, ) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let record = Entity::find() .filter(Column::OrgId.eq(org_id)) .filter(Column::Name.eq(name)) @@ -296,10 +291,9 @@ pub async fn get_by_name( /// Rotate (regenerate) a token's value. Returns the new token. pub async fn rotate_token(org_id: &str, name: &str) -> Result { - let _lock = get_lock().await; let new_token = generate_token(); let now = chrono::Utc::now().timestamp_micros(); - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::update_many() .col_expr(Column::Token, Expr::value(new_token.clone())) @@ -316,8 +310,7 @@ pub async fn rotate_token(org_id: &str, name: &str) -> Result Result<(), errors::Error> { - let _lock = get_lock().await; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::delete_many() .filter(Column::OrgId.eq(org_id)) @@ -334,8 +327,7 @@ pub async fn delete_by_org(org_id: &str) -> Result<(), errors::Error> { /// Used by super-cluster sync to remove a token (e.g. the old value after a /// rotation) on a receiving cluster. pub async fn remove_by_token(org_id: &str, token: &str) -> Result<(), errors::Error> { - let _lock = get_lock().await; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::delete_many() .filter(Column::OrgId.eq(org_id)) @@ -350,7 +342,7 @@ pub async fn remove_by_token(org_id: &str, token: &str) -> Result<(), errors::Er /// Count tokens for an org. pub async fn count_by_org(org_id: &str) -> Result { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let count = Entity::find() .filter(Column::OrgId.eq(org_id)) .count(client) @@ -362,9 +354,8 @@ pub async fn count_by_org(org_id: &str) -> Result { /// Enable or disable a token by org_id and name. pub async fn set_enabled(org_id: &str, name: &str, enabled: bool) -> Result<(), errors::Error> { - let _lock = get_lock().await; let now = chrono::Utc::now().timestamp_micros(); - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::update_many() .col_expr(Column::Enabled, Expr::value(enabled)) @@ -381,7 +372,7 @@ pub async fn set_enabled(org_id: &str, name: &str, enabled: bool) -> Result<(), /// List all enabled tokens (for cache bootstrapping). Returns (org_id, token, name) tuples. pub async fn list_all_enabled() -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let records = Entity::find() .filter(Column::Enabled.eq(true)) .select_only() diff --git a/src/infra/src/table/org_storage_providers.rs b/src/infra/src/table/org_storage_providers.rs index eb25e5fd30..ac4a634398 100644 --- a/src/infra/src/table/org_storage_providers.rs +++ b/src/infra/src/table/org_storage_providers.rs @@ -19,9 +19,9 @@ use hashbrown::HashMap; use sea_orm::{ColumnTrait, EntityTrait, QueryFilter, Set, sea_query::OnConflict}; use serde::{Deserialize, Serialize}; -use super::{entity::org_storage_providers::*, get_lock}; +use super::entity::org_storage_providers::*; use crate::{ - db::{ORM_CLIENT, connect_to_orm}, + db::{get_orm_client_ro, get_orm_client_rw}, errors, table::cipher, }; @@ -143,8 +143,7 @@ pub async fn add(entry: OrgStorageProvider) -> Result<(), errors::Error> { data: Set(encrypted_data), }; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; - let _lock = get_lock().await; + let client = get_orm_client_rw().await; Entity::insert(model) .on_conflict( @@ -171,15 +170,12 @@ pub async fn get_for_org(org_id: &str) -> Result, err } } - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; - let _lock = get_lock().await; + let client = get_orm_client_rw().await; let res = Entity::find() .filter(Column::OrgId.eq(org_id)) .one(client) .await?; - // we must drop lock here, as the get_dek itself can attempt lock, which results in deadlock - drop(_lock); let dek = cipher::get_dek(org_id).await?; let ret = res @@ -225,7 +221,7 @@ pub fn remove_from_cache(org_id: &str) { } pub async fn list_all() -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let res = Entity::find().all(client).await?; let temp: Vec = res.into_iter().map(|v| v.into()).collect(); @@ -252,8 +248,7 @@ pub async fn prime_cache() -> Result<(), anyhow::Error> { /// Deletes the storage provider entry for the given org. pub async fn delete_by_org(org_id: &str) -> Result<(), errors::Error> { - let _lock = get_lock().await; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::delete_many() .filter(Column::OrgId.eq(org_id)) .exec(client) diff --git a/src/infra/src/table/org_users.rs b/src/infra/src/table/org_users.rs index eea3435c4b..5f0e10449c 100644 --- a/src/infra/src/table/org_users.rs +++ b/src/infra/src/table/org_users.rs @@ -28,15 +28,12 @@ use sea_orm::{ }; use serde::{Deserialize, Serialize}; -use super::{ - entity::{ - org_users::{ActiveModel, Column, Entity, Model}, - organizations, users, - }, - get_lock, +use super::entity::{ + org_users::{ActiveModel, Column, Entity, Model}, + organizations, users, }; use crate::{ - db::{ORM_CLIENT, connect_to_orm}, + db::{get_orm_client_ro, get_orm_client_rw}, errors::{self, DbError, Error}, }; @@ -182,7 +179,7 @@ impl FromQueryResult for OrgUserExpandedRecord { } pub async fn create_table() -> Result<(), errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let builder = client.get_database_backend(); let schema = Schema::new(builder); @@ -243,10 +240,7 @@ pub async fn add_with_flags( allow_static_token: Set(allow_static_token), }; - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; match Entity::insert(record).exec(client).await { Ok(_) => Ok(()), Err(e) => match e.sql_err() { @@ -257,10 +251,7 @@ pub async fn add_with_flags( } pub async fn remove(org_id: &str, email: &str) -> Result<(), errors::Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::delete_many() .filter(Column::OrgId.eq(org_id)) .filter(Expr::expr(Func::lower(Expr::col(Column::Email))).eq(email.to_lowercase())) @@ -272,10 +263,7 @@ pub async fn remove(org_id: &str, email: &str) -> Result<(), errors::Error> { } pub async fn remove_by_user(email: &str) -> Result<(), errors::Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::delete_many() .filter(Expr::expr(Func::lower(Expr::col(Column::Email))).eq(email.to_lowercase())) .exec(client) @@ -292,10 +280,7 @@ pub async fn update( token: &str, rum_token: Option, ) -> Result<(), errors::Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; // There can be only one record with one org_id and email. // Hence the below updates only one record. Entity::update_many() @@ -320,10 +305,7 @@ pub async fn update_rum_token( email: &str, rum_token: &str, ) -> Result<(), errors::Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; // There can be only one record with one org_id and email. // Hence the below updates only one record. Entity::update_many() @@ -342,10 +324,7 @@ pub async fn update_rum_token( } pub async fn update_token(org_id: &str, email: &str, token: &str) -> Result<(), errors::Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; // There can be only one record with one org_id and email. // Hence the below updates only one record. Entity::update_many() @@ -364,7 +343,7 @@ pub async fn update_token(org_id: &str, email: &str, token: &str) -> Result<(), } pub async fn get(org_id: &str, email: &str) -> Result { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let record = Entity::find() .filter(Column::OrgId.eq(org_id)) .filter(Expr::expr(Func::lower(Expr::col(Column::Email))).eq(email.to_lowercase())) @@ -377,7 +356,7 @@ pub async fn get(org_id: &str, email: &str) -> Result Result { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let user = Entity::find() .filter(Column::OrgId.eq(org_id)) .filter(Column::Role.eq(UserRole::Admin as i16)) @@ -393,7 +372,7 @@ pub async fn get_expanded_user_org( org_id: &str, email: &str, ) -> Result { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let record = Entity::find() .filter(Column::OrgId.eq(org_id)) .filter( @@ -431,7 +410,7 @@ pub async fn get_user_by_rum_token( org_id: &str, rum_token: &str, ) -> Result { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let record = Entity::find() .filter(Column::RumToken.eq(rum_token)) .filter(Column::OrgId.eq(org_id)) @@ -463,7 +442,7 @@ pub async fn get_user_by_rum_token( } pub async fn list_users_by_org(org_id: &str) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let records = Entity::find() .filter(Column::OrgId.eq(org_id)) .all(client) @@ -477,7 +456,7 @@ pub async fn list_users_by_org(org_id: &str) -> Result, error } pub async fn list_orgs_by_user(email: &str) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let records = Entity::find() .filter( Expr::expr(Func::lower(Expr::col((Entity, Column::Email)))).eq(email.to_lowercase()), @@ -503,7 +482,7 @@ pub async fn list_orgs_by_user(email: &str) -> Result } pub async fn list(limit: Option) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let mut res = Entity::find().order_by(Column::CreatedAt, Order::Desc); if let Some(limit) = limit { res = res.limit(limit as u64); @@ -520,7 +499,7 @@ pub async fn list(limit: Option) -> Result, errors::Erro } pub async fn len() -> usize { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let len = Entity::find().count(client).await; match len { @@ -533,10 +512,7 @@ pub async fn len() -> usize { } pub async fn clear() -> Result<(), errors::Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::delete_many() .exec(client) .await diff --git a/src/infra/src/table/organizations.rs b/src/infra/src/table/organizations.rs index f2f2dd8958..2e63a1bd93 100644 --- a/src/infra/src/table/organizations.rs +++ b/src/infra/src/table/organizations.rs @@ -24,12 +24,9 @@ use sea_orm::{ QuerySelect, Schema, Set, entity::prelude::*, }; -use super::{ - entity::organizations::{ActiveModel, Column, Entity, Model}, - get_lock, -}; +use super::entity::organizations::{ActiveModel, Column, Entity, Model}; use crate::{ - db::{ORM_CLIENT, connect_to_orm}, + db::{get_orm_client_ro, get_orm_client_rw}, errors::{self, DbError, Error}, }; @@ -106,7 +103,7 @@ pub struct OrgId { } pub async fn create_table() -> Result<(), errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let builder = client.get_database_backend(); let schema = Schema::new(builder); @@ -153,10 +150,7 @@ pub async fn add( trial_ends_at: now + day_micros(15), }; - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; match Entity::insert(record).exec(client).await { Ok(_) => { let mut cache = CACHE.write().await; @@ -172,12 +166,9 @@ pub async fn add( #[cfg(feature = "cloud")] pub async fn set_trial_period_end(org_id: &str, new_end: i64) -> Result<(), errors::Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - let update_time = chrono::Utc::now().timestamp_micros(); - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::update_many() .col_expr(Column::TrialEndsAt, Expr::value(new_end)) .col_expr(Column::UpdatedAt, Expr::value(update_time)) @@ -198,12 +189,9 @@ pub async fn set_trial_period_end(org_id: &str, new_end: i64) -> Result<(), erro } pub async fn rename(org_id: &str, new_name: &str) -> Result<(), errors::Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - let update_time = chrono::Utc::now().timestamp_micros(); - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::update_many() .col_expr(Column::OrgName, Expr::value(new_name.to_string())) .col_expr( @@ -226,10 +214,7 @@ pub async fn rename(org_id: &str, new_name: &str) -> Result<(), errors::Error> { } pub async fn remove(org_id: &str) -> Result<(), errors::Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::delete_many() .filter(Column::Identifier.eq(org_id)) .exec(client) @@ -245,7 +230,7 @@ pub async fn remove(org_id: &str) -> Result<(), errors::Error> { } pub async fn get(org_id: &str) -> Result { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; { let cache = CACHE.read().await; if let Some(v) = cache.get(org_id) { @@ -271,7 +256,7 @@ pub async fn get(org_id: &str) -> Result { } pub async fn list(filter: ListFilter) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let mut res = Entity::find().order_by(Column::CreatedAt, Order::Desc); if let Some(limit) = filter.limit { res = res.limit(limit as u64); @@ -294,7 +279,7 @@ pub async fn list(filter: ListFilter) -> Result, errors: } pub async fn get_by_name(org_name: &str) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let records = Entity::find() .filter(Column::OrgName.eq(org_name)) .all(client) @@ -308,7 +293,7 @@ pub async fn get_by_name(org_name: &str) -> Result, erro } pub async fn len() -> usize { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let len = Entity::find().count(client).await; match len { @@ -321,10 +306,7 @@ pub async fn len() -> usize { } pub async fn clear() -> Result<(), errors::Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::delete_many() .exec(client) .await @@ -338,10 +320,7 @@ pub async fn is_empty() -> bool { } pub async fn batch_remove(org_ids: Vec) -> Result<(), errors::Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::delete_many() .filter(Column::Identifier.is_in(org_ids.clone())) .exec(client) @@ -359,10 +338,7 @@ pub async fn batch_remove(org_ids: Vec) -> Result<(), errors::Error> { } pub async fn set_status(org_id: &str, status: &str) -> Result<(), errors::Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let now = config::utils::time::now_micros(); Entity::update_many() .col_expr(Column::Status, Expr::value(status)) @@ -383,10 +359,7 @@ pub async fn set_status_if( expected_status: &str, new_status: &str, ) -> Result { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let now = config::utils::time::now_micros(); let result = Entity::update_many() .col_expr(Column::Status, Expr::value(new_status)) @@ -413,10 +386,7 @@ pub async fn set_status_if_with_deleted_at( new_status: &str, deleted_at: Option, ) -> Result { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let now = config::utils::time::now_micros(); let result = Entity::update_many() .col_expr(Column::Status, Expr::value(new_status)) diff --git a/src/infra/src/table/providers.rs b/src/infra/src/table/providers.rs index e1b52cc018..e77424f462 100644 --- a/src/infra/src/table/providers.rs +++ b/src/infra/src/table/providers.rs @@ -22,9 +22,8 @@ use sea_orm::{ColumnTrait, ConnectionTrait, EntityTrait, QueryFilter, QueryOrder, Schema, Set}; use serde::{Deserialize, Serialize}; -use super::get_lock; use crate::{ - db::{ORM_CLIENT, connect_to_orm}, + db::{get_orm_client_ro, get_orm_client_rw}, errors, table::{ cipher, @@ -93,7 +92,7 @@ fn model_to_provider(model: Model, dek: &[u8]) -> Result Result<(), errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let builder = client.get_database_backend(); let schema = Schema::new(builder); @@ -117,8 +116,6 @@ pub async fn add(provider: &Provider) -> Result<(), errors::Error> { Some(encrypt_data(&dek, &plaintext)?) }; - let _lock = get_lock().await; - let record = ActiveModel { id: Set(provider.id.clone()), org_id: Set(provider.org_id.clone()), @@ -133,7 +130,7 @@ pub async fn add(provider: &Provider) -> Result<(), errors::Error> { updated_at: Set(provider.updated_at), }; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::insert(record).exec(client).await?; Ok(()) @@ -149,8 +146,6 @@ pub async fn update(provider: &Provider) -> Result<(), errors::Error> { Some(encrypt_data(&dek, &plaintext)?) }; - let _lock = get_lock().await; - let record = ActiveModel { id: Set(provider.id.clone()), org_id: Set(provider.org_id.clone()), @@ -165,14 +160,14 @@ pub async fn update(provider: &Provider) -> Result<(), errors::Error> { updated_at: Set(provider.updated_at), }; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::update(record).exec(client).await?; Ok(()) } pub async fn get(id: &str) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let model = Entity::find().filter(Column::Id.eq(id)).one(client).await?; @@ -186,7 +181,7 @@ pub async fn get(id: &str) -> Result, errors::Error> { } pub async fn get_all_by_org(org_id: &str) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let rows = Entity::find() .filter(Column::OrgId.eq(org_id)) @@ -205,7 +200,7 @@ pub async fn get_all_by_org(org_id: &str) -> Result, errors::Error } pub async fn get_default(org_id: &str) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let model = Entity::find() .filter(Column::OrgId.eq(org_id)) @@ -223,9 +218,7 @@ pub async fn get_default(org_id: &str) -> Result, errors::Error } pub async fn delete(id: &str) -> Result<(), errors::Error> { - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::delete_many() .filter(Column::Id.eq(id)) .exec(client) @@ -235,9 +228,7 @@ pub async fn delete(id: &str) -> Result<(), errors::Error> { } pub async fn delete_all_by_org(org_id: &str) -> Result<(), errors::Error> { - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::delete_many() .filter(Column::OrgId.eq(org_id)) .exec(client) @@ -247,7 +238,7 @@ pub async fn delete_all_by_org(org_id: &str) -> Result<(), errors::Error> { } pub async fn exists(id: &str) -> Result { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let record = Entity::find().filter(Column::Id.eq(id)).one(client).await?; Ok(record.is_some()) diff --git a/src/infra/src/table/ratelimit.rs b/src/infra/src/table/ratelimit.rs index 7da74ae92b..c28e415198 100644 --- a/src/infra/src/table/ratelimit.rs +++ b/src/infra/src/table/ratelimit.rs @@ -25,9 +25,8 @@ use sea_orm::{ }; use serde::{Deserialize, Serialize}; -use super::get_lock; use crate::{ - db::{ORM_CLIENT, connect_to_orm}, + db::{get_orm_client_ro, get_orm_client_rw}, orm_err, table::{ entity::rate_limit_rules::{ActiveModel, Column, Entity}, @@ -58,7 +57,7 @@ pub async fn fetch_rules( org_id: Option, user_role: Option, ) -> Result, anyhow::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let mut res = Entity::find() .select_only() .column(Column::Org) @@ -91,7 +90,7 @@ pub async fn fetch_rules( } pub async fn fetch_rules_by_id(rule_id: &str) -> Result, anyhow::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let res = Entity::find() .select_only() .column(Column::Org) @@ -113,9 +112,6 @@ pub async fn fetch_rules_by_id(rule_id: &str) -> Result, a } pub async fn add(rule: RuleEntry) -> Result<(), anyhow::Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - match rule { RuleEntry::Single(rule) => add_single(rule).await, RuleEntry::Batch(rules) => add_batch(rules).await, @@ -128,7 +124,7 @@ async fn add_batch(rules: Vec) -> Result<(), anyhow::Error> { return Ok(()); } - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; // Collect all rule IDs and validate let rule_ids: Vec = rules @@ -212,7 +208,7 @@ async fn add_upsert_batch(rules: Vec) -> Result<(), anyhow::Error return Ok(()); } - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let current_time = config::utils::time::now_micros(); // Start transaction @@ -320,7 +316,7 @@ async fn add_single(rule: RatelimitRule) -> Result<(), anyhow::Error> { let rule_id = rule.rule_id.clone(); match fetch_rules_by_id(rule_id.unwrap().as_str()).await { Ok(None) => { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let active_rule = ActiveModel { org: Set(rule.org), rule_id: Set(rule.rule_id.unwrap()), @@ -354,9 +350,6 @@ async fn add_single(rule: RatelimitRule) -> Result<(), anyhow::Error> { } pub async fn update(rule: RuleEntry) -> Result<(), anyhow::Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - match rule { RuleEntry::Single(rule) => update_single(rule).await, RuleEntry::Batch(rules) => update_batch(rules).await, @@ -372,7 +365,7 @@ async fn update_single(rule: RatelimitRule) -> Result<(), anyhow::Error> { match fetch_rules_by_id(rule_id.as_str()).await { Ok(Some(_)) => { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; match Entity::update_many() .col_expr(Column::Org, Expr::value(&rule.org)) .col_expr( @@ -418,7 +411,7 @@ async fn update_batch(rules: Vec) -> Result<(), anyhow::Error> { return Ok(()); } - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let txn = client .begin() .await @@ -506,12 +499,9 @@ async fn update_batch(rules: Vec) -> Result<(), anyhow::Error> { } pub async fn delete(rule_id: String) -> Result<(), anyhow::Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - match fetch_rules_by_id(rule_id.as_str()).await { Ok(Some(_)) => { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; match Entity::delete_many() .filter(Column::RuleId.eq(rule_id)) .exec(client) @@ -536,7 +526,7 @@ pub async fn list( api: Option<&str>, user_id: Option<&str>, ) -> Result, anyhow::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let mut query = Entity::find() .select_only() .column(Column::Org) @@ -611,7 +601,6 @@ mod tests { .append_query_results(vec![vec![create_test_model()]]) .into_connection(); - // Override the ORM_CLIENT for testing let rules = Entity::find() .select_only() .column(Column::Org) @@ -766,9 +755,9 @@ mod tests { async fn setup_test_db(mock_db: sea_orm::DatabaseConnection) { INIT.call_once(|| { // Another test in the shared process may already have initialized - // ORM_CLIENT; the validation paths under test never reach a query, - // so losing the race is fine. - let _ = ORM_CLIENT.set(mock_db); + // the write client; the validation paths under test never reach a + // query, so losing the race is fine. + let _ = crate::db::ORM_CLIENT_RW.set(mock_db); }); } diff --git a/src/infra/src/table/re_pattern.rs b/src/infra/src/table/re_pattern.rs index c0af1c6395..c469444f7c 100644 --- a/src/infra/src/table/re_pattern.rs +++ b/src/infra/src/table/re_pattern.rs @@ -17,9 +17,9 @@ use config::ider; use sea_orm::{ActiveValue::Unchanged, ColumnTrait, EntityTrait, QueryFilter, Set, SqlErr}; use serde::{Deserialize, Serialize}; -use super::{entity::re_patterns::*, get_lock}; +use super::entity::re_patterns::*; use crate::{ - db::{ORM_CLIENT, connect_to_orm}, + db::{get_orm_client_ro, get_orm_client_rw}, errors, }; @@ -78,24 +78,18 @@ pub async fn add(entry: PatternEntry) -> Result<(), errors::Error> { pattern: Set(entry.pattern), }; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; + let client = get_orm_client_rw().await; match Entity::insert(record).exec(client).await { Ok(_) => {} - Err(e) => { - drop(_lock); - match e.sql_err() { - Some(SqlErr::UniqueConstraintViolation(_)) => { - return Err(errors::Error::DbError(errors::DbError::UniqueViolation)); - } - _ => { - return Err(e.into()); - } + Err(e) => match e.sql_err() { + Some(SqlErr::UniqueConstraintViolation(_)) => { + return Err(errors::Error::DbError(errors::DbError::UniqueViolation)); } - } + _ => { + return Err(e.into()); + } + }, } - drop(_lock); Ok(()) } @@ -109,35 +103,27 @@ pub async fn update_pattern(id: &str, new_pattern: &str) -> Result<(), errors::E ..Default::default() }; - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::update(record) .filter(Column::Id.eq(id)) .exec(client) .await?; - drop(_lock); Ok(()) } pub async fn remove(id: &str) -> Result<(), errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; + let client = get_orm_client_rw().await; Entity::delete_many() .filter(Column::Id.eq(id)) .exec(client) .await?; - drop(_lock); - Ok(()) } pub async fn get(id: &str) -> Result, errors::DbError> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; Entity::find() .filter(Column::Id.eq(id)) .into_model::() @@ -148,7 +134,7 @@ pub async fn get(id: &str) -> Result, errors::DbError> { } pub async fn list_by_org(org: &str) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let res = Entity::find() .filter(Column::Org.eq(org)) @@ -164,7 +150,7 @@ pub async fn list_by_org(org: &str) -> Result, errors::Error> } pub async fn list_all() -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let records = Entity::find().into_model::().all(client).await?; @@ -176,9 +162,7 @@ pub async fn list_all() -> Result, errors::Error> { } pub async fn clear() -> Result<(), errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; + let client = get_orm_client_rw().await; Entity::delete_many().exec(client).await?; Ok(()) @@ -186,8 +170,7 @@ pub async fn clear() -> Result<(), errors::Error> { /// Deletes all re_pattern entries belonging to the given org. pub async fn delete_by_org(org: &str) -> Result<(), errors::Error> { - let _lock = get_lock().await; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::delete_many() .filter(Column::Org.eq(org)) .exec(client) diff --git a/src/infra/src/table/re_pattern_stream_map.rs b/src/infra/src/table/re_pattern_stream_map.rs index f9316d76a3..d9c06ce871 100644 --- a/src/infra/src/table/re_pattern_stream_map.rs +++ b/src/infra/src/table/re_pattern_stream_map.rs @@ -17,9 +17,9 @@ use config::meta::stream::StreamType; use sea_orm::{ColumnTrait, DbErr, EntityTrait, QueryFilter, Set, SqlErr, TransactionTrait}; use serde::{Deserialize, Serialize}; -use super::{entity::re_pattern_stream_map::*, get_lock}; +use super::entity::re_pattern_stream_map::*; use crate::{ - db::{ORM_CLIENT, connect_to_orm}, + db::{get_orm_client_ro, get_orm_client_rw}, errors, }; @@ -124,24 +124,18 @@ pub async fn add(entry: PatternAssociationEntry) -> Result<(), errors::Error> { ..Default::default() }; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; + let client = get_orm_client_rw().await; match Entity::insert(record).exec(client).await { Ok(_) => {} - Err(e) => { - drop(_lock); - match e.sql_err() { - Some(SqlErr::UniqueConstraintViolation(_)) => { - return Err(errors::Error::DbError(errors::DbError::UniqueViolation)); - } - _ => { - return Err(e.into()); - } + Err(e) => match e.sql_err() { + Some(SqlErr::UniqueConstraintViolation(_)) => { + return Err(errors::Error::DbError(errors::DbError::UniqueViolation)); } - } + _ => { + return Err(e.into()); + } + }, } - drop(_lock); Ok(()) } @@ -150,9 +144,7 @@ pub async fn batch_process( added: Vec, removed: Vec, ) -> Result<(), errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; + let client = get_orm_client_rw().await; let txn = client.begin().await?; // we MUST first remove the entries and then add. This is because @@ -205,7 +197,7 @@ pub async fn batch_process( pub async fn get_by_pattern_id( pattern_id: &str, ) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let res = Entity::find() .filter(Column::PatternId.eq(pattern_id)) .into_model::() @@ -219,7 +211,7 @@ pub async fn get_by_pattern_id( } pub async fn list_all() -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let records = Entity::find().into_model::().all(client).await?; @@ -235,9 +227,7 @@ pub async fn remove_associations_by_stream( stream: &str, stype: StreamType, ) -> Result<(), errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; + let client = get_orm_client_rw().await; Entity::delete_many() .filter(Column::Org.eq(org)) .filter(Column::Stream.eq(stream)) @@ -245,15 +235,11 @@ pub async fn remove_associations_by_stream( .exec(client) .await?; - drop(_lock); - Ok(()) } pub async fn clear() -> Result<(), errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; + let client = get_orm_client_rw().await; Entity::delete_many().exec(client).await?; Ok(()) @@ -261,8 +247,7 @@ pub async fn clear() -> Result<(), errors::Error> { /// Deletes all re_pattern stream map entries belonging to the given org. pub async fn delete_by_org(org: &str) -> Result<(), errors::Error> { - let _lock = get_lock().await; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::delete_many() .filter(Column::Org.eq(org)) .exec(client) diff --git a/src/infra/src/table/reports/mod.rs b/src/infra/src/table/reports/mod.rs index 88b183b7c9..47095ac211 100644 --- a/src/infra/src/table/reports/mod.rs +++ b/src/infra/src/table/reports/mod.rs @@ -60,10 +60,9 @@ pub async fn get_by_name( folder_snowflake_id: &str, report_name: &str, ) -> Result, Error> { - let _lock = super::get_lock().await; let txn = conn.begin().await?; let Some(models) = queries::SelectReportAndJoinRelationsResult::get( - conn, + &txn, org_id, folder_snowflake_id, report_name, @@ -81,10 +80,9 @@ pub async fn get_by_id( conn: &C, report_id: &str, ) -> Result, Error> { - let _lock = super::get_lock().await; let txn = conn.begin().await?; let Some(models) = - queries::SelectReportAndJoinRelationsResult::get_by_id(conn, report_id).await? + queries::SelectReportAndJoinRelationsResult::get_by_id(&txn, report_id).await? else { return Ok(None); }; @@ -105,7 +103,6 @@ pub async fn create_report( report: MetaReport, report_id: Option, ) -> Result<(String, MetaReport), Error> { - let _lock = super::get_lock().await; let txn = conn.begin().await?; let report_id = report_id @@ -201,7 +198,6 @@ pub async fn update_report( new_folder_snowflake_id: Option<&str>, report: MetaReport, ) -> Result { - let _lock = super::get_lock().await; let txn = conn.begin().await?; let Some(models) = queries::SelectReportAndJoinRelationsResult::get( @@ -304,7 +300,6 @@ pub async fn delete_by_id( conn: &C, report_id: &str, ) -> Result { - let _lock = super::get_lock().await; let txn = conn.begin().await?; let Some(report_model) = reports::Entity::find() .filter(reports::Column::Id.eq(report_id)) @@ -328,11 +323,10 @@ pub async fn update_by_id( new_folder_snowflake_id: Option<&str>, report: MetaReport, ) -> Result { - let _lock = super::get_lock().await; let txn = conn.begin().await?; let Some(models) = - queries::SelectReportAndJoinRelationsResult::get_by_id(conn, report_id).await? + queries::SelectReportAndJoinRelationsResult::get_by_id(&txn, report_id).await? else { return Err(Error::ReportNotFound); }; @@ -423,10 +417,9 @@ pub async fn delete_by_name( folder_snowflake_id: &str, report_name: &str, ) -> Result { - let _lock = super::get_lock().await; let txn = conn.begin().await?; let Some((_folder_model, Some(report_model))) = - queries::get_report_from_folder(conn, org_id, folder_snowflake_id, report_name).await? + queries::get_report_from_folder(&txn, org_id, folder_snowflake_id, report_name).await? else { return Ok(String::new()); }; @@ -438,17 +431,13 @@ pub async fn delete_by_name( /// Delete all reports. pub async fn delete_all(conn: &C) -> Result<(), Error> { - let _lock = super::get_lock().await; reports::Entity::delete_many().exec(conn).await?; Ok(()) } /// Deletes all reports belonging to the given org. pub async fn delete_by_org(org_id: &str) -> Result<(), Error> { - let _lock = super::get_lock().await; - let client = crate::db::ORM_CLIENT - .get_or_init(crate::db::connect_to_orm) - .await; + let client = crate::db::get_orm_client_rw().await; reports::Entity::delete_many() .filter(reports::Column::Org.eq(org_id)) .exec(client) @@ -461,7 +450,6 @@ pub async fn list_reports( conn: &C, params: &ListReportsParams, ) -> Result, Error> { - let _lock = super::get_lock().await; let reports = queries::ListReportsQueryResult::get(conn, params).await?; Ok(reports) } diff --git a/src/infra/src/table/score_configs.rs b/src/infra/src/table/score_configs.rs index 94f7b7b3fd..78b5f7546c 100644 --- a/src/infra/src/table/score_configs.rs +++ b/src/infra/src/table/score_configs.rs @@ -18,9 +18,8 @@ use sea_orm::{ }; use serde::{Deserialize, Serialize}; -use super::get_lock; use crate::{ - db::{ORM_CLIENT, connect_to_orm}, + db::{get_orm_client_ro, get_orm_client_rw}, errors, table::entity::score_configs::{ActiveModel, Column, Entity, Model}, }; @@ -101,7 +100,7 @@ impl From for ScoreConfig { } pub async fn create_table() -> Result<(), errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let builder = client.get_database_backend(); let schema = Schema::new(builder); @@ -116,8 +115,6 @@ pub async fn create_table() -> Result<(), errors::Error> { } pub async fn add(config: &ScoreConfig) -> Result<(), errors::Error> { - let _lock = get_lock().await; - let record = ActiveModel { id: Set(config.id.clone()), org_id: Set(config.org_id.clone()), @@ -134,15 +131,13 @@ pub async fn add(config: &ScoreConfig) -> Result<(), errors::Error> { updated_at: Set(config.updated_at), }; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::insert(record).exec(client).await?; Ok(()) } pub async fn update(config: &ScoreConfig) -> Result<(), errors::Error> { - let _lock = get_lock().await; - let record = ActiveModel { id: Set(config.id.clone()), org_id: Set(config.org_id.clone()), @@ -159,14 +154,14 @@ pub async fn update(config: &ScoreConfig) -> Result<(), errors::Error> { updated_at: Set(config.updated_at), }; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::update(record).exec(client).await?; Ok(()) } pub async fn get(id: &str) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let record = Entity::find() .filter(Column::Id.eq(id)) @@ -181,7 +176,7 @@ pub async fn get_by_entity_id( org_id: &str, entity_id: &str, ) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let record = Entity::find() .filter(Column::OrgId.eq(org_id)) @@ -200,7 +195,7 @@ pub async fn get_by_entity_id_and_version( entity_id: &str, version: i32, ) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let record = Entity::find() .filter(Column::OrgId.eq(org_id)) @@ -217,7 +212,7 @@ pub async fn get_by_active_name( org_id: &str, name: &str, ) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let record = Entity::find() .filter(Column::OrgId.eq(org_id)) @@ -235,7 +230,7 @@ pub async fn get_versions( org_id: &str, entity_id: &str, ) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let records = Entity::find() .filter(Column::OrgId.eq(org_id)) @@ -251,7 +246,7 @@ pub async fn get_versions( } pub async fn get_all_by_org(org_id: &str) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; // Get the latest version per logical entity. let records: Vec = Entity::find() @@ -274,9 +269,7 @@ pub async fn get_all_by_org(org_id: &str) -> Result, errors::Er } pub async fn delete(entity_id: &str, org_id: &str) -> Result<(), errors::Error> { - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; // Deactivate all versions let records = Entity::find() .filter(Column::OrgId.eq(org_id)) @@ -297,9 +290,7 @@ pub async fn delete(entity_id: &str, org_id: &str) -> Result<(), errors::Error> } pub async fn delete_all_by_org(org_id: &str) -> Result<(), errors::Error> { - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::delete_many() .filter(Column::OrgId.eq(org_id)) .exec(client) @@ -309,7 +300,7 @@ pub async fn delete_all_by_org(org_id: &str) -> Result<(), errors::Error> { } pub async fn exists(id: &str) -> Result { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let record = Entity::find().filter(Column::Id.eq(id)).one(client).await?; Ok(record.is_some()) @@ -317,7 +308,7 @@ pub async fn exists(id: &str) -> Result { /// Returns the latest version number for a given entity, or 0 if none exist. pub async fn get_latest_version(org_id: &str, entity_id: &str) -> Result { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let model = Entity::find() .filter(Column::OrgId.eq(org_id)) diff --git a/src/infra/src/table/scorers.rs b/src/infra/src/table/scorers.rs index 86fa60c88f..153edca520 100644 --- a/src/infra/src/table/scorers.rs +++ b/src/infra/src/table/scorers.rs @@ -19,10 +19,9 @@ use sea_orm::{ }; use serde::{Deserialize, Serialize}; -use super::get_lock; pub use crate::table::entity::scorers::ScorerType; use crate::{ - db::{ORM_CLIENT, connect_to_orm}, + db::{get_orm_client_ro, get_orm_client_rw}, errors, table::entity::scorers::{ActiveModel, Column, Entity, Model}, }; @@ -69,7 +68,7 @@ impl From for Scorer { } pub async fn create_table() -> Result<(), errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let builder = client.get_database_backend(); let schema = Schema::new(builder); @@ -84,8 +83,6 @@ pub async fn create_table() -> Result<(), errors::Error> { } pub async fn add(scorer: &Scorer) -> Result<(), errors::Error> { - let _lock = get_lock().await; - let record = ActiveModel { id: Set(scorer.id.clone()), org_id: Set(scorer.org_id.clone()), @@ -104,15 +101,13 @@ pub async fn add(scorer: &Scorer) -> Result<(), errors::Error> { updated_at: Set(scorer.updated_at), }; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::insert(record).exec(client).await?; Ok(()) } pub async fn update(scorer: &Scorer) -> Result<(), errors::Error> { - let _lock = get_lock().await; - let record = ActiveModel { id: Set(scorer.id.clone()), org_id: Set(scorer.org_id.clone()), @@ -131,14 +126,14 @@ pub async fn update(scorer: &Scorer) -> Result<(), errors::Error> { updated_at: Set(scorer.updated_at), }; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::update(record).exec(client).await?; Ok(()) } pub async fn get(id: &str) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let record = Entity::find() .filter(Column::Id.eq(id)) @@ -153,7 +148,7 @@ pub async fn get_by_entity_id( org_id: &str, entity_id: &str, ) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let record = Entity::find() .filter(Column::OrgId.eq(org_id)) @@ -172,7 +167,7 @@ pub async fn get_by_entity_id_and_version( entity_id: &str, version: i32, ) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let record = Entity::find() .filter(Column::OrgId.eq(org_id)) @@ -186,7 +181,7 @@ pub async fn get_by_entity_id_and_version( } pub async fn get_by_active_name(org_id: &str, name: &str) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let record = Entity::find() .filter(Column::OrgId.eq(org_id)) @@ -201,7 +196,7 @@ pub async fn get_by_active_name(org_id: &str, name: &str) -> Result Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let records = Entity::find() .filter(Column::OrgId.eq(org_id)) @@ -217,7 +212,7 @@ pub async fn get_versions(org_id: &str, entity_id: &str) -> Result, } pub async fn get_all_by_org(org_id: &str) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let records: Vec = Entity::find() .filter(Column::OrgId.eq(org_id)) @@ -242,7 +237,7 @@ pub async fn get_by_type( org_id: &str, scorer_type: &ScorerType, ) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let records: Vec = Entity::find() .filter(Column::OrgId.eq(org_id)) @@ -268,7 +263,7 @@ pub async fn has_active_by_score_config_id( org_id: &str, score_config_entity_id: &str, ) -> Result { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let count = Entity::find() .filter(Column::OrgId.eq(org_id)) @@ -281,9 +276,7 @@ pub async fn has_active_by_score_config_id( } pub async fn delete(entity_id: &str, org_id: &str) -> Result<(), errors::Error> { - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; // Deactivate all versions let records = Entity::find() .filter(Column::OrgId.eq(org_id)) @@ -304,9 +297,7 @@ pub async fn delete(entity_id: &str, org_id: &str) -> Result<(), errors::Error> } pub async fn delete_all_by_org(org_id: &str) -> Result<(), errors::Error> { - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::delete_many() .filter(Column::OrgId.eq(org_id)) .exec(client) @@ -316,7 +307,7 @@ pub async fn delete_all_by_org(org_id: &str) -> Result<(), errors::Error> { } pub async fn exists(id: &str) -> Result { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let record = Entity::find().filter(Column::Id.eq(id)).one(client).await?; Ok(record.is_some()) @@ -324,7 +315,7 @@ pub async fn exists(id: &str) -> Result { /// Returns the latest version number for a given entity, or 0 if none exist. pub async fn get_latest_version(org_id: &str, entity_id: &str) -> Result { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let model = Entity::find() .filter(Column::OrgId.eq(org_id)) diff --git a/src/infra/src/table/search_job/search_job_partitions.rs b/src/infra/src/table/search_job/search_job_partitions.rs index 6c4559dbfc..665302fccc 100644 --- a/src/infra/src/table/search_job/search_job_partitions.rs +++ b/src/infra/src/table/search_job/search_job_partitions.rs @@ -21,11 +21,11 @@ use sea_orm::{ use serde::{Deserialize, Serialize}; use super::{ - super::{entity::search_job_partitions::*, get_lock}, + super::entity::search_job_partitions::*, common::{OperatorType, Value}, }; use crate::{ - db::{ORM_CLIENT, connect_to_orm}, + db::{get_orm_client_ro, get_orm_client_rw}, errors, orm_err, }; @@ -127,9 +127,7 @@ pub async fn submit_partitions(job_id: &str, jobs: Vec) -> Result<(), err return orm_err!("partitions array cannot be empty"); } - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let tx = match client.begin().await { Ok(tx) => tx, Err(e) => return orm_err!(format!("submit partition job start transaction error: {e}")), @@ -178,7 +176,7 @@ pub async fn submit_partitions(job_id: &str, jobs: Vec) -> Result<(), err } pub async fn get_partition_jobs(job_id: &str) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; // sql: select * from search_job_partitions where job_id = job_id let res = Entity::find() @@ -194,10 +192,7 @@ pub async fn get_partition_jobs(job_id: &str) -> Result, errors::Erro } pub async fn set(operator: SetOperator) -> Result<(), errors::Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let mut query = Entity::update_many(); @@ -229,10 +224,7 @@ pub async fn set(operator: SetOperator) -> Result<(), errors::Error> { } pub async fn clean_deleted_partition_job(job_id: &str) -> Result<(), errors::Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let res = Entity::delete_many() .filter(Column::JobId.eq(job_id)) diff --git a/src/infra/src/table/search_job/search_job_results.rs b/src/infra/src/table/search_job/search_job_results.rs index 4c288d8d16..e6d85d9e09 100644 --- a/src/infra/src/table/search_job/search_job_results.rs +++ b/src/infra/src/table/search_job/search_job_results.rs @@ -16,9 +16,9 @@ use sea_orm::{ColumnTrait, EntityTrait, Order, QueryFilter, QueryOrder}; use serde::{Deserialize, Serialize}; -use super::super::{entity::search_job_results::*, get_lock}; +use super::super::entity::search_job_results::*; use crate::{ - db::{ORM_CLIENT, connect_to_orm}, + db::{get_orm_client_ro, get_orm_client_rw}, errors, orm_err, }; @@ -28,10 +28,7 @@ pub enum JobResultOperator { } pub async fn get(job_id: &str) -> Result, errors::Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let res = Entity::find() .filter(Column::JobId.eq(job_id)) @@ -46,10 +43,7 @@ pub async fn get(job_id: &str) -> Result, errors::Error> { } pub async fn clean_deleted_job_result(job_id: &str) -> Result<(), errors::Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let res = Entity::delete_many() .filter(Column::JobId.eq(job_id)) diff --git a/src/infra/src/table/search_job/search_jobs.rs b/src/infra/src/table/search_job/search_jobs.rs index 63f05156a4..e6dcd4a09c 100644 --- a/src/infra/src/table/search_job/search_jobs.rs +++ b/src/infra/src/table/search_job/search_jobs.rs @@ -23,20 +23,17 @@ use sea_orm::{ use serde::{Deserialize, Serialize}; use super::{ - super::{ - entity::{ - search_job_partitions::{Column as PartitionJobColumn, Entity as PartitionJobEntity}, - search_job_results::{ - ActiveModel as JobResultModel, Column as JobResultColumn, Entity as JobResultEntity, - }, - search_jobs::*, + super::entity::{ + search_job_partitions::{Column as PartitionJobColumn, Entity as PartitionJobEntity}, + search_job_results::{ + ActiveModel as JobResultModel, Column as JobResultColumn, Entity as JobResultEntity, }, - get_lock, + search_jobs::*, }, common::{OperatorType, Value}, }; use crate::{ - db::{ORM_CLIENT, connect_to_orm}, + db::{get_orm_client_ro, get_orm_client_rw}, errors, orm_err, }; @@ -167,10 +164,7 @@ pub struct SetOperator { #[allow(clippy::too_many_arguments)] pub async fn submit(job: ActiveModel) -> Result<(), errors::Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let _res = match Entity::insert(job).exec(client).await { Ok(res) => res, Err(e) => return orm_err!(format!("submit search job error: {e}")), @@ -181,10 +175,7 @@ pub async fn submit(job: ActiveModel) -> Result<(), errors::Error> { // get the job and update status pub async fn get_job(updated_at: i64) -> Result, errors::Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let tx = match client.begin().await { Ok(tx) => tx, @@ -234,10 +225,7 @@ pub async fn get_job(updated_at: i64) -> Result, errors::Error> { } pub async fn cancel_job(job_id: &str, update_at: i64) -> Result { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let tx = match client.begin().await { Ok(tx) => tx, @@ -298,10 +286,7 @@ pub async fn cancel_job(job_id: &str, update_at: i64) -> Result Result { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let mut query = Entity::update_many(); @@ -333,10 +318,7 @@ pub async fn set(operator: SetOperator) -> Result { } pub async fn clean_deleted_job(job_id: &str) -> Result<(), errors::Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let res = Entity::delete_many() .filter(Column::Id.eq(job_id)) @@ -354,10 +336,7 @@ pub async fn clean_deleted_job(job_id: &str) -> Result<(), errors::Error> { /// Child rows are deleted explicitly because SQLite does not enforce FK cascades /// unless `PRAGMA foreign_keys = ON` is set. pub async fn delete_by_org(org_id: &str) -> Result<(), errors::Error> { - // make sure only one client is writing to the database (only for SQLite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; // Collect job IDs for this org let job_ids: Vec = Entity::find() @@ -422,10 +401,7 @@ pub async fn retry_search_job( new_trace_id: &str, updated_at: i64, ) -> Result<(), errors::Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let tx = match client.begin().await { Ok(tx) => tx, @@ -523,7 +499,7 @@ pub async fn retry_search_job( } pub async fn get(job_id: &str, org_id: &str) -> Result { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let res = Entity::find() .filter(Column::Id.eq(job_id)) .filter(Column::OrgId.eq(org_id)) @@ -538,7 +514,7 @@ pub async fn get(job_id: &str, org_id: &str) -> Result { } pub async fn list_status_by_org_id(org_id: &str) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let res = Entity::find() .filter(Column::OrgId.eq(org_id)) .filter(Column::Status.ne(4)) @@ -555,7 +531,7 @@ pub async fn list_status_by_org_id(org_id: &str) -> Result, errors::E } pub async fn get_deleted_jobs() -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let res = Entity::find() .filter(Column::Status.eq(4)) @@ -601,10 +577,7 @@ pub async fn set_job_start( node: &str, updated_at: i64, ) -> Result<(), errors::Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let res = Entity::update_many() .col_expr(Column::Status, Expr::value(1)) diff --git a/src/infra/src/table/search_queue.rs b/src/infra/src/table/search_queue.rs index 829272a4c7..826b1060fc 100644 --- a/src/infra/src/table/search_queue.rs +++ b/src/infra/src/table/search_queue.rs @@ -19,9 +19,9 @@ use sea_orm::{ Statement, }; -use super::{entity::search_queue::*, get_lock}; +use super::entity::search_queue::*; use crate::{ - db::{ORM_CLIENT, connect_to_orm}, + db::{get_orm_client_ro, get_orm_client_rw}, errors, }; @@ -40,20 +40,14 @@ pub async fn add( ..Default::default() }; - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::insert(record).exec(client).await?; Ok(()) } pub async fn delete(id: i64) -> Result<(), errors::Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::delete_many() .filter(Column::Id.eq(id)) .exec(client) @@ -63,10 +57,7 @@ pub async fn delete(id: i64) -> Result<(), errors::Error> { } pub async fn delete_by_trace_id(trace_id: &str) -> Result<(), errors::Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::delete_many() .filter(Column::TraceId.eq(trace_id)) .exec(client) @@ -76,7 +67,7 @@ pub async fn delete_by_trace_id(trace_id: &str) -> Result<(), errors::Error> { } pub async fn count(work_group: &str, user_id: Option<&str>) -> Result { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let mut query = Entity::find().filter(Column::WorkGroup.eq(work_group)); if let Some(user_id) = user_id { query = query.filter(Column::UserId.eq(user_id)); @@ -103,7 +94,7 @@ pub async fn count_all_levels( org_id: Option<&str>, user_id: Option<&str>, ) -> Result<(usize, usize, usize), errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let backend = client.get_database_backend(); // Build the SQL query using window functions for efficient counting @@ -152,20 +143,14 @@ pub async fn count_all_levels( } pub async fn clear() -> Result<(), errors::Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::delete_many().exec(client).await?; Ok(()) } pub async fn clean_incomplete(expired: i64) -> Result<(), errors::Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::delete_many() .filter(Column::CreatedAt.lt(expired)) .exec(client) @@ -176,8 +161,7 @@ pub async fn clean_incomplete(expired: i64) -> Result<(), errors::Error> { /// Deletes all search queue entries belonging to the given org. pub async fn delete_by_org(org_id: &str) -> Result<(), errors::Error> { - let _lock = get_lock().await; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::delete_many() .filter(Column::OrgId.eq(org_id)) .exec(client) diff --git a/src/infra/src/table/service_streams.rs b/src/infra/src/table/service_streams.rs index 068ed66b6b..6f6a4d368e 100644 --- a/src/infra/src/table/service_streams.rs +++ b/src/infra/src/table/service_streams.rs @@ -26,9 +26,8 @@ use sea_orm::{ use serde::{Deserialize, Serialize}; use svix_ksuid::KsuidLike; -use super::get_lock; use crate::{ - db::{ORM_CLIENT, ORM_CLIENT_DDL, connect_to_orm, connect_to_orm_ddl}, + db::{get_orm_client_ddl, get_orm_client_ro, get_orm_client_rw}, errors::{self, DbError, Error}, }; @@ -146,7 +145,7 @@ impl ServiceRecord { } pub async fn create_table() -> Result<(), errors::Error> { - let client = ORM_CLIENT_DDL.get_or_init(connect_to_orm_ddl).await; + let client = get_orm_client_ddl().await; let builder = client.get_database_backend(); let schema = Schema::new(builder); @@ -192,8 +191,7 @@ pub async fn put( // are stable regardless of insertion order of the original object. record.disambiguation = normalize_json_object(record.disambiguation); - let _lock = get_lock().await; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; // Look up the exact row via (org_id, service_name, set_id, disambiguation). // Scoping by set_id ensures records from different identity sets never merge. @@ -508,7 +506,7 @@ fn union_stream_array( /// List all services for an organization pub async fn list(org_id: &str) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let records = Entity::find() .filter(Column::OrgId.eq(org_id)) @@ -524,7 +522,7 @@ pub async fn list_by_name( org_id: &str, service_name: &str, ) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let records = Entity::find() .filter(Column::OrgId.eq(org_id)) @@ -542,7 +540,7 @@ pub async fn list_filtered_by_service( org_id: &str, service_name: Option<&str>, ) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let query = match service_name { Some(name) => Entity::find() @@ -562,8 +560,7 @@ pub async fn list_filtered_by_service( /// Delete all service records for a specific identity set within an organization. /// Called when a set is removed from the config to clean up stale data. pub async fn delete_by_set_id(org_id: &str, set_id: &str) -> Result<(), errors::Error> { - let _lock = get_lock().await; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::delete_many() .filter(Column::OrgId.eq(org_id)) @@ -577,8 +574,7 @@ pub async fn delete_by_set_id(org_id: &str, set_id: &str) -> Result<(), errors:: /// Delete all services for an organization, returning the number of deleted rows pub async fn delete_all(org_id: &str) -> Result { - let _lock = get_lock().await; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let result = Entity::delete_many() .filter(Column::OrgId.eq(org_id)) @@ -591,7 +587,7 @@ pub async fn delete_all(org_id: &str) -> Result { /// List distinct organization IDs that have service_streams records pub async fn list_distinct_orgs() -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let backend = client.get_database_backend(); let sql = "SELECT DISTINCT org_id FROM service_streams ORDER BY org_id"; @@ -612,8 +608,7 @@ pub async fn list_distinct_orgs() -> Result, errors::Error> { /// Delete stale records (last_seen older than threshold) pub async fn delete_stale(org_id: &str, older_than_micros: i64) -> Result { - let _lock = get_lock().await; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let result = Entity::delete_many() .filter(Column::OrgId.eq(org_id)) @@ -627,7 +622,7 @@ pub async fn delete_stale(org_id: &str, older_than_micros: i64) -> Result Result { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; Entity::find() .filter(Column::OrgId.eq(org_id)) diff --git a/src/infra/src/table/sessions.rs b/src/infra/src/table/sessions.rs index 29aa14a4be..054c1ae21b 100644 --- a/src/infra/src/table/sessions.rs +++ b/src/infra/src/table/sessions.rs @@ -15,18 +15,15 @@ use sea_orm::{ColumnTrait, EntityTrait, QueryFilter, Set, sea_query::OnConflict}; -use super::{ - entity::sessions::{ActiveModel, Column, Entity, Model}, - get_lock, -}; +use super::entity::sessions::{ActiveModel, Column, Entity, Model}; use crate::{ - db::{ORM_CLIENT, connect_to_orm}, + db::{get_orm_client_ro, get_orm_client_rw}, 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 client = get_orm_client_ro().await; let session = Entity::find() .filter(Column::SessionId.eq(session_id)) .one(client) @@ -46,10 +43,7 @@ pub async fn set_with_expiry( access_token: &str, expires_at: i64, ) -> Result<(), errors::Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let now = chrono::Utc::now().timestamp_micros(); // Use atomic upsert to avoid race conditions @@ -81,10 +75,7 @@ pub async fn set_with_expiry( /// Deletes a session by session_id pub async fn delete(session_id: &str) -> Result<(), errors::Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::delete_many() .filter(Column::SessionId.eq(session_id)) .exec(client) @@ -94,7 +85,7 @@ pub async fn delete(session_id: &str) -> Result<(), errors::Error> { /// Lists all sessions pub async fn list() -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let sessions = Entity::find().all(client).await?; Ok(sessions) } @@ -102,10 +93,7 @@ pub async fn list() -> Result, errors::Error> { /// Deletes all expired sessions from the database /// This is more efficient than deleting one at a time pub async fn delete_expired() -> Result { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let now = chrono::Utc::now().timestamp(); let result = Entity::delete_many() diff --git a/src/infra/src/table/short_urls.rs b/src/infra/src/table/short_urls.rs index 3bcddb3084..83adfc5c59 100644 --- a/src/infra/src/table/short_urls.rs +++ b/src/infra/src/table/short_urls.rs @@ -22,11 +22,9 @@ use sea_orm::{ }; use serde::{Deserialize, Serialize}; -use super::get_lock; use crate::{ db::{ - IndexStatement, ORM_CLIENT, ORM_CLIENT_DDL, connect_to_orm, connect_to_orm_ddl, postgres, - sqlite, + IndexStatement, get_orm_client_ddl, get_orm_client_ro, get_orm_client_rw, postgres, sqlite, }, errors::{self, DbError, Error}, }; @@ -90,7 +88,7 @@ pub async fn init() -> Result<(), errors::Error> { } pub async fn create_table() -> Result<(), errors::Error> { - let client = ORM_CLIENT_DDL.get_or_init(connect_to_orm_ddl).await; + let client = get_orm_client_ddl().await; let builder = client.get_database_backend(); let schema = Schema::new(builder); @@ -113,7 +111,7 @@ pub async fn create_table_index() -> Result<(), errors::Error> { &["created_ts"], ); - let client = ORM_CLIENT_DDL.get_or_init(connect_to_orm_ddl).await; + let client = get_orm_client_ddl().await; match client.get_database_backend() { DatabaseBackend::Postgres => { postgres::create_index(index1).await?; @@ -137,10 +135,7 @@ pub async fn add(short_id: &str, original_url: &str, org_id: &str) -> Result Result Result<(), errors::Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::delete_many() .filter(Column::ShortId.eq(short_id)) .exec(client) @@ -167,7 +159,7 @@ pub async fn remove(short_id: &str) -> Result<(), errors::Error> { } pub async fn get(short_id: &str, org_id: &str) -> Result { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let mut query = Entity::find() .select_only() .column(Column::ShortId) @@ -195,7 +187,7 @@ pub async fn get(short_id: &str, org_id: &str) -> Result) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let mut res = Entity::find() .select_only() .column(Column::ShortId) @@ -211,7 +203,7 @@ pub async fn list(limit: Option) -> Result, errors::Err } pub async fn len() -> usize { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let len = Entity::find().count(client).await; match len { @@ -224,10 +216,7 @@ pub async fn len() -> usize { } pub async fn delete_by_org(org_id: &str) -> Result<(), errors::Error> { - // make sure only one client is writing to the database (only for SQLite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::delete_many() .filter(Column::OrgId.eq(org_id)) .exec(client) @@ -237,10 +226,7 @@ pub async fn delete_by_org(org_id: &str) -> Result<(), errors::Error> { } pub async fn clear() -> Result<(), errors::Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::delete_many().exec(client).await?; Ok(()) @@ -254,7 +240,7 @@ pub async fn get_expired( expired_before: i64, limit: Option, ) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let mut res = Entity::find() .select_only() .column(Column::ShortId) @@ -267,10 +253,7 @@ pub async fn get_expired( } pub async fn batch_remove(short_ids: Vec) -> Result<(), errors::Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::delete_many() .filter(Column::ShortId.is_in(short_ids)) .exec(client) diff --git a/src/infra/src/table/slo.rs b/src/infra/src/table/slo.rs index 8998a8a422..3d52fa3600 100644 --- a/src/infra/src/table/slo.rs +++ b/src/infra/src/table/slo.rs @@ -44,7 +44,7 @@ use sea_orm::{ sea_query::{Expr, Func, SimpleExpr}, }; -use super::{entity::slo_status, get_lock}; +use super::entity::slo_status; use crate::errors; /// `COALESCE(col, 0) + delta` — the increment applied in SQL rather than @@ -122,9 +122,6 @@ pub async fn apply_status( db: &DatabaseConnection, write: &StatusWrite, ) -> Result { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - let txn = db.begin().await?; let outcome = match apply_status_in_txn(&txn, write).await { Ok(outcome) => outcome, @@ -297,9 +294,6 @@ pub async fn init_generation( slo_id: &str, definition_generation: i32, ) -> Result<(), errors::Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - slo_status::ActiveModel { slo_id: Set(slo_id.to_string()), group_key: Set(slo_status::ROLLUP_GROUP_KEY.to_string()), @@ -318,9 +312,6 @@ pub async fn bump_generation( slo_id: &str, new_generation: i32, ) -> Result<(), errors::Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - let txn = db.begin().await?; // Delete rather than null out: the group set itself belongs to the old // definition, so a group that no longer exists under the new one must not @@ -354,9 +345,6 @@ pub async fn reconcile_from_slices( ) -> Result<(), errors::Error> { let (good, total, covered) = recomputed; - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - // Assignment, not increment: this is a rebuild from the source of truth, // and the watermark is deliberately untouched — reconciliation repairs the // aggregate, not the read clamp. @@ -381,9 +369,6 @@ pub async fn delete_by_org(db: &DatabaseConnection, org: &str) -> Result<(), err return Ok(()); } - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - slo_status::Entity::delete_many() .filter(slo_status::Column::SloId.is_in(slo_ids)) .exec(db) diff --git a/src/infra/src/table/slo_backfill_jobs.rs b/src/infra/src/table/slo_backfill_jobs.rs index 4a317f4cfb..f17e255bcb 100644 --- a/src/infra/src/table/slo_backfill_jobs.rs +++ b/src/infra/src/table/slo_backfill_jobs.rs @@ -25,7 +25,7 @@ use sea_orm::{ Set, }; -use super::{entity::slo_backfill_jobs, get_lock}; +use super::entity::slo_backfill_jobs; use crate::errors::Error; pub const STATE_QUEUED: i32 = 1; @@ -60,9 +60,6 @@ pub async fn queue( return Ok(()); } - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - slo_backfill_jobs::ActiveModel { slo_id: Set(slo_id.to_string()), definition_generation: Set(generation), @@ -92,9 +89,6 @@ pub async fn record_progress( return Ok(()); }; - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - let previous = model.rows_written; let mut active = model.into_active_model(); active.state = Set(STATE_RUNNING); @@ -136,9 +130,6 @@ async fn set_state( return Ok(()); }; - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - let mut active = model.into_active_model(); active.state = Set(state); if let Some(e) = error { @@ -150,9 +141,6 @@ async fn set_state( /// Drop every job for an SLO — used when the SLO itself is deleted. pub async fn delete_all(db: &DatabaseConnection, slo_id: &str) -> Result<(), Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - slo_backfill_jobs::Entity::delete_many() .filter(slo_backfill_jobs::Column::SloId.eq(slo_id)) .exec(db) @@ -171,9 +159,6 @@ pub async fn delete_by_org(db: &DatabaseConnection, org: &str) -> Result<(), Err return Ok(()); } - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - slo_backfill_jobs::Entity::delete_many() .filter(slo_backfill_jobs::Column::SloId.is_in(slo_ids)) .exec(db) diff --git a/src/infra/src/table/slo_budget.rs b/src/infra/src/table/slo_budget.rs index 09e84883d1..609265d87a 100644 --- a/src/infra/src/table/slo_budget.rs +++ b/src/infra/src/table/slo_budget.rs @@ -31,10 +31,7 @@ use sea_orm::{ TransactionTrait, }; -use super::{ - entity::{slo_budget, slo_budget_charges}, - get_lock, -}; +use super::entity::{slo_budget, slo_budget_charges}; use crate::errors::Error; /// Charge states, as stored. @@ -105,9 +102,6 @@ pub async fn charge( rows: i64, cap: i64, ) -> Result<(), ChargeError> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - for _ in 0..CAS_RETRIES { let current = get(db, org).await?; let (version, active, residual) = match ¤t { @@ -204,9 +198,6 @@ pub async fn retire( generation: i32, expires_at: i64, ) -> Result<(), ChargeError> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - for _ in 0..CAS_RETRIES { let Some(budget) = get(db, org).await? else { return Ok(()); @@ -265,9 +256,6 @@ pub async fn expire_residuals( org: &str, now: i64, ) -> Result { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - for _ in 0..CAS_RETRIES { let Some(budget) = get(db, org).await? else { return Ok(0); @@ -321,9 +309,6 @@ pub async fn expire_residuals( /// No CAS: the org is going away, so there is no concurrent charge whose /// arithmetic this could invalidate. pub async fn delete_by_org(db: &DatabaseConnection, org: &str) -> Result<(), Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - let txn = db.begin().await?; slo_budget_charges::Entity::delete_many() .filter(slo_budget_charges::Column::Org.eq(org)) diff --git a/src/infra/src/table/slos.rs b/src/infra/src/table/slos.rs index 31fc428d03..28575712e7 100644 --- a/src/infra/src/table/slos.rs +++ b/src/infra/src/table/slos.rs @@ -32,10 +32,7 @@ use sea_orm::{ PaginatorTrait, QueryFilter, QueryOrder, QuerySelect, Set, TransactionTrait, sea_query::Expr, }; -use super::{ - entity::{slo_status, slos}, - get_lock, -}; +use super::entity::{slo_status, slos}; use crate::errors::{DbError, Error}; /// What an update did to the writing epoch. @@ -141,9 +138,6 @@ pub async fn create( now: i64, editor: Option<&str>, ) -> Result<(), Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - let txn = db.begin().await?; let mut model = to_model(slo, now)?; model.created_at = Set(now); @@ -174,9 +168,6 @@ pub async fn update( now: i64, editor: Option<&str>, ) -> Result { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - let txn = db.begin().await?; let Some(existing) = slos::Entity::find_by_id(slo.id.clone()) @@ -294,9 +285,6 @@ pub async fn bump_generation( id: &str, now: i64, ) -> Result, Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - let txn = db.begin().await?; let Some(existing) = slos::Entity::find_by_id(id.to_string()) .filter(slos::Column::Org.eq(org)) @@ -341,9 +329,6 @@ pub async fn bump_generation( /// does not instantly free the storage its slices occupy, and pretending /// otherwise would let an org delete-and-recreate its way past the cap. pub async fn delete(db: &DatabaseConnection, org: &str, id: &str) -> Result { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - let txn = db.begin().await?; let Some(model) = slos::Entity::find_by_id(id.to_string()) .filter(slos::Column::Org.eq(org)) @@ -382,9 +367,6 @@ pub(crate) async fn ids_in_org(db: &DatabaseConnection, org: &str) -> Result Result<(), Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - slos::Entity::delete_many() .filter(slos::Column::Org.eq(org)) .exec(db) @@ -408,9 +390,6 @@ pub async fn set_enabled( return Ok(false); }; - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - let mut active = model.into_active_model(); active.enabled = Set(enabled); active.updated_at = Set(now); @@ -445,9 +424,6 @@ pub async fn move_to_folder( return Ok(0); } - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - let res = slos::Entity::update_many() .col_expr(slos::Column::FolderId, Expr::value(dst_folder_id)) .col_expr(slos::Column::UpdatedAt, Expr::value(now)) diff --git a/src/infra/src/table/source_maps.rs b/src/infra/src/table/source_maps.rs index 7ea9b003c5..0e28e4ab5d 100644 --- a/src/infra/src/table/source_maps.rs +++ b/src/infra/src/table/source_maps.rs @@ -19,9 +19,9 @@ use sea_orm::{ }; use serde::{Deserialize, Serialize}; -use super::{entity::source_maps::*, get_lock}; +use super::entity::source_maps::*; use crate::{ - db::{ORM_CLIENT, connect_to_orm}, + db::{get_orm_client_ro, get_orm_client_rw}, errors, }; @@ -154,8 +154,7 @@ pub async fn add_many(entries: Vec) -> Result<(), errors::Error> { }) .collect::>(); - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; - let _lock = get_lock().await; + let client = get_orm_client_rw().await; let txn = client.begin().await?; @@ -186,8 +185,7 @@ pub async fn delete_group( env: Option, version: Option, ) -> Result<(), errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; - let _lock = get_lock().await; + let client = get_orm_client_rw().await; let mut stmt = Entity::delete_many().filter(Column::Org.eq(org)); @@ -220,8 +218,7 @@ pub async fn get_sourcemap_file( env: &Option, version: &Option, ) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; - let _lock = get_lock().await; + let client = get_orm_client_ro().await; let mut stmt = Entity::find() .filter(Column::Org.eq(org)) @@ -249,8 +246,7 @@ pub async fn list_files( env: Option, version: Option, ) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; - let _lock = get_lock().await; + let client = get_orm_client_ro().await; let mut stmt = Entity::find().filter(Column::Org.eq(org)); @@ -276,8 +272,7 @@ pub async fn list_files( } pub async fn update_cluster(entry: SourceMap) -> Result<(), errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; - let _lock = get_lock().await; + let client = get_orm_client_rw().await; Entity::update_many() .filter(Column::Org.eq(entry.org)) @@ -294,8 +289,7 @@ pub async fn get_values( org_id: &str, limit: u64, ) -> Result<(Vec, Vec, Vec), errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; - let _lock = get_lock().await; + let client = get_orm_client_ro().await; let services = Entity::find() .filter(Column::Org.eq(org_id)) diff --git a/src/infra/src/table/synthetics_agents.rs b/src/infra/src/table/synthetics_agents.rs index 65dbb4fad8..ce4d38355b 100644 --- a/src/infra/src/table/synthetics_agents.rs +++ b/src/infra/src/table/synthetics_agents.rs @@ -29,12 +29,9 @@ use sea_orm::{ ColumnTrait, EntityTrait, QueryFilter, QueryOrder, QuerySelect, Set, sea_query::Expr, }; -use super::{ - entity::synthetics_agents::{ActiveModel, Column, Entity, Model}, - get_lock, -}; +use super::entity::synthetics_agents::{ActiveModel, Column, Entity, Model}; use crate::{ - db::{ORM_CLIENT, connect_to_orm}, + db::{get_orm_client_ro, get_orm_client_rw}, errors::{self, DbError, Error}, }; @@ -103,8 +100,7 @@ async fn invalidate_and_publish(agent_id: &str) { /// Insert an agent row, or refresh version/capabilities/last_seen_at when the /// id already exists (idempotent re-register after restart). pub async fn register(record: &SyntheticsAgentRecord) -> Result<(), errors::Error> { - let _lock = get_lock().await; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let existing = Entity::find_by_id(&record.id) .one(client) .await @@ -143,11 +139,6 @@ pub async fn register(record: &SyntheticsAgentRecord) -> Result<(), errors::Erro .map_err(|e| Error::DbError(DbError::SeaORMError(e.to_string())))?; } } - // Release the SQLite write mutex BEFORE emitting. On a sqlite meta_store - // the coordinator's put() takes the *same* CLIENT_RW lock `get_lock()` - // returns, so emitting while holding it deadlocks the process — and the - // mutex is then held forever, hanging every later synthetics query. - drop(_lock); // Register is where capabilities change, so the cached copy is now stale. invalidate_and_publish(&record.id).await; Ok(()) @@ -159,8 +150,7 @@ pub async fn register(record: &SyntheticsAgentRecord) -> Result<(), errors::Erro /// lease, which is exactly the rate of the reads the cache serves. Callers that /// need a current `last_seen_at` use [`get`], not [`get_cached`]. pub async fn touch(agent_id: &str, now_us: i64) -> Result<(), errors::Error> { - let _lock = get_lock().await; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::update_many() .col_expr(Column::LastSeenAt, Expr::value(now_us)) .filter(Column::Id.eq(agent_id)) @@ -174,7 +164,7 @@ pub async fn touch(agent_id: &str, now_us: i64) -> Result<(), errors::Error> { pub async fn list_by_location( location_id: &str, ) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let rows = Entity::find() .filter(Column::LocationId.eq(location_id)) .order_by_desc(Column::LastSeenAt) @@ -197,7 +187,7 @@ pub async fn list_by_locations( if location_ids.is_empty() { return Ok(std::collections::HashMap::new()); } - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let rows = Entity::find() .filter(Column::LocationId.is_in(location_ids.to_vec())) .order_by_desc(Column::LastSeenAt) @@ -226,7 +216,7 @@ pub async fn find_by_identity( location_id: &str, name: &str, ) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let row = Entity::find() .filter(Column::OrgId.eq(org_id)) .filter(Column::LocationId.eq(location_id)) @@ -245,7 +235,7 @@ pub async fn find_by_identity( pub async fn count_by_token( org_id: &str, ) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let rows: Vec> = Entity::find() .filter(Column::OrgId.eq(org_id)) .select_only() @@ -279,7 +269,7 @@ pub async fn get_cached(agent_id: &str) -> Result, /// Reads an agent straight from the database, `last_seen_at` included. pub async fn get(agent_id: &str) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let row = Entity::find_by_id(agent_id) .one(client) .await diff --git a/src/infra/src/table/synthetics_checks.rs b/src/infra/src/table/synthetics_checks.rs index aede7cae8a..789483025a 100644 --- a/src/infra/src/table/synthetics_checks.rs +++ b/src/infra/src/table/synthetics_checks.rs @@ -180,7 +180,6 @@ pub async fn get( org_id: &str, id: &str, ) -> Result, errors::Error> { - let _lock = super::get_lock().await; let maybe = get_model(conn, org_id, id).await?; maybe.map(Synthetic::try_from).transpose() } @@ -259,7 +258,6 @@ pub async fn list( org_id: &str, params: &ListSyntheticsParams, ) -> Result, errors::Error> { - let _lock = super::get_lock().await; let models = list_models(conn, org_id, params).await?; // A single unreadable row used to 500 the whole org's list — the UI renders // that as "no checks yet", so a total outage looked like an empty state. @@ -276,7 +274,6 @@ pub async fn count( org_id: &str, params: &ListSyntheticsParams, ) -> Result { - let _lock = super::get_lock().await; let q = Entity::find() .filter(Column::OrgId.eq(org_id)) .apply_filters(params); @@ -315,7 +312,6 @@ pub async fn list_referencing_location( org_id: &str, location_id: &str, ) -> Result, errors::Error> { - let _lock = super::get_lock().await; let models = Entity::find() .filter(Column::OrgId.eq(org_id)) .all(conn) @@ -354,7 +350,6 @@ pub async fn create( check: Synthetic, use_given_id: bool, ) -> Result { - let _lock = super::get_lock().await; let txn = conn.begin().await?; let now = config::utils::time::now_micros(); let id = new_check_id(&check, use_given_id); @@ -372,11 +367,6 @@ pub async fn create( let model = am.insert(&txn).await?.try_into_model()?; let result = Synthetic::try_from(model)?; txn.commit().await?; - // Release the SQLite write mutex BEFORE emitting. On a sqlite meta_store - // the coordinator's put() takes the *same* CLIENT_RW lock `get_lock()` - // returns, so emitting while holding it deadlocks the process — and the - // mutex is then held forever, hanging every later synthetics query. - drop(_lock); invalidate_and_publish(&result.org_id, &result.id).await; Ok(result) } @@ -387,7 +377,6 @@ pub async fn update( id: &str, check: Synthetic, ) -> Result { - let _lock = super::get_lock().await; let txn = conn.begin().await?; let Some(m) = get_model(&txn, org_id, id).await? else { @@ -401,11 +390,6 @@ pub async fn update( let model = am.update(&txn).await?.try_into_model()?; let result = Synthetic::try_from(model)?; txn.commit().await?; - // Release the SQLite write mutex BEFORE emitting. On a sqlite meta_store - // the coordinator's put() takes the *same* CLIENT_RW lock `get_lock()` - // returns, so emitting while holding it deadlocks the process — and the - // mutex is then held forever, hanging every later synthetics query. - drop(_lock); invalidate_and_publish(&result.org_id, &result.id).await; Ok(result) } @@ -415,7 +399,6 @@ pub async fn put( org_id: &str, check: Synthetic, ) -> Result { - let _lock = super::get_lock().await; let txn = conn.begin().await?; let now = config::utils::time::now_micros(); @@ -441,11 +424,6 @@ pub async fn put( }; txn.commit().await?; - // Release the SQLite write mutex BEFORE emitting. On a sqlite meta_store - // the coordinator's put() takes the *same* CLIENT_RW lock `get_lock()` - // returns, so emitting while holding it deadlocks the process — and the - // mutex is then held forever, hanging every later synthetics query. - drop(_lock); invalidate_and_publish(&result.org_id, &result.id).await; Ok(result) } @@ -455,17 +433,11 @@ pub async fn delete( org_id: &str, id: &str, ) -> Result { - let _lock = super::get_lock().await; let res = Entity::delete_many() .filter(Column::OrgId.eq(org_id)) .filter(Column::Id.eq(id)) .exec(conn) .await?; - // Release the SQLite write mutex BEFORE emitting. On a sqlite meta_store - // the coordinator's put() takes the *same* CLIENT_RW lock `get_lock()` - // returns, so emitting while holding it deadlocks the process — and the - // mutex is then held forever, hanging every later synthetics query. - drop(_lock); invalidate_and_publish_delete(org_id, id).await; Ok(res.rows_affected > 0) } @@ -477,7 +449,6 @@ pub async fn move_to_folder( ids: &[String], dst_folder_id: &str, ) -> Result { - let _lock = super::get_lock().await; if ids.is_empty() { return Ok(0); } @@ -491,11 +462,6 @@ pub async fn move_to_folder( .filter(Column::Id.is_in(ids.to_vec())) .exec(conn) .await?; - // Release the SQLite write mutex BEFORE emitting. On a sqlite meta_store - // the coordinator's put() takes the *same* CLIENT_RW lock `get_lock()` - // returns, so emitting while holding it deadlocks the process — and the - // mutex is then held forever, hanging every later synthetics query. - drop(_lock); for id in ids { invalidate_and_publish(org_id, id).await; } @@ -509,7 +475,6 @@ pub async fn set_enabled( id: &str, enabled: bool, ) -> Result { - let _lock = super::get_lock().await; let res = Entity::update_many() .col_expr(Column::Enabled, Expr::value(enabled)) .col_expr( @@ -520,11 +485,6 @@ pub async fn set_enabled( .filter(Column::Id.eq(id)) .exec(conn) .await?; - // Release the SQLite write mutex BEFORE emitting. On a sqlite meta_store - // the coordinator's put() takes the *same* CLIENT_RW lock `get_lock()` - // returns, so emitting while holding it deadlocks the process — and the - // mutex is then held forever, hanging every later synthetics query. - drop(_lock); invalidate_and_publish(org_id, id).await; Ok(res.rows_affected > 0) } @@ -621,7 +581,6 @@ pub async fn fetch_due( now_us: i64, limit: u64, ) -> Result, errors::Error> { - let _lock = super::get_lock().await; let models = due_checks_query(now_us, limit).all(conn).await?; let (due, _skipped) = convert_batch(models, "fetch_due"); @@ -698,7 +657,6 @@ pub async fn fetch_overdue( now_us: i64, limit: u64, ) -> Result, errors::Error> { - let _lock = super::get_lock().await; let rows = due_checks_query(now_us, limit) .select_only() .column(Column::Id) @@ -754,7 +712,6 @@ pub async fn claim_due( where C: ConnectionTrait + TransactionTrait, { - let _lock = super::get_lock().await; let txn = conn.begin().await?; let mut query = due_checks_query(now_us, limit); @@ -800,7 +757,6 @@ pub async fn advance_schedule( last_triggered_at: i64, next_run_at: i64, ) -> Result<(), errors::Error> { - let _lock = super::get_lock().await; Entity::update_many() .col_expr(Column::LastTriggeredAt, Expr::value(last_triggered_at)) .col_expr(Column::NextRunAt, Expr::value(next_run_at)) @@ -830,7 +786,6 @@ pub async fn update_last_check_status( id: &str, status: i32, ) -> Result { - let _lock = super::get_lock().await; let res = Entity::update_many() .col_expr(Column::LastCheckStatus, Expr::value(status)) .filter(Column::Id.eq(id)) @@ -907,7 +862,6 @@ pub async fn update_alert_state_if( expected: AlertState, state: AlertState, ) -> Result { - let _lock = super::get_lock().await; let res = Entity::update_many() .col_expr( Column::ConsecutiveFailures, diff --git a/src/infra/src/table/synthetics_jobs.rs b/src/infra/src/table/synthetics_jobs.rs index cb3caff3fa..116f8f8b84 100644 --- a/src/infra/src/table/synthetics_jobs.rs +++ b/src/infra/src/table/synthetics_jobs.rs @@ -27,10 +27,7 @@ use sea_orm::{ use serde::Serialize; use svix_ksuid::KsuidLike as _; -use super::{ - entity::synthetics_jobs::{Column, Entity}, - get_lock, -}; +use super::entity::synthetics_jobs::{Column, Entity}; use crate::errors; // ── Types ───────────────────────────────────────────────────────────────────── @@ -154,9 +151,6 @@ pub async fn enqueue( ON CONFLICT (synthetics_id, location, scheduled_ts) DO NOTHING "#; - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - conn.execute(Statement::from_sql_and_values( conn.get_database_backend(), sql, @@ -227,9 +221,6 @@ pub async fn drain_check( conn: &C, synthetics_id: &str, ) -> Result { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - let sql = "DELETE FROM synthetics_jobs WHERE synthetics_id = $1"; let res = conn .execute(Statement::from_sql_and_values( @@ -283,10 +274,7 @@ pub async fn lease_batch( // trusted: a probe does not get to decide how much of the queue it may take. let limit = limit.clamp(1, MAX_LEASE_BATCH); - // make sure only one client is writing to the database(only for sqlite). - // One lock across the candidate pick and the claim; released before the // read-back, which only reads rows pinned by our own stamp. - let _lock = get_lock().await; // Step 1: pick candidate IDs. // @@ -342,7 +330,6 @@ pub async fn lease_batch( .filter(Column::Status.eq(0i32)) .exec(conn) .await?; - drop(_lock); // Step 3: return ONLY the rows THIS call actually won — pinned by our // (claimed_by, claimed_at) stamp. A racing agent that lost the guard above @@ -450,9 +437,6 @@ pub async fn ack_complete( WHERE id = $4 AND status = 1{owner_clause}" ); - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - let applied = conn .execute(Statement::from_sql_and_values( conn.get_database_backend(), @@ -511,9 +495,6 @@ pub async fn requeue_expired( AND dispatch_attempts < $2 "#; - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - let res = conn .execute(Statement::from_sql_and_values( conn.get_database_backend(), @@ -554,10 +535,7 @@ pub async fn dead_letter_expired( // the remainder on its next tick. const MAX_DEAD_LETTER_BATCH: i64 = 500; - // make sure only one client is writing to the database(only for sqlite). - // One lock across the bounded SELECT and the per-row CAS updates — drops // at fn end. - let _lock = get_lock().await; // Step 1: find candidates before marking them dead. `status` comes back so // the CAS can require the row not to have moved under us. @@ -668,9 +646,6 @@ pub async fn fail_dispatch( current_attempts: i32, max_attempts: i32, ) -> Result { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - if current_attempts < max_attempts { // `AND status = 1` is what keeps this from resurrecting finished work. A // dispatch can be judged failed after the probe has already acked — a @@ -869,9 +844,6 @@ pub async fn pending_backlog( /// count here is a bug signal, not routine housekeeping — the reaper logs it at /// warn for that reason. pub async fn prune_stale(conn: &C, now_us: i64) -> Result { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - let sql = r#" DELETE FROM synthetics_jobs WHERE status = 0 AND valid_until < $1 diff --git a/src/infra/src/table/synthetics_locations.rs b/src/infra/src/table/synthetics_locations.rs index 3791d52bac..055a248089 100644 --- a/src/infra/src/table/synthetics_locations.rs +++ b/src/infra/src/table/synthetics_locations.rs @@ -28,12 +28,9 @@ use std::{ use sea_orm::{ColumnTrait, EntityTrait, QueryFilter, QueryOrder, Set, sea_query::Expr}; use tokio::sync::RwLock; -use super::{ - entity::synthetics_locations::{ActiveModel, Column, Entity, Model}, - get_lock, -}; +use super::entity::synthetics_locations::{ActiveModel, Column, Entity, Model}; use crate::{ - db::{ORM_CLIENT, connect_to_orm}, + db::{get_orm_client_ro, get_orm_client_rw}, errors::{self, DbError, Error}, }; @@ -96,7 +93,7 @@ async fn ensure_fresh() -> Result<(), errors::Error> { return Ok(()); } - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let rows: Vec = Entity::find() .order_by_asc(Column::Kind) .order_by_asc(Column::Label) @@ -159,8 +156,7 @@ impl From for SyntheticsLocationRecord { /// Insert a new location row. pub async fn add(record: &SyntheticsLocationRecord) -> Result<(), errors::Error> { - let _lock = get_lock().await; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let model = ActiveModel { id: Set(record.id.clone()), org_id: Set(record.org_id.clone()), @@ -179,11 +175,6 @@ pub async fn add(record: &SyntheticsLocationRecord) -> Result<(), errors::Error> .exec(client) .await .map_err(|e| Error::DbError(DbError::SeaORMError(e.to_string())))?; - // Release the SQLite write mutex BEFORE emitting. On a sqlite meta_store - // the coordinator's put() takes the *same* CLIENT_RW lock `get_lock()` - // returns, so emitting while holding it deadlocks the process — and the - // mutex is then held forever, hanging every later synthetics query. - drop(_lock); invalidate_and_publish().await; Ok(()) } @@ -208,7 +199,7 @@ pub async fn list_visible(org_id: &str) -> Result, /// All private rows across orgs — used by the staleness watcher. pub async fn list_private() -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let rows = Entity::find() .filter(Column::Kind.eq(KIND_PRIVATE)) .all(client) @@ -247,35 +238,23 @@ fn update_stmt(id: &str, label: &str, enabled: bool, now: i64) -> sea_orm::Updat /// Update label/enabled on a location. pub async fn update(id: &str, label: &str, enabled: bool) -> Result<(), errors::Error> { - let _lock = get_lock().await; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let now = chrono::Utc::now().timestamp_micros(); update_stmt(id, label, enabled, now) .exec(client) .await .map_err(|e| Error::DbError(DbError::SeaORMError(e.to_string())))?; - // Release the SQLite write mutex BEFORE emitting. On a sqlite meta_store - // the coordinator's put() takes the *same* CLIENT_RW lock `get_lock()` - // returns, so emitting while holding it deadlocks the process — and the - // mutex is then held forever, hanging every later synthetics query. - drop(_lock); invalidate_and_publish().await; Ok(()) } /// Delete a location row. pub async fn remove(id: &str) -> Result<(), errors::Error> { - let _lock = get_lock().await; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::delete_by_id(id) .exec(client) .await .map_err(|e| Error::DbError(DbError::SeaORMError(e.to_string())))?; - // Release the SQLite write mutex BEFORE emitting. On a sqlite meta_store - // the coordinator's put() takes the *same* CLIENT_RW lock `get_lock()` - // returns, so emitting while holding it deadlocks the process — and the - // mutex is then held forever, hanging every later synthetics query. - drop(_lock); invalidate_and_publish().await; Ok(()) } @@ -296,8 +275,7 @@ pub async fn remove(id: &str) -> Result<(), errors::Error> { /// Deliberately NOT cached — `LOCATIONS_CACHE` serves definition reads, and a /// stale `down_notified_at` would reintroduce exactly the duplicate this fixes. pub async fn try_claim_down_notification(id: &str, now_us: i64) -> Result { - let _lock = get_lock().await; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let res = Entity::update_many() .col_expr(Column::DownNotifiedAt, Expr::value(now_us)) .filter(Column::Id.eq(id)) @@ -314,8 +292,7 @@ pub async fn try_claim_down_notification(id: &str, now_us: i64) -> Result Result<(), errors::Error> { - let _lock = get_lock().await; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::update_many() .col_expr(Column::DownNotifiedAt, Expr::value(0i64)) .filter(Column::Id.eq(id)) diff --git a/src/infra/src/table/synthetics_probe_tokens.rs b/src/infra/src/table/synthetics_probe_tokens.rs index 6866949c51..833285d1e0 100644 --- a/src/infra/src/table/synthetics_probe_tokens.rs +++ b/src/infra/src/table/synthetics_probe_tokens.rs @@ -29,12 +29,9 @@ use sea_orm::{ ColumnTrait, EntityTrait, QueryFilter, Set, SqlErr, TransactionTrait, sea_query::Expr, }; -use super::{ - entity::synthetics_probe_tokens::{ActiveModel, Column, Entity, Model}, - get_lock, -}; +use super::entity::synthetics_probe_tokens::{ActiveModel, Column, Entity, Model}; use crate::{ - db::{ORM_CLIENT, connect_to_orm}, + db::{get_orm_client_ro, get_orm_client_rw}, errors::{self, DbError, Error}, }; @@ -129,7 +126,6 @@ pub fn generate_token() -> String { /// Insert a new probe token row. Fails with a clear message if a token with the /// same `(org_id, name)` already exists (the unique constraint). pub async fn add(record: &SyntheticsProbeTokenRecord) -> Result<(), errors::Error> { - let _lock = get_lock().await; let model = ActiveModel { id: Set(record.id.clone()), org_id: Set(record.org_id.clone()), @@ -146,14 +142,9 @@ pub async fn add(record: &SyntheticsProbeTokenRecord) -> Result<(), errors::Erro created_at: Set(record.created_at), updated_at: Set(record.updated_at), }; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; match Entity::insert(model).exec(client).await { Ok(_) => { - // Release the SQLite write mutex BEFORE emitting. On a sqlite meta_store - // the coordinator's put() takes the *same* CLIENT_RW lock `get_lock()` - // returns, so emitting while holding it deadlocks the process — and the - // mutex is then held forever, hanging every later synthetics query. - drop(_lock); // A new token can become the org default, so both caches are stale. invalidate_and_publish(&record.org_id).await; Ok(()) @@ -182,7 +173,7 @@ pub async fn find_global(token: &str) -> Result Result, errors::Error> { use sea_orm::{Order, QueryOrder}; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let records = Entity::find() .filter(Column::OrgId.eq(org_id)) .order_by(Column::IsDefault, Order::Desc) @@ -243,7 +234,7 @@ pub async fn get_by_name( org_id: &str, name: &str, ) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let record = Entity::find() .filter(Column::OrgId.eq(org_id)) .filter(Column::Name.eq(name)) @@ -270,9 +261,8 @@ pub async fn get_by_name( /// middleware on every probe request, and an uncached validator is a cheap way /// to generate load from outside. pub async fn set_enabled(org_id: &str, name: &str, enabled: bool) -> Result<(), errors::Error> { - let _lock = get_lock().await; let now = chrono::Utc::now().timestamp_micros(); - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::update_many() .col_expr(Column::Enabled, Expr::value(enabled)) .col_expr(Column::UpdatedAt, Expr::value(now)) @@ -281,11 +271,6 @@ pub async fn set_enabled(org_id: &str, name: &str, enabled: bool) -> Result<(), .exec(client) .await .map_err(|e| Error::DbError(DbError::SeaORMError(e.to_string())))?; - // Release the SQLite write mutex BEFORE emitting. On a sqlite meta_store - // the coordinator's put() takes the *same* CLIENT_RW lock `get_lock()` - // returns, so emitting while holding it deadlocks the process — and the - // mutex is then held forever, hanging every later synthetics query. - drop(_lock); invalidate_and_publish(org_id).await; Ok(()) } @@ -295,9 +280,8 @@ pub async fn set_enabled(org_id: &str, name: &str, enabled: bool) -> Result<(), /// previous default stays valid during a rotation overlap window until it is /// explicitly disabled. pub async fn set_default(org_id: &str, name: &str) -> Result<(), errors::Error> { - let _lock = get_lock().await; let now = chrono::Utc::now().timestamp_micros(); - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let txn = client .begin() .await @@ -320,11 +304,6 @@ pub async fn set_default(org_id: &str, name: &str) -> Result<(), errors::Error> txn.commit() .await .map_err(|e| Error::DbError(DbError::SeaORMError(e.to_string())))?; - // Release the SQLite write mutex BEFORE emitting. On a sqlite meta_store - // the coordinator's put() takes the *same* CLIENT_RW lock `get_lock()` - // returns, so emitting while holding it deadlocks the process — and the - // mutex is then held forever, hanging every later synthetics query. - drop(_lock); invalidate_and_publish(org_id).await; Ok(()) } diff --git a/src/infra/src/table/synthetics_runs.rs b/src/infra/src/table/synthetics_runs.rs index 7d5fed4f3e..6dab1ca9d7 100644 --- a/src/infra/src/table/synthetics_runs.rs +++ b/src/infra/src/table/synthetics_runs.rs @@ -21,7 +21,6 @@ use sea_orm::{ConnectionTrait, Statement, Value}; -use super::get_lock; use crate::errors; // ── Types ───────────────────────────────────────────────────────────────────── @@ -69,9 +68,6 @@ pub async fn insert_run( VALUES ($1, $2, $3, $4, $5, $6, 0, NULL, $7, NULL) "#; - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - conn.execute(Statement::from_sql_and_values( conn.get_database_backend(), sql, @@ -268,10 +264,6 @@ pub async fn increment_jobs_done( job_result: i32, now_us: i64, ) -> Result, errors::Error> { - // make sure only one client is writing to the database(only for sqlite). - // One lock across all three steps — the guard drops at fn end. - let _lock = get_lock().await; - // Step 1: count this job and fold its severity in. // // `jobs_done < job_count` stops a duplicate or replayed ack from pushing the diff --git a/src/infra/src/table/system_prompts.rs b/src/infra/src/table/system_prompts.rs index f4dfe17432..5a3f88b022 100644 --- a/src/infra/src/table/system_prompts.rs +++ b/src/infra/src/table/system_prompts.rs @@ -19,9 +19,9 @@ use sea_orm::{ QuerySelect, Schema, Set, }; -use super::{entity::system_prompts::Model, get_lock}; +use super::entity::system_prompts::Model; use crate::{ - db::{ORM_CLIENT, connect_to_orm}, + db::{get_orm_client_ro, get_orm_client_rw}, errors, table::entity::system_prompts::{ActiveModel, Column, Entity}, }; @@ -38,7 +38,7 @@ impl TryFrom for AIPrompt { } } pub async fn create_table() -> Result<(), errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let builder = client.get_database_backend(); let schema = Schema::new(builder); @@ -58,10 +58,8 @@ pub async fn add(prompt: &AIPrompt) -> Result<(), errors::Error> { content: Set(prompt.content.clone()), updated_at: Set(prompt.updated_at), }; - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::insert(record).exec(client).await?; Ok(()) @@ -74,19 +72,13 @@ pub async fn update(prompt: &AIPrompt) -> Result<(), errors::Error> { updated_at: Set(prompt.updated_at), }; - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::update(record).exec(client).await?; Ok(()) } pub async fn remove(_org_id: &str, r#type: PromptType) -> Result<(), errors::Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let _ = Entity::delete_many() .filter(Column::Type.eq(r#type.to_string())) .exec(client) @@ -96,7 +88,7 @@ pub async fn remove(_org_id: &str, r#type: PromptType) -> Result<(), errors::Err } pub async fn get(r#type: PromptType) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let record = Entity::find() .filter(Column::Type.eq(r#type.to_string())) @@ -112,7 +104,7 @@ pub async fn get(r#type: PromptType) -> Result, errors::Error> /// /// If the user prompt is not found, it returns None pub async fn get_all() -> Result<(Option, Option), errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let res = Entity::find() // This limit doesn't matter much, but still .limit(2) @@ -135,7 +127,7 @@ pub async fn get_all() -> Result<(Option, Option), errors::E } pub async fn exists(r#type: PromptType) -> Result { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let record = Entity::find() .filter(Column::Type.eq(r#type.to_string())) .one(client) @@ -145,16 +137,13 @@ pub async fn exists(r#type: PromptType) -> Result { } pub async fn len() -> Result { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let len = Entity::find().count(client).await?; Ok(len as usize) } pub async fn clear() -> Result<(), errors::Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::delete_many().exec(client).await?; Ok(()) diff --git a/src/infra/src/table/system_settings.rs b/src/infra/src/table/system_settings.rs index a3a196e274..5df3173a00 100644 --- a/src/infra/src/table/system_settings.rs +++ b/src/infra/src/table/system_settings.rs @@ -29,9 +29,8 @@ use sea_orm::{ ColumnTrait, EntityTrait, QueryFilter, QueryOrder, }; -use super::get_lock; use crate::{ - db::{ORM_CLIENT, connect_to_orm}, + db::{get_orm_client_ro, get_orm_client_rw}, table::entity::system_settings::{ActiveModel, Column, Entity, Model}, }; @@ -42,7 +41,7 @@ pub async fn get( user_id: Option<&str>, key: &str, ) -> Result, anyhow::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let mut query = Entity::find() .filter(Column::Scope.eq(scope.as_str())) @@ -116,7 +115,7 @@ pub async fn list( user_id: Option<&str>, category: Option<&str>, ) -> Result, anyhow::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let mut query = Entity::find().order_by_asc(Column::SettingKey); @@ -182,10 +181,7 @@ pub async fn list_resolved( pub async fn set(setting: &SystemSetting) -> Result { setting.validate().map_err(|e| anyhow!(e))?; - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let now = chrono::Utc::now().timestamp_micros(); // Check if setting already exists @@ -253,7 +249,7 @@ pub async fn delete( user_id: Option<&str>, key: &str, ) -> Result { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let mut query = Entity::delete_many() .filter(Column::Scope.eq(scope.as_str())) @@ -287,9 +283,6 @@ pub async fn delete( } } - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - let result = query .exec(client) .await @@ -300,10 +293,7 @@ pub async fn delete( /// Delete all settings for an organization pub async fn delete_org_settings(org_id: &str) -> Result { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let result = Entity::delete_many() .filter(Column::OrgId.eq(org_id)) @@ -316,10 +306,7 @@ pub async fn delete_org_settings(org_id: &str) -> Result { /// Delete all settings for a user in an organization pub async fn delete_user_settings(org_id: &str, user_id: &str) -> Result { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let result = Entity::delete_many() .filter(Column::OrgId.eq(org_id)) diff --git a/src/infra/src/table/templates.rs b/src/infra/src/table/templates.rs index f4de448b6c..2490a0ddc9 100644 --- a/src/infra/src/table/templates.rs +++ b/src/infra/src/table/templates.rs @@ -25,12 +25,9 @@ use sea_orm::{ }; use crate::{ - db::{ORM_CLIENT, connect_to_orm}, + db::{get_orm_client_ro, get_orm_client_rw}, errors::{Error, TemplateError}, - table::{ - entity::templates::{ActiveModel, Column, Entity, Model}, - get_lock, - }, + table::entity::templates::{ActiveModel, Column, Entity, Model}, }; const DEFAULT_ORG: &str = "default"; @@ -61,10 +58,7 @@ impl TryFrom for Template { } pub async fn put(template: Template) -> Result { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let title = match &template.template_type { TemplateType::Email { title } => Some(title.to_string()), _ => None, @@ -93,7 +87,7 @@ pub async fn put(template: Template) -> Result { } pub async fn get(org_id: &str, name: &str) -> Result, Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; match get_model(client, org_id, name).await? { Some(model) => Ok(Some(Template::try_from(model)?)), None => Ok(None), @@ -101,7 +95,7 @@ pub async fn get(org_id: &str, name: &str) -> Result, Error> { } pub async fn list(org_id: &str) -> Result, Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let templates = list_models(client, Some(org_id)) .await? .into_iter() @@ -111,7 +105,7 @@ pub async fn list(org_id: &str) -> Result, Error> { } pub async fn list_all() -> Result, Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let templates = list_models(client, None) .await? .into_iter() @@ -121,10 +115,7 @@ pub async fn list_all() -> Result, Error> { } pub async fn delete(org_id: &str, name: &str) -> Result<(), Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let model = get_model(client, org_id, name).await?; if let Some(model) = model { @@ -162,8 +153,7 @@ async fn list_models( /// Deletes all templates belonging to the given org. pub async fn delete_by_org(org_id: &str) -> Result<(), Error> { - let _lock = get_lock().await; - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::delete_many() .filter(Column::Org.eq(org_id)) .exec(client) diff --git a/src/infra/src/table/timed_annotation_panels.rs b/src/infra/src/table/timed_annotation_panels.rs index f1537d4b73..2726668890 100644 --- a/src/infra/src/table/timed_annotation_panels.rs +++ b/src/infra/src/table/timed_annotation_panels.rs @@ -16,17 +16,14 @@ use config::ider; use sea_orm::{ColumnTrait, EntityTrait, QueryFilter, Set, TransactionTrait}; -use super::{entity::timed_annotation_panels, get_lock}; +use super::entity::timed_annotation_panels; use crate::{ - db::{ORM_CLIENT, connect_to_orm}, + db::{get_orm_client_ro, get_orm_client_rw}, errors, }; pub async fn get_panels(timed_annotation_id: &str) -> Result, errors::Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let panels = timed_annotation_panels::Entity::find() .filter(timed_annotation_panels::Column::TimedAnnotationId.eq(timed_annotation_id)) .all(client) @@ -45,10 +42,7 @@ pub async fn insert_many_panels( return Ok(()); } - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let txn = client.begin().await?; for panel_id in panel_ids { @@ -76,10 +70,7 @@ pub async fn delete_many_panels( return Ok(()); } - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let txn = client.begin().await?; timed_annotation_panels::Entity::delete_many() diff --git a/src/infra/src/table/timed_annotations.rs b/src/infra/src/table/timed_annotations.rs index 60dbb90ab3..49e1cbbb85 100644 --- a/src/infra/src/table/timed_annotations.rs +++ b/src/infra/src/table/timed_annotations.rs @@ -22,15 +22,9 @@ use sea_orm::{ TransactionTrait, prelude::Expr, }; -use super::{ - entity::{ - dashboards, folders, timed_annotation_panels, - timed_annotations::{self}, - }, - get_lock, -}; +use super::entity::{dashboards, folders, timed_annotation_panels, timed_annotations}; use crate::{ - db::{ORM_CLIENT, connect_to_orm}, + db::{get_orm_client_ro, get_orm_client_rw}, errors, }; @@ -40,10 +34,7 @@ pub async fn get( start_time: i64, end_time: i64, ) -> Result, errors::Error> { - // make sure only one client is writing to the database (only for SQLite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; // Step 1: Resolve the user-facing `dashboard_id` to the primary key (KSUID) let dashboard_record = dashboards::Entity::find() @@ -189,10 +180,7 @@ pub async fn get( } pub async fn delete(dashboard_id: &str, timed_annotation_id: &str) -> Result<(), errors::Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let dashboard_record = dashboards::Entity::find() .filter(dashboards::Column::DashboardId.eq(dashboard_id)) .one(client) @@ -224,10 +212,7 @@ pub async fn delete_many( dashboard_id: &str, timed_annotation_ids: &[String], ) -> Result<(), errors::Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let dashboard_record = dashboards::Entity::find() .filter(dashboards::Column::DashboardId.eq(dashboard_id)) .one(client) @@ -268,10 +253,7 @@ pub async fn delete_many( /// timed_annotation_panels are deleted first because SQLite does not enforce FK cascades /// unless `PRAGMA foreign_keys = ON` is set. pub async fn delete_by_org(org_id: &str) -> Result<(), errors::Error> { - // make sure only one client is writing to the database (only for SQLite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; // Step 1: Collect folder PKs for this org let folder_ids: Vec = folders::Entity::find() @@ -331,11 +313,8 @@ pub async fn get_one( dashboard_id: &str, timed_annotation_id: &str, ) -> Result { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - // Initialize the ORM client - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; // Step 1: Resolve the user-facing `dashboard_id` to the primary key let dashboard_record = dashboards::Entity::find() @@ -399,11 +378,8 @@ pub async fn update( dashboard_id: &str, timed_annotation: TimedAnnotation, ) -> Result<(), errors::Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - // Initialize the ORM client - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let txn = client.begin().await?; @@ -552,10 +528,7 @@ pub async fn add_many( timed_annotations: Vec, use_given_id: bool, ) -> Result, errors::Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let txn = client.begin().await?; let mut inserted_annotations = Vec::new(); diff --git a/src/infra/src/table/trial_quota_usage.rs b/src/infra/src/table/trial_quota_usage.rs index 34f6cbf6d1..b592b6bb18 100644 --- a/src/infra/src/table/trial_quota_usage.rs +++ b/src/infra/src/table/trial_quota_usage.rs @@ -19,7 +19,7 @@ use sea_orm::{ }; use crate::{ - db::{ORM_CLIENT, connect_to_orm}, + db::{get_orm_client_ro, get_orm_client_rw}, table::entity::trial_quota_usage, }; @@ -30,8 +30,7 @@ pub async fn batch_increment(records: Vec<(String, String, i64)>) -> Result<(), if records.is_empty() { return Ok(()); } - let db = ORM_CLIENT.get_or_init(connect_to_orm).await; - let _lock = super::get_lock().await; + let db = get_orm_client_rw().await; // one transaction for the whole batch so the write lock is held for a // single commit instead of one autocommit round-trip per record let txn = db.begin().await?; @@ -74,14 +73,14 @@ pub async fn batch_increment(records: Vec<(String, String, i64)>) -> Result<(), /// Load all quota records (all features, all orgs). /// Called once on node startup to populate the in-memory DashMap. pub async fn load_all() -> Result, sea_orm::DbErr> { - let db = ORM_CLIENT.get_or_init(connect_to_orm).await; + let db = get_orm_client_ro().await; trial_quota_usage::Entity::find().all(db).await } /// Get total usage across all features for an org (sum of usage_count). /// Note: PostgreSQL SUM(bigint) returns NUMERIC, so we cast to BIGINT for Rust i64 compat. pub async fn get_total_usage_for_org(org_id: &str) -> Result { - let db = ORM_CLIENT.get_or_init(connect_to_orm).await; + let db = get_orm_client_ro().await; let result: Option> = trial_quota_usage::Entity::find() .filter(trial_quota_usage::Column::OrgId.eq(org_id)) .select_only() @@ -101,7 +100,7 @@ pub async fn get_total_usage_for_org(org_id: &str) -> Result Result, sea_orm::DbErr> { - let db = ORM_CLIENT.get_or_init(connect_to_orm).await; + let db = get_orm_client_ro().await; let result = trial_quota_usage::Entity::find() .filter(trial_quota_usage::Column::OrgId.eq(org_id)) .select_only() @@ -113,7 +112,7 @@ pub async fn get_usage_limit_for_org(org_id: &str) -> Result, sea_or } pub async fn load_all_usage_limits() -> Result, sea_orm::DbErr> { - let db = ORM_CLIENT.get_or_init(connect_to_orm).await; + let db = get_orm_client_ro().await; let results: Vec<(String, Option)> = trial_quota_usage::Entity::find() .select_only() .column(trial_quota_usage::Column::OrgId) @@ -132,8 +131,7 @@ pub async fn load_all_usage_limits() -> Result, sea_orm::DbEr /// Set one shared limit on every feature row for an organization. The AI chat /// row is upserted first so organizations without prior usage can be configured. pub async fn set_usage_limit_for_org(org_id: &str, usage_limit: i64) -> Result<(), sea_orm::DbErr> { - let db = ORM_CLIENT.get_or_init(connect_to_orm).await; - let _lock = super::get_lock().await; + let db = get_orm_client_rw().await; let txn = db.begin().await?; let now = config::utils::time::now_micros(); let active_model = trial_quota_usage::ActiveModel { @@ -180,7 +178,7 @@ pub async fn get_for_org_feature( org_id: &str, feature: &str, ) -> Result, sea_orm::DbErr> { - let db = ORM_CLIENT.get_or_init(connect_to_orm).await; + let db = get_orm_client_ro().await; trial_quota_usage::Entity::find() .filter(trial_quota_usage::Column::OrgId.eq(org_id)) .filter(trial_quota_usage::Column::Feature.eq(feature)) @@ -190,7 +188,7 @@ pub async fn get_for_org_feature( /// Get the highest notified checkpoint for an org (across all feature rows). pub async fn get_notified_checkpoint(org_id: &str) -> Result { - let db = ORM_CLIENT.get_or_init(connect_to_orm).await; + let db = get_orm_client_ro().await; let result = trial_quota_usage::Entity::find() .filter(trial_quota_usage::Column::OrgId.eq(org_id)) .column_as( @@ -210,8 +208,7 @@ pub async fn update_notified_checkpoint( org_id: &str, checkpoint: i16, ) -> Result { - let db = ORM_CLIENT.get_or_init(connect_to_orm).await; - let _lock = super::get_lock().await; + let db = get_orm_client_rw().await; let result = trial_quota_usage::Entity::update_many() .filter(trial_quota_usage::Column::OrgId.eq(org_id)) .filter(trial_quota_usage::Column::NotifiedCheckpoint.lt(checkpoint)) @@ -226,8 +223,7 @@ pub async fn update_notified_checkpoint( } pub async fn reset_notified_checkpoint(org_id: &str) -> Result<(), sea_orm::DbErr> { - let db = ORM_CLIENT.get_or_init(connect_to_orm).await; - let _lock = super::get_lock().await; + let db = get_orm_client_rw().await; trial_quota_usage::Entity::update_many() .filter(trial_quota_usage::Column::OrgId.eq(org_id)) .col_expr( @@ -241,7 +237,7 @@ pub async fn reset_notified_checkpoint(org_id: &str) -> Result<(), sea_orm::DbEr /// Load all notified checkpoints (one per org, max across features). pub async fn load_all_checkpoints() -> Result, sea_orm::DbErr> { - let db = ORM_CLIENT.get_or_init(connect_to_orm).await; + let db = get_orm_client_ro().await; let results: Vec<(String, Option)> = trial_quota_usage::Entity::find() .select_only() .column(trial_quota_usage::Column::OrgId) @@ -263,8 +259,7 @@ pub async fn load_all_checkpoints() -> Result, sea_orm::DbErr /// Deletes all trial quota usage entries for the given org. pub async fn delete_by_org(org_id: &str) -> Result<(), sea_orm::DbErr> { use sea_orm::{ColumnTrait, EntityTrait, QueryFilter}; - let db = ORM_CLIENT.get_or_init(connect_to_orm).await; - let _lock = super::get_lock().await; + let db = get_orm_client_rw().await; trial_quota_usage::Entity::delete_many() .filter(trial_quota_usage::Column::OrgId.eq(org_id)) .exec(db) diff --git a/src/infra/src/table/users.rs b/src/infra/src/table/users.rs index a34bb3a94a..10ba6f633b 100644 --- a/src/infra/src/table/users.rs +++ b/src/infra/src/table/users.rs @@ -23,12 +23,9 @@ use sea_orm::{ }; use serde::{Deserialize, Serialize}; -use super::{ - entity::users::{ActiveModel, Column, Entity, Model}, - get_lock, -}; +use super::entity::users::{ActiveModel, Column, Entity, Model}; use crate::{ - db::{ORM_CLIENT, connect_to_orm}, + db::{get_orm_client_ro, get_orm_client_rw}, errors::{self, DbError, Error}, }; @@ -114,7 +111,7 @@ impl From<&UserRecord> for DBUser { } pub async fn create_table() -> Result<(), errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let builder = client.get_database_backend(); let schema = Schema::new(builder); @@ -147,10 +144,7 @@ pub async fn add(user: UserRecord) -> Result<(), errors::Error> { id: Set(ider::uuid()), }; - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; match Entity::insert(record).exec(client).await { Ok(_) => Ok(()), Err(e) => match e.sql_err() { @@ -167,10 +161,7 @@ pub async fn update( password: &str, password_ext: Option, ) -> Result { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let result = Entity::update_many() .col_expr(Column::FirstName, Expr::value(first_name)) @@ -190,10 +181,7 @@ pub async fn update( } pub async fn remove(email: &str) -> Result<(), errors::Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::delete_many() .filter(Expr::expr(Func::lower(Expr::col(Column::Email))).eq(email.to_lowercase())) .exec(client) @@ -204,7 +192,7 @@ pub async fn remove(email: &str) -> Result<(), errors::Error> { } pub async fn get(email: &str) -> Result { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let record = Entity::find() .filter(Expr::expr(Func::lower(Expr::col(Column::Email))).eq(email.to_lowercase())) .one(client) @@ -216,7 +204,7 @@ pub async fn get(email: &str) -> Result { } pub async fn get_root_user() -> Result { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let record = Entity::find() .filter(Column::IsRoot.eq(true)) .one(client) @@ -228,7 +216,7 @@ pub async fn get_root_user() -> Result { } pub async fn list(limit: Option) -> Result, errors::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let mut res = Entity::find().order_by(Column::CreatedAt, Order::Desc); if let Some(limit) = limit { res = res.limit(limit as u64); @@ -245,7 +233,7 @@ pub async fn list(limit: Option) -> Result, errors::Error> } pub async fn len() -> usize { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let len = Entity::find().count(client).await; match len { @@ -258,10 +246,7 @@ pub async fn len() -> usize { } pub async fn clear() -> Result<(), errors::Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; Entity::delete_many() .exec(client) .await @@ -275,10 +260,7 @@ pub async fn is_empty() -> bool { } pub async fn batch_remove(emails: Vec) -> Result<(), errors::Error> { - // make sure only one client is writing to the database(only for sqlite) - let _lock = get_lock().await; - - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_rw().await; let lowered_emails: Vec = emails.iter().map(|e| e.to_lowercase()).collect(); Entity::delete_many() .filter(Expr::expr(Func::lower(Expr::col(Column::Email))).is_in(lowered_emails)) diff --git a/src/infra/src/table/workflows.rs b/src/infra/src/table/workflows.rs index a5e964f9c0..355776b1f2 100644 --- a/src/infra/src/table/workflows.rs +++ b/src/infra/src/table/workflows.rs @@ -17,12 +17,9 @@ use config::meta::pipeline::components::{Edge, Node}; use sea_orm::{ColumnTrait, EntityTrait, QueryFilter, Set, TransactionTrait, prelude::Expr}; use serde::{Deserialize, Serialize}; -use super::{ - entity::{workflow_drafts, workflow_errors, workflow_run_data, workflows}, - get_lock, -}; +use super::entity::{workflow_drafts, workflow_errors, workflow_run_data, workflows}; use crate::{ - db::{ORM_CLIENT, connect_to_orm}, + db::{get_orm_client_ro, get_orm_client_rw}, table::entity::workflow_associations, }; @@ -176,7 +173,7 @@ impl std::fmt::Display for WorkflowTriggerEntity { } pub async fn list_all() -> Result, anyhow::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let entities = workflows::Entity::find().all(client).await?; let mut ret = Vec::with_capacity(entities.len()); for e in entities { @@ -186,7 +183,7 @@ pub async fn list_all() -> Result, anyhow::Error> { } pub async fn list_by_org(org_id: &str) -> Result, anyhow::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let entities = workflows::Entity::find() .filter(workflows::Column::OrgId.eq(org_id)) .all(client) @@ -202,7 +199,7 @@ pub async fn get_by_org_wid( org_id: &str, workflow_id: &str, ) -> Result, anyhow::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let res = workflows::Entity::find() .filter(workflows::Column::Id.eq(workflow_id)) .filter(workflows::Column::OrgId.eq(org_id)) @@ -218,7 +215,7 @@ pub async fn get_errors_for_run( wid: &str, run_id: &str, ) -> Result, anyhow::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let res = workflow_errors::Entity::find() .filter(workflow_errors::Column::OrgId.eq(org_id)) .filter(workflow_errors::Column::WorkflowId.eq(wid)) @@ -236,7 +233,7 @@ pub async fn list_errors_for_workflow_run( wid: &str, run_id: &str, ) -> Result, anyhow::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let res = workflow_errors::Entity::find() .filter(workflow_errors::Column::OrgId.eq(org_id)) .filter(workflow_errors::Column::WorkflowId.eq(wid)) @@ -249,8 +246,7 @@ pub async fn list_errors_for_workflow_run( } pub async fn save_workflow(workflow: Workflow) -> Result<(), anyhow::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; - let _lock = get_lock().await; + let client = get_orm_client_rw().await; let now = chrono::Utc::now().timestamp_millis(); let model = workflows::ActiveModel { @@ -270,8 +266,7 @@ pub async fn save_workflow(workflow: Workflow) -> Result<(), anyhow::Error> { } pub async fn update_workflow(workflow: Workflow) -> Result<(), anyhow::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; - let _lock = get_lock().await; + let client = get_orm_client_rw().await; let now = chrono::Utc::now().timestamp_millis(); let model = workflows::ActiveModel { @@ -292,8 +287,7 @@ pub async fn update_workflow(workflow: Workflow) -> Result<(), anyhow::Error> { } pub async fn save_workflow_errors(errors: WorkflowRunErrors) -> Result<(), anyhow::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; - let _lock = get_lock().await; + let client = get_orm_client_rw().await; let model = workflow_errors::ActiveModel { cluster: Set(errors.cluster), org_id: Set(errors.org_id), @@ -311,8 +305,7 @@ pub async fn save_workflow_errors(errors: WorkflowRunErrors) -> Result<(), anyho pub async fn update_error_input_file_cluster_data( errors: WorkflowRunErrors, ) -> Result<(), anyhow::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; - let _lock = get_lock().await; + let client = get_orm_client_rw().await; workflow_errors::Entity::update_many() .filter(workflow_errors::Column::OrgId.eq(errors.org_id)) @@ -333,8 +326,7 @@ pub async fn update_error_input_file_cluster_data( } pub async fn delete_workflow(id: &str) -> Result<(), anyhow::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; - let _lock = get_lock().await; + let client = get_orm_client_rw().await; workflows::Entity::delete_many() .filter(workflows::Column::Id.eq(id)) @@ -348,8 +340,7 @@ pub async fn delete_workflow(id: &str) -> Result<(), anyhow::Error> { pub async fn delete_all_errors_older_than( ts: i64, ) -> Result, anyhow::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; - let _lock = get_lock().await; + let client = get_orm_client_rw().await; // ideally this could be delete returning, but sqlite does not // support that with seaorm, so transaction instead, which is still not @@ -385,8 +376,7 @@ pub async fn delete_all_errors_older_than( } pub async fn delete_all_runs_older_than(ts: i64) -> Result, anyhow::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; - let _lock = get_lock().await; + let client = get_orm_client_rw().await; // ideally this could be delete returning, but sqlite does not // support that with seaorm, so transaction instead, which is still not @@ -428,7 +418,7 @@ pub async fn get_run_data( workflow_id: &str, run_id: &str, ) -> Result, anyhow::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let ret = workflow_run_data::Entity::find() .filter(workflow_run_data::Column::OrgId.eq(org_id)) @@ -445,8 +435,7 @@ pub async fn delete_run_data( workflow_id: &str, run_id: &str, ) -> Result<(), anyhow::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; - let _lock = get_lock().await; + let client = get_orm_client_rw().await; workflow_run_data::Entity::delete_many() .filter(workflow_run_data::Column::OrgId.eq(org_id)) @@ -459,8 +448,7 @@ pub async fn delete_run_data( } pub async fn save_workflow_run_data(entry: WorkflowRunData) -> Result<(), anyhow::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; - let _lock = get_lock().await; + let client = get_orm_client_rw().await; let model = workflow_run_data::ActiveModel { org_id: Set(entry.org_id), @@ -480,7 +468,7 @@ pub async fn get_all_associations_for_workflow( org_id: &str, workflow_id: &str, ) -> Result, anyhow::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let ret = workflow_associations::Entity::find() .filter(workflow_associations::Column::OrgId.eq(org_id)) @@ -496,7 +484,7 @@ pub async fn get_all_associations_for_entity( entity_id: &str, entity_type: &str, ) -> Result, anyhow::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let ret = workflow_associations::Entity::find() .filter(workflow_associations::Column::OrgId.eq(org_id)) @@ -512,7 +500,7 @@ pub async fn get_all_associations_for_trigger_type( org_id: &str, trigger: &str, ) -> Result, anyhow::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let ret = workflow_associations::Entity::find() .filter(workflow_associations::Column::OrgId.eq(org_id)) @@ -524,8 +512,7 @@ pub async fn get_all_associations_for_trigger_type( } pub async fn add_workflow_association(entry: WorkflowAssociation) -> Result<(), anyhow::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; - let _lock = get_lock().await; + let client = get_orm_client_rw().await; let model = workflow_associations::ActiveModel { id: Default::default(), @@ -547,8 +534,7 @@ pub async fn delete_workflow_association( workflow_id: &str, entity_id: &str, ) -> Result<(), anyhow::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; - let _lock = get_lock().await; + let client = get_orm_client_rw().await; workflow_associations::Entity::delete_many() .filter(workflow_associations::Column::OrgId.eq(org_id)) @@ -560,8 +546,7 @@ pub async fn delete_workflow_association( } pub async fn delete_association_by_workflow(org_id: &str, id: &str) -> Result<(), anyhow::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; - let _lock = get_lock().await; + let client = get_orm_client_rw().await; workflow_associations::Entity::delete_many() .filter(workflow_associations::Column::OrgId.eq(org_id)) @@ -575,8 +560,7 @@ pub async fn delete_association_by_trigger( org_id: &str, trigger: &str, ) -> Result<(), anyhow::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; - let _lock = get_lock().await; + let client = get_orm_client_rw().await; workflow_associations::Entity::delete_many() .filter(workflow_associations::Column::OrgId.eq(org_id)) @@ -591,8 +575,7 @@ pub async fn delete_association_by_trigger_workflow( trigger: &str, workflow_id: &str, ) -> Result<(), anyhow::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; - let _lock = get_lock().await; + let client = get_orm_client_rw().await; workflow_associations::Entity::delete_many() .filter(workflow_associations::Column::OrgId.eq(org_id)) @@ -604,8 +587,7 @@ pub async fn delete_association_by_trigger_workflow( } pub async fn delete_association_by_entity(org_id: &str, entity: &str) -> Result<(), anyhow::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; - let _lock = get_lock().await; + let client = get_orm_client_rw().await; workflow_associations::Entity::delete_many() .filter(workflow_associations::Column::OrgId.eq(org_id)) @@ -616,7 +598,7 @@ pub async fn delete_association_by_entity(org_id: &str, entity: &str) -> Result< } pub async fn list_drafts_by_org(org_id: &str) -> Result, anyhow::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let entities = workflow_drafts::Entity::find() .filter(workflow_drafts::Column::OrgId.eq(org_id)) .all(client) @@ -632,7 +614,7 @@ pub async fn get_draft_by_org_draft_id( org_id: &str, id: &str, ) -> Result, anyhow::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; + let client = get_orm_client_ro().await; let res = workflow_drafts::Entity::find() .filter(workflow_drafts::Column::Id.eq(id)) .filter(workflow_drafts::Column::OrgId.eq(org_id)) @@ -644,8 +626,7 @@ pub async fn get_draft_by_org_draft_id( } pub async fn save_draft(draft: Workflow) -> Result<(), anyhow::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; - let _lock = get_lock().await; + let client = get_orm_client_rw().await; let now = chrono::Utc::now().timestamp_millis(); let model = workflow_drafts::ActiveModel { @@ -664,8 +645,7 @@ pub async fn save_draft(draft: Workflow) -> Result<(), anyhow::Error> { } pub async fn update_draft(draft: Workflow) -> Result<(), anyhow::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; - let _lock = get_lock().await; + let client = get_orm_client_rw().await; let now = chrono::Utc::now().timestamp_millis(); let model = workflow_drafts::ActiveModel { @@ -685,8 +665,7 @@ pub async fn update_draft(draft: Workflow) -> Result<(), anyhow::Error> { } pub async fn delete_draft(id: &str) -> Result<(), anyhow::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; - let _lock = get_lock().await; + let client = get_orm_client_rw().await; workflow_drafts::Entity::delete_many() .filter(workflow_drafts::Column::Id.eq(id)) @@ -696,8 +675,7 @@ pub async fn delete_draft(id: &str) -> Result<(), anyhow::Error> { } pub async fn promote_draft_to_workflow(org_id: &str, draft: Workflow) -> Result<(), anyhow::Error> { - let client = ORM_CLIENT.get_or_init(connect_to_orm).await; - let _lock = get_lock().await; + let client = get_orm_client_rw().await; let id = draft.id.clone(); diff --git a/src/jobs/src/job/scheduler.rs b/src/jobs/src/job/scheduler.rs index 57e13a9687..66c784a89c 100644 --- a/src/jobs/src/job/scheduler.rs +++ b/src/jobs/src/job/scheduler.rs @@ -15,6 +15,8 @@ use config::{cluster::LOCAL_NODE, get_config, spawn_pausable_job}; #[cfg(feature = "enterprise")] +use infra::db::get_orm_client_rw; +#[cfg(feature = "enterprise")] use o2_enterprise::enterprise::common::config::get_config as get_o2_config; use crate::service; @@ -145,14 +147,7 @@ async fn run_schedule_jobs() -> Result<(), anyhow::Error> { async fn cleanup_alert_dedup_state() -> Result<(), anyhow::Error> { log::debug!("[ALERT DEDUP CLEANUP] Starting cleanup of old deduplication state"); - // Get database connection - let db = match infra::db::ORM_CLIENT.get() { - Some(db) => db, - None => { - log::warn!("[ALERT DEDUP CLEANUP] ORM client not available, skipping cleanup"); - return Ok(()); - } - }; + let db = get_orm_client_rw().await; // Default cleanup: Remove records older than 24 hours // This is conservative - most alerts have shorter time windows diff --git a/src/jobs/src/job/slo_maintenance.rs b/src/jobs/src/job/slo_maintenance.rs index 408bc6c710..ce5364e3e0 100644 --- a/src/jobs/src/job/slo_maintenance.rs +++ b/src/jobs/src/job/slo_maintenance.rs @@ -33,6 +33,7 @@ //! each other. use config::{cluster::LOCAL_NODE, get_config, spawn_pausable_job, utils::time::now_micros}; +use infra::db::get_orm_client_rw; pub fn run() { if !LOCAL_NODE.is_scheduler() { @@ -67,9 +68,7 @@ pub fn run() { } async fn sweep() -> Result<(), anyhow::Error> { - let db = infra::db::ORM_CLIENT - .get() - .ok_or_else(|| anyhow::anyhow!("database not initialized"))?; + let db = get_orm_client_rw().await; let slos = infra::table::slos::list_enabled(db) .await diff --git a/src/migration/mod.rs b/src/migration/mod.rs index 14f34d5204..60ddaed8cb 100644 --- a/src/migration/mod.rs +++ b/src/migration/mod.rs @@ -21,7 +21,7 @@ //! - `migrate-file-list`: Migrate file_list related tables only use ::config::DB_SCHEMA_VERSION; -use infra::db::{ORM_CLIENT, ORM_CLIENT_DDL, connect_to_orm, connect_to_orm_ddl}; +use infra::db::{get_orm_client_ddl, get_orm_client_ro, get_orm_client_rw}; mod adapter; mod config; @@ -32,8 +32,9 @@ pub use config::MigrationConfig; pub use migrator::{run_file_list, run_meta}; pub async fn init_db() -> std::result::Result<(), anyhow::Error> { - // we init client here to avoid deadlocks - ORM_CLIENT.get_or_init(connect_to_orm).await; + // warm both pools before the migration starts hitting them + get_orm_client_ro().await; + get_orm_client_rw().await; // a missing version is Ok(0), so an error here means the db is // unreachable or overloaded: never assume a fresh install, retry then abort const MAX_RETRIES: usize = 5; @@ -79,7 +80,7 @@ pub async fn init_db() -> std::result::Result<(), anyhow::Error> { } // we initialize both clients here to avoid potential deadlock afterwards - ORM_CLIENT_DDL.get_or_init(connect_to_orm_ddl).await; + get_orm_client_ddl().await; // migrate infra_sea_orm if let Err(e) = infra::table::migrate().await { diff --git a/src/sourcemaps/src/lib.rs b/src/sourcemaps/src/lib.rs index d4c609e19b..ece4f7004d 100644 --- a/src/sourcemaps/src/lib.rs +++ b/src/sourcemaps/src/lib.rs @@ -544,9 +544,7 @@ mod tests { use sea_orm::{ConnectionTrait, Schema}; infra::db::create_table().await.unwrap(); - let client = infra::db::ORM_CLIENT - .get_or_init(infra::db::connect_to_orm) - .await; + let client = infra::db::get_orm_client_rw().await; let backend = client.get_database_backend(); let mut stmt = Schema::new(backend) .create_table_from_entity(infra::table::entity::source_maps::Entity); diff --git a/src/super_cluster_queue/src/alerts.rs b/src/super_cluster_queue/src/alerts.rs index 8df6633047..3c972429fb 100644 --- a/src/super_cluster_queue/src/alerts.rs +++ b/src/super_cluster_queue/src/alerts.rs @@ -13,11 +13,7 @@ // You should have received a copy of the GNU Affero General Public License // along with this program. If not, see . -use infra::{ - db::{ORM_CLIENT, connect_to_orm}, - errors::Result, - table, -}; +use infra::{db::get_orm_client_rw, errors::Result, table}; use o2_enterprise::enterprise::super_cluster::queue::{AlertMessage, Message, MessageType}; pub(crate) async fn process(msg: Message) -> Result<()> { @@ -38,7 +34,7 @@ pub(crate) async fn process(msg: Message) -> Result<()> { } pub(crate) async fn process_msg(msg: AlertMessage) -> Result<()> { - let conn = ORM_CLIENT.get_or_init(connect_to_orm).await; + let conn = get_orm_client_rw().await; match msg { AlertMessage::Create { org_id, diff --git a/src/super_cluster_queue/src/anomaly_detection.rs b/src/super_cluster_queue/src/anomaly_detection.rs index 278cb5eead..10306ba879 100644 --- a/src/super_cluster_queue/src/anomaly_detection.rs +++ b/src/super_cluster_queue/src/anomaly_detection.rs @@ -22,11 +22,7 @@ //! Only updates the local database for API reads. Training and detection run //! exclusively on the primary region (where the scheduler runs). -use infra::{ - db::{ORM_CLIENT, connect_to_orm}, - errors::Result, - table::anomaly_detection::config as table, -}; +use infra::{db::get_orm_client_rw, errors::Result, table::anomaly_detection::config as table}; use o2_enterprise::enterprise::super_cluster::queue::{AnomalyDetectionMessage, Message}; pub(crate) async fn process(msg: Message) -> Result<()> { @@ -41,7 +37,7 @@ pub(crate) async fn process(msg: Message) -> Result<()> { } pub(crate) async fn process_msg(msg: AnomalyDetectionMessage) -> Result<()> { - let db = ORM_CLIENT.get_or_init(connect_to_orm).await; + let db = get_orm_client_rw().await; match msg { AnomalyDetectionMessage::ConfigCreate { config } => { diff --git a/src/super_cluster_queue/src/eval_annotation_queues.rs b/src/super_cluster_queue/src/eval_annotation_queues.rs index 0374791148..7dc15ba3b8 100644 --- a/src/super_cluster_queue/src/eval_annotation_queues.rs +++ b/src/super_cluster_queue/src/eval_annotation_queues.rs @@ -6,7 +6,7 @@ // (at your option) any later version. use infra::{ - db::{ORM_CLIENT, connect_to_orm}, + db::get_orm_client_rw, errors::{Error, Result}, table::entity::{ llm_annotation_queue_bindings as queue_bindings, llm_annotation_queue_items as queue_items, @@ -54,7 +54,7 @@ async fn apply_put(queue: queues::Model, bindings: Vec) - )); } - let db = ORM_CLIENT.get_or_init(connect_to_orm).await; + let db = get_orm_client_rw().await; let txn = db.begin().await?; match queues::Entity::find_by_id(&queue.id).one(&txn).await? { Some(current) if current.org_id != queue.org_id => { @@ -93,7 +93,7 @@ async fn apply_put(queue: queues::Model, bindings: Vec) - } async fn apply_delete(org_id: &str, queue_id: &str) -> Result<()> { - let db = ORM_CLIENT.get_or_init(connect_to_orm).await; + let db = get_orm_client_rw().await; let txn = db.begin().await?; queue_bindings::Entity::delete_many() .filter(queue_bindings::Column::OrgId.eq(org_id)) diff --git a/src/super_cluster_queue/src/eval_datasets.rs b/src/super_cluster_queue/src/eval_datasets.rs index 23eee196d3..b08f7e0db9 100644 --- a/src/super_cluster_queue/src/eval_datasets.rs +++ b/src/super_cluster_queue/src/eval_datasets.rs @@ -6,7 +6,7 @@ // (at your option) any later version. use infra::{ - db::{ORM_CLIENT, connect_to_orm}, + db::get_orm_client_rw, errors::{Error, Result}, table::entity::{llm_dataset_items as dataset_items, llm_datasets as datasets}, }; @@ -50,7 +50,7 @@ async fn apply_put(dataset: datasets::Model, items: Vec) - )); } - let db = ORM_CLIENT.get_or_init(connect_to_orm).await; + let db = get_orm_client_rw().await; let txn = db.begin().await?; match datasets::Entity::find_by_id(&dataset.id).one(&txn).await? { Some(current) if current.org_id != dataset.org_id => { @@ -108,7 +108,7 @@ async fn apply_put(dataset: datasets::Model, items: Vec) - } async fn apply_delete(org_id: &str, dataset_id: &str) -> Result<()> { - let db = ORM_CLIENT.get_or_init(connect_to_orm).await; + let db = get_orm_client_rw().await; let txn = db.begin().await?; dataset_items::Entity::delete_many() .filter(dataset_items::Column::OrgId.eq(org_id)) diff --git a/src/super_cluster_queue/src/reports.rs b/src/super_cluster_queue/src/reports.rs index b3b52ec38f..6cc70cc648 100644 --- a/src/super_cluster_queue/src/reports.rs +++ b/src/super_cluster_queue/src/reports.rs @@ -15,11 +15,7 @@ use std::str::FromStr; -use infra::{ - db::{ORM_CLIENT, connect_to_orm}, - errors::Result, - table, -}; +use infra::{db::get_orm_client_rw, errors::Result, table}; use o2_enterprise::enterprise::super_cluster::queue::{Message, MessageType, ReportMessage}; use svix_ksuid::Ksuid; @@ -41,7 +37,7 @@ pub(crate) async fn process(msg: Message) -> Result<()> { } pub(crate) async fn process_msg(msg: ReportMessage) -> Result<()> { - let conn = ORM_CLIENT.get_or_init(connect_to_orm).await; + let conn = get_orm_client_rw().await; match msg { ReportMessage::Create { org_id, diff --git a/src/super_cluster_queue/src/scheduler.rs b/src/super_cluster_queue/src/scheduler.rs index e9bf315438..45ac6e1a65 100644 --- a/src/super_cluster_queue/src/scheduler.rs +++ b/src/super_cluster_queue/src/scheduler.rs @@ -22,7 +22,7 @@ use config::{ #[allow(clippy::single_component_path_imports)] use db; use infra::{ - db::{ORM_CLIENT, connect_to_orm}, + db::get_orm_client_ro, errors::{Error, Result}, scheduler, }; @@ -105,7 +105,7 @@ async fn update(msg: Message) -> Result<()> { } // First check if the module record exists in this region, to verify that the module record // is not deleted. - let conn = ORM_CLIENT.get_or_init(connect_to_orm).await; + let conn = get_orm_client_ro().await; match trigger.module { TriggerModule::Alert => { let Ok(alert_id) = svix_ksuid::Ksuid::from_str(&trigger.module_key) else { diff --git a/src/super_cluster_queue/src/synthetics.rs b/src/super_cluster_queue/src/synthetics.rs index 3065cb1bb6..62765acdd0 100644 --- a/src/super_cluster_queue/src/synthetics.rs +++ b/src/super_cluster_queue/src/synthetics.rs @@ -47,7 +47,7 @@ use config::meta::{ synthetics::Synthetic, }; use infra::{ - db::{ORM_CLIENT, connect_to_orm}, + db::get_orm_client_rw, errors::{Error, Result}, table, }; @@ -261,7 +261,7 @@ fn encrypt_with(dek: &[u8], check: &mut Synthetic) -> Result<()> { } async fn process_msg(msg: SyntheticsMessage) -> Result<()> { - let conn = ORM_CLIENT.get_or_init(connect_to_orm).await; + let conn = get_orm_client_rw().await; match msg { SyntheticsMessage::Create { org_id, payload } => { let folder_slug = payload.folder_slug.clone(); diff --git a/src/synthetics/src/dispatcher.rs b/src/synthetics/src/dispatcher.rs index acf59736db..ff4ce5490c 100644 --- a/src/synthetics/src/dispatcher.rs +++ b/src/synthetics/src/dispatcher.rs @@ -30,7 +30,7 @@ use aws_config::{BehaviorVersion, Region}; use aws_sdk_lambda::{Client as LambdaClient, primitives::Blob}; use dashmap::DashMap; use infra::{ - db::{ORM_CLIENT, connect_to_orm}, + db::get_orm_client_rw, table::{synthetics_checks, synthetics_jobs, synthetics_runs}, }; @@ -102,7 +102,7 @@ pub async fn run() { loop { tokio::time::sleep(TICK).await; - let db = ORM_CLIENT.get_or_init(connect_to_orm).await; + let db = get_orm_client_rw().await; let now_us = config::utils::time::now_micros(); @@ -258,10 +258,7 @@ async fn mint_default_token(org_id: &str) -> Option { } async fn mark_failure(row: &synthetics_jobs::LeasedRow, reason: Option<&str>) -> bool { - let db = match ORM_CLIENT.get() { - Some(db) => db, - None => return false, - }; + let db = get_orm_client_rw().await; let now_us = config::utils::time::now_micros(); // SyntheticStatus Error = DB int 4; job status Error = 6. // @@ -631,17 +628,15 @@ async fn handle_dispatch_failure( ingest_token: &str, error_msg: &str, ) { - let outcome = match ORM_CLIENT.get() { - Some(db) => synthetics_jobs::fail_dispatch( - db, - &row.id, - row.dispatch_attempts, - crate::MAX_DISPATCH_ATTEMPTS, - ) - .await - .unwrap_or(synthetics_jobs::DispatchFailureOutcome::DeadLettered), - None => synthetics_jobs::DispatchFailureOutcome::DeadLettered, - }; + let db = get_orm_client_rw().await; + let outcome = synthetics_jobs::fail_dispatch( + db, + &row.id, + row.dispatch_attempts, + crate::MAX_DISPATCH_ATTEMPTS, + ) + .await + .unwrap_or(synthetics_jobs::DispatchFailureOutcome::DeadLettered); match outcome { synthetics_jobs::DispatchFailureOutcome::Requeued => { diff --git a/src/synthetics/src/job_api.rs b/src/synthetics/src/job_api.rs index c12ebc6c65..373f001a5e 100644 --- a/src/synthetics/src/job_api.rs +++ b/src/synthetics/src/job_api.rs @@ -24,7 +24,7 @@ use std::collections::HashMap; use config::meta::synthetics::{Synthetic, SyntheticAuth, for_each_string_at_path}; use infra::{ - db::ORM_CLIENT, + db::{get_orm_client_ro, get_orm_client_rw}, table::{ org_ingestion_tokens, synthetics_checks, synthetics_jobs, synthetics_locations, synthetics_runs, @@ -265,9 +265,7 @@ pub async fn artifact_urls( req: ArtifactUrlsRequest, token_org: &str, ) -> anyhow::Result { - let conn = ORM_CLIENT - .get() - .ok_or_else(|| anyhow::anyhow!("Database not initialized"))?; + let conn = get_orm_client_ro().await; let check = synthetics_jobs::get_by_id(conn, &req.job_id) .await @@ -572,9 +570,7 @@ pub enum AlertDecision { /// `check.auth` is redacted in the response — the probe reads credentials from /// `env_inject` only so the encrypted blob never leaves the backend. pub async fn resolve(req: ResolveRequest, token_org: &str) -> anyhow::Result { - let conn = ORM_CLIENT - .get() - .ok_or_else(|| anyhow::anyhow!("Database not initialized"))?; + let conn = get_orm_client_rw().await; let check = synthetics_jobs::get_by_id(conn, &req.job_id) .await @@ -821,9 +817,7 @@ fn redact_auth(auth: SyntheticAuth) -> SyntheticAuth { /// Marks the job complete, increments the run counter, and returns context for notifications. /// Returns `run_complete = true` when all jobs in the run have acked. pub async fn ack(req: AckRequest, token_org: &str) -> anyhow::Result { - let conn = ORM_CLIENT - .get() - .ok_or_else(|| anyhow::anyhow!("Database not initialized"))?; + let conn = get_orm_client_rw().await; // Fetch the leased row first for context (location, check_id, org_id). let check = synthetics_jobs::get_by_id(conn, &req.job_id) diff --git a/src/synthetics/src/reaper/mod.rs b/src/synthetics/src/reaper/mod.rs index 747a87199b..4c0156b290 100644 --- a/src/synthetics/src/reaper/mod.rs +++ b/src/synthetics/src/reaper/mod.rs @@ -40,7 +40,7 @@ pub mod orphan; use std::time::Duration; use infra::{ - db::{ORM_CLIENT, connect_to_orm}, + db::get_orm_client_rw, table::{org_ingestion_tokens, synthetics_checks, synthetics_jobs, synthetics_runs}, }; @@ -56,7 +56,7 @@ pub async fn run() { loop { tokio::time::sleep(TICK).await; - let db = ORM_CLIENT.get_or_init(connect_to_orm).await; + let db = get_orm_client_rw().await; let now_us = config::utils::time::now_micros(); diff --git a/src/synthetics/src/reaper/orphan.rs b/src/synthetics/src/reaper/orphan.rs index f5555d8d70..2adb1cb17a 100644 --- a/src/synthetics/src/reaper/orphan.rs +++ b/src/synthetics/src/reaper/orphan.rs @@ -97,7 +97,7 @@ use std::{ use config::{cluster::LOCAL_NODE, meta::synthetics::SyntheticFrequency}; use infra::{ - db::{ORM_CLIENT, connect_to_orm}, + db::get_orm_client_ro, table::{org_ingestion_tokens, synthetics_checks, synthetics_checks::OrphanCandidate}, }; @@ -346,7 +346,7 @@ pub async fn run() { continue; } - let db = ORM_CLIENT.get_or_init(connect_to_orm).await; + let db = get_orm_client_ro().await; detect( db, &client, diff --git a/src/synthetics/src/scheduler.rs b/src/synthetics/src/scheduler.rs index 39efb91873..39dbcc1735 100644 --- a/src/synthetics/src/scheduler.rs +++ b/src/synthetics/src/scheduler.rs @@ -31,7 +31,7 @@ use config::{ utils::hash::{Sum64, fnv}, }; use infra::{ - db::{ORM_CLIENT, connect_to_orm}, + db::get_orm_client_rw, table::{synthetics_checks, synthetics_jobs, synthetics_locations, synthetics_runs}, }; use serde::Serialize; @@ -55,7 +55,7 @@ pub async fn run() { loop { tokio::time::sleep(TICK).await; - let db = ORM_CLIENT.get_or_init(connect_to_orm).await; + let db = get_orm_client_rw().await; let now_us = config::utils::time::now_micros(); diff --git a/src/synthetics/src/service/checks.rs b/src/synthetics/src/service/checks.rs index 0ac639631d..312c8210d4 100644 --- a/src/synthetics/src/service/checks.rs +++ b/src/synthetics/src/service/checks.rs @@ -1,3 +1,5 @@ +use infra::db::{get_orm_client_ro, get_orm_client_rw}; + use super::*; // ── Synthetics CRUD ────────────────────────────────────────────────────────────── @@ -52,7 +54,7 @@ pub async fn create_synthetic( body = encrypt_synthetic_auth(org_id, body).await?; body.owner = Some(created_by.to_owned()); - let conn = db()?; + let conn = get_orm_client_rw().await; let mut result = synthetics_checks::create(conn, org_id, body, false) .await .map_err(|e| anyhow::anyhow!(e.to_string()))?; @@ -109,7 +111,7 @@ pub async fn create_synthetic( } pub async fn get_synthetic(org_id: &str, id: &str) -> anyhow::Result> { - let conn = db()?; + let conn = get_orm_client_ro().await; let mut check = synthetics_checks::get(conn, org_id, id) .await .map_err(|e| anyhow::anyhow!(e.to_string()))?; @@ -133,7 +135,7 @@ pub async fn update_synthetic( id: &str, mut body: Synthetic, ) -> anyhow::Result { - let conn = db()?; + let conn = get_orm_client_rw().await; // Normalise locations exactly like create — a bare region stored on update // would never dispatch (region is derived from the "aws-" prefix). @@ -241,7 +243,7 @@ pub async fn update_synthetic( } pub async fn delete_synthetic(org_id: &str, id: &str) -> anyhow::Result { - let conn = db()?; + let conn = get_orm_client_rw().await; // Drain any queued checks before deleting. synthetics_jobs::drain_check(conn, id) .await @@ -271,7 +273,7 @@ pub async fn list_synthetics( org_id: &str, params: &ListSyntheticsParams, ) -> anyhow::Result { - let conn = db()?; + let conn = get_orm_client_rw().await; // synthetics_checks.folder_id stores the KSUID PK (folders.id), not the slug // (folders.folder_id). Resolve any slug → PK before filtering. Falls back to the @@ -351,7 +353,7 @@ pub async fn list_synthetics( /// (they expire via the reaper). Re-enabling resets next_run_at to 0 so the /// scheduler fires immediately on the next tick. pub async fn set_synthetic_enabled(org_id: &str, id: &str, enabled: bool) -> anyhow::Result { - let conn = db()?; + let conn = get_orm_client_rw().await; let changed = synthetics_checks::set_enabled(conn, org_id, id, enabled) .await .map_err(|e| anyhow::anyhow!(e.to_string()))?; @@ -394,7 +396,7 @@ pub async fn delete_synthetics_bulk( ids: &[String], _folder_id: Option<&str>, ) -> anyhow::Result<()> { - let conn = db()?; + let conn = get_orm_client_rw().await; let ofga = ofga_enabled(); // Hoisted so the loop does not re-read the config per id. #[cfg(feature = "enterprise")] @@ -432,7 +434,7 @@ pub async fn move_synthetics( ids: &[String], dst_folder_id: &str, ) -> anyhow::Result<()> { - let conn = db()?; + let conn = get_orm_client_rw().await; // Resolve slug → KSUID PK. Auto-create the default folder when necessary, // matching the behaviour of create_synthetic. @@ -531,7 +533,7 @@ pub async fn move_synthetics( /// Triggers an immediate run by resetting `next_run_at` to 0. The scheduler /// picks it up on the next tick (within 5 seconds) and inserts jobs. pub async fn run_synthetic_now(org_id: &str, id: &str) -> anyhow::Result<()> { - let conn = db()?; + let conn = get_orm_client_rw().await; let check = synthetics_checks::get(conn, org_id, id) .await .map_err(|e| anyhow::anyhow!(e.to_string()))? diff --git a/src/synthetics/src/service/locations.rs b/src/synthetics/src/service/locations.rs index 4d22b844fe..d3a992489e 100644 --- a/src/synthetics/src/service/locations.rs +++ b/src/synthetics/src/service/locations.rs @@ -1,3 +1,5 @@ +use infra::db::{get_orm_client_ro, get_orm_client_rw}; + use super::*; // ── Location registry (synthetics_locations table) ─────────────────────────── @@ -108,9 +110,13 @@ pub async fn list_locations_for_org(org_id: &str) -> anyhow::Result anyhow::Result anyhow::Result<()> { location_for_write(org_id, is_root, id).await?; - let conn = db()?; + let conn = get_orm_client_ro().await; let refs = synthetics_checks::count_referencing_location(conn, id) .await .map_err(|e| anyhow::anyhow!(e.to_string()))?; diff --git a/src/synthetics/src/service/mod.rs b/src/synthetics/src/service/mod.rs index cdbbf16640..94ffca232f 100644 --- a/src/synthetics/src/service/mod.rs +++ b/src/synthetics/src/service/mod.rs @@ -33,12 +33,9 @@ use config::meta::{ for_each_string_at_path, take_strings_at_path, }, }; -use infra::{ - db::ORM_CLIENT, - table::{ - cipher, folders, synthetics_agents, synthetics_checks, synthetics_jobs, - synthetics_locations, synthetics_runs, - }, +use infra::table::{ + cipher, folders, synthetics_agents, synthetics_checks, synthetics_jobs, synthetics_locations, + synthetics_runs, }; // ── OpenFGA ─────────────────────────────────────────────────────────────────── // @@ -96,12 +93,6 @@ pub use tokens::*; // ── DB helper ───────────────────────────────────────────────────────────────── -fn db() -> anyhow::Result<&'static sea_orm::DatabaseConnection> { - ORM_CLIENT - .get() - .ok_or_else(|| anyhow::anyhow!("Database not initialized")) -} - /// Mints this org's default synthetics folder, in THIS region only. /// /// It deliberately does not broadcast, and broadcasting would not help. diff --git a/src/synthetics/src/service/runs.rs b/src/synthetics/src/service/runs.rs index f8b1af2885..474a6f8207 100644 --- a/src/synthetics/src/service/runs.rs +++ b/src/synthetics/src/service/runs.rs @@ -1,3 +1,5 @@ +use infra::db::get_orm_client_ro; + use super::*; // ── Runs API ────────────────────────────────────────────────────────────────── @@ -72,7 +74,7 @@ pub async fn list_runs( page: i64, page_size: i64, ) -> anyhow::Result { - let conn = db()?; + let conn = get_orm_client_ro().await; let (rows, total) = synthetics_runs::list_runs( conn, synthetics_runs::ListRunsParams { @@ -98,7 +100,7 @@ pub async fn get_run_detail( synthetics_id: &str, run_id: &str, ) -> anyhow::Result> { - let conn = db()?; + let conn = get_orm_client_ro().await; let row = synthetics_runs::get_run(conn, run_id) .await .map_err(|e| anyhow::anyhow!(e.to_string()))?;