diff --git a/litellm-rust/crates/db/.sqlx/query-085edfb9af2a914466c75b447bc7bb2d97532a9cd461cc5d8da1e1090560ca1f.json b/litellm-rust/crates/db/.sqlx/query-085edfb9af2a914466c75b447bc7bb2d97532a9cd461cc5d8da1e1090560ca1f.json new file mode 100644 index 00000000000..a415666d0be --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-085edfb9af2a914466c75b447bc7bb2d97532a9cd461cc5d8da1e1090560ca1f.json @@ -0,0 +1,23 @@ +{ + "db_name": "PostgreSQL", + "query": "SELECT pg_advisory_xact_lock(hashtext($1)) IS NULL AS \"locked!\"\n", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "locked!", + "type_info": "Bool", + "origin": "Expression" + } + ], + "parameters": { + "Left": [ + "Text" + ] + }, + "nullable": [ + null + ] + }, + "hash": "085edfb9af2a914466c75b447bc7bb2d97532a9cd461cc5d8da1e1090560ca1f" +} diff --git a/litellm-rust/crates/db/.sqlx/query-5c74f47a2b84574133c520d235192bf33e16fea1dbb0c32cf89a376a874927f4.json b/litellm-rust/crates/db/.sqlx/query-5c74f47a2b84574133c520d235192bf33e16fea1dbb0c32cf89a376a874927f4.json new file mode 100644 index 00000000000..7eef5278cb5 --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-5c74f47a2b84574133c520d235192bf33e16fea1dbb0c32cf89a376a874927f4.json @@ -0,0 +1,21 @@ +{ + "db_name": "PostgreSQL", + "query": "SELECT set_config('lock_timeout', '5000ms', true)", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "set_config", + "type_info": "Text", + "origin": "Expression" + } + ], + "parameters": { + "Left": [] + }, + "nullable": [ + null + ] + }, + "hash": "5c74f47a2b84574133c520d235192bf33e16fea1dbb0c32cf89a376a874927f4" +} diff --git a/litellm-rust/crates/db/queries/spend/lock_team.sql b/litellm-rust/crates/db/queries/spend/lock_team.sql new file mode 100644 index 00000000000..06cdb806ef0 --- /dev/null +++ b/litellm-rust/crates/db/queries/spend/lock_team.sql @@ -0,0 +1 @@ +SELECT pg_advisory_xact_lock(hashtext($1)) IS NULL AS "locked!" diff --git a/litellm-rust/crates/db/src/lib.rs b/litellm-rust/crates/db/src/lib.rs index 3a9032d73dc..4f009b411fd 100644 --- a/litellm-rust/crates/db/src/lib.rs +++ b/litellm-rust/crates/db/src/lib.rs @@ -18,4 +18,7 @@ pub use daily_spend::{DailyEndUserRow, DailyTeamRow, DailyUserRow}; pub use end_user::{EndUser, RateLimits, find_end_user}; pub use error::Error; pub use failure::DbFailure; -pub use spend::{EntitySpend, add_end_user_spend, add_key_spend, add_team_spend, add_user_spend}; +pub use spend::{ + EntitySpend, add_end_user_spend, add_key_spend, add_team_spend, add_user_spend, begin, + lock_team, +}; diff --git a/litellm-rust/crates/db/src/spend/mod.rs b/litellm-rust/crates/db/src/spend/mod.rs index 5191da2adb5..0459e63aaf8 100644 --- a/litellm-rust/crates/db/src/spend/mod.rs +++ b/litellm-rust/crates/db/src/spend/mod.rs @@ -1,9 +1,11 @@ mod end_user_spend; mod key_spend; +mod team_lock; mod team_spend; mod user_spend; pub use end_user_spend::add_end_user_spend; pub use key_spend::{EntitySpend, add_key_spend}; +pub use team_lock::{begin, lock_team}; pub use team_spend::add_team_spend; pub use user_spend::add_user_spend; diff --git a/litellm-rust/crates/db/src/spend/team_lock.rs b/litellm-rust/crates/db/src/spend/team_lock.rs new file mode 100644 index 00000000000..7a6b9c7f7d0 --- /dev/null +++ b/litellm-rust/crates/db/src/spend/team_lock.rs @@ -0,0 +1,98 @@ +use sqlx::{PgPool, Postgres, Transaction}; + +use crate::Error; + +pub async fn begin(pool: &PgPool) -> Result, Error> { + let mut transaction = pool.begin().await?; + sqlx::query!("SELECT set_config('lock_timeout', '5000ms', true)") + .fetch_one(&mut *transaction) + .await?; + Ok(transaction) +} + +pub async fn lock_team( + transaction: &mut Transaction<'_, Postgres>, + team_id: &str, +) -> Result<(), Error> { + sqlx::query_file!("queries/spend/lock_team.sql", team_id) + .fetch_one(&mut **transaction) + .await?; + Ok(()) +} + +#[cfg(test)] +mod tests { + use std::error::Error as StdError; + + use super::{begin, lock_team}; + use crate::test_support::database_pool; + use rstest::rstest; + use sqlx::{PgPool, postgres::PgPoolOptions}; + + #[rstest] + #[case::transaction_scoped("team-lock-scope")] + #[tokio::test] + #[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"] + async fn team_lock_is_held_until_rollback( + #[case] team_id: &str, + #[future(awt)] database_pool: Result>, + ) -> Result<(), Box> { + let database_pool = database_pool?; + let pool = PgPoolOptions::new() + .max_connections(2) + .connect_with(database_pool.connect_options().as_ref().clone()) + .await?; + let mut first = begin(&pool).await?; + let mut second = pool.begin().await?; + #[expect( + clippy::disallowed_methods, + reason = "Test assertions are outside the compile-checked query set" + )] + let original_timeout: String = sqlx::query_scalar("SELECT current_setting('lock_timeout')") + .fetch_one(&mut *second) + .await?; + lock_team(&mut first, team_id).await?; + #[expect( + clippy::disallowed_methods, + reason = "Test assertions are outside the compile-checked query set" + )] + let (backend, timeout): (i32, String) = + sqlx::query_as("SELECT pg_backend_pid(), current_setting('lock_timeout')") + .fetch_one(&mut *first) + .await?; + assert_eq!(timeout, "5s"); + #[expect( + clippy::disallowed_methods, + reason = "Test assertions are outside the compile-checked query set" + )] + let held: bool = sqlx::query_scalar("SELECT pg_try_advisory_xact_lock(hashtext($1))") + .bind(team_id) + .fetch_one(&mut *second) + .await?; + assert!(!held); + first.rollback().await?; + #[expect( + clippy::disallowed_methods, + reason = "Test assertions are outside the compile-checked query set" + )] + let released: bool = sqlx::query_scalar("SELECT pg_try_advisory_xact_lock(hashtext($1))") + .bind(team_id) + .fetch_one(&mut *second) + .await?; + assert!(released); + let mut next = pool.begin().await?; + #[expect( + clippy::disallowed_methods, + reason = "Test assertions are outside the compile-checked query set" + )] + let (next_backend, next_timeout): (i32, String) = + sqlx::query_as("SELECT pg_backend_pid(), current_setting('lock_timeout')") + .fetch_one(&mut *next) + .await?; + assert_eq!(next_backend, backend); + assert_eq!(next_timeout, original_timeout); + next.rollback().await?; + second.rollback().await?; + Ok(()) + } +}