diff --git a/litellm-rust/crates/db/.sqlx/query-5c74f47a2b84574133c520d235192bf33e16fea1dbb0c32cf89a376a874927f4.json b/litellm-rust/crates/db/.sqlx/query-805611b4198a8f11e66e29b77afa021f3a07befaedd77d4949e9401b120a5322.json similarity index 59% rename from litellm-rust/crates/db/.sqlx/query-5c74f47a2b84574133c520d235192bf33e16fea1dbb0c32cf89a376a874927f4.json rename to litellm-rust/crates/db/.sqlx/query-805611b4198a8f11e66e29b77afa021f3a07befaedd77d4949e9401b120a5322.json index 7eef5278cb5..7fb7d2aed5f 100644 --- a/litellm-rust/crates/db/.sqlx/query-5c74f47a2b84574133c520d235192bf33e16fea1dbb0c32cf89a376a874927f4.json +++ b/litellm-rust/crates/db/.sqlx/query-805611b4198a8f11e66e29b77afa021f3a07befaedd77d4949e9401b120a5322.json @@ -1,6 +1,6 @@ { "db_name": "PostgreSQL", - "query": "SELECT set_config('lock_timeout', '5000ms', true)", + "query": "SELECT set_config('lock_timeout', $1::text, true)", "describe": { "columns": [ { @@ -11,11 +11,13 @@ } ], "parameters": { - "Left": [] + "Left": [ + "Text" + ] }, "nullable": [ null ] }, - "hash": "5c74f47a2b84574133c520d235192bf33e16fea1dbb0c32cf89a376a874927f4" + "hash": "805611b4198a8f11e66e29b77afa021f3a07befaedd77d4949e9401b120a5322" } diff --git a/litellm-rust/crates/db/src/error.rs b/litellm-rust/crates/db/src/error.rs index af0ee030896..1bd1416d0ee 100644 --- a/litellm-rust/crates/db/src/error.rs +++ b/litellm-rust/crates/db/src/error.rs @@ -8,21 +8,64 @@ pub enum Error { InvalidTimestamp(#[from] chrono::ParseError), } +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ErrorKind { + Unreachable, + Capacity, + Deadlock, + SerializationFailure, + LockTimeout, + ReadOnly, + DataRejected, + Other, +} + impl Error { - pub fn is_retryable(&self) -> bool { + pub fn kind(&self) -> ErrorKind { match self { - Self::Query(sqlx::Error::Database(error)) => { - matches!( - error.code().as_deref(), - Some("40P01" | "55P03" | "40001" | "53300") - ) - } - Self::Query(sqlx::Error::Io(_) | sqlx::Error::PoolTimedOut) => true, + Self::Query(sqlx::Error::Io(_)) => ErrorKind::Unreachable, + Self::Query(sqlx::Error::PoolTimedOut) => ErrorKind::Capacity, + Self::Query(sqlx::Error::Database(error)) => sqlstate_kind(error.code().as_deref()), + _ => ErrorKind::Other, + } + } + + pub fn is_retryable(&self) -> bool { + matches!( + self.kind(), + ErrorKind::Unreachable + | ErrorKind::Capacity + | ErrorKind::Deadlock + | ErrorKind::SerializationFailure + | ErrorKind::LockTimeout + ) + } + + pub fn is_unavailable(&self) -> bool { + matches!(self.kind(), ErrorKind::Unreachable | ErrorKind::Capacity) + } + + pub fn is_requeue_safe(&self) -> bool { + match self { + Self::Query(sqlx::Error::PoolTimedOut | sqlx::Error::PoolClosed) => true, + Self::Query(sqlx::Error::Database(_)) => self.kind() != ErrorKind::DataRejected, _ => false, } } } +fn sqlstate_kind(code: Option<&str>) -> ErrorKind { + match code { + Some("53300") => ErrorKind::Capacity, + Some("40P01") => ErrorKind::Deadlock, + Some("40001") => ErrorKind::SerializationFailure, + Some("55P03") => ErrorKind::LockTimeout, + Some("25006") => ErrorKind::ReadOnly, + Some(code) if code.starts_with("22") || code.starts_with("23") => ErrorKind::DataRejected, + _ => ErrorKind::Other, + } +} + #[cfg(test)] mod tests { use std::error::Error as StdError; @@ -30,29 +73,52 @@ mod tests { use rstest::rstest; use sqlx::PgPool; - use super::Error; + use super::{Error, ErrorKind}; use crate::test_support::database_pool; #[rstest] - #[case::closed_pool(sqlx::Error::PoolClosed, false)] - #[case::pool_timeout(sqlx::Error::PoolTimedOut, true)] - #[case::io(sqlx::Error::Io(std::io::ErrorKind::ConnectionReset.into()), true)] - #[case::missing_row(sqlx::Error::RowNotFound, false)] - fn classifies_client_errors(#[case] error: sqlx::Error, #[case] retryable: bool) { - assert_eq!(Error::from(error).is_retryable(), retryable); + #[case::closed_pool(sqlx::Error::PoolClosed, ErrorKind::Other, false, false, true)] + #[case::pool_timeout(sqlx::Error::PoolTimedOut, ErrorKind::Capacity, true, true, true)] + #[case::io( + sqlx::Error::Io(std::io::ErrorKind::ConnectionReset.into()), + ErrorKind::Unreachable, + true, + true, + false + )] + #[case::missing_row(sqlx::Error::RowNotFound, ErrorKind::Other, false, false, false)] + fn classifies_client_errors( + #[case] error: sqlx::Error, + #[case] kind: ErrorKind, + #[case] retryable: bool, + #[case] unavailable: bool, + #[case] requeue_safe: bool, + ) { + let error = Error::from(error); + + assert_eq!(error.kind(), kind); + assert_eq!(error.is_retryable(), retryable); + assert_eq!(error.is_unavailable(), unavailable); + assert_eq!(error.is_requeue_safe(), requeue_safe); } #[rstest] - #[case::connection_capacity("53300", true)] - #[case::deadlock("40P01", true)] - #[case::lock_unavailable("55P03", true)] - #[case::serialization_failure("40001", true)] - #[case::unique_violation("23505", false)] + #[case::connection_capacity("53300", ErrorKind::Capacity, true, true, true)] + #[case::deadlock("40P01", ErrorKind::Deadlock, true, false, true)] + #[case::lock_unavailable("55P03", ErrorKind::LockTimeout, true, false, true)] + #[case::serialization_failure("40001", ErrorKind::SerializationFailure, true, false, true)] + #[case::read_only_after_failover("25006", ErrorKind::ReadOnly, false, false, true)] + #[case::unique_violation("23505", ErrorKind::DataRejected, false, false, false)] + #[case::numeric_overflow("22003", ErrorKind::DataRejected, false, false, false)] + #[case::undefined_table("42P01", ErrorKind::Other, false, false, true)] #[tokio::test] #[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"] async fn classifies_server_errors( #[case] code: &str, + #[case] kind: ErrorKind, #[case] retryable: bool, + #[case] unavailable: bool, + #[case] requeue_safe: bool, #[future(awt)] database_pool: Result>, ) -> Result<(), Box> { let pool = database_pool?; @@ -72,7 +138,11 @@ mod tests { error.as_database_error().and_then(|error| error.code()), Some(code.into()) ); - assert_eq!(Error::from(error).is_retryable(), retryable); + let error = Error::from(error); + assert_eq!(error.kind(), kind); + assert_eq!(error.is_retryable(), retryable); + assert_eq!(error.is_unavailable(), unavailable); + assert_eq!(error.is_requeue_safe(), requeue_safe); tx.rollback().await?; Ok(()) } diff --git a/litellm-rust/crates/db/src/lib.rs b/litellm-rust/crates/db/src/lib.rs index 541f97c7db9..3a29c1625ba 100644 --- a/litellm-rust/crates/db/src/lib.rs +++ b/litellm-rust/crates/db/src/lib.rs @@ -16,7 +16,7 @@ mod test_support; pub use daily_spend::{DailyEndUserRow, DailyTeamRow, DailyUserRow}; pub use end_user::{EndUser, RateLimits, find_end_user}; -pub use error::Error; +pub use error::{Error, ErrorKind}; pub use failure::DbFailure; pub use spend::{ DailySpend, DailyTable, EntitySpend, SpendLog, SpendLogRow, TeamMemberSpend, diff --git a/litellm-rust/crates/db/src/spend/team_lock.rs b/litellm-rust/crates/db/src/spend/team_lock.rs index 7a6b9c7f7d0..17f854d23c2 100644 --- a/litellm-rust/crates/db/src/spend/team_lock.rs +++ b/litellm-rust/crates/db/src/spend/team_lock.rs @@ -1,12 +1,20 @@ +use std::time::Duration; + use sqlx::{PgPool, Postgres, Transaction}; use crate::Error; -pub async fn begin(pool: &PgPool) -> Result, Error> { +pub async fn begin( + pool: &PgPool, + lock_timeout: Duration, +) -> Result, Error> { let mut transaction = pool.begin().await?; - sqlx::query!("SELECT set_config('lock_timeout', '5000ms', true)") - .fetch_one(&mut *transaction) - .await?; + sqlx::query!( + "SELECT set_config('lock_timeout', $1::text, true)", + format!("{}ms", lock_timeout.as_millis()) + ) + .fetch_one(&mut *transaction) + .await?; Ok(transaction) } @@ -22,7 +30,7 @@ pub async fn lock_team( #[cfg(test)] mod tests { - use std::error::Error as StdError; + use std::{error::Error as StdError, time::Duration}; use super::{begin, lock_team}; use crate::test_support::database_pool; @@ -42,7 +50,7 @@ mod tests { .max_connections(2) .connect_with(database_pool.connect_options().as_ref().clone()) .await?; - let mut first = begin(&pool).await?; + let mut first = begin(&pool, Duration::from_millis(1500)).await?; let mut second = pool.begin().await?; #[expect( clippy::disallowed_methods, @@ -60,7 +68,7 @@ mod tests { sqlx::query_as("SELECT pg_backend_pid(), current_setting('lock_timeout')") .fetch_one(&mut *first) .await?; - assert_eq!(timeout, "5s"); + assert_eq!(timeout, "1500ms"); #[expect( clippy::disallowed_methods, reason = "Test assertions are outside the compile-checked query set" diff --git a/litellm-rust/crates/db/src/spend/team_member_spend.rs b/litellm-rust/crates/db/src/spend/team_member_spend.rs index 965a39ada4f..d4120d15d6e 100644 --- a/litellm-rust/crates/db/src/spend/team_member_spend.rs +++ b/litellm-rust/crates/db/src/spend/team_member_spend.rs @@ -49,7 +49,7 @@ pub async fn add_team_member_spend( #[cfg(test)] mod tests { - use std::error::Error as StdError; + use std::{error::Error as StdError, time::Duration}; use super::super::begin; use super::{TeamMemberSpend, add_team_member_spend}; @@ -90,7 +90,7 @@ mod tests { #[future(awt)] database_pool: Result>, ) -> Result<(), Box> { let pool = database_pool?; - let mut tx = begin(&pool).await?; + let mut tx = begin(&pool, Duration::from_secs(5)).await?; #[expect( clippy::disallowed_methods, reason = "Fixture writes are outside the compile-checked query set"