diff --git a/litellm-rust/crates/db/.sqlx/query-43f2e12a3ffb1de29b5b031a1173b87a20a2664b0bef3616631d2c24487c1e54.json b/litellm-rust/crates/db/.sqlx/query-43f2e12a3ffb1de29b5b031a1173b87a20a2664b0bef3616631d2c24487c1e54.json new file mode 100644 index 00000000000..a758ae5126a --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-43f2e12a3ffb1de29b5b031a1173b87a20a2664b0bef3616631d2c24487c1e54.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 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" +} diff --git a/litellm-rust/crates/db/queries/spend/add_user_spend.sql b/litellm-rust/crates/db/queries/spend/add_user_spend.sql new file mode 100644 index 00000000000..27fc12e82f0 --- /dev/null +++ b/litellm-rust/crates/db/queries/spend/add_user_spend.sql @@ -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) diff --git a/litellm-rust/crates/db/src/lib.rs b/litellm-rust/crates/db/src/lib.rs index b3abf5f5173..6589b46b24a 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}; +pub use spend::{EntitySpend, add_key_spend, add_user_spend}; diff --git a/litellm-rust/crates/db/src/spend/mod.rs b/litellm-rust/crates/db/src/spend/mod.rs index e34886c6130..98530805fdd 100644 --- a/litellm-rust/crates/db/src/spend/mod.rs +++ b/litellm-rust/crates/db/src/spend/mod.rs @@ -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; diff --git a/litellm-rust/crates/db/src/spend/user_spend.rs b/litellm-rust/crates/db/src/spend/user_spend.rs new file mode 100644 index 00000000000..2c89c04c1ef --- /dev/null +++ b/litellm-rust/crates/db/src/spend/user_spend.rs @@ -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>, + ) -> 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_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, + #[case] expected_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 _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, + #[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_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(()) + } +}