starred/openobserve
1
0
Fork 0
mirror of https://github.com/openobserve/openobserve.git synced 2026-11-04 08:09:27 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor(infra): split ORM read/write clients, retire the sqlite write lock (#13927)

Browse Source 
Paired with o2-enterprise `refactor/orm-rw-ro-client` (same branch
name).

## Problem

Every ORM caller went through a single client backed by the sqlite
**write** pool, so writes were kept apart only by `table::get_lock()` —
a convention 365 call sites had to remember. Forgetting it surfaced as
`(code: 517) database is locked`: a deferred transaction that reads and
then upgrades to a write finds the snapshot moved under it, and sqlite
reports that **without ever consulting `busy_timeout`**, which is why
the failures were immediate rather than retried.

## Approach

Serialization now comes from the pool, not from a lock. The sqlite write
pool holds **one connection** — all sqlite allows anyway — so writers
queue on acquire and there is nothing to forget.

That is only affordable because reads move to the read-only pool.
Measured, same load, fresh data dir each run:

| | reads on the write connection | reads on the RO pool |
|---|---|---|
| 200 concurrent reads | p50 **135 ms** | p50 **1 ms** |
| writes during mixed load | p50 **189 ms** | p50 **14 ms** |

Both were 100% success with zero lock errors, so this is purely
throughput — but without the split, reads queue behind the single writer
and drag writes down with them.

## API

`infra::db` now exposes three functions and nothing else:

```rust
pub async fn get_orm_client_ro()  -> &'static DatabaseConnection
pub async fn get_orm_client_rw()  -> &'static DatabaseConnection
pub async fn get_orm_client_ddl() -> &'static DatabaseConnection
```

The statics behind them, every `connect_to_orm_*`, and the raw
sqlite/postgres pools are now private or `pub(crate)`. `ORM_CLIENT`,
`connect_to_orm` and `table::get_lock()` are gone, as is the raw
`CLIENT_RW` mutex (56 lock acquisitions, plus the 52 guard derefs that
existed only to reach through it).

Final split: **RO 247 / RW 360 / DDL 9**; `get_lock` 0, `ORM_CLIENT` 0,
`CLIENT_RW.lock()` 0.

## How read/write was decided

Per call site, then verified by resolving each call through **its own
file's imports** (bare-name matching is useless here — `get_by_id`,
`list` and `update` collide across modules) and propagating write-ness
along the edges that hand a connection to a callee.

A wrong write client only costs throughput; a wrong read client fails in
production. So the audit is zero-miss in the read direction, and a site
is only downgraded when *every* connection-passing callee resolves.

## Behaviour changes on sqlite

All three make sqlite match what postgres already does today — postgres
has never had this lock:

- check-then-write pairs like `folders::put` are no longer atomic; they
rely on their unique index (`folders_org_folder_id_idx`)
- `db::put` runs two sequential transactions, so concurrent puts to one
key are last-write-wins
- `db::delete` runs several statements outside a transaction

`write_pool_is_single_connection_on_sqlite` is now the **only** guard
against reintroducing 517 — please don't drop it.

## Review guide

The mutex removal is separable if you want to look at it on its own:
`src/infra/src/{db,file_list,scheduler,pipeline,schema/history}/sqlite.rs`.
Everything else is the client split.

## Verification

- `cargo check --workspace --all-targets` and the CI clippy command:
clean
- Tests match a HEAD baseline built with the same `.env`, exactly: infra
1410/2, core 1629/1, db 431/0, api-management 179/0. The three failures
are pre-existing and environmental (two need a postgres that isn't
running locally)
- Static scans, all zero: queries on the pool inside an open
transaction; transactions spanning a second connection acquisition;
ro/rw misclassification
- Runtime smoke on a fresh data dir: 200 concurrent reads, then 100
reads + 100 writes simultaneously — 400/400 HTTP 200, rows persisted,
**0** lock errors, 0 ERROR/WARN lines
- Enterprise verified via the `Cargo.toml.openobserve` swap (`cargo
check --all-targets`): clean

Also fixes `users::tests::set_up`, which truncated shared tables but
released its lock before the test body ran, letting a concurrent test
clear rows mid-test.

