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 user spend query (#48229)

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:04 -07:00 • committed by GitHub
parent 35cd5fc05f
commit 8e61c22fe0
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 205 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-43f2e12a3ffb1de29b5b031a1173b87a20a2664b0bef3616631d2c24487c1e54.json generated Normal file
View file

@ -0,0 +1,15 @@
{
"db_name": "PostgreSQL",
"query": "WITH added AS (\n SELECT id, SUM(cost) AS cost\n FROM UNNEST($1::text[], $2::float8[]) AS a(id, cost)\n GROUP BY id\n), locked AS (\n SELECT u.user_id FROM \"LiteLLM_UserTable\" u\n WHERE u.user_id IN (SELECT id FROM added)\n ORDER BY u.user_id FOR UPDATE\n)\nUPDATE \"LiteLLM_UserTable\" u\nSET spend = u.spend + added.cost,\n updated_at = (NOW() AT TIME ZONE 'UTC')\nFROM added\nWHERE u.user_id = added.id AND u.user_id IN (SELECT user_id FROM locked)\n",
"describe": {
"columns": [],
"parameters": {
"Left": [
"TextArray",
"Float8Array"
]
},
"nullable": []
},
"hash": "43f2e12a3ffb1de29b5b031a1173b87a20a2664b0bef3616631d2c24487c1e54"
}

14
litellm-rust/crates/db/queries/spend/add_user_spend.sql Normal file
View file

@ -0,0 +1,14 @@
WITH added AS (
SELECT id, SUM(cost) AS cost
FROM UNNEST($1::text[], $2::float8[]) AS a(id, cost)
GROUP BY id
), locked AS (
SELECT u.user_id FROM "LiteLLM_UserTable" u
WHERE u.user_id IN (SELECT id FROM added)
ORDER BY u.user_id FOR UPDATE
)
UPDATE "LiteLLM_UserTable" u
SET spend = u.spend + added.cost,
updated_at = (NOW() AT TIME ZONE 'UTC')
FROM added
WHERE u.user_id = added.id AND u.user_id IN (SELECT user_id FROM locked)

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};
pub use spend::{EntitySpend, add_key_spend, add_user_spend};

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

@ -1,3 +1,5 @@
mod key_spend;
mod user_spend;
pub use key_spend::{EntitySpend, add_key_spend};
pub use user_spend::add_user_spend;

173
litellm-rust/crates/db/src/spend/user_spend.rs Normal file
View file

@ -0,0 +1,173 @@
use sqlx::PgExecutor;
use super::EntitySpend;
use crate::Error;
pub async fn add_user_spend(
executor: impl PgExecutor<'_>,
spend: &EntitySpend,
) -> Result<(), Error> {
spend.validate()?;
if spend.is_empty() {
return Ok(());
}
sqlx::query_file!("queries/spend/add_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_user_spend};
use crate::{Error, test_support::database_pool};
use rstest::rstest;
use sqlx::PgPool;
#[rstest]
#[case::fractional_costs(0.25, 2.5, 4.0)]
#[case::zero_costs(0.0, 2.0, 3.0)]
#[tokio::test]
#[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"]
async fn user_spend_matches_python_update_many(
#[case] cost: f64,
#[case] expected_a_spend: f64,
#[case] expected_z_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 _inserted = sqlx::query(
r#"INSERT INTO "LiteLLM_UserTable" (user_id, spend, updated_at)
VALUES ('user-spend-a', 2, '2000-01-01'),
('user-spend-z', 3, '2000-01-01'),
('user-spend-untouched', 100, '2000-01-01')"#,
)
.execute(&mut *tx)
.await?;
let spend = EntitySpend {
ids: vec![
"user-spend-z".into(),
"user-spend-a".into(),
"user-spend-missing".into(),
],
costs: vec![2.0 * cost, cost, 50.0],
};
add_user_spend(&mut *tx, &EntitySpend::default()).await?;
add_user_spend(&mut *tx, &spend).await?;
add_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)> = sqlx::query_as(
r#"SELECT user_id, spend, updated_at > TIMESTAMP '2000-01-01'
FROM "LiteLLM_UserTable"
WHERE user_id = ANY($1) ORDER BY user_id"#,
)
.bind(vec![
"user-spend-a",
"user-spend-z",
"user-spend-missing",
"user-spend-untouched",
])
.fetch_all(&mut *tx)
.await?;
assert_eq!(
rows,
vec![
("user-spend-a".into(), expected_a_spend, true),
("user-spend-untouched".into(), 100.0, false),
("user-spend-z".into(), expected_z_spend, true),
]
);
Ok(())
}
#[rstest]
#[case::direct_duplicate_batch(vec![0.25, 0.75], 3.0)]
#[tokio::test]
#[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"]
async fn user_spend_sums_duplicate_ids(
#[case] costs: Vec<f64>,
#[case] expected_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 _inserted = sqlx::query(
r#"INSERT INTO "LiteLLM_UserTable" (user_id, spend, updated_at)
VALUES ('user-spend-duplicate', 2, '2000-01-01')"#,
)
.execute(&mut *tx)
.await?;
let spend = EntitySpend {
ids: vec!["user-spend-duplicate".into(), "user-spend-duplicate".into()],
costs,
};
add_user_spend(&mut *tx, &spend).await?;
#[expect(
clippy::disallowed_methods,
reason = "Test assertions are outside the compile-checked query set"
)]
let row: (f64, bool) = sqlx::query_as(
r#"SELECT spend, updated_at > TIMESTAMP '2000-01-01'
FROM "LiteLLM_UserTable" WHERE user_id = 'user-spend-duplicate'"#,
)
.fetch_one(&mut *tx)
.await?;
assert_eq!(row, (expected_spend, true));
Ok(())
}
#[rstest]
#[case::missing_cost(vec![])]
#[case::extra_cost(vec![0.25, 0.75])]
#[tokio::test]
#[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"]
async fn user_spend_rejects_malformed_batch_without_writes(
#[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 _inserted = sqlx::query(
r#"INSERT INTO "LiteLLM_UserTable" (user_id, spend, updated_at)
VALUES ('user-spend-malformed', 2, '2000-01-01')"#,
)
.execute(&mut *tx)
.await?;
let spend = EntitySpend {
ids: vec!["user-spend-malformed".into()],
costs,
};
let result = add_user_spend(&mut *tx, &spend).await;
assert!(matches!(result, Err(Error::InvalidSpendBatch { .. })));
#[expect(
clippy::disallowed_methods,
reason = "Test assertions are outside the compile-checked query set"
)]
let row: (f64, bool) = sqlx::query_as(
r#"SELECT spend, updated_at = TIMESTAMP '2000-01-01'
FROM "LiteLLM_UserTable" WHERE user_id = 'user-spend-malformed'"#,
)
.fetch_one(&mut *tx)
.await?;
assert_eq!(row, (2.0, true));
Ok(())
}
}