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 team-member spend upsert (#48233)

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:06 -07:00 • committed by GitHub
parent c531692976
commit 46410c8da9
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 197 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
litellm-rust/Cargo.lock generated
View file

@ -3699,6 +3699,7 @@ dependencies = [
"macro_rules_attribute",
"rstest",
"serde",
"serde_json",
"sqlx",
"thiserror 2.0.19",
"tokio",

14
litellm-rust/crates/db/.sqlx/query-b0eb48b795bdaf728c79e986d2027df60ff71d8fed05b53be4458b335af88147.json generated Normal file
View file

@ -0,0 +1,14 @@
{
"db_name": "PostgreSQL",
"query": "INSERT INTO \"LiteLLM_TeamMembership\" (user_id, team_id, spend, total_spend)\nSELECT member.user_id, member.team_id, member.cost, member.cost\nFROM jsonb_to_recordset($1::jsonb) AS member(user_id text, team_id text, cost float8)\nWHERE EXISTS (\n SELECT 1 FROM \"LiteLLM_TeamTable\" t\n WHERE t.team_id = member.team_id\n AND t.members_with_roles @> jsonb_build_array(jsonb_build_object('user_id', member.user_id))\n)\n OR EXISTS (\n SELECT 1 FROM \"LiteLLM_TeamMembership\" m\n WHERE m.user_id = member.user_id AND m.team_id = member.team_id\n)\nON CONFLICT (user_id, team_id) DO UPDATE\nSET spend = \"LiteLLM_TeamMembership\".spend + EXCLUDED.spend,\n total_spend = \"LiteLLM_TeamMembership\".total_spend + EXCLUDED.total_spend\n",
"describe": {
"columns": [],
"parameters": {
"Left": [
"Jsonb"
]
},
"nullable": []
},
"hash": "b0eb48b795bdaf728c79e986d2027df60ff71d8fed05b53be4458b335af88147"
}

1
litellm-rust/crates/db/Cargo.toml
View file

@ -9,6 +9,7 @@ repository.workspace = true
chrono.workspace = true
macro_rules_attribute.workspace = true
serde.workspace = true
serde_json.workspace = true
sqlx.workspace = true
thiserror.workspace = true

15
litellm-rust/crates/db/queries/spend/upsert_team_member_spend.sql Normal file
View file

@ -0,0 +1,15 @@
INSERT INTO "LiteLLM_TeamMembership" (user_id, team_id, spend, total_spend)
SELECT member.user_id, member.team_id, member.cost, member.cost
FROM jsonb_to_recordset($1::jsonb) AS member(user_id text, team_id text, cost float8)
WHERE EXISTS (
SELECT 1 FROM "LiteLLM_TeamTable" t
WHERE t.team_id = member.team_id
AND t.members_with_roles @> jsonb_build_array(jsonb_build_object('user_id', member.user_id))
)
OR EXISTS (
SELECT 1 FROM "LiteLLM_TeamMembership" m
WHERE m.user_id = member.user_id AND m.team_id = member.team_id
)
ON CONFLICT (user_id, team_id) DO UPDATE
SET spend = "LiteLLM_TeamMembership".spend + EXCLUDED.spend,
total_spend = "LiteLLM_TeamMembership".total_spend + EXCLUDED.total_spend

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

@ -19,6 +19,6 @@ pub use end_user::{EndUser, RateLimits, find_end_user};
pub use error::Error;
pub use failure::DbFailure;
pub use spend::{
EntitySpend, add_end_user_spend, add_key_spend, add_team_spend, add_user_spend, begin,
lock_team,
EntitySpend, TeamMemberSpend, add_end_user_spend, add_key_spend, add_team_member_spend,
add_team_spend, add_user_spend, begin, lock_team,
};

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

@ -1,11 +1,13 @@
mod end_user_spend;
mod key_spend;
mod team_lock;
mod team_member_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_lock::{begin, lock_team};
pub use team_member_spend::{TeamMemberSpend, add_team_member_spend};
pub use team_spend::add_team_spend;
pub use user_spend::add_user_spend;

162
litellm-rust/crates/db/src/spend/team_member_spend.rs Normal file
View file

@ -0,0 +1,162 @@
use sqlx::{Postgres, Transaction};
use super::lock_team;
use crate::Error;
#[macro_rules_attribute::apply(crate::db_input)]
pub struct TeamMemberSpend {
pub team_id: String,
pub user_id: String,
pub cost: f64,
}
pub async fn add_team_member_spend(
transaction: &mut Transaction<'_, Postgres>,
members: &[TeamMemberSpend],
) -> Result<(), Error> {
if members.is_empty() {
return Ok(());
}
let mut members = members.to_vec();
members.sort_by(|left, right| {
(&left.team_id, &left.user_id).cmp(&(&right.team_id, &right.user_id))
});
let mut teams: Vec<&str> = members
.iter()
.map(|member| member.team_id.as_str())
.collect();
teams.dedup();
for team in teams {
lock_team(transaction, team).await?;
}
let rows = serde_json::Value::Array(
members
.iter()
.map(|member| {
serde_json::json!({
"user_id": member.user_id,
"team_id": member.team_id,
"cost": member.cost,
})
})
.collect(),
);
sqlx::query_file!("queries/spend/upsert_team_member_spend.sql", rows)
.execute(&mut **transaction)
.await?;
Ok(())
}
#[cfg(test)]
mod tests {
use std::error::Error as StdError;
use super::super::begin;
use super::{TeamMemberSpend, add_team_member_spend};
use crate::test_support::database_pool;
use rstest::rstest;
use sqlx::PgPool;
fn members(cost: f64) -> [TeamMemberSpend; 3] {
[
TeamMemberSpend {
team_id: "member-spend-z".into(),
user_id: "off-roster-existing".into(),
cost: 2.0 * cost,
},
TeamMemberSpend {
team_id: "member-spend-a".into(),
user_id: "on-roster-missing".into(),
cost,
},
TeamMemberSpend {
team_id: "member-spend-a".into(),
user_id: "off-roster-missing".into(),
cost: 50.0 * cost,
},
]
}
#[rstest]
#[case::zero_then_quarter(0.25, 0.5, 3.0, 8.0)]
#[case::zero_then_eighth(0.125, 0.25, 2.5, 7.5)]
#[tokio::test]
#[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"]
async fn membership_spend_preserves_roster_guard_and_fractional_costs(
#[case] cost: f64,
#[case] expected_created: f64,
#[case] expected_spend: f64,
#[case] expected_total: f64,
#[future(awt)] database_pool: Result<PgPool, Box<dyn StdError>>,
) -> Result<(), Box<dyn StdError>> {
let pool = database_pool?;
let mut tx = begin(&pool).await?;
#[expect(
clippy::disallowed_methods,
reason = "Fixture writes are outside the compile-checked query set"
)]
let _teams = sqlx::query(
r#"INSERT INTO "LiteLLM_TeamTable" (team_id, members_with_roles)
VALUES ('member-spend-a', '[{"user_id":"on-roster-missing","role":"user"}]'),
('member-spend-z', '[]')"#,
)
.execute(&mut *tx)
.await?;
#[expect(
clippy::disallowed_methods,
reason = "Fixture writes are outside the compile-checked query set"
)]
let _members = sqlx::query(
r#"INSERT INTO "LiteLLM_TeamMembership" (team_id, user_id, spend, total_spend)
VALUES ('member-spend-z', 'off-roster-existing', 2, 7),
('member-spend-a', 'untouched', 100, 200)"#,
)
.execute(&mut *tx)
.await?;
add_team_member_spend(&mut tx, &[]).await?;
add_team_member_spend(&mut tx, &members(0.0)).await?;
add_team_member_spend(&mut tx, &members(cost)).await?;
add_team_member_spend(&mut tx, &members(cost)).await?;
#[expect(
clippy::disallowed_methods,
reason = "Test assertions are outside the compile-checked query set"
)]
let rows: Vec<(String, String, f64, f64)> = sqlx::query_as(
r#"SELECT team_id, user_id, spend, total_spend FROM "LiteLLM_TeamMembership"
WHERE team_id IN ('member-spend-a', 'member-spend-z') ORDER BY team_id, user_id"#,
)
.fetch_all(&mut *tx)
.await?;
assert_eq!(
rows,
vec![
(
"member-spend-a".into(),
"on-roster-missing".into(),
expected_created,
expected_created
),
("member-spend-a".into(), "untouched".into(), 100.0, 200.0),
(
"member-spend-z".into(),
"off-roster-existing".into(),
expected_spend,
expected_total
),
]
);
tx.rollback().await?;
#[expect(
clippy::disallowed_methods,
reason = "Test assertions are outside the compile-checked query set"
)]
let remaining: i64 = sqlx::query_scalar(
r#"SELECT COUNT(*) FROM "LiteLLM_TeamMembership"
WHERE team_id IN ('member-spend-a', 'member-spend-z')"#,
)
.fetch_one(&pool)
.await?;
assert_eq!(remaining, 0);
Ok(())
}
}