https://claude.ai/code/session_016ouf1RxiMpiRw5W7GvCUDh
This commit is contained in:
Hengfei Yang
2026-08-25 08:37:50 +00:00
committed by GitHub
parent abc0fdc42d
commit ccd254ad00
119 changed files with 1022 additions and 2134 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/api/common/src/auth/validator.rs
+3 -3
View File
@@ -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;
src/api/management/src/request/alerts/dedup_stats.rs
+3 -3
View File
@@ -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()
src/api/management/src/request/alerts/history.rs
+1 -3
View File
@@ -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(_) => {
src/api/management/src/request/alerts/mod.rs
+23 -23
View File
@@ -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<UserEmail>,
) -> 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<HashMap<String, String>>,
#[cfg(feature = "enterprise")] Headers(user_email): Headers<UserEmail>,
) -> 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<String> = 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<String, Trigger> =
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<String>)>,
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<UserEmail>,
Json(req_body): Json<MoveAlertsRequestBody>,
) -> 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
src/cli/basic/cli.rs
+2 -5
View File
@@ -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<bool, anyhow::Error> {
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" => {
src/core/src/alerts/alert.rs
+8 -18
View File
@@ -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<C: ConnectionTrait>(
}
pub async fn get_by_id_db(org_id: &str, alert_id: Ksuid) -> Result<Alert, AlertError> {
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<Vec<(Folder, Alert)>, 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
src/core/src/alerts/composite/service.rs
+9 -17
View File
@@ -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<Option<alert_composites::CompositeWithChildren>, 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.
src/core/src/alerts/deduplication.rs
+1 -7
View File
@@ -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<alert_dedup_state::Model, sea_orm::DbErr> {
// 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));
src/core/src/alerts/scheduler/handlers.rs
+12 -33
View File
@@ -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()))?
src/core/src/anomaly_detection.rs
+13 -48
View File
@@ -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<Vec<serde_json::Value>> {
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<Option<serde_json::Value>> {
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<serde_json::Value> {
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<String>,
folder_id: Option<String>,
) -> Result<serde_json::Value> {
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<serde_json::Value> {
// 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<serde_json::V
/// Run detection for a configuration
pub async fn detect_anomalies(org_id: &str, anomaly_id: &str) -> Result<serde_json::Value> {
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,
src/core/src/dashboards/reports.rs
+12 -12
View File
@@ -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<Report, ReportError> {
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<Vec<String>>,
) -> Result<Vec<table::reports::ListReportsQueryResult>, 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, &params)
.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?;
src/core/src/ofga/migrations/mod.rs
+5 -5
View File
@@ -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
}
src/core/src/org_cleanup.rs
+9 -13
View File
@@ -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}"))?;
src/core/src/slo/backfill.rs
+6 -7
View File
@@ -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<ChunkOutcome, anyhow::Error> {
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<bool, 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 status = slo_table::load_status(db, &slo.id, "").await?;
Ok(status.is_some_and(|s| s.definition_generation == slo.definition_generation))
}
src/core/src/slo/evaluate.rs
+5 -4
View File
@@ -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);
src/core/src/slo/job.rs
+6 -7
View File
@@ -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<PassOutcome, anyhow::Error> {
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<usize, anyhow::Error> {
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 {
src/core/src/slo/reconcile.rs
+2 -4
View File
@@ -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<Vec<Rebuilt>, 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());
src/core/src/slo/service.rs
+18 -40
View File
@@ -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<Vec<Slo>, 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<Vec<SloEligibleAlert>, 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<u64, SloError> {
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<bool, SloError> {
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<bool, SloError> {
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<Option<(Slo, Option<config::meta::slo::SloStatusView>)>, 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<Vec<(Slo, Option<config::meta::slo::SloStatusView>)>, 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<Vec<config::meta::slo::SloStatusView>, 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).
src/core/src/stream.rs
+1 -3
View File
@@ -63,9 +63,7 @@ async fn cleanup_related_resources(
.iter()
.filter_map(|alert| alert.id)
.collect::<Vec<_>>();
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)?;
src/core/src/synthetics.rs
+1 -3
View File
@@ -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
src/core/src/users.rs
+21 -17
View File
@@ -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",
src/db/src/alerts/alert.rs
+8 -8
View File
@@ -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<Option<Alert>, 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<Alert, infra::errors::Error> {
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<String>,
) -> 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(());
src/db/src/folders.rs
+2 -5
View File
@@ -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);
src/db/src/pipeline_errors.rs
+8 -19
View File
@@ -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<Vec<pipeline_last_errors::Model>, 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<Option<pipeline_last_errors::Model>, 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<u64, infra::errors::Error> {
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))
src/infra/src/db/mod.rs
+35 -22
View File
@@ -40,25 +40,44 @@ static SUPER_CLUSTER: OnceCell<Box<dyn Db>> = OnceCell::const_new();
pub const SQLITE_STORE: &str = "sqlite";
pub static ORM_CLIENT: OnceCell<DatabaseConnection> = OnceCell::const_new();
pub static ORM_CLIENT_DDL: OnceCell<DatabaseConnection> = OnceCell::const_new();
static ORM_CLIENT_RO: OnceCell<DatabaseConnection> = OnceCell::const_new();
pub(crate) static ORM_CLIENT_RW: OnceCell<DatabaseConnection> = OnceCell::const_new();
static ORM_CLIENT_DDL: OnceCell<DatabaseConnection> = 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);
}
}
src/infra/src/db/postgres.rs
+8 -8
View File
@@ -36,9 +36,9 @@ use tokio::sync::{OnceCell, mpsc};
use super::{DBIndex, IndexStatement};
use crate::errors::*;
pub static CLIENT: Lazy<Pool<Postgres>> = Lazy::new(|| connect(false, false));
pub static CLIENT_RO: Lazy<Pool<Postgres>> = Lazy::new(|| connect(true, false));
pub static CLIENT_DDL: Lazy<Pool<Postgres>> = Lazy::new(|| connect(false, true));
pub(crate) static CLIENT_RO: Lazy<Pool<Postgres>> = Lazy::new(|| connect(true, false));
pub(crate) static CLIENT_RW: Lazy<Pool<Postgres>> = Lazy::new(|| connect(false, false));
pub(crate) static CLIENT_DDL: Lazy<Pool<Postgres>> = Lazy::new(|| connect(false, true));
static INDICES: OnceCell<HashSet<DBIndex>> = OnceCell::const_new();
/// The three pools come from separately configured DSNs whose roles can have
@@ -47,7 +47,7 @@ static INDICES: OnceCell<HashSet<DBIndex>> = 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<Postgres> {
}
async fn cache_indices() -> HashSet<DBIndex> {
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<i64>,
) -> 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<super::UpdateFn>,
) -> 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 {
src/infra/src/db/sqlite.rs
+20 -28
View File
@@ -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<Pool<Sqlite>> = 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<Pool<Sqlite>> = Lazy::new(connect_rw);
pub static CLIENT_RW: Lazy<Arc<Mutex<Pool<Sqlite>>>> =
Lazy::new(|| Arc::new(Mutex::new(CLIENT_RW_POOL.clone())));
pub(crate) static CLIENT_RO: Lazy<Pool<Sqlite>> = Lazy::new(connect_ro);
/// Holds a single connection, so every writer serializes on it.
pub(crate) static CLIENT_RW: Lazy<Pool<Sqlite>> = Lazy::new(connect_rw);
static INDICES: OnceCell<HashSet<DBIndex>> = OnceCell::const_new();
pub static CHANNEL: Lazy<SqliteDbChannel> = Lazy::new(SqliteDbChannel::new);
@@ -75,9 +72,12 @@ fn connect_rw() -> Pool<Sqlite> {
.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(())
}
src/infra/src/file_list/postgres.rs
+30 -30
View File
@@ -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<i64> {
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<u64> {
let pool = CLIENT.clone();
let pool = CLIENT_RW.clone();
let mut conditions: Vec<String> = 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<stdHashMap<String, stdHashMap<String, i64>>> {
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<Vec<(i64, String, i64)>> {
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<i64> {
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);
}
src/infra/src/file_list/sqlite.rs
+25 -53
View File
@@ -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<Sqlite> = 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<i64> {
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<Vec<super::MergeJobRecord>> {
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<u64> {
let client = CLIENT_RW.clone();
let client = client.lock().await;
let mut conditions: Vec<String> = 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<Vec<(i64, String, i64)>> {
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::<Vec<_>>()
.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::<Vec<_>>();
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::<String, &str>("stream");
@@ -2070,7 +2043,7 @@ pub async fn create_table_index() -> Result<()> {
let id = r.get::<i64, &str>("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(())
src/infra/src/pipeline/postgres.rs
+4 -4
View File
@@ -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<Pipeline> {
let pool = CLIENT.clone();
let pool = CLIENT_RW.clone();
let pipeline = sqlx::query_as::<_, Pipeline>("SELECT * FROM pipeline WHERE id = $1;")
.bind(pipeline_id)
src/infra/src/pipeline/sqlite.rs
+3 -9
View File
@@ -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<Pipeline> {
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;")
src/infra/src/scheduler/postgres.rs
+14 -14
View File
@@ -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::<Vec<_>>()
.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<TriggerModule>,
) -> Result<Vec<Trigger>> {
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"])
src/infra/src/scheduler/sqlite.rs
+11 -23
View File
@@ -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<bool> {
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<TriggerModule>,
) -> Result<Vec<Trigger>> {
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<Trigger> = q.fetch_all(&*client).await?;
let jobs: Vec<Trigger> = 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"),
src/infra/src/schema/history/postgres.rs
+2 -2
View File
@@ -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();
src/infra/src/schema/history/sqlite.rs
+2 -4
View File
@@ -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(())
src/infra/src/table/action_scripts.rs
+12 -28
View File
@@ -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<Model> 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<String, errors::Error> {
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<String, errors::Error> {
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<Action, errors::Error> {
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<Action, errors::Error> {
pub async fn list(org_id: &str, limit: Option<i64>) -> Result<Vec<Action>, 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<i64>) -> Result<Vec<Action>, error
}
pub async fn contains(id: &str, org_id: &str) -> Result<bool, 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))
.filter(Column::OrgId.eq(org_id))
@@ -204,16 +192,13 @@ pub async fn contains(id: &str, org_id: &str) -> Result<bool, errors::Error> {
}
pub async fn len() -> Result<usize, errors::Error> {
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<bool, errors::Error> {
/// 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)
src/infra/src/table/alert_composites.rs
+1 -19
View File
@@ -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<C: ConnectionTrait>(
org: &str,
id: &str,
) -> Result<i64, sea_orm::DbErr> {
// 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<C>(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<C>(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()
src/infra/src/table/alert_eval_intervals.rs
+6 -12
View File
@@ -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<Vec<AlertEvalInterval>, 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<C: sea_orm::ConnectionTrait>(
/// 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<Option<i64>, 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<C: sea_orm::ConnectionTrait>(
/// 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<u64, 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_before_with(client, cutoff_us).await
}
@@ -263,10 +260,7 @@ pub async fn delete_before_with<C: sea_orm::ConnectionTrait>(
/// 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
}
src/infra/src/table/alert_incidents.rs
+21 -56
View File
@@ -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<Option<alert_incidents::Model>, 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<String>,
) -> Result<alert_incidents::Model, 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();
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<bool, 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;
// Use transaction for atomic update
let txn = client
.begin()
@@ -172,12 +163,9 @@ pub async fn update_status(
id: &str,
status: &str,
) -> Result<alert_incidents::Model, 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())))?;
@@ -202,12 +190,9 @@ pub async fn update_title(
id: &str,
title: &str,
) -> Result<alert_incidents::Model, 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())))?;
@@ -228,12 +213,9 @@ pub async fn update_severity(
id: &str,
severity: &str,
) -> Result<alert_incidents::Model, 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())))?;
@@ -255,7 +237,7 @@ pub async fn list(
limit: u64,
offset: u64,
) -> Result<Vec<alert_incidents::Model>, 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<Vec<alert_incident_alerts::Model>, 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<Option<alert_incidents::Model>, 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<Option<alert_incidents::Model>, 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<u64, errors::Error> {
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<u64, errors::Error> {
/// Count incidents with optional status filter
pub async fn count(org_id: &str, status: Option<&str>) -> Result<u64, 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));
@@ -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<serde_json::Value>,
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<i64>,
limit: Option<u64>,
) -> Result<Vec<alert_incidents::Model>, 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)
src/infra/src/table/alert_states.rs
+16 -40
View File
@@ -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<alert_states::Model> for AlertState {
/// Fetch the state row for one `(alert_id, group_key)`.
pub async fn get(alert_id: &str, group_key: &str) -> Result<Option<AlertState>, 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<Vec<AlertState>, 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<Vec<AlertState>, errors
/// All per-group rows for one alert.
pub async fn list_groups(alert_id: &str) -> Result<Vec<AlertState>, 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<C: sea_orm::ConnectionTrait + TransactionTrait>(
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<C: sea_orm::ConnectionTrait + TransactionTrait>(
/// 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<bool, errors::Error> {
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<C: sea_orm::ConnectionTrait + TransactionTr
plan: &GroupPlan,
alert_id: &str,
) -> Result<bool, 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?;
// §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<bool, errors::Error> {
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<C: sea_orm::ConnectionTrait + Transacti
episode: DeliveryEpisode,
outcome: DeliveryOutcome,
) -> Result<bool, 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?;
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<Vec<String>, 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<C: sea_orm::ConnectionTrait>(
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<C: sea_orm::ConnectionTrait>(
/// 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<u64, errors::Error> {
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<C: sea_orm::ConnectionTrait>(
conn: &C,
alert_id: &str,
) -> Result<u64, errors::Error> {
// 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<Vec<StateTransition>, 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<Vec<StateTransition>, 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<C: sea_orm::ConnectionTrait + TransactionTrait
conn: &C,
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 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<C: sea_orm::ConnectionTrait + TransactionTrait
/// Retention for the append-only transition log. Governed by audit needs, set
/// independently of the `triggers` stream retention.
pub async fn delete_transitions_before(cutoff: 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;
alert_state_transitions::Entity::delete_many()
.filter(alert_state_transitions::Column::At.lt(cutoff))
.exec(client)
src/infra/src/table/alerts/mod.rs
+3 -18
View File
@@ -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<C: ConnectionTrait>(
org_id: &str,
alert_id: Ksuid,
) -> Result<Option<(MetaFolder, MetaAlert)>, 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<C: ConnectionTrait>(
stream_name: &str,
alert_name: &str,
) -> Result<Option<(MetaFolder, MetaAlert)>, 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<C: TransactionTrait>(
folder_id: &str,
alert: MetaAlert,
) -> Result<MetaAlert, errors::Error> {
let _lock = super::get_lock().await;
let txn = conn.begin().await?;
let rslt: Result<alerts::Model, errors::Error> = match get_model_by_name(
&txn,
@@ -341,7 +338,6 @@ pub async fn create<C: TransactionTrait>(
alert: MetaAlert,
use_given_id: bool,
) -> Result<MetaAlert, errors::Error> {
let _lock = super::get_lock().await;
let txn = conn.begin().await?;
// Get the destination folder.
@@ -396,7 +392,6 @@ pub async fn update<C: TransactionTrait + ConnectionTrait>(
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<C: TransactionTrait + ConnectionTrait>(
};
// 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<C: ConnectionTrait>(
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<C: ConnectionTrait>(
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<C: ConnectionTrait>(
conn: &C,
params: ListAlertsParams,
) -> Result<Vec<(MetaFolder, MetaAlert)>, 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<C: ConnectionTrait>(
slo_id: &str,
exclude_alert_id: Option<&str>,
) -> Result<Vec<(i64, i64)>, 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<C: ConnectionTrait>(
org: &str,
slo_id: &str,
) -> Result<Vec<(String, String)>, 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<C: ConnectionTrait>(
org: &str,
alert_id: &str,
) -> Result<Option<i64>, 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<C: ConnectionTrait>(
/// Lists all alerts.
pub async fn list_all<C: ConnectionTrait>(conn: &C) -> Result<Vec<MetaAlert>, 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)
src/infra/src/table/anomaly_detection/config.rs
-8
View File
@@ -39,7 +39,6 @@ pub async fn get_by_id<C: ConnectionTrait>(
org_id: &str,
anomaly_id: &str,
) -> Result<Option<Model>> {
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<C: ConnectionTrait>(
/// Lists all configs for an organisation, ordered by `created_at` descending.
pub async fn list_by_org<C: ConnectionTrait>(conn: &C, org_id: &str) -> Result<Vec<Model>> {
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<C: ConnectionTrait>(conn: &C, org_id: &str) -> Result<V
/// Lists all enabled configs across all organisations.
/// Used by startup recovery to re-create any missing detection triggers.
pub async fn list_all_enabled<C: ConnectionTrait>(conn: &C) -> Result<Vec<Model>> {
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<C: TransactionTrait + ConnectionTrait>(
conn: &C,
config: Model,
) -> Result<Model> {
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<C: TransactionTrait + ConnectionTrait>(
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<C: TransactionTrait + ConnectionTrait>(
conn: &C,
incoming: Model,
) -> Result<Model> {
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<C: TransactionTrait + ConnectionTrait>(
conn: &C,
incoming: Model,
) -> Result<Model> {
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<C: TransactionTrait + ConnectionTrait>(
/// Deletes a config by `anomaly_id` + `org_id`.
pub async fn delete<C: ConnectionTrait>(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))
src/infra/src/table/anomaly_detection/models.rs
-3
View File
@@ -36,7 +36,6 @@ type Result<T> = std::result::Result<T, Error>;
/// object storage. Returns the inserted row so callers can inspect the stored
/// values.
pub async fn insert<C: ConnectionTrait>(conn: &C, model: Model) -> Result<Model> {
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<C: ConnectionTrait>(
conn: &C,
anomaly_id: &str,
) -> Result<Vec<Model>> {
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<C: ConnectionTrait>(
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))
src/infra/src/table/backfill_jobs.rs
+9 -19
View File
@@ -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<Model> for BackfillJob {
}
pub async fn get(org: &str, job_id: &str) -> Result<BackfillJob, 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::Id.eq(job_id))
@@ -69,7 +68,7 @@ pub async fn get(org: &str, job_id: &str) -> Result<BackfillJob, errors::Error>
}
pub async fn list_by_org(org: &str) -> Result<Vec<BackfillJob>, 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<Vec<BackfillJob>, 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()
src/infra/src/table/cipher.rs
+17 -37
View File
@@ -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<Option<Model>, 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<Option<Model>, er
}
pub async fn list_all(limit: Option<i64>) -> Result<Vec<CipherEntry>, 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<i64>,
) -> Result<Vec<CipherEntry>, 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<Vec<u8>, 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)
src/infra/src/table/compactor_manual_jobs.rs
+9 -24
View File
@@ -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<CompactorManualJob, 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::Id.eq(ksuid))
.one(client)
@@ -110,7 +107,7 @@ pub async fn get_by_key(
key: &str,
status: Option<Status>,
) -> Result<CompactorManualJob, errors::Error> {
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<Vec<CompactorManualJob>, 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<CompactorManualJob>) -> 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<CompactorManualJob>) -> 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))
src/infra/src/table/dashboards.rs
+9 -23
View File
@@ -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<Option<Dashboard>, 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<Option<(Folder, Dashboard)>, 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<Vec<(Folder, Dashboard)>, 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<Vec<(Folder, Dashboard
/// Lists all existing dashboards
pub async fn list_all() -> Result<Vec<(String, Dashboard)>, 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<Dashboard, errors::Error> {
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<JsonValue, errors::Error>
/// 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<String> = folders::Entity::find()
.filter(folders::Column::Org.eq(org_id))
src/infra/src/table/destinations.rs
+10 -19
View File
@@ -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<String>) =
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<Option<destinations::Destination>, 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<Vec<destinations::Destination>, 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<Vec<destinations::Destination>, 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<Vec<destinations::Destination>, 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)
src/infra/src/table/distinct_values.rs
+8 -22
View File
@@ -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<Vec<DistinctFieldRecord>, 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<u64, errors::Error> {
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
src/infra/src/table/enrichment_table_urls.rs
+17 -28
View File
@@ -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<Model> 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<Option<EnrichmentTableUrlRecord>, 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<Vec<EnrichmentTableUrlRecord>, errors::Error>` - List of records or error
pub async fn list_by_org(org: &str) -> Result<Vec<EnrichmentTableUrlRecord>, 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<Vec<EnrichmentTableUrlRecord>, 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<Option<EnrichmentTableUrlRecord>, errors::Error>` - The record or None
pub async fn get_by_id(job_id: &str) -> Result<Option<EnrichmentTableUrlRecord>, 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<Option<EnrichmentTableUrlRecord>,
/// # 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<bool, errors::Error>` - true if any job is processing
pub async fn has_processing_jobs(org: &str, table_name: &str) -> Result<bool, errors::Error> {
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<bool, er
/// # Returns
/// * `Result<(), errors::Error>` - 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<Vec<EnrichmentTableUrlRecord>, 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<Model> = 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)
src/infra/src/table/enrichment_tables.rs
+12 -23
View File
@@ -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<Vec<EnrichmentTableRecord>, 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<i64>,
) -> Result<Vec<EnrichmentTableRecord>, 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<EnrichmentTab
/// # Returns
/// * `Result<Vec<EnrichmentTableRecord>, errors::Error>` - List of records or error
pub async fn get_by_org(org: &str) -> Result<Vec<EnrichmentTableRecord>, 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<Vec<EnrichmentTableRecord>, errors:
/// # Returns
/// * `Result<bool, errors::Error>` - True if records exist, false otherwise
pub async fn contains(org: &str, table_name: &str) -> Result<bool, 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))
@@ -297,7 +290,7 @@ pub async fn contains(org: &str, table_name: &str) -> Result<bool, errors::Error
/// # Returns
/// * `Result<usize, errors::Error>` - Count of records
pub async fn count_by_org_and_name(org: &str, table_name: &str) -> Result<usize, errors::Error> {
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<usize,
/// # Returns
/// * `Result<usize, errors::Error>` - Total count of records
pub async fn count() -> Result<usize, errors::Error> {
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<usize, errors::Error> {
/// # 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<bool, errors::Error> {
/// # Returns
/// * `Result<Vec<(String, String)>, errors::Error>` - List of (org_id, table_name) pairs
pub async fn list() -> Result<Vec<(String, String)>, 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<Vec<(String, String)>, 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)
src/infra/src/table/external_alerts.rs
+6 -10
View File
@@ -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<ExternalAlertRecord, errors::Error> {
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<Option<ExternalAlertRecord>, 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<Vec<ExternalAlertRecord>, 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()))
src/infra/src/table/folders.rs
+11 -20
View File
@@ -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<Option<Folder>, 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<Option<Folder>, 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<Vec<Folder>, 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<Option<String>, 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<Option<String>, 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<Option<String>, errors::Error> {
pub async fn get_name_and_display_name_by_pk(
pk: &str,
) -> Result<Option<(String, String)>, 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)
src/infra/src/table/gen_ai_agents.rs
+6 -9
View File
@@ -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<u64, errors::Error> {
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<AgentRecord>) -> Result<usize, errors::Err
return Ok(0);
}
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 stored = 0;
for record in records {
@@ -138,7 +136,7 @@ pub async fn upsert_many(records: Vec<AgentRecord>) -> Result<usize, errors::Err
}
pub async fn list(org_id: &str, filter: &AgentListFilter) -> Result<Vec<Model>, 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<u64, 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 mut query = Entity::delete_many().filter(Column::OrgId.eq(org_id));
if let Some(source_stream) = source_stream {
src/infra/src/table/incident_events.rs
+7 -23
View File
@@ -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<Vec<IncidentEvent>, 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<Option<usize>, 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)
src/infra/src/table/incident_integrations.rs
+17 -26
View File
@@ -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<Option<IncidentIntegrationRecord>, 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<Vec<IncidentIntegrationRecord>, 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<IncidentIntegrationRecord, errors::Error> {
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<String, errors::Error> {
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<bool, 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))
@@ -479,7 +470,7 @@ pub async fn delete(org_id: &str, id: &str) -> Result<bool, errors::Error> {
/// List all observed senders for an integration.
pub async fn list_senders(integration_id: &str) -> Result<Vec<SenderRecord>, 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)
src/infra/src/table/kv_store.rs
+8 -18
View File
@@ -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<Option<Model>, 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<Option<Model>, 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<Vec<String>, 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<Vec<String>, 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)
src/infra/src/table/mod.rs
+4 -34
View File
@@ -13,14 +13,10 @@
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see <http://www.gnu.org/licenses/>.
use config::{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<u32, anyhow::Error> {
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<u32, anyhow::Error> {
}
pub async fn down(steps: Option<u32>) -> 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<tokio::sync::MutexGuard<'static, sqlx::Pool<sqlx::Sqlite>>> {
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) => {
src/infra/src/table/model_pricing.rs
+9 -14
View File
@@ -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<Model> for ModelPricingDefinition {
@@ -62,8 +59,7 @@ impl TryFrom<Model> for ModelPricingDefinition {
}
pub async fn put(item: ModelPricingDefinition) -> Result<ModelPricingDefinition, 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();
let tiers_json = json::to_value(&item.tiers)?;
@@ -140,7 +136,7 @@ pub async fn put(item: ModelPricingDefinition) -> Result<ModelPricingDefinition,
}
pub async fn get(org_id: &str, name: &str) -> Result<Option<ModelPricingDefinition>, 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<Option<ModelPricingDefiniti
}
pub async fn get_by_id(id: &str) -> Result<Option<ModelPricingDefinition>, 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<Option<ModelPricingDefinition>, 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<Vec<String>, 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<Vec<String>, Error> {
}
pub async fn list(org_id: &str) -> Result<Vec<ModelPricingDefinition>, 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<Vec<ModelPricingDefinition>, 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<bool, Error> {
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<bool, Error> {
/// 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<Vec<String>, 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)
src/infra/src/table/online_eval_jobs/mod.rs
+15 -26
View File
@@ -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<Model> 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<Option<OnlineEvalJob>, 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<Option<OnlineEvalJob>, errors::Error> {
}
pub async fn get_by_org(id: &str, org_id: &str) -> Result<Option<OnlineEvalJob>, 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<Option<OnlineEvalJob>,
}
pub async fn get_all_by_org(org_id: &str) -> Result<Vec<OnlineEvalJob>, 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<Vec<OnlineEvalJob>, 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<Vec<OnlineEvalJob>, 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<Vec<OnlineEvalJob>, 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<bool, errors::Error> {
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<bool, 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)).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)
src/infra/src/table/org_ai_toolsets.rs
+13 -22
View File
@@ -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<Option<OrgToolset>, 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<Option<OrgToolset>, 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<Option<OrgToolset>, 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<i64>,
) -> Result<Vec<OrgToolset>, 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(())
}
src/infra/src/table/org_cleanup_tasks.rs
+10 -21
View File
@@ -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<ActiveModel> = 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<Vec<CleanupTask>, 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<Vec<CleanupTask>, 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<bool, 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();
let result = Entity::update_many()
.col_expr(Column::Status, Expr::value("running"))
@@ -107,8 +100,7 @@ pub async fn mark_running(id: &str) -> Result<bool, errors::Error> {
}
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<u64, 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();
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<Vec<CleanupTask>, 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<u64, 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 result = Entity::delete_many()
.filter(Column::OrgId.eq(org_id))
.exec(client)
src/infra/src/table/org_ingestion_tokens.rs
+15 -24
View File
@@ -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<Option<OrgIngestionTokenRecord>, 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<Option<OrgIngestionTokenRecord>, 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<Vec<OrgIngestionTokenListRecord>, 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<Option<OrgIngestionTokenRecord>, 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<String, errors::Error> {
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<String, errors::Er
/// Delete all tokens for an org (cascade on org deletion).
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))
@@ -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<u64, errors::Error> {
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<u64, errors::Error> {
/// 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<Vec<(String, String, String)>, 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()
src/infra/src/table/org_storage_providers.rs
+6 -11
View File
@@ -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<Option<OrgStorageProvider>, 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<Vec<OrgStorageProvider>, 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<OrgStorageProvider> = 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)
src/infra/src/table/org_users.rs
+20 -44
View File
@@ -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<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;
// 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<OrgUserRecord, 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(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<OrgUserRecord, errors::Err
}
pub async fn get_admin(org_id: &str) -> Result<OrgUserRecord, errors::Error> {
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<OrgUserExpandedRecord, 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(
@@ -431,7 +410,7 @@ pub async fn get_user_by_rum_token(
org_id: &str,
rum_token: &str,
) -> Result<OrgUserExpandedRecord, 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::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<Vec<OrgUserRecord>, 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<Vec<OrgUserRecord>, error
}
pub async fn list_orgs_by_user(email: &str) -> Result<Vec<UserOrgExpandedRecord>, 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<Vec<UserOrgExpandedRecord>
}
pub async fn list(limit: Option<i64>) -> Result<Vec<OrgUserRecord>, 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<i64>) -> Result<Vec<OrgUserRecord>, 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
src/infra/src/table/organizations.rs
+16 -46
View File
@@ -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<OrganizationRecord, errors::Error> {
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<OrganizationRecord, errors::Error> {
}
pub async fn list(filter: ListFilter) -> Result<Vec<OrganizationRecord>, 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<Vec<OrganizationRecord>, errors:
}
pub async fn get_by_name(org_name: &str) -> Result<Vec<OrganizationRecord>, 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<Vec<OrganizationRecord>, 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<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;
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<String>) -> 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<bool, 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();
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<i64>,
) -> Result<bool, 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();
let result = Entity::update_many()
.col_expr(Column::Status, Expr::value(new_status))
src/infra/src/table/providers.rs
+10 -19
View File
@@ -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<Provider, errors::Error
}
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);
@@ -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<Option<Provider>, 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<Option<Provider>, errors::Error> {
}
pub async fn get_all_by_org(org_id: &str) -> Result<Vec<Provider>, 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<Vec<Provider>, errors::Error
}
pub async fn get_default(org_id: &str) -> Result<Option<Provider>, 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<Option<Provider>, 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<bool, 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)).one(client).await?;
Ok(record.is_some())
src/infra/src/table/ratelimit.rs
+13 -24
View File
@@ -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<String>,
user_role: Option<String>,
) -> Result<Vec<RatelimitRule>, 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<Option<RatelimitRule>, 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<Option<RatelimitRule>, 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<RatelimitRule>) -> 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<String> = rules
@@ -212,7 +208,7 @@ async fn add_upsert_batch(rules: Vec<RatelimitRule>) -> 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<RatelimitRule>) -> 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<RatelimitRule>) -> 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<Vec<RatelimitRule>, 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);
});
}
src/infra/src/table/re_pattern.rs
+17 -34
View File
@@ -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<Option<PatternEntry>, 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::<Model>()
@@ -148,7 +134,7 @@ pub async fn get(id: &str) -> Result<Option<PatternEntry>, errors::DbError> {
}
pub async fn list_by_org(org: &str) -> Result<Vec<PatternEntry>, 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<Vec<PatternEntry>, errors::Error>
}
pub async fn list_all() -> Result<Vec<PatternEntry>, 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::<Model>().all(client).await?;
@@ -176,9 +162,7 @@ pub async fn list_all() -> Result<Vec<PatternEntry>, 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)
src/infra/src/table/re_pattern_stream_map.rs
+16 -31
View File
@@ -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<PatternAssociationEntry>,
removed: Vec<PatternAssociationEntry>,
) -> 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<Vec<PatternAssociationEntry>, 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::<Model>()
@@ -219,7 +211,7 @@ pub async fn get_by_pattern_id(
}
pub async fn list_all() -> Result<Vec<PatternAssociationEntry>, 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::<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)
src/infra/src/table/reports/mod.rs
+5 -17
View File
@@ -60,10 +60,9 @@ pub async fn get_by_name<C: ConnectionTrait + TransactionTrait>(
folder_snowflake_id: &str,
report_name: &str,
) -> Result<Option<(MetaFolder, MetaReport)>, 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<C: ConnectionTrait + TransactionTrait>(
conn: &C,
report_id: &str,
) -> Result<Option<(MetaFolder, MetaReport)>, 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<C: ConnectionTrait + TransactionTrait>(
report: MetaReport,
report_id: Option<Ksuid>,
) -> 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<C: ConnectionTrait + TransactionTrait>(
new_folder_snowflake_id: Option<&str>,
report: MetaReport,
) -> Result<String, Error> {
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<C: ConnectionTrait + TransactionTrait>(
conn: &C,
report_id: &str,
) -> Result<String, Error> {
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<C: ConnectionTrait + TransactionTrait>(
new_folder_snowflake_id: Option<&str>,
report: MetaReport,
) -> Result<String, 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 Err(Error::ReportNotFound);
};
@@ -423,10 +417,9 @@ pub async fn delete_by_name<C: ConnectionTrait + TransactionTrait>(
folder_snowflake_id: &str,
report_name: &str,
) -> Result<String, Error> {
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<C: ConnectionTrait + TransactionTrait>(
/// Delete all reports.
pub async fn delete_all<C: ConnectionTrait>(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<C: ConnectionTrait>(
conn: &C,
params: &ListReportsParams,
) -> Result<Vec<ListReportsQueryResult>, Error> {
let _lock = super::get_lock().await;
let reports = queries::ListReportsQueryResult::get(conn, params).await?;
Ok(reports)
}
src/infra/src/table/score_configs.rs
+14 -23
View File
@@ -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<Model> 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<Option<ScoreConfig>, 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<Option<ScoreConfig>, 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<Option<ScoreConfig>, 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<Option<ScoreConfig>, 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<Vec<ScoreConfig>, 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<Vec<ScoreConfig>, 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<Model> = Entity::find()
@@ -274,9 +269,7 @@ pub async fn get_all_by_org(org_id: &str) -> Result<Vec<ScoreConfig>, 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<bool, 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)).one(client).await?;
Ok(record.is_some())
@@ -317,7 +308,7 @@ pub async fn exists(id: &str) -> Result<bool, errors::Error> {
/// 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<i32, 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))
src/infra/src/table/scorers.rs
+16 -25
View File
@@ -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<Model> 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<Option<Scorer>, 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<Option<Scorer>, 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<Option<Scorer>, 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<Option<Scorer>, 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<Option<Score
}
pub async fn get_versions(org_id: &str, entity_id: &str) -> Result<Vec<Scorer>, 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<Vec<Scorer>,
}
pub async fn get_all_by_org(org_id: &str) -> Result<Vec<Scorer>, errors::Error> {
let client = ORM_CLIENT.get_or_init(connect_to_orm).await;
let client = get_orm_client_rw().await;
let records: Vec<Model> = 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<Vec<Scorer>, errors::Error> {
let client = ORM_CLIENT.get_or_init(connect_to_orm).await;
let client = get_orm_client_rw().await;
let records: Vec<Model> = 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<bool, errors::Error> {
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<bool, 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)).one(client).await?;
Ok(record.is_some())
@@ -324,7 +315,7 @@ pub async fn exists(id: &str) -> Result<bool, errors::Error> {
/// 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<i32, 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))
src/infra/src/table/search_job/search_job_partitions.rs
+6 -14
View File
@@ -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<Model>) -> 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<Model>) -> Result<(), err
}
pub async fn get_partition_jobs(job_id: &str) -> Result<Vec<Model>, 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<Vec<Model>, 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))
src/infra/src/table/search_job/search_job_results.rs
+4 -10
View File
@@ -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<Vec<Model>, 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<Vec<Model>, 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))
src/infra/src/table/search_job/search_jobs.rs
+17 -44
View File
@@ -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<Option<Model>, 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<Option<Model>, errors::Error> {
}
pub async fn cancel_job(job_id: &str, update_at: i64) -> Result<i64, 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,
@@ -298,10 +286,7 @@ pub async fn cancel_job(job_id: &str, update_at: i64) -> Result<i64, errors::Err
}
pub async fn set(operator: SetOperator) -> Result<UpdateResult, 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();
@@ -333,10 +318,7 @@ pub async fn set(operator: SetOperator) -> Result<UpdateResult, errors::Error> {
}
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<String> = 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<Model, 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::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<Model, errors::Error> {
}
pub async fn list_status_by_org_id(org_id: &str) -> Result<Vec<Model>, 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<Vec<Model>, errors::E
}
pub async fn get_deleted_jobs() -> Result<Vec<Model>, 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))
src/infra/src/table/search_queue.rs
+10 -26
View File
@@ -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<usize, 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::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)
src/infra/src/table/service_streams.rs
+11 -16
View File
@@ -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<Vec<ServiceRecord>, 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<Vec<ServiceRecord>, 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<Vec<ServiceRecord>, 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<u64, 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 result = Entity::delete_many()
.filter(Column::OrgId.eq(org_id))
@@ -591,7 +587,7 @@ pub async fn delete_all(org_id: &str) -> Result<u64, errors::Error> {
/// List distinct organization IDs that have service_streams records
pub async fn list_distinct_orgs() -> Result<Vec<String>, 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<Vec<String>, errors::Error> {
/// Delete stale records (last_seen older than threshold)
pub async fn delete_stale(org_id: &str, older_than_micros: i64) -> Result<u64, 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 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<u64, e
/// Get total row count for an organization
pub async fn count(org_id: &str) -> Result<u64, errors::Error> {
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))
src/infra/src/table/sessions.rs
+7 -19
View File
@@ -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<Option<Model>, 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<Vec<Model>, 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<Vec<Model>, 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<u64, 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();
let result = Entity::delete_many()
src/infra/src/table/short_urls.rs
+12 -29
View File
@@ -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<boo
..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;
let inserted = Entity::insert(record)
.on_conflict(
sea_orm::sea_query::OnConflict::column(Column::ShortId)
@@ -154,10 +149,7 @@ pub async fn add(short_id: &str, original_url: &str, org_id: &str) -> Result<boo
}
pub async fn remove(short_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::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<ShortUrlRecord, 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::ShortId)
@@ -195,7 +187,7 @@ pub async fn get(short_id: &str, org_id: &str) -> Result<ShortUrlRecord, errors:
}
pub async fn list(limit: Option<i64>) -> Result<Vec<ShortUrlRecord>, 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<i64>) -> Result<Vec<ShortUrlRecord>, 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<i64>,
) -> Result<Vec<String>, 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<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;
Entity::delete_many()
.filter(Column::ShortId.is_in(short_ids))
.exec(client)
src/infra/src/table/slo.rs
+1 -16
View File
@@ -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<WriteOutcome, 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?;
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)
src/infra/src/table/slo_backfill_jobs.rs
+1 -16
View File
@@ -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)
src/infra/src/table/slo_budget.rs
+1 -16
View File
@@ -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 &current {
@@ -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<i64, 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(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))
src/infra/src/table/slos.rs
+1 -25
View File
@@ -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<GenerationEffect, 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(slo.id.clone())
@@ -294,9 +285,6 @@ pub async fn bump_generation(
id: &str,
now: i64,
) -> Result<Option<(i32, i32)>, 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<bool, 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(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<Vec
/// modules, which resolve their org through `ids_in_org` and therefore run
/// BEFORE this (`org_cleanup::step_delete_db_resources`).
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;
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))
src/infra/src/table/source_maps.rs
+8 -14
View File
@@ -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<SourceMap>) -> Result<(), errors::Error> {
})
.collect::<Vec<_>>();
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<String>,
version: Option<String>,
) -> 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<String>,
version: &Option<String>,
) -> Result<Option<SourceMap>, 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<String>,
version: Option<String>,
) -> Result<Vec<SourceMap>, 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<String>, Vec<String>, Vec<String>), 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))
src/infra/src/table/synthetics_agents.rs
+9 -19
View File
@@ -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<Vec<SyntheticsAgentRecord>, 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<Option<SyntheticsAgentRecord>, 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<std::collections::HashMap<String, i64>, errors::Error> {
let client = ORM_CLIENT.get_or_init(connect_to_orm).await;
let client = get_orm_client_ro().await;
let rows: Vec<Option<String>> = Entity::find()
.filter(Column::OrgId.eq(org_id))
.select_only()
@@ -279,7 +269,7 @@ pub async fn get_cached(agent_id: &str) -> Result<Option<SyntheticsAgentRecord>,
/// Reads an agent straight from the database, `last_seen_at` included.
pub async fn get(agent_id: &str) -> Result<Option<SyntheticsAgentRecord>, 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
src/infra/src/table/synthetics_checks.rs
-46
View File
@@ -180,7 +180,6 @@ pub async fn get<C: ConnectionTrait>(
org_id: &str,
id: &str,
) -> Result<Option<Synthetic>, 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<C: ConnectionTrait>(
org_id: &str,
params: &ListSyntheticsParams,
) -> Result<Vec<Synthetic>, 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<C: ConnectionTrait>(
org_id: &str,
params: &ListSyntheticsParams,
) -> Result<u64, errors::Error> {
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<C: ConnectionTrait>(
org_id: &str,
location_id: &str,
) -> Result<Vec<Synthetic>, 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<C: TransactionTrait>(
check: Synthetic,
use_given_id: bool,
) -> Result<Synthetic, errors::Error> {
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<C: TransactionTrait>(
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<C: TransactionTrait>(
id: &str,
check: Synthetic,
) -> Result<Synthetic, errors::Error> {
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<C: TransactionTrait>(
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<C: TransactionTrait>(
org_id: &str,
check: Synthetic,
) -> Result<Synthetic, errors::Error> {
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<C: TransactionTrait>(
};
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<C: ConnectionTrait>(
org_id: &str,
id: &str,
) -> Result<bool, errors::Error> {
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<C: ConnectionTrait>(
ids: &[String],
dst_folder_id: &str,
) -> Result<u64, errors::Error> {
let _lock = super::get_lock().await;
if ids.is_empty() {
return Ok(0);
}
@@ -491,11 +462,6 @@ pub async fn move_to_folder<C: ConnectionTrait>(
.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<C: ConnectionTrait>(
id: &str,
enabled: bool,
) -> Result<bool, errors::Error> {
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<C: ConnectionTrait>(
.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<C: ConnectionTrait>(
now_us: i64,
limit: u64,
) -> Result<Vec<DueCheck>, 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<C: ConnectionTrait>(
now_us: i64,
limit: u64,
) -> Result<Vec<OrphanCandidate>, 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<C>(
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<C: ConnectionTrait>(
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<C: ConnectionTrait>(
id: &str,
status: i32,
) -> Result<bool, errors::Error> {
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<C: ConnectionTrait>(
expected: AlertState,
state: AlertState,
) -> Result<bool, errors::Error> {
let _lock = super::get_lock().await;
let res = Entity::update_many()
.col_expr(
Column::ConsecutiveFailures,
src/infra/src/table/synthetics_jobs.rs
+1 -29
View File
@@ -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<C: ConnectionTrait>(
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<C: ConnectionTrait>(
conn: &C,
synthetics_id: &str,
) -> Result<u64, errors::Error> {
// 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<C: ConnectionTrait>(
// 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<C: ConnectionTrait>(
.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<C: ConnectionTrait>(
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<C: ConnectionTrait>(
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<C: ConnectionTrait>(
// 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<C: ConnectionTrait>(
current_attempts: i32,
max_attempts: i32,
) -> Result<DispatchFailureOutcome, errors::Error> {
// 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<C: ConnectionTrait>(
/// count here is a bug signal, not routine housekeeping — the reaper logs it at
/// warn for that reason.
pub async fn prune_stale<C: ConnectionTrait>(conn: &C, now_us: i64) -> Result<u64, errors::Error> {
// 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
src/infra/src/table/synthetics_locations.rs
+9 -32
View File
@@ -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<SyntheticsLocationRecord> = Entity::find()
.order_by_asc(Column::Kind)
.order_by_asc(Column::Label)
@@ -159,8 +156,7 @@ impl From<Model> 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<Vec<SyntheticsLocationRecord>,
/// All private rows across orgs — used by the staleness watcher.
pub async fn list_private() -> Result<Vec<SyntheticsLocationRecord>, 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<bool, 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::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<bool,
/// the result is the same. Guarding it would leave the flag set if the one node
/// that "won" the clear died before the next tick, silencing the next outage.
pub async fn clear_down_notification(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::update_many()
.col_expr(Column::DownNotifiedAt, Expr::value(0i64))
.filter(Column::Id.eq(id))
src/infra/src/table/synthetics_probe_tokens.rs
+9 -30
View File
@@ -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<Option<SyntheticsProbeTokenRecor
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::Token.eq(token))
.filter(Column::Enabled.eq(true))
@@ -206,7 +197,7 @@ pub async fn find_default(
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::IsDefault.eq(true))
@@ -224,7 +215,7 @@ pub async fn find_default(
/// Callers mask the token value before returning it to a UI.
pub async fn list_by_org(org_id: &str) -> Result<Vec<SyntheticsProbeTokenRecord>, 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<Option<SyntheticsProbeTokenRecord>, 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(())
}
src/infra/src/table/synthetics_runs.rs
-8
View File
@@ -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<C: ConnectionTrait>(
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<C: ConnectionTrait>(
job_result: i32,
now_us: i64,
) -> Result<Option<(i32, i32)>, 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
src/infra/src/table/system_prompts.rs
+11 -22
View File
@@ -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<Model> 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<Option<AIPrompt>, 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<Option<AIPrompt>, errors::Error>
///
/// If the user prompt is not found, it returns None
pub async fn get_all() -> Result<(Option<AIPrompt>, Option<AIPrompt>), 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<AIPrompt>, Option<AIPrompt>), errors::E
}
pub async fn exists(r#type: PromptType) -> Result<bool, 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()))
.one(client)
@@ -145,16 +137,13 @@ pub async fn exists(r#type: PromptType) -> Result<bool, errors::Error> {
}
pub async fn len() -> Result<usize, errors::Error> {
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(())
src/infra/src/table/system_settings.rs
+7 -20
View File
@@ -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<Option<SystemSetting>, 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<Vec<SystemSetting>, 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<SystemSetting, anyhow::Error> {
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<bool, anyhow::Error> {
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<u64, anyhow::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 result = Entity::delete_many()
.filter(Column::OrgId.eq(org_id))
@@ -316,10 +306,7 @@ pub async fn delete_org_settings(org_id: &str) -> Result<u64, anyhow::Error> {
/// Delete all settings for a user in an organization
pub async fn delete_user_settings(org_id: &str, user_id: &str) -> Result<u64, anyhow::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 result = Entity::delete_many()
.filter(Column::OrgId.eq(org_id))
src/infra/src/table/templates.rs
+8 -18
View File
@@ -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<Model> for Template {
}
pub async fn put(template: Template) -> Result<Template, 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 title = match &template.template_type {
TemplateType::Email { title } => Some(title.to_string()),
_ => None,
@@ -93,7 +87,7 @@ pub async fn put(template: Template) -> Result<Template, Error> {
}
pub async fn get(org_id: &str, name: &str) -> Result<Option<Template>, 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<Option<Template>, Error> {
}
pub async fn list(org_id: &str) -> Result<Vec<Template>, 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<Vec<Template>, Error> {
}
pub async fn list_all() -> Result<Vec<(String, Template)>, 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<Vec<(String, Template)>, 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)
src/infra/src/table/timed_annotation_panels.rs
+5 -14
View File
@@ -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<Vec<String>, 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()
src/infra/src/table/timed_annotations.rs
+9 -36
View File
@@ -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<Vec<TimedAnnotation>, 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<String> = folders::Entity::find()
@@ -331,11 +313,8 @@ pub async fn get_one(
dashboard_id: &str,
timed_annotation_id: &str,
) -> Result<TimedAnnotation, 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_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<TimedAnnotation>,
use_given_id: bool,
) -> Result<Vec<TimedAnnotation>, 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();
src/infra/src/table/trial_quota_usage.rs
+13 -18
View File
@@ -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<Vec<trial_quota_usage::Model>, 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<i64, sea_orm::DbErr> {
let db = ORM_CLIENT.get_or_init(connect_to_orm).await;
let db = get_orm_client_ro().await;
let result: Option<Option<i64>> = 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<i64, sea_orm::DbErr
/// Get the explicitly configured shared usage limit for an organization.
/// A missing value means the deployment-wide default still applies.
pub async fn get_usage_limit_for_org(org_id: &str) -> Result<Option<i64>, 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<Option<i64>, sea_or
}
pub async fn load_all_usage_limits() -> Result<Vec<(String, i64)>, 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<i64>)> = 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<Vec<(String, i64)>, 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<Option<trial_quota_usage::Model>, 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<i16, 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))
.column_as(
@@ -210,8 +208,7 @@ pub async fn update_notified_checkpoint(
org_id: &str,
checkpoint: i16,
) -> Result<bool, 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 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<Vec<(String, i16)>, 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<i16>)> = trial_quota_usage::Entity::find()
.select_only()
.column(trial_quota_usage::Column::OrgId)
@@ -263,8 +259,7 @@ pub async fn load_all_checkpoints() -> Result<Vec<(String, i16)>, 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)
src/infra/src/table/users.rs
+12 -30
View File
@@ -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<String>,
) -> Result<u64, 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 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<UserRecord, 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(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<UserRecord, errors::Error> {
}
pub async fn get_root_user() -> Result<UserRecord, 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::IsRoot.eq(true))
.one(client)
@@ -228,7 +216,7 @@ pub async fn get_root_user() -> Result<UserRecord, errors::Error> {
}
pub async fn list(limit: Option<i64>) -> Result<Vec<UserRecord>, 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<i64>) -> Result<Vec<UserRecord>, 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<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 lowered_emails: Vec<String> = emails.iter().map(|e| e.to_lowercase()).collect();
Entity::delete_many()
.filter(Expr::expr(Func::lower(Expr::col(Column::Email))).is_in(lowered_emails))
src/infra/src/table/workflows.rs
+32 -54
View File
@@ -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<Vec<Workflow>, 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<Vec<Workflow>, anyhow::Error> {
}
pub async fn list_by_org(org_id: &str) -> Result<Vec<Workflow>, 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<Option<Workflow>, 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<Option<WorkflowRunErrors>, 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<Option<WorkflowRunErrors>, 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<Vec<WorkflowRunErrors>, 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<Vec<WorkflowRunData>, 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<Option<String>, 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<Vec<WorkflowAssociation>, 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<Vec<WorkflowAssociation>, 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<Vec<WorkflowAssociation>, 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<Vec<Workflow>, 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<Option<Workflow>, 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();
src/jobs/src/job/scheduler.rs
+3 -8
View File
@@ -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

Some files were not shown because too many files have changed in this diff Show More