From 8f456eb7d51f52d62621e5672ee47ae9079e078d Mon Sep 17 00:00:00 2001 From: yujonglee Date: Mon, 26 Oct 2026 17:05:05 -0700 Subject: [PATCH] feat(db): add typed end-user spend query (#48231) Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...2288e0f2c1f12a2a79ed53f0002826724cd0a.json | 15 ++ .../db/queries/spend/add_end_user_spend.sql | 6 + litellm-rust/crates/db/src/lib.rs | 2 +- .../crates/db/src/spend/end_user_spend.rs | 197 ++++++++++++++++++ litellm-rust/crates/db/src/spend/mod.rs | 2 + 5 files changed, 221 insertions(+), 1 deletion(-) create mode 100644 litellm-rust/crates/db/.sqlx/query-345c02595f55a7f0274ef7214b02288e0f2c1f12a2a79ed53f0002826724cd0a.json create mode 100644 litellm-rust/crates/db/queries/spend/add_end_user_spend.sql create mode 100644 litellm-rust/crates/db/src/spend/end_user_spend.rs diff --git a/litellm-rust/crates/db/.sqlx/query-345c02595f55a7f0274ef7214b02288e0f2c1f12a2a79ed53f0002826724cd0a.json b/litellm-rust/crates/db/.sqlx/query-345c02595f55a7f0274ef7214b02288e0f2c1f12a2a79ed53f0002826724cd0a.json new file mode 100644 index 00000000000..123ff15ade2 --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-345c02595f55a7f0274ef7214b02288e0f2c1f12a2a79ed53f0002826724cd0a.json @@ -0,0 +1,15 @@ +{ + "db_name": "PostgreSQL", + "query": "INSERT INTO \"LiteLLM_EndUserTable\" (user_id, spend, blocked)\nSELECT id, SUM(cost), false FROM UNNEST($1::text[], $2::float8[]) AS a(id, cost)\nGROUP BY id\nORDER BY id\nON CONFLICT (user_id) DO UPDATE\nSET spend = \"LiteLLM_EndUserTable\".spend + EXCLUDED.spend\n", + "describe": { + "columns": [], + "parameters": { + "Left": [ + "TextArray", + "Float8Array" + ] + }, + "nullable": [] + }, + "hash": "345c02595f55a7f0274ef7214b02288e0f2c1f12a2a79ed53f0002826724cd0a" +} diff --git a/litellm-rust/crates/db/queries/spend/add_end_user_spend.sql b/litellm-rust/crates/db/queries/spend/add_end_user_spend.sql new file mode 100644 index 00000000000..06a5b96d254 --- /dev/null +++ b/litellm-rust/crates/db/queries/spend/add_end_user_spend.sql @@ -0,0 +1,6 @@ +INSERT INTO "LiteLLM_EndUserTable" (user_id, spend, blocked) +SELECT id, SUM(cost), false FROM UNNEST($1::text[], $2::float8[]) AS a(id, cost) +GROUP BY id +ORDER BY id +ON CONFLICT (user_id) DO UPDATE +SET spend = "LiteLLM_EndUserTable".spend + EXCLUDED.spend diff --git a/litellm-rust/crates/db/src/lib.rs b/litellm-rust/crates/db/src/lib.rs index 1dda3a987b0..3a9032d73dc 100644 --- a/litellm-rust/crates/db/src/lib.rs +++ b/litellm-rust/crates/db/src/lib.rs @@ -18,4 +18,4 @@ 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_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}; diff --git a/litellm-rust/crates/db/src/spend/end_user_spend.rs b/litellm-rust/crates/db/src/spend/end_user_spend.rs new file mode 100644 index 00000000000..ca327ac7e18 --- /dev/null +++ b/litellm-rust/crates/db/src/spend/end_user_spend.rs @@ -0,0 +1,197 @@ +use sqlx::PgExecutor; + +use super::EntitySpend; +use crate::Error; + +pub async fn add_end_user_spend( + executor: impl PgExecutor<'_>, + spend: &EntitySpend, +) -> Result<(), Error> { + spend.validate()?; + if spend.is_empty() { + return Ok(()); + } + sqlx::query_file!( + "queries/spend/add_end_user_spend.sql", + &spend.ids, + &spend.costs + ) + .execute(executor) + .await?; + Ok(()) +} + +#[cfg(test)] +mod tests { + use std::error::Error as StdError; + + use super::{EntitySpend, add_end_user_spend}; + use crate::test_support::database_pool; + use rstest::rstest; + use sqlx::PgPool; + + #[rstest] + #[case::fractional_costs(0.25, 0.75, false, 2.5, 1.5)] + #[case::zero_costs(0.0, 0.0, false, 2.0, 0.0)] + #[case::duplicate_ids(0.25, 0.75, true, 3.0, 3.0)] + #[tokio::test] + #[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"] + async fn end_user_spend_matches_python_upsert( + #[case] existing_cost: f64, + #[case] new_cost: f64, + #[case] duplicate_ids: bool, + #[case] expected_existing_spend: f64, + #[case] expected_new_spend: f64, + #[future(awt)] database_pool: Result>, + ) -> Result<(), Box> { + let pool = database_pool?; + let mut tx = pool.begin().await?; + #[expect( + clippy::disallowed_methods, + reason = "Fixture writes are outside the compile-checked query set" + )] + let _budget = sqlx::query( + r#"INSERT INTO "LiteLLM_BudgetTable" (budget_id, created_by, updated_by) + VALUES ('end-user-spend-budget', 'db-query-test', 'db-query-test')"#, + ) + .execute(&mut *tx) + .await?; + #[expect( + clippy::disallowed_methods, + reason = "Fixture writes are outside the compile-checked query set" + )] + let _users = sqlx::query( + r#"INSERT INTO "LiteLLM_EndUserTable" (user_id, spend, blocked, budget_id) + VALUES ('end-user-spend-z', 2, true, 'end-user-spend-budget'), + ('end-user-spend-untouched', 100, true, NULL)"#, + ) + .execute(&mut *tx) + .await?; + let spend = if duplicate_ids { + EntitySpend { + ids: vec![ + "end-user-spend-z".into(), + "end-user-spend-a".into(), + "end-user-spend-z".into(), + "end-user-spend-a".into(), + ], + costs: vec![existing_cost, new_cost, existing_cost, new_cost], + } + } else { + EntitySpend { + ids: vec!["end-user-spend-z".into(), "end-user-spend-a".into()], + costs: vec![existing_cost, new_cost], + } + }; + add_end_user_spend(&mut *tx, &spend).await?; + add_end_user_spend(&mut *tx, &spend).await?; + #[expect( + clippy::disallowed_methods, + reason = "Test assertions are outside the compile-checked query set" + )] + let rows: Vec<(String, f64, bool, Option)> = sqlx::query_as( + r#"SELECT user_id, spend, blocked, budget_id FROM "LiteLLM_EndUserTable" + WHERE user_id = ANY($1) ORDER BY user_id"#, + ) + .bind(vec![ + "end-user-spend-a", + "end-user-spend-z", + "end-user-spend-untouched", + ]) + .fetch_all(&mut *tx) + .await?; + assert_eq!( + rows, + vec![ + ("end-user-spend-a".into(), expected_new_spend, false, None), + ("end-user-spend-untouched".into(), 100.0, true, None), + ( + "end-user-spend-z".into(), + expected_existing_spend, + true, + Some("end-user-spend-budget".into()), + ), + ] + ); + Ok(()) + } + + #[rstest] + #[case::missing_cost(vec!["end-user-invalid-existing".into(), "end-user-invalid-new".into()], vec![1.0])] + #[case::extra_cost(vec!["end-user-invalid-existing".into()], vec![1.0, 2.0])] + #[case::empty_ids(vec![], vec![1.0])] + #[case::empty_costs(vec!["end-user-invalid-existing".into()], vec![])] + #[tokio::test] + #[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"] + async fn malformed_end_user_spend_is_rejected_before_writes( + #[case] ids: Vec, + #[case] costs: Vec, + #[future(awt)] database_pool: Result>, + ) -> Result<(), Box> { + let pool = database_pool?; + let mut tx = pool.begin().await?; + #[expect( + clippy::disallowed_methods, + reason = "Fixture writes are outside the compile-checked query set" + )] + let _user = sqlx::query( + r#"INSERT INTO "LiteLLM_EndUserTable" (user_id, spend, blocked) + VALUES ('end-user-invalid-existing', 8.5, true)"#, + ) + .execute(&mut *tx) + .await?; + let spend = EntitySpend { ids, costs }; + let result = add_end_user_spend(&mut *tx, &spend).await; + assert!(matches!( + result, + Err(crate::Error::InvalidSpendBatch { ids, costs }) + if ids == spend.ids.len() && costs == spend.costs.len() + )); + #[expect( + clippy::disallowed_methods, + reason = "Test assertions are outside the compile-checked query set" + )] + let rows: Vec<(String, f64, bool)> = sqlx::query_as( + r#"SELECT user_id, spend, blocked FROM "LiteLLM_EndUserTable" + WHERE user_id IN ('end-user-invalid-existing', 'end-user-invalid-new')"#, + ) + .fetch_all(&mut *tx) + .await?; + assert_eq!(rows, vec![("end-user-invalid-existing".into(), 8.5, true)]); + Ok(()) + } + + #[rstest] + #[case::empty_batch()] + #[tokio::test] + #[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"] + async fn empty_end_user_spend_is_a_noop( + #[future(awt)] database_pool: Result>, + ) -> Result<(), Box> { + let pool = database_pool?; + let mut tx = pool.begin().await?; + #[expect( + clippy::disallowed_methods, + reason = "Fixture writes are outside the compile-checked query set" + )] + let _user = sqlx::query( + r#"INSERT INTO "LiteLLM_EndUserTable" (user_id, spend, blocked) + VALUES ('end-user-empty-spend', 8.5, true)"#, + ) + .execute(&mut *tx) + .await?; + add_end_user_spend(&mut *tx, &EntitySpend::default()).await?; + #[expect( + clippy::disallowed_methods, + reason = "Test assertions are outside the compile-checked query set" + )] + let row: (f64, bool, Option) = sqlx::query_as( + r#"SELECT spend, blocked, budget_id FROM "LiteLLM_EndUserTable" + WHERE user_id = 'end-user-empty-spend'"#, + ) + .fetch_one(&mut *tx) + .await?; + assert_eq!(row, (8.5, true, None)); + Ok(()) + } +} diff --git a/litellm-rust/crates/db/src/spend/mod.rs b/litellm-rust/crates/db/src/spend/mod.rs index 775be335150..5191da2adb5 100644 --- a/litellm-rust/crates/db/src/spend/mod.rs +++ b/litellm-rust/crates/db/src/spend/mod.rs @@ -1,7 +1,9 @@ +mod end_user_spend; mod key_spend; 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_spend::add_team_spend; pub use user_spend::add_user_spend;