scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-29 06:33:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 200

feat(db): add typed end-user spend query (#48231)

Browse source 
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yujonglee 2026-10-26 17:05:05 -07:00 • committed by GitHub
parent 3b72c08c4f
commit 8f456eb7d5
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 221 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

15
litellm-rust/crates/db/.sqlx/query-345c02595f55a7f0274ef7214b02288e0f2c1f12a2a79ed53f0002826724cd0a.json generated Normal file
View file

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

6
litellm-rust/crates/db/queries/spend/add_end_user_spend.sql Normal file
View file

@ -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

2
litellm-rust/crates/db/src/lib.rs
View file

@ -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};

197
litellm-rust/crates/db/src/spend/end_user_spend.rs Normal file
View file

@ -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<PgPool, Box<dyn StdError>>,
) -> Result<(), Box<dyn StdError>> {
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<String>)> = 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<String>,
#[case] costs: Vec<f64>,
#[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 = "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<PgPool, Box<dyn StdError>>,
) -> Result<(), Box<dyn StdError>> {
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<String>) = 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(())
}
}

2
litellm-rust/crates/db/src/spend/mod.rs
View file

@ -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;