From 35cd5fc05f7e91c7bc9c91637e14452a3ed55360 Mon Sep 17 00:00:00 2001 From: yujonglee Date: Mon, 26 Oct 2026 17:05:04 -0700 Subject: [PATCH] feat(db): add typed key spend query (#48217) Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...10731a0cacc1ae5b92eec32967f539948f563.json | 15 + .../crates/db/queries/spend/add_key_spend.sql | 16 + litellm-rust/crates/db/src/lib.rs | 2 + litellm-rust/crates/db/src/spend/key_spend.rs | 309 ++++++++++++++++++ litellm-rust/crates/db/src/spend/mod.rs | 3 + 5 files changed, 345 insertions(+) create mode 100644 litellm-rust/crates/db/.sqlx/query-3829a383d0ae50578f0fd75da3910731a0cacc1ae5b92eec32967f539948f563.json create mode 100644 litellm-rust/crates/db/queries/spend/add_key_spend.sql create mode 100644 litellm-rust/crates/db/src/spend/key_spend.rs create mode 100644 litellm-rust/crates/db/src/spend/mod.rs diff --git a/litellm-rust/crates/db/.sqlx/query-3829a383d0ae50578f0fd75da3910731a0cacc1ae5b92eec32967f539948f563.json b/litellm-rust/crates/db/.sqlx/query-3829a383d0ae50578f0fd75da3910731a0cacc1ae5b92eec32967f539948f563.json new file mode 100644 index 00000000000..e85c0986a60 --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-3829a383d0ae50578f0fd75da3910731a0cacc1ae5b92eec32967f539948f563.json @@ -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 t.token FROM \"LiteLLM_VerificationToken\" t\n WHERE t.token IN (SELECT id FROM added)\n ORDER BY t.token FOR UPDATE\n)\nUPDATE \"LiteLLM_VerificationToken\" t\nSET spend = t.spend + added.cost,\n total_spend = t.total_spend + added.cost,\n last_active = (clock_timestamp() AT TIME ZONE 'UTC'),\n updated_at = (clock_timestamp() AT TIME ZONE 'UTC')\nFROM added\nWHERE t.token = added.id AND t.token IN (SELECT token FROM locked)\n", + "describe": { + "columns": [], + "parameters": { + "Left": [ + "TextArray", + "Float8Array" + ] + }, + "nullable": [] + }, + "hash": "3829a383d0ae50578f0fd75da3910731a0cacc1ae5b92eec32967f539948f563" +} diff --git a/litellm-rust/crates/db/queries/spend/add_key_spend.sql b/litellm-rust/crates/db/queries/spend/add_key_spend.sql new file mode 100644 index 00000000000..61850714acd --- /dev/null +++ b/litellm-rust/crates/db/queries/spend/add_key_spend.sql @@ -0,0 +1,16 @@ +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 t.token FROM "LiteLLM_VerificationToken" t + WHERE t.token IN (SELECT id FROM added) + ORDER BY t.token FOR UPDATE +) +UPDATE "LiteLLM_VerificationToken" t +SET spend = t.spend + added.cost, + total_spend = t.total_spend + added.cost, + last_active = (clock_timestamp() AT TIME ZONE 'UTC'), + updated_at = (clock_timestamp() AT TIME ZONE 'UTC') +FROM added +WHERE t.token = added.id AND t.token IN (SELECT token FROM locked) diff --git a/litellm-rust/crates/db/src/lib.rs b/litellm-rust/crates/db/src/lib.rs index ca5cd366eec..b3abf5f5173 100644 --- a/litellm-rust/crates/db/src/lib.rs +++ b/litellm-rust/crates/db/src/lib.rs @@ -10,6 +10,7 @@ mod daily_spend; mod end_user; mod error; mod failure; +mod spend; #[cfg(test)] mod test_support; @@ -17,3 +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}; diff --git a/litellm-rust/crates/db/src/spend/key_spend.rs b/litellm-rust/crates/db/src/spend/key_spend.rs new file mode 100644 index 00000000000..5a6ba13307c --- /dev/null +++ b/litellm-rust/crates/db/src/spend/key_spend.rs @@ -0,0 +1,309 @@ +use sqlx::PgExecutor; +use std::collections::BTreeMap; + +use crate::Error; + +#[macro_rules_attribute::apply(crate::db_input)] +#[derive(Default)] +pub struct EntitySpend { + pub ids: Vec, + pub costs: Vec, +} + +impl EntitySpend { + pub(crate) fn validate(&self) -> Result<(), Error> { + if self.ids.len() != self.costs.len() { + return Err(Error::InvalidSpendBatch { + ids: self.ids.len(), + costs: self.costs.len(), + }); + } + Ok(()) + } + + pub fn is_empty(&self) -> bool { + self.ids.is_empty() + } +} + +impl FromIterator<(String, f64)> for EntitySpend { + fn from_iter>(entries: I) -> Self { + let entries = entries + .into_iter() + .fold(BTreeMap::new(), |mut sums, (id, cost)| { + *sums.entry(id).or_insert(0.0) += cost; + sums + }); + let (ids, costs) = entries.into_iter().unzip(); + Self { ids, costs } + } +} + +pub async fn add_key_spend( + executor: impl PgExecutor<'_>, + spend: &EntitySpend, +) -> Result<(), Error> { + spend.validate()?; + if spend.is_empty() { + return Ok(()); + } + sqlx::query_file!("queries/spend/add_key_spend.sql", &spend.ids, &spend.costs) + .execute(executor) + .await?; + Ok(()) +} + +#[cfg(test)] +mod tests { + use std::error::Error as StdError; + + use super::{EntitySpend, add_key_spend}; + use crate::{Error, test_support::database_pool}; + use rstest::rstest; + use sqlx::PgPool; + + #[rstest] + #[case::direct(false)] + #[case::collected(true)] + #[tokio::test] + #[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"] + async fn duplicate_key_spend_accumulates( + #[case] collected: bool, + #[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" + )] + sqlx::query( + r#"INSERT INTO "LiteLLM_VerificationToken" (token, spend, total_spend, updated_at) + VALUES ('key-spend-duplicate', 2, 7, '2000-01-01')"#, + ) + .execute(&mut *tx) + .await?; + let spend = if collected { + [ + ("key-spend-duplicate".into(), 0.25), + ("key-spend-duplicate".into(), 0.75), + ] + .into_iter() + .collect() + } else { + EntitySpend { + ids: vec!["key-spend-duplicate".into(), "key-spend-duplicate".into()], + costs: vec![0.25, 0.75], + } + }; + add_key_spend(&mut *tx, &spend).await?; + #[expect( + clippy::disallowed_methods, + reason = "Test assertions are outside the compile-checked query set" + )] + let row: (f64, f64) = sqlx::query_as( + r#"SELECT spend, total_spend FROM "LiteLLM_VerificationToken" + WHERE token = 'key-spend-duplicate'"#, + ) + .fetch_one(&mut *tx) + .await?; + assert_eq!(row, (3.0, 8.0)); + Ok(()) + } + + #[rstest] + #[case::extra_id(vec!["key-spend-invalid".into(), "key-spend-invalid".into()], vec![0.25])] + #[case::extra_cost(vec!["key-spend-invalid".into()], vec![0.25, 0.75])] + #[case::empty_ids(vec![], vec![0.25])] + #[tokio::test] + #[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"] + async fn malformed_key_spend_leaves_transaction_usable( + #[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" + )] + sqlx::query( + r#"INSERT INTO "LiteLLM_VerificationToken" + (token, spend, total_spend, last_active, updated_at) + VALUES ('key-spend-invalid', 2, 7, NULL, '2000-01-01')"#, + ) + .execute(&mut *tx) + .await?; + let spend = EntitySpend { ids, costs }; + let result = add_key_spend(&mut *tx, &spend).await; + assert!(matches!( + result, + Err(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 unchanged: (f64, f64, bool, bool) = sqlx::query_as( + r#"SELECT spend, total_spend, last_active IS NULL, + updated_at = TIMESTAMP '2000-01-01' + FROM "LiteLLM_VerificationToken" WHERE token = 'key-spend-invalid'"#, + ) + .fetch_one(&mut *tx) + .await?; + assert_eq!(unchanged, (2.0, 7.0, true, true)); + let valid = EntitySpend { + ids: vec!["key-spend-invalid".into()], + costs: vec![0.5], + }; + add_key_spend(&mut *tx, &valid).await?; + #[expect( + clippy::disallowed_methods, + reason = "Test assertions are outside the compile-checked query set" + )] + let updated: (f64, f64) = sqlx::query_as( + r#"SELECT spend, total_spend FROM "LiteLLM_VerificationToken" + WHERE token = 'key-spend-invalid'"#, + ) + .fetch_one(&mut *tx) + .await?; + assert_eq!(updated, (2.5, 7.5)); + Ok(()) + } + + #[rstest] + #[tokio::test] + #[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"] + async fn key_activity_uses_update_time_in_old_transaction( + #[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" + )] + sqlx::query( + r#"INSERT INTO "LiteLLM_VerificationToken" (token, spend, total_spend, updated_at) + VALUES ('key-spend-old-transaction', 2, 7, '2000-01-01')"#, + ) + .execute(&mut *tx) + .await?; + #[expect( + clippy::disallowed_methods, + reason = "Transaction timing setup is outside the compile-checked query set" + )] + sqlx::query("SELECT pg_sleep(0.02)") + .execute(&mut *tx) + .await?; + #[expect( + clippy::disallowed_methods, + reason = "Test assertions are outside the compile-checked query set" + )] + let boundary: chrono::NaiveDateTime = sqlx::query_scalar( + "SELECT date_trunc('milliseconds', clock_timestamp() AT TIME ZONE 'UTC')", + ) + .fetch_one(&mut *tx) + .await?; + let spend = EntitySpend { + ids: vec!["key-spend-old-transaction".into()], + costs: vec![0.25], + }; + add_key_spend(&mut *tx, &spend).await?; + #[expect( + clippy::disallowed_methods, + reason = "Test assertions are outside the compile-checked query set" + )] + let row: (chrono::NaiveDateTime, chrono::NaiveDateTime) = sqlx::query_as( + r#"SELECT last_active, updated_at FROM "LiteLLM_VerificationToken" + WHERE token = 'key-spend-old-transaction'"#, + ) + .fetch_one(&mut *tx) + .await?; + assert!(row.0 >= boundary, "last_active predates update boundary"); + assert!(row.1 >= boundary, "updated_at predates update boundary"); + Ok(()) + } + + #[rstest] + #[case::fractional_costs(0.25, 2.5, 7.5, 4.0, 12.0)] + #[case::zero_costs(0.0, 2.0, 7.0, 3.0, 11.0)] + #[tokio::test] + #[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"] + async fn key_spend_matches_python_update_many( + #[case] cost: f64, + #[case] expected_a_spend: f64, + #[case] expected_a_total: f64, + #[case] expected_z_spend: f64, + #[case] expected_z_total: 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 _inserted = sqlx::query( + r#"INSERT INTO "LiteLLM_VerificationToken" + (token, spend, total_spend, last_active, updated_at) + VALUES ('key-spend-a', 2, 7, NULL, '2000-01-01'), + ('key-spend-z', 3, 11, NULL, '2000-01-01'), + ('key-spend-untouched', 100, 200, NULL, '2000-01-01')"#, + ) + .execute(&mut *tx) + .await?; + let spend = EntitySpend { + ids: vec![ + "key-spend-z".into(), + "key-spend-a".into(), + "key-spend-missing".into(), + ], + costs: vec![2.0 * cost, cost, 50.0], + }; + add_key_spend(&mut *tx, &EntitySpend::default()).await?; + add_key_spend(&mut *tx, &spend).await?; + add_key_spend(&mut *tx, &spend).await?; + #[expect( + clippy::disallowed_methods, + reason = "Test assertions are outside the compile-checked query set" + )] + let rows: Vec<(String, f64, f64, bool, bool)> = sqlx::query_as( + r#"SELECT token, spend, total_spend, last_active IS NOT NULL, + updated_at > TIMESTAMP '2000-01-01' + FROM "LiteLLM_VerificationToken" + WHERE token = ANY($1) ORDER BY token"#, + ) + .bind(vec![ + "key-spend-a", + "key-spend-z", + "key-spend-missing", + "key-spend-untouched", + ]) + .fetch_all(&mut *tx) + .await?; + assert_eq!( + rows, + vec![ + ( + "key-spend-a".into(), + expected_a_spend, + expected_a_total, + true, + true + ), + ("key-spend-untouched".into(), 100.0, 200.0, false, false), + ( + "key-spend-z".into(), + expected_z_spend, + expected_z_total, + true, + true + ), + ] + ); + Ok(()) + } +} diff --git a/litellm-rust/crates/db/src/spend/mod.rs b/litellm-rust/crates/db/src/spend/mod.rs new file mode 100644 index 00000000000..e34886c6130 --- /dev/null +++ b/litellm-rust/crates/db/src/spend/mod.rs @@ -0,0 +1,3 @@ +mod key_spend; + +pub use key_spend::{EntitySpend, add_key_spend};