From 554b011787043a4cfc6a48d7d011bbfca05b6e22 Mon Sep 17 00:00:00 2001 From: yujonglee Date: Mon, 26 Oct 2026 17:05:04 -0700 Subject: [PATCH] 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 | 106 ++++++ litellm-rust/Cargo.lock | 5 + litellm-rust/Cargo.toml | 1 + ...64311d25428fd393bac4250891f148597e92e.json | 78 ++++ litellm-rust/crates/db/AGENTS.md | 25 ++ litellm-rust/crates/db/Cargo.toml | 7 + .../db/queries/end_user/find_end_user.sql | 10 + litellm-rust/crates/db/src/daily_spend.rs | 76 ++++ litellm-rust/crates/db/src/end_user.rs | 359 ++++++++++++++++++ litellm-rust/crates/db/src/error.rs | 79 ++++ litellm-rust/crates/db/src/lib.rs | 18 +- litellm-rust/crates/db/src/test_support.rs | 10 + 12 files changed, 773 insertions(+), 1 deletion(-) create mode 100644 .github/workflows/test-rust-db.yml create mode 100644 litellm-rust/crates/db/.sqlx/query-377d063cd0fc3a265d07d01f5aa64311d25428fd393bac4250891f148597e92e.json create mode 100644 litellm-rust/crates/db/AGENTS.md create mode 100644 litellm-rust/crates/db/queries/end_user/find_end_user.sql create mode 100644 litellm-rust/crates/db/src/daily_spend.rs create mode 100644 litellm-rust/crates/db/src/end_user.rs create mode 100644 litellm-rust/crates/db/src/error.rs create mode 100644 litellm-rust/crates/db/src/test_support.rs diff --git a/.github/workflows/test-rust-db.yml b/.github/workflows/test-rust-db.yml new file mode 100644 index 00000000000..bb88b593d5b --- /dev/null +++ b/.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 diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index cb50324e1dd..d58ba9ecb5f 100644 --- a/litellm-rust/Cargo.lock +++ b/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]] diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 9e5e0a96b73..fc65355c047 100644 --- a/litellm-rust/Cargo.toml +++ b/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" diff --git a/litellm-rust/crates/db/.sqlx/query-377d063cd0fc3a265d07d01f5aa64311d25428fd393bac4250891f148597e92e.json b/litellm-rust/crates/db/.sqlx/query-377d063cd0fc3a265d07d01f5aa64311d25428fd393bac4250891f148597e92e.json new file mode 100644 index 00000000000..734935466ff --- /dev/null +++ b/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" +} diff --git a/litellm-rust/crates/db/AGENTS.md b/litellm-rust/crates/db/AGENTS.md new file mode 100644 index 00000000000..1e08317230e --- /dev/null +++ b/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 diff --git a/litellm-rust/crates/db/Cargo.toml b/litellm-rust/crates/db/Cargo.toml index ca9d6f555b0..bb6d5655b60 100644 --- a/litellm-rust/crates/db/Cargo.toml +++ b/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 diff --git a/litellm-rust/crates/db/queries/end_user/find_end_user.sql b/litellm-rust/crates/db/queries/end_user/find_end_user.sql new file mode 100644 index 00000000000..c10ef8b877d --- /dev/null +++ b/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) diff --git a/litellm-rust/crates/db/src/daily_spend.rs b/litellm-rust/crates/db/src/daily_spend.rs new file mode 100644 index 00000000000..9db07514aab --- /dev/null +++ b/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, + pub api_key: String, + pub model: Option, + pub custom_llm_provider: Option, + pub endpoint: Option, + pub model_group: Option, + pub mcp_namespaced_tool_name: Option, + 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, + pub api_key: String, + pub model: Option, + pub custom_llm_provider: Option, + pub endpoint: Option, + pub model_group: Option, + pub mcp_namespaced_tool_name: Option, + 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, + pub api_key: String, + pub model: Option, + pub custom_llm_provider: Option, + pub endpoint: Option, + pub model_group: Option, + pub mcp_namespaced_tool_name: Option, + 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, +} diff --git a/litellm-rust/crates/db/src/end_user.rs b/litellm-rust/crates/db/src/end_user.rs new file mode 100644 index 00000000000..66661b51640 --- /dev/null +++ b/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, + pub tpm: Option, +} + +#[derive(Clone, Debug, Default, PartialEq)] +pub struct EndUser { + pub spend: Option, + pub blocked: bool, + pub max_budget: Option, + pub limits: RateLimits, +} + +#[macro_rules_attribute::apply(crate::db_row)] +struct EndUserRow { + spend: Option, + blocked: Option, + max_budget: Option, + rpm_limit: Option, + tpm_limit: Option, +} + +pub async fn find_end_user( + executor: impl PgExecutor<'_>, + end_user_id: &str, + default_budget_id: Option<&str>, + key_budget_id: Option<&str>, +) -> Result { + 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, + #[future(awt)] database_pool: Result>, + ) -> Result<(), Box> { + 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, + #[case] default_budget: BudgetFixture, + #[case] override_budget: Option, + #[case] expected_budget: Option, + #[case] insert_override_budget: bool, + #[future(awt)] database_pool: Result>, + ) -> Result<(), Box> { + 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(()) + } +} diff --git a/litellm-rust/crates/db/src/error.rs b/litellm-rust/crates/db/src/error.rs new file mode 100644 index 00000000000..af0ee030896 --- /dev/null +++ b/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>, + ) -> Result<(), Box> { + 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(()) + } +} diff --git a/litellm-rust/crates/db/src/lib.rs b/litellm-rust/crates/db/src/lib.rs index b7ce211d906..ca5cd366eec 100644 --- a/litellm-rust/crates/db/src/lib.rs +++ b/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; diff --git a/litellm-rust/crates/db/src/test_support.rs b/litellm-rust/crates/db/src/test_support.rs new file mode 100644 index 00000000000..992be4a4447 --- /dev/null +++ b/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> { + let database_url = std::env::var("DATABASE_URL")?; + Ok(PgPool::connect(&database_url).await?) +}