feat(db): add transaction-scoped team lock query (#48232)

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

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"
+}

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"
+}

litellm-rust/crates/db/queries/spend/lock_team.sql

@@ -0,0 +1 @@
+SELECT pg_advisory_xact_lock(hashtext($1)) IS NULL AS "locked!"

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,
+};

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;

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<Transaction<'static, Postgres>, 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<PgPool, Box<dyn StdError>>,
+    ) -> Result<(), Box<dyn StdError>> {
+        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(())
+    }
+}