feat(rust): add typed database query foundation (#48214)

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

.github/workflows/test-rust-db.yml

@@ -0,0 +1,106 @@
+name: LiteLLM Rust DB
+
+on:
+  push:
+    paths:
+      - "litellm-rust/crates/db/**"
+      - "litellm-rust/Cargo.toml"
+      - "litellm-rust/Cargo.lock"
+      - "litellm-rust/clippy.toml"
+      - "litellm-proxy-extras/litellm_proxy_extras/migrations/**"
+      - "**/schema.prisma"
+      - ".cargo/**"
+      - "rust-toolchain.toml"
+      - "pyproject.toml"
+      - "uv.lock"
+      - ".github/actions/setup-uv-with-retries/**"
+      - ".github/workflows/test-rust-db.yml"
+  pull_request:
+    branches:
+      - main
+      - "litellm_**"
+    paths:
+      - "litellm-rust/crates/db/**"
+      - "litellm-rust/Cargo.toml"
+      - "litellm-rust/Cargo.lock"
+      - "litellm-rust/clippy.toml"
+      - "litellm-proxy-extras/litellm_proxy_extras/migrations/**"
+      - "**/schema.prisma"
+      - ".cargo/**"
+      - "rust-toolchain.toml"
+      - "pyproject.toml"
+      - "uv.lock"
+      - ".github/actions/setup-uv-with-retries/**"
+      - ".github/workflows/test-rust-db.yml"
+
+permissions:
+  contents: read
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
+  cancel-in-progress: true
+
+env:
+  CARGO_TERM_COLOR: always
+
+jobs:
+  db-queries:
+    runs-on: ubuntu-latest
+    timeout-minutes: 30
+    services:
+      postgres:
+        image: postgres:16.6@sha256:557fea37a744d5f4c8faab304b0a90858b53ab119735a88c131fd19dab802f36
+        env:
+          POSTGRES_DB: litellm
+          POSTGRES_PASSWORD: postgres
+          POSTGRES_USER: postgres
+        ports:
+          - 5432:5432
+        options: >-
+          --health-cmd "pg_isready -U postgres -d litellm"
+          --health-interval 10s
+          --health-timeout 5s
+          --health-retries 5
+    env:
+      DATABASE_URL: postgresql://postgres:postgres@localhost:5432/litellm
+      SQLX_OFFLINE: "false"
+    defaults:
+      run:
+        working-directory: litellm-rust
+    steps:
+      - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
+        with:
+          persist-credentials: false
+
+      - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
+        with:
+          python-version: "3.12"
+
+      - uses: ./.github/actions/setup-uv-with-retries
+        with:
+          version: "0.10.9"
+
+      - run: rustup toolchain install --no-self-update
+
+      - uses: Swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2.9.2
+        with:
+          workspaces: litellm-rust
+          cache-on-failure: true
+
+      - name: Install Python dependencies
+        working-directory: .
+        run: uv sync --frozen --extra extra-proxy
+
+      - name: Deploy Prisma migrations
+        working-directory: .
+        run: uv run --no-sync prisma migrate deploy --schema litellm-proxy-extras/litellm_proxy_extras/schema.prisma
+
+      - name: Install SQLx CLI
+        run: cargo install sqlx-cli --version 0.9.0 --locked --no-default-features -F postgres
+
+      - name: Check SQLx query metadata
+        working-directory: litellm-rust/crates/db
+        run: cargo sqlx prepare --check
+
+      - name: Test database queries
+        run: cargo test -p litellm-db --locked -- --include-ignored

litellm-rust/Cargo.lock

@@ -3695,8 +3695,13 @@ dependencies = [
 name = "litellm-db"
 version = "0.1.0"
 dependencies = [
+ "chrono",
+ "macro_rules_attribute",
+ "rstest",
  "serde",
  "sqlx",
+ "thiserror 2.0.19",
+ "tokio",
 ]
 
 [[package]]

litellm-rust/Cargo.toml

@@ -104,6 +104,7 @@ serde_with = { version = "=3.21.0", default-features = false, features = ["std",
 sha2 = "0.10"
 syn = { version = "2", default-features = false }
 sqlx = { version = "0.9.0", default-features = false, features = ["json", "macros", "postgres", "runtime-tokio", "chrono", "tls-rustls-ring-native-roots"] }
+chrono = { version = "0.4.45", default-features = false, features = ["std"] }
 proc-macro2 = "1"
 quote = "1"
 subtle = "2"

litellm-rust/crates/db/.sqlx/query-377d063cd0fc3a265d07d01f5aa64311d25428fd393bac4250891f148597e92e.json

@@ -0,0 +1,78 @@
+{
+  "db_name": "PostgreSQL",
+  "query": "SELECT e.spend AS \"spend?\",\n       e.blocked AS \"blocked?\",\n       b.max_budget AS \"max_budget?\",\n       b.rpm_limit AS \"rpm_limit?\",\n       b.tpm_limit AS \"tpm_limit?\"\nFROM (SELECT 1) AS one\nLEFT JOIN \"LiteLLM_EndUserTable\" AS e ON e.user_id = $1\nLEFT JOIN \"LiteLLM_BudgetTable\" AS assigned ON assigned.budget_id = e.budget_id\nLEFT JOIN \"LiteLLM_BudgetTable\" AS key_budget ON key_budget.budget_id = $3\nLEFT JOIN \"LiteLLM_BudgetTable\" AS b ON b.budget_id = COALESCE(assigned.budget_id, key_budget.budget_id, $2)\n",
+  "describe": {
+    "columns": [
+      {
+        "ordinal": 0,
+        "name": "spend?",
+        "type_info": "Float8",
+        "origin": {
+          "Table": {
+            "table": "\"LiteLLM_EndUserTable\"",
+            "name": "spend"
+          }
+        }
+      },
+      {
+        "ordinal": 1,
+        "name": "blocked?",
+        "type_info": "Bool",
+        "origin": {
+          "Table": {
+            "table": "\"LiteLLM_EndUserTable\"",
+            "name": "blocked"
+          }
+        }
+      },
+      {
+        "ordinal": 2,
+        "name": "max_budget?",
+        "type_info": "Float8",
+        "origin": {
+          "Table": {
+            "table": "\"LiteLLM_BudgetTable\"",
+            "name": "max_budget"
+          }
+        }
+      },
+      {
+        "ordinal": 3,
+        "name": "rpm_limit?",
+        "type_info": "Int8",
+        "origin": {
+          "Table": {
+            "table": "\"LiteLLM_BudgetTable\"",
+            "name": "rpm_limit"
+          }
+        }
+      },
+      {
+        "ordinal": 4,
+        "name": "tpm_limit?",
+        "type_info": "Int8",
+        "origin": {
+          "Table": {
+            "table": "\"LiteLLM_BudgetTable\"",
+            "name": "tpm_limit"
+          }
+        }
+      }
+    ],
+    "parameters": {
+      "Left": [
+        "Text",
+        "Text",
+        "Text"
+      ]
+    },
+    "nullable": [
+      true,
+      true,
+      true,
+      true,
+      true
+    ]
+  },
+  "hash": "377d063cd0fc3a265d07d01f5aa64311d25428fd393bac4250891f148597e92e"
+}

litellm-rust/crates/db/AGENTS.md

@@ -0,0 +1,25 @@
+# Database crate
+
+Keep each static query in `queries/` and call it through a typed SQLx query function. Rust does not create or migrate the database
+
+- Own reusable row projections in `src/` and export those consumed by other crates or integration tests
+- Apply `#[macro_rules_attribute::apply(crate::db_row)]` to row types for shared derives, without adding serialization or schema derives automatically
+- Apply `#[macro_rules_attribute::apply(crate::db_input)]` to write inputs, adding `Default` only where supported
+- Keep function tests inline in their owning query module and reuse the strict database fixture in `test_support`
+- Transaction-local setup in `begin` is the approved exception to static query files
+- Keep write inputs separate from stored rows and keep test-only assertions out of public types
+- `FromRow` supports runtime reads, but production static reads still use `query_file_as!` for compile-time checking
+
+Database integration tests are ignored by default, but still compile in workspace checks. Run them explicitly against a local database already migrated by Prisma. Missing or invalid database configuration must fail the selected tests:
+
+```sh
+DATABASE_URL=postgresql://... cargo test -p litellm-db --locked -- --include-ignored
+```
+
+After changing a query, refresh the checked query metadata from this directory:
+
+```sh
+SQLX_OFFLINE=false cargo sqlx prepare
+```
+
+Tests use transactions and leave them uncommitted so fixture rows roll back when each test finishes

litellm-rust/crates/db/Cargo.toml

@@ -6,5 +6,12 @@ license.workspace = true
 repository.workspace = true
 
 [dependencies]
+chrono.workspace = true
+macro_rules_attribute.workspace = true
 serde.workspace = true
 sqlx.workspace = true
+thiserror.workspace = true
+
+[dev-dependencies]
+rstest.workspace = true
+tokio.workspace = true

litellm-rust/crates/db/queries/end_user/find_end_user.sql

@@ -0,0 +1,10 @@
+SELECT e.spend AS "spend?",
+       e.blocked AS "blocked?",
+       b.max_budget AS "max_budget?",
+       b.rpm_limit AS "rpm_limit?",
+       b.tpm_limit AS "tpm_limit?"
+FROM (SELECT 1) AS one
+LEFT JOIN "LiteLLM_EndUserTable" AS e ON e.user_id = $1
+LEFT JOIN "LiteLLM_BudgetTable" AS assigned ON assigned.budget_id = e.budget_id
+LEFT JOIN "LiteLLM_BudgetTable" AS key_budget ON key_budget.budget_id = $3
+LEFT JOIN "LiteLLM_BudgetTable" AS b ON b.budget_id = COALESCE(assigned.budget_id, key_budget.budget_id, $2)

litellm-rust/crates/db/src/daily_spend.rs

@@ -0,0 +1,76 @@
+use sqlx::types::chrono::NaiveDateTime;
+
+#[macro_rules_attribute::apply(crate::db_row)]
+pub struct DailyUserRow {
+    pub id: String,
+    pub user_id: Option<String>,
+    pub api_key: String,
+    pub model: Option<String>,
+    pub custom_llm_provider: Option<String>,
+    pub endpoint: Option<String>,
+    pub model_group: Option<String>,
+    pub mcp_namespaced_tool_name: Option<String>,
+    pub prompt_tokens: i64,
+    pub completion_tokens: i64,
+    pub api_requests: i64,
+    pub successful_requests: i64,
+    pub failed_requests: i64,
+    pub cache_read_input_tokens: i64,
+    pub cache_creation_input_tokens: i64,
+    pub total_response_time_ms: i64,
+    pub timed_requests: i64,
+    pub spend: f64,
+    pub date: String,
+    pub created_at: NaiveDateTime,
+    pub updated_at: NaiveDateTime,
+}
+
+#[macro_rules_attribute::apply(crate::db_row)]
+pub struct DailyTeamRow {
+    pub id: String,
+    pub team_id: Option<String>,
+    pub api_key: String,
+    pub model: Option<String>,
+    pub custom_llm_provider: Option<String>,
+    pub endpoint: Option<String>,
+    pub model_group: Option<String>,
+    pub mcp_namespaced_tool_name: Option<String>,
+    pub prompt_tokens: i64,
+    pub completion_tokens: i64,
+    pub api_requests: i64,
+    pub successful_requests: i64,
+    pub failed_requests: i64,
+    pub cache_read_input_tokens: i64,
+    pub cache_creation_input_tokens: i64,
+    pub total_response_time_ms: i64,
+    pub timed_requests: i64,
+    pub spend: f64,
+    pub date: String,
+    pub created_at: NaiveDateTime,
+    pub updated_at: NaiveDateTime,
+}
+
+#[macro_rules_attribute::apply(crate::db_row)]
+pub struct DailyEndUserRow {
+    pub id: String,
+    pub end_user_id: Option<String>,
+    pub api_key: String,
+    pub model: Option<String>,
+    pub custom_llm_provider: Option<String>,
+    pub endpoint: Option<String>,
+    pub model_group: Option<String>,
+    pub mcp_namespaced_tool_name: Option<String>,
+    pub prompt_tokens: i64,
+    pub completion_tokens: i64,
+    pub api_requests: i64,
+    pub successful_requests: i64,
+    pub failed_requests: i64,
+    pub cache_read_input_tokens: i64,
+    pub cache_creation_input_tokens: i64,
+    pub total_response_time_ms: i64,
+    pub timed_requests: i64,
+    pub spend: f64,
+    pub date: String,
+    pub created_at: NaiveDateTime,
+    pub updated_at: NaiveDateTime,
+}

litellm-rust/crates/db/src/end_user.rs

@@ -0,0 +1,359 @@
+use sqlx::PgExecutor;
+
+use crate::Error;
+
+#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
+pub struct RateLimits {
+    pub rpm: Option<i64>,
+    pub tpm: Option<i64>,
+}
+
+#[derive(Clone, Debug, Default, PartialEq)]
+pub struct EndUser {
+    pub spend: Option<f64>,
+    pub blocked: bool,
+    pub max_budget: Option<f64>,
+    pub limits: RateLimits,
+}
+
+#[macro_rules_attribute::apply(crate::db_row)]
+struct EndUserRow {
+    spend: Option<f64>,
+    blocked: Option<bool>,
+    max_budget: Option<f64>,
+    rpm_limit: Option<i64>,
+    tpm_limit: Option<i64>,
+}
+
+pub async fn find_end_user(
+    executor: impl PgExecutor<'_>,
+    end_user_id: &str,
+    default_budget_id: Option<&str>,
+    key_budget_id: Option<&str>,
+) -> Result<EndUser, Error> {
+    let row = sqlx::query_file_as!(
+        EndUserRow,
+        "queries/end_user/find_end_user.sql",
+        end_user_id,
+        default_budget_id,
+        key_budget_id,
+    )
+    .fetch_one(executor)
+    .await?;
+    Ok(EndUser {
+        spend: row.spend,
+        blocked: row.blocked.unwrap_or(false),
+        max_budget: row.max_budget,
+        limits: RateLimits {
+            rpm: row.rpm_limit,
+            tpm: row.tpm_limit,
+        },
+    })
+}
+
+#[cfg(test)]
+mod tests {
+    use std::error::Error as StdError;
+
+    use super::{EndUser, RateLimits, find_end_user};
+    use crate::test_support::database_pool;
+    use rstest::rstest;
+    use sqlx::{PgPool, Postgres, Transaction};
+
+    #[derive(Clone, Copy)]
+    struct BudgetFixture {
+        budget_id: &'static str,
+        max_budget: f64,
+        rpm_limit: i64,
+        tpm_limit: i64,
+    }
+
+    async fn insert_budget(
+        tx: &mut Transaction<'_, Postgres>,
+        budget: BudgetFixture,
+    ) -> Result<(), sqlx::Error> {
+        #[expect(
+            clippy::disallowed_methods,
+            reason = "Fixture inserts are outside the compile-checked query set"
+        )]
+        let result = sqlx::query(
+            r#"
+            INSERT INTO "LiteLLM_BudgetTable"
+                (budget_id, max_budget, rpm_limit, tpm_limit, created_by, updated_by)
+            VALUES ($1, $2, $3, $4, 'db-query-test', 'db-query-test')
+            "#,
+        )
+        .bind(budget.budget_id)
+        .bind(budget.max_budget)
+        .bind(budget.rpm_limit)
+        .bind(budget.tpm_limit)
+        .execute(&mut **tx)
+        .await;
+        result.map(|_| ())
+    }
+
+    async fn insert_end_user(
+        tx: &mut Transaction<'_, Postgres>,
+        user_id: &str,
+        spend: f64,
+        blocked: bool,
+        budget_id: Option<&str>,
+    ) -> Result<(), sqlx::Error> {
+        #[expect(
+            clippy::disallowed_methods,
+            reason = "Fixture inserts are outside the compile-checked query set"
+        )]
+        let result = sqlx::query(
+            r#"
+            INSERT INTO "LiteLLM_EndUserTable" (user_id, spend, blocked, budget_id)
+            VALUES ($1, $2, $3, $4)
+            "#,
+        )
+        .bind(user_id)
+        .bind(spend)
+        .bind(blocked)
+        .bind(budget_id)
+        .execute(&mut **tx)
+        .await;
+        result.map(|_| ())
+    }
+
+    #[rstest]
+    #[case::without_default_budget(None)]
+    #[case::with_default_budget(Some(BudgetFixture {
+        budget_id: "missing-user-default-budget",
+        max_budget: 11.5,
+        rpm_limit: 12,
+        tpm_limit: 1300,
+    }))]
+    #[tokio::test]
+    #[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"]
+    async fn missing_end_user_maps_to_defaults(
+        #[case] default_budget: Option<BudgetFixture>,
+        #[future(awt)] database_pool: Result<PgPool, Box<dyn StdError>>,
+    ) -> Result<(), Box<dyn StdError>> {
+        let pool = database_pool?;
+        let mut tx = pool.begin().await?;
+        if let Some(budget) = default_budget {
+            insert_budget(&mut tx, budget).await?;
+        }
+
+        let actual = find_end_user(
+            &mut *tx,
+            "missing-end-user",
+            default_budget.map(|budget| budget.budget_id),
+            None,
+        )
+        .await?;
+
+        assert_eq!(
+            actual,
+            EndUser {
+                spend: None,
+                blocked: false,
+                max_budget: default_budget.map(|budget| budget.max_budget),
+                limits: RateLimits {
+                    rpm: default_budget.map(|budget| budget.rpm_limit),
+                    tpm: default_budget.map(|budget| budget.tpm_limit),
+                },
+            }
+        );
+        Ok(())
+    }
+
+    #[rstest]
+    #[case::row_budget_precedes_default(
+        "row-budget-priority-user",
+        Some(BudgetFixture {
+            budget_id: "row-budget-priority-row",
+            max_budget: 10.0,
+            rpm_limit: 101,
+            tpm_limit: 1001,
+        }),
+        BudgetFixture {
+            budget_id: "row-budget-priority-default",
+            max_budget: 20.0,
+            rpm_limit: 201,
+            tpm_limit: 2001,
+        },
+        None,
+        Some(BudgetFixture {
+            budget_id: "row-budget-priority-row",
+            max_budget: 10.0,
+            rpm_limit: 101,
+            tpm_limit: 1001,
+        }),
+        false
+    )]
+    #[case::default_applies_without_row_budget(
+        "default-budget-user",
+        None,
+        BudgetFixture {
+            budget_id: "default-budget-fallback",
+            max_budget: 30.0,
+            rpm_limit: 301,
+            tpm_limit: 3001,
+        },
+        None,
+        Some(BudgetFixture {
+            budget_id: "default-budget-fallback",
+            max_budget: 30.0,
+            rpm_limit: 301,
+            tpm_limit: 3001,
+        }),
+        false
+    )]
+    #[case::assigned_budget_precedes_key_when_all_ids_differ(
+        "override-priority-user",
+        Some(BudgetFixture {
+            budget_id: "override-priority-row",
+            max_budget: 40.0,
+            rpm_limit: 401,
+            tpm_limit: 4001,
+        }),
+        BudgetFixture {
+            budget_id: "override-priority-default",
+            max_budget: 50.0,
+            rpm_limit: 501,
+            tpm_limit: 5001,
+        },
+        Some(BudgetFixture {
+            budget_id: "override-priority-override",
+            max_budget: 60.0,
+            rpm_limit: 601,
+            tpm_limit: 6001,
+        }),
+        Some(BudgetFixture {
+            budget_id: "override-priority-row",
+            max_budget: 40.0,
+            rpm_limit: 401,
+            tpm_limit: 4001,
+        }),
+        true
+    )]
+    #[case::missing_key_budget_keeps_assigned_budget(
+        "dangling-override-user",
+        Some(BudgetFixture {
+            budget_id: "dangling-override-row",
+            max_budget: 70.0,
+            rpm_limit: 701,
+            tpm_limit: 7001,
+        }),
+        BudgetFixture {
+            budget_id: "dangling-override-default",
+            max_budget: 80.0,
+            rpm_limit: 801,
+            tpm_limit: 8001,
+        },
+        Some(BudgetFixture {
+            budget_id: "dangling-override-missing",
+            max_budget: 90.0,
+            rpm_limit: 901,
+            tpm_limit: 9001,
+        }),
+        Some(BudgetFixture {
+            budget_id: "dangling-override-row",
+            max_budget: 70.0,
+            rpm_limit: 701,
+            tpm_limit: 7001,
+        }),
+        false
+    )]
+    #[case::key_budget_precedes_proxy_default(
+        "key-default-user",
+        None,
+        BudgetFixture {
+            budget_id: "key-default-proxy",
+            max_budget: 80.0,
+            rpm_limit: 801,
+            tpm_limit: 8001,
+        },
+        Some(BudgetFixture {
+            budget_id: "key-default-key",
+            max_budget: 90.0,
+            rpm_limit: 901,
+            tpm_limit: 9001,
+        }),
+        Some(BudgetFixture {
+            budget_id: "key-default-key",
+            max_budget: 90.0,
+            rpm_limit: 901,
+            tpm_limit: 9001,
+        }),
+        true
+    )]
+    #[case::missing_key_budget_falls_back_to_proxy_default(
+        "missing-key-default-user",
+        None,
+        BudgetFixture {
+            budget_id: "missing-key-default-proxy",
+            max_budget: 80.0,
+            rpm_limit: 801,
+            tpm_limit: 8001,
+        },
+        Some(BudgetFixture {
+            budget_id: "missing-key-default-key",
+            max_budget: 90.0,
+            rpm_limit: 901,
+            tpm_limit: 9001,
+        }),
+        Some(BudgetFixture {
+            budget_id: "missing-key-default-proxy",
+            max_budget: 80.0,
+            rpm_limit: 801,
+            tpm_limit: 8001,
+        }),
+        false
+    )]
+    #[tokio::test]
+    #[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"]
+    async fn budget_resolution_matches_python_precedence(
+        #[case] user_id: &'static str,
+        #[case] user_budget: Option<BudgetFixture>,
+        #[case] default_budget: BudgetFixture,
+        #[case] override_budget: Option<BudgetFixture>,
+        #[case] expected_budget: Option<BudgetFixture>,
+        #[case] insert_override_budget: bool,
+        #[future(awt)] database_pool: Result<PgPool, Box<dyn StdError>>,
+    ) -> Result<(), Box<dyn StdError>> {
+        let pool = database_pool?;
+        let mut tx = pool.begin().await?;
+        insert_budget(&mut tx, default_budget).await?;
+        if let Some(budget) = user_budget {
+            insert_budget(&mut tx, budget).await?;
+        }
+        if let (Some(budget), true) = (override_budget, insert_override_budget) {
+            insert_budget(&mut tx, budget).await?;
+        }
+        insert_end_user(
+            &mut tx,
+            user_id,
+            7.5,
+            true,
+            user_budget.map(|budget| budget.budget_id),
+        )
+        .await?;
+
+        let actual = find_end_user(
+            &mut *tx,
+            user_id,
+            Some(default_budget.budget_id),
+            override_budget.map(|budget| budget.budget_id),
+        )
+        .await?;
+
+        assert_eq!(
+            actual,
+            EndUser {
+                spend: Some(7.5),
+                blocked: true,
+                max_budget: expected_budget.map(|budget| budget.max_budget),
+                limits: RateLimits {
+                    rpm: expected_budget.map(|budget| budget.rpm_limit),
+                    tpm: expected_budget.map(|budget| budget.tpm_limit),
+                },
+            }
+        );
+        Ok(())
+    }
+}

litellm-rust/crates/db/src/error.rs

@@ -0,0 +1,79 @@
+#[derive(Debug, thiserror::Error)]
+pub enum Error {
+    #[error("database query failed")]
+    Query(#[from] sqlx::Error),
+    #[error("spend batch has {ids} ids and {costs} costs")]
+    InvalidSpendBatch { ids: usize, costs: usize },
+    #[error("invalid spend timestamp")]
+    InvalidTimestamp(#[from] chrono::ParseError),
+}
+
+impl Error {
+    pub fn is_retryable(&self) -> bool {
+        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,
+            _ => false,
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::error::Error as StdError;
+
+    use rstest::rstest;
+    use sqlx::PgPool;
+
+    use super::Error;
+    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);
+    }
+
+    #[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)]
+    #[tokio::test]
+    #[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"]
+    async fn classifies_server_errors(
+        #[case] code: &str,
+        #[case] retryable: bool,
+        #[future(awt)] database_pool: Result<PgPool, Box<dyn StdError>>,
+    ) -> Result<(), Box<dyn StdError>> {
+        let pool = database_pool?;
+        let mut tx = pool.begin().await?;
+        #[expect(
+            clippy::disallowed_methods,
+            reason = "Test query must raise the requested SQLSTATE"
+        )]
+        let error = sqlx::query(sqlx::AssertSqlSafe(format!(
+            "DO $$ BEGIN RAISE EXCEPTION USING ERRCODE = '{code}'; END $$"
+        )))
+        .execute(&mut *tx)
+        .await
+        .expect_err("PostgreSQL must raise the requested error");
+
+        assert_eq!(
+            error.as_database_error().and_then(|error| error.code()),
+            Some(code.into())
+        );
+        assert_eq!(Error::from(error).is_retryable(), retryable);
+        tx.rollback().await?;
+        Ok(())
+    }
+}

litellm-rust/crates/db/src/lib.rs

@@ -1,3 +1,19 @@
-mod failure;
+macro_rules_attribute::attribute_alias! {
+    #[apply(db_input)] =
+        #[derive(Clone, Debug, PartialEq)];
 
+    #[apply(db_row)] =
+        #[derive(Clone, Debug, PartialEq, sqlx::FromRow)];
+}
+
+mod daily_spend;
+mod end_user;
+mod error;
+mod failure;
+#[cfg(test)]
+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 failure::DbFailure;

litellm-rust/crates/db/src/test_support.rs

@@ -0,0 +1,10 @@
+use std::error::Error as StdError;
+
+use rstest::fixture;
+use sqlx::PgPool;
+
+#[fixture]
+pub async fn database_pool() -> Result<PgPool, Box<dyn StdError>> {
+    let database_url = std::env::var("DATABASE_URL")?;
+    Ok(PgPool::connect(&database_url).await?)
+}