Merge fbbce9725c into ff2709f5c3

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

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<PgPool, Box<dyn StdError>>,
     ) -> Result<(), Box<dyn StdError>> {
         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(())
     }

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,

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<Transaction<'static, Postgres>, Error> {
+pub async fn begin(
+    pool: &PgPool,
+    lock_timeout: Duration,
+) -> 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?;
+    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"

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<PgPool, Box<dyn StdError>>,
     ) -> Result<(), Box<dyn StdError>> {
         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"