feat(db): add typed daily team spend upsert (#48236)

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:07 -07:00 • committed by GitHub
parent 1b874ceb85
commit a2545134f3
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 559 additions and 1 deletions

View file

@ -0,0 +1,31 @@
{
"db_name": "PostgreSQL",
"query": "INSERT INTO \"LiteLLM_DailyTeamSpend\" (\n id, team_id, date, api_key, model, custom_llm_provider,\n mcp_namespaced_tool_name, endpoint, model_group, prompt_tokens,\n completion_tokens, api_requests, successful_requests, failed_requests,\n cache_read_input_tokens, cache_creation_input_tokens, total_response_time_ms,\n timed_requests, spend, updated_at\n)\nSELECT gen_random_uuid()::text, r.entity_id, r.date, r.api_key, r.model,\n r.custom_llm_provider, r.mcp_namespaced_tool_name, r.endpoint, r.model_group, r.prompt_tokens,\n r.completion_tokens, r.api_requests, r.successful_requests,\n r.failed_requests, r.cache_read, r.cache_creation, r.response_time,\n r.timed_requests, r.spend, (NOW() AT TIME ZONE 'UTC')\nFROM UNNEST(\n $1::text[], $2::text[], $3::text[], $4::text[], $5::text[],\n $6::text[], $7::text[], $8::int8[], $9::int8[], $10::int8[],\n $11::int8[], $12::int8[], $13::int8[], $14::int8[], $15::int8[],\n $16::int8[], $17::float8[], $18::text[]\n) AS r(\n entity_id, date, api_key, model, custom_llm_provider, endpoint,\n model_group, prompt_tokens, completion_tokens, api_requests,\n successful_requests, failed_requests, cache_read, cache_creation,\n response_time, timed_requests, spend, mcp_namespaced_tool_name\n)\nORDER BY r.date, r.entity_id, r.api_key, r.model, r.custom_llm_provider,\n r.mcp_namespaced_tool_name, r.endpoint\nON CONFLICT (\n team_id, date, api_key, model, custom_llm_provider,\n mcp_namespaced_tool_name, endpoint\n) DO UPDATE SET\n prompt_tokens = \"LiteLLM_DailyTeamSpend\".prompt_tokens + EXCLUDED.prompt_tokens,\n completion_tokens = \"LiteLLM_DailyTeamSpend\".completion_tokens + EXCLUDED.completion_tokens,\n api_requests = \"LiteLLM_DailyTeamSpend\".api_requests + EXCLUDED.api_requests,\n successful_requests = \"LiteLLM_DailyTeamSpend\".successful_requests + EXCLUDED.successful_requests,\n failed_requests = \"LiteLLM_DailyTeamSpend\".failed_requests + EXCLUDED.failed_requests,\n cache_read_input_tokens = \"LiteLLM_DailyTeamSpend\".cache_read_input_tokens + EXCLUDED.cache_read_input_tokens,\n cache_creation_input_tokens = \"LiteLLM_DailyTeamSpend\".cache_creation_input_tokens + EXCLUDED.cache_creation_input_tokens,\n total_response_time_ms = \"LiteLLM_DailyTeamSpend\".total_response_time_ms + EXCLUDED.total_response_time_ms,\n timed_requests = \"LiteLLM_DailyTeamSpend\".timed_requests + EXCLUDED.timed_requests,\n spend = \"LiteLLM_DailyTeamSpend\".spend + EXCLUDED.spend,\n updated_at = (NOW() AT TIME ZONE 'UTC')\n",
"describe": {
"columns": [],
"parameters": {
"Left": [
"TextArray",
"TextArray",
"TextArray",
"TextArray",
"TextArray",
"TextArray",
"TextArray",
"Int8Array",
"Int8Array",
"Int8Array",
"Int8Array",
"Int8Array",
"Int8Array",
"Int8Array",
"Int8Array",
"Int8Array",
"Float8Array",
"TextArray"
]
},
"nullable": []
},
"hash": "9c51401a8907f818e96da28050b74dca981ff954452b35ffbace3e022035b416"
}

View file

@ -0,0 +1,40 @@
INSERT INTO "LiteLLM_DailyTeamSpend" (
id, team_id, date, api_key, model, custom_llm_provider,
mcp_namespaced_tool_name, endpoint, model_group, prompt_tokens,
completion_tokens, api_requests, successful_requests, failed_requests,
cache_read_input_tokens, cache_creation_input_tokens, total_response_time_ms,
timed_requests, spend, updated_at
)
SELECT gen_random_uuid()::text, r.entity_id, r.date, r.api_key, r.model,
r.custom_llm_provider, r.mcp_namespaced_tool_name, r.endpoint, r.model_group, r.prompt_tokens,
r.completion_tokens, r.api_requests, r.successful_requests,
r.failed_requests, r.cache_read, r.cache_creation, r.response_time,
r.timed_requests, r.spend, (NOW() AT TIME ZONE 'UTC')
FROM UNNEST(
$1::text[], $2::text[], $3::text[], $4::text[], $5::text[],
$6::text[], $7::text[], $8::int8[], $9::int8[], $10::int8[],
$11::int8[], $12::int8[], $13::int8[], $14::int8[], $15::int8[],
$16::int8[], $17::float8[], $18::text[]
) AS r(
entity_id, date, api_key, model, custom_llm_provider, endpoint,
model_group, prompt_tokens, completion_tokens, api_requests,
successful_requests, failed_requests, cache_read, cache_creation,
response_time, timed_requests, spend, mcp_namespaced_tool_name
)
ORDER BY r.date, r.entity_id, r.api_key, r.model, r.custom_llm_provider,
r.mcp_namespaced_tool_name, r.endpoint
ON CONFLICT (
team_id, date, api_key, model, custom_llm_provider,
mcp_namespaced_tool_name, endpoint
) DO UPDATE SET
prompt_tokens = "LiteLLM_DailyTeamSpend".prompt_tokens + EXCLUDED.prompt_tokens,
completion_tokens = "LiteLLM_DailyTeamSpend".completion_tokens + EXCLUDED.completion_tokens,
api_requests = "LiteLLM_DailyTeamSpend".api_requests + EXCLUDED.api_requests,
successful_requests = "LiteLLM_DailyTeamSpend".successful_requests + EXCLUDED.successful_requests,
failed_requests = "LiteLLM_DailyTeamSpend".failed_requests + EXCLUDED.failed_requests,
cache_read_input_tokens = "LiteLLM_DailyTeamSpend".cache_read_input_tokens + EXCLUDED.cache_read_input_tokens,
cache_creation_input_tokens = "LiteLLM_DailyTeamSpend".cache_creation_input_tokens + EXCLUDED.cache_creation_input_tokens,
total_response_time_ms = "LiteLLM_DailyTeamSpend".total_response_time_ms + EXCLUDED.total_response_time_ms,
timed_requests = "LiteLLM_DailyTeamSpend".timed_requests + EXCLUDED.timed_requests,
spend = "LiteLLM_DailyTeamSpend".spend + EXCLUDED.spend,
updated_at = (NOW() AT TIME ZONE 'UTC')

View file

@ -21,5 +21,5 @@ pub use failure::DbFailure;
pub use spend::{
DailySpend, EntitySpend, SpendLog, SpendLogRow, TeamMemberSpend, add_end_user_spend,
add_key_spend, add_team_member_spend, add_team_spend, add_user_spend, begin, insert_spend_logs,
lock_team, upsert_daily_user_spend,
lock_team, upsert_daily_team_spend, upsert_daily_user_spend,
};

View file

@ -0,0 +1,485 @@
use sqlx::PgExecutor;
use super::DailySpend;
use crate::Error;
pub async fn upsert_daily_team_spend(
executor: impl PgExecutor<'_>,
rows: &[DailySpend],
) -> Result<(), Error> {
if rows.is_empty() {
return Ok(());
}
macro_rules! column {
($field:ident) => {
rows.iter()
.map(|row| row.$field.clone())
.collect::<Vec<_>>()
};
}
sqlx::query_file!(
"queries/spend/upsert_daily_team_spend.sql",
&column!(entity_id),
&column!(date),
&column!(api_key),
&column!(model),
&column!(custom_llm_provider),
&column!(endpoint),
&column!(model_group) as &[Option<String>],
&column!(prompt_tokens),
&column!(completion_tokens),
&column!(api_requests),
&column!(successful_requests),
&column!(failed_requests),
&column!(cache_read_input_tokens),
&column!(cache_creation_input_tokens),
&column!(total_response_time_ms),
&column!(timed_requests),
&column!(spend),
&column!(mcp_namespaced_tool_name),
)
.execute(executor)
.await?;
Ok(())
}
#[cfg(test)]
mod tests {
use std::error::Error as StdError;
use super::{DailySpend, upsert_daily_team_spend};
use crate::{DailyTeamRow, test_support::database_pool};
use rstest::rstest;
use sqlx::{PgPool, types::chrono::NaiveDateTime};
#[rstest]
#[case::distinct_tools("server.fetch", "server.search")]
#[case::default_and_named_tool("", "server.search")]
#[tokio::test]
#[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"]
async fn daily_team_spend_separates_tools_and_accumulates(
#[case] first_tool: &str,
#[case] second_tool: &str,
#[future(awt)] database_pool: Result<PgPool, Box<dyn StdError>>,
) -> Result<(), Box<dyn StdError>> {
let pool = database_pool?;
let mut tx = pool.begin().await?;
let first = DailySpend {
entity_id: "daily-team-tools".to_owned(),
date: "2026-01-15".to_owned(),
api_key: "daily-key-tools".to_owned(),
model: "daily-model".to_owned(),
custom_llm_provider: "daily-provider".to_owned(),
endpoint: "/mcp".to_owned(),
mcp_namespaced_tool_name: first_tool.to_owned(),
api_requests: 1,
spend: 0.125,
..Default::default()
};
let second = DailySpend {
mcp_namespaced_tool_name: second_tool.to_owned(),
api_requests: 2,
spend: 0.25,
..first.clone()
};
upsert_daily_team_spend(&mut *tx, &[second.clone(), first.clone()]).await?;
upsert_daily_team_spend(
&mut *tx,
&[
DailySpend {
api_requests: 3,
spend: 0.5,
..first
},
DailySpend {
api_requests: 5,
spend: 1.0,
..second
},
],
)
.await?;
#[expect(
clippy::disallowed_methods,
reason = "The integration test reads stored tool attribution and counters"
)]
let rows = sqlx::query_as::<_, DailyTeamRow>(
r#"SELECT * FROM "LiteLLM_DailyTeamSpend"
WHERE team_id = $1 AND date = $2
ORDER BY mcp_namespaced_tool_name"#,
)
.bind("daily-team-tools")
.bind("2026-01-15")
.fetch_all(&mut *tx)
.await?;
assert_eq!(rows.len(), 2);
assert_eq!(
rows[0].mcp_namespaced_tool_name.as_deref(),
Some(first_tool)
);
assert_eq!(
rows[1].mcp_namespaced_tool_name.as_deref(),
Some(second_tool)
);
assert_eq!((rows[0].api_requests, rows[0].spend), (4, 0.625));
assert_eq!((rows[1].api_requests, rows[1].spend), (7, 1.25));
Ok(())
}
#[rstest]
#[case::full_keys(
"daily-team-base",
"daily-key-base",
"daily-model",
"daily-provider",
"/chat/completions",
Some("first-group")
)]
#[case::initial_model_group_none(
"daily-team-base",
"daily-key-base",
"daily-model",
"daily-provider",
"/chat/completions",
None
)]
#[case::empty_optional_keys("", "", "", "", "", Some(""))]
#[tokio::test]
#[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"]
async fn daily_team_spend_upsert_preserves_keys_and_accumulates(
#[case] team_id: &'static str,
#[case] api_key: &'static str,
#[case] model: &'static str,
#[case] custom_llm_provider: &'static str,
#[case] endpoint: &'static str,
#[case] first_model_group: Option<&'static str>,
#[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 = "The integration test exercises UTC timestamp behavior under a non-UTC session"
)]
sqlx::query("SET LOCAL TIME ZONE 'America/Los_Angeles'")
.execute(&mut *tx)
.await?;
let first_batch = [
DailySpend {
entity_id: team_id.to_owned(),
date: "2026-01-15".to_owned(),
api_key: api_key.to_owned(),
model: model.to_owned(),
custom_llm_provider: custom_llm_provider.to_owned(),
mcp_namespaced_tool_name: String::new(),
endpoint: endpoint.to_owned(),
model_group: first_model_group.map(str::to_owned),
prompt_tokens: 1,
completion_tokens: 2,
api_requests: 3,
successful_requests: 4,
failed_requests: 5,
cache_read_input_tokens: 6,
cache_creation_input_tokens: 7,
total_response_time_ms: 8,
timed_requests: 9,
spend: 0.125,
},
DailySpend {
entity_id: team_id.to_owned(),
date: "2026-01-15".to_owned(),
api_key: api_key.to_owned(),
model: model.to_owned(),
custom_llm_provider: custom_llm_provider.to_owned(),
mcp_namespaced_tool_name: String::new(),
endpoint: "endpoint-sentinel".to_owned(),
model_group: Some("sentinel-endpoint-group".to_owned()),
prompt_tokens: 101,
completion_tokens: 102,
api_requests: 103,
successful_requests: 104,
failed_requests: 105,
cache_read_input_tokens: 106,
cache_creation_input_tokens: 107,
total_response_time_ms: 108,
timed_requests: 109,
spend: 1.25,
},
DailySpend {
entity_id: team_id.to_owned(),
date: "2026-01-15".to_owned(),
api_key: "api-key-sentinel".to_owned(),
model: model.to_owned(),
custom_llm_provider: custom_llm_provider.to_owned(),
mcp_namespaced_tool_name: String::new(),
endpoint: endpoint.to_owned(),
model_group: Some("sentinel-key-group".to_owned()),
prompt_tokens: 201,
completion_tokens: 202,
api_requests: 203,
successful_requests: 204,
failed_requests: 205,
cache_read_input_tokens: 206,
cache_creation_input_tokens: 207,
total_response_time_ms: 208,
timed_requests: 209,
spend: 2.5,
},
DailySpend {
entity_id: "entity-sentinel".to_owned(),
date: "2026-01-15".to_owned(),
api_key: api_key.to_owned(),
model: model.to_owned(),
custom_llm_provider: custom_llm_provider.to_owned(),
mcp_namespaced_tool_name: String::new(),
endpoint: endpoint.to_owned(),
model_group: Some("sentinel-entity-group".to_owned()),
prompt_tokens: 301,
completion_tokens: 302,
api_requests: 303,
successful_requests: 304,
failed_requests: 305,
cache_read_input_tokens: 306,
cache_creation_input_tokens: 307,
total_response_time_ms: 308,
timed_requests: 309,
spend: 3.75,
},
];
upsert_daily_team_spend(&mut *tx, &first_batch).await?;
#[expect(
clippy::disallowed_methods,
reason = "The integration test captures the inserted row identity before conflict handling"
)]
let (initial_id, initial_created_at): (String, NaiveDateTime) = sqlx::query_as(
r#"SELECT id, created_at
FROM "LiteLLM_DailyTeamSpend"
WHERE team_id = $1 AND date = $2 AND api_key = $3 AND model = $4
AND custom_llm_provider = $5 AND endpoint = $6"#,
)
.bind(team_id)
.bind("2026-01-15")
.bind(api_key)
.bind(model)
.bind(custom_llm_provider)
.bind(endpoint)
.fetch_one(&mut *tx)
.await?;
assert!(!initial_id.is_empty());
#[expect(
clippy::disallowed_methods,
reason = "The integration test seeds updated_at before exercising the conflict update"
)]
let seeded_updated_at: NaiveDateTime = sqlx::query_scalar(
r#"UPDATE "LiteLLM_DailyTeamSpend"
SET updated_at = TIMESTAMP '2000-01-01'
WHERE id = $1
RETURNING updated_at"#,
)
.bind(&initial_id)
.fetch_one(&mut *tx)
.await?;
let increment = [DailySpend {
entity_id: team_id.to_owned(),
date: "2026-01-15".to_owned(),
api_key: api_key.to_owned(),
model: model.to_owned(),
custom_llm_provider: custom_llm_provider.to_owned(),
mcp_namespaced_tool_name: String::new(),
endpoint: endpoint.to_owned(),
model_group: Some("second-group".to_owned()),
prompt_tokens: 10,
completion_tokens: 20,
api_requests: 30,
successful_requests: 40,
failed_requests: 50,
cache_read_input_tokens: 60,
cache_creation_input_tokens: 70,
total_response_time_ms: 80,
timed_requests: 90,
spend: 0.625,
}];
upsert_daily_team_spend(&mut *tx, &increment).await?;
upsert_daily_team_spend(&mut *tx, &[]).await?;
#[expect(
clippy::disallowed_methods,
reason = "The integration test captures the transaction timestamp for assertions"
)]
let expected_updated_at: NaiveDateTime =
sqlx::query_scalar("SELECT (NOW() AT TIME ZONE 'UTC')::timestamp(3)")
.fetch_one(&mut *tx)
.await?;
#[expect(
clippy::disallowed_methods,
reason = "The integration test reads all rows to verify attribution and accumulation"
)]
let rows: Vec<DailyTeamRow> = sqlx::query_as(
r#"SELECT id, team_id, api_key, model, custom_llm_provider, endpoint,
model_group, mcp_namespaced_tool_name, prompt_tokens,
completion_tokens, api_requests, successful_requests,
failed_requests, cache_read_input_tokens,
cache_creation_input_tokens, total_response_time_ms,
timed_requests, spend, date, created_at, updated_at
FROM "LiteLLM_DailyTeamSpend"
WHERE date = $1 AND model = $2 AND custom_llm_provider = $3
AND mcp_namespaced_tool_name = '' AND team_id = ANY($4)
ORDER BY CASE
WHEN team_id = $5 AND api_key = $6 AND endpoint = $7 THEN 0
WHEN team_id = $5 AND api_key = $6 AND endpoint = $9 THEN 1
WHEN team_id = $5 AND api_key = $8 AND endpoint = $7 THEN 2
WHEN team_id = $10 AND api_key = $6 AND endpoint = $7 THEN 3
ELSE 4
END"#,
)
.bind("2026-01-15")
.bind(model)
.bind(custom_llm_provider)
.bind(vec![team_id, "entity-sentinel"])
.bind(team_id)
.bind(api_key)
.bind(endpoint)
.bind("api-key-sentinel")
.bind("endpoint-sentinel")
.bind("entity-sentinel")
.fetch_all(&mut *tx)
.await?;
assert_eq!(rows.len(), 4);
assert_eq!(rows[0].id, initial_id);
assert_eq!(rows[0].team_id.as_deref(), Some(team_id));
assert_eq!(rows[0].api_key, api_key);
assert_eq!(rows[0].model.as_deref(), Some(model));
assert_eq!(
rows[0].custom_llm_provider.as_deref(),
Some(custom_llm_provider)
);
assert_eq!(rows[0].endpoint.as_deref(), Some(endpoint));
assert_eq!(rows[0].model_group.as_deref(), first_model_group);
assert_eq!(rows[0].mcp_namespaced_tool_name.as_deref(), Some(""));
assert_eq!(
(
rows[0].prompt_tokens,
rows[0].completion_tokens,
rows[0].api_requests,
rows[0].successful_requests,
rows[0].failed_requests,
rows[0].cache_read_input_tokens,
rows[0].cache_creation_input_tokens,
rows[0].total_response_time_ms,
rows[0].timed_requests,
rows[0].spend
),
(11, 22, 33, 44, 55, 66, 77, 88, 99, 0.75)
);
assert_eq!(rows[0].date, "2026-01-15");
assert_eq!(rows[0].created_at, initial_created_at);
assert!(rows[0].updated_at > seeded_updated_at);
assert_eq!(rows[0].updated_at, expected_updated_at);
assert!(!rows[1].id.is_empty());
assert_eq!(rows[1].team_id.as_deref(), Some(team_id));
assert_eq!(rows[1].api_key, api_key);
assert_eq!(rows[1].model.as_deref(), Some(model));
assert_eq!(
rows[1].custom_llm_provider.as_deref(),
Some(custom_llm_provider)
);
assert_eq!(rows[1].endpoint.as_deref(), Some("endpoint-sentinel"));
assert_eq!(
rows[1].model_group.as_deref(),
Some("sentinel-endpoint-group")
);
assert_eq!(rows[1].mcp_namespaced_tool_name.as_deref(), Some(""));
assert_eq!(
(
rows[1].prompt_tokens,
rows[1].completion_tokens,
rows[1].api_requests,
rows[1].successful_requests,
rows[1].failed_requests,
rows[1].cache_read_input_tokens,
rows[1].cache_creation_input_tokens,
rows[1].total_response_time_ms,
rows[1].timed_requests,
rows[1].spend
),
(101, 102, 103, 104, 105, 106, 107, 108, 109, 1.25)
);
assert_eq!(rows[1].date, "2026-01-15");
assert!(rows[1].updated_at > seeded_updated_at);
assert_eq!(rows[1].updated_at, expected_updated_at);
assert!(!rows[2].id.is_empty());
assert_eq!(rows[2].team_id.as_deref(), Some(team_id));
assert_eq!(rows[2].api_key, "api-key-sentinel");
assert_eq!(rows[2].model.as_deref(), Some(model));
assert_eq!(
rows[2].custom_llm_provider.as_deref(),
Some(custom_llm_provider)
);
assert_eq!(rows[2].endpoint.as_deref(), Some(endpoint));
assert_eq!(rows[2].model_group.as_deref(), Some("sentinel-key-group"));
assert_eq!(rows[2].mcp_namespaced_tool_name.as_deref(), Some(""));
assert_eq!(
(
rows[2].prompt_tokens,
rows[2].completion_tokens,
rows[2].api_requests,
rows[2].successful_requests,
rows[2].failed_requests,
rows[2].cache_read_input_tokens,
rows[2].cache_creation_input_tokens,
rows[2].total_response_time_ms,
rows[2].timed_requests,
rows[2].spend
),
(201, 202, 203, 204, 205, 206, 207, 208, 209, 2.5)
);
assert_eq!(rows[2].date, "2026-01-15");
assert!(rows[2].updated_at > seeded_updated_at);
assert_eq!(rows[2].updated_at, expected_updated_at);
assert!(!rows[3].id.is_empty());
assert_eq!(rows[3].team_id.as_deref(), Some("entity-sentinel"));
assert_eq!(rows[3].api_key, api_key);
assert_eq!(rows[3].model.as_deref(), Some(model));
assert_eq!(
rows[3].custom_llm_provider.as_deref(),
Some(custom_llm_provider)
);
assert_eq!(rows[3].endpoint.as_deref(), Some(endpoint));
assert_eq!(
rows[3].model_group.as_deref(),
Some("sentinel-entity-group")
);
assert_eq!(rows[3].mcp_namespaced_tool_name.as_deref(), Some(""));
assert_eq!(
(
rows[3].prompt_tokens,
rows[3].completion_tokens,
rows[3].api_requests,
rows[3].successful_requests,
rows[3].failed_requests,
rows[3].cache_read_input_tokens,
rows[3].cache_creation_input_tokens,
rows[3].total_response_time_ms,
rows[3].timed_requests,
rows[3].spend
),
(301, 302, 303, 304, 305, 306, 307, 308, 309, 3.75)
);
assert_eq!(rows[3].date, "2026-01-15");
assert!(rows[3].updated_at > seeded_updated_at);
assert_eq!(rows[3].updated_at, expected_updated_at);
Ok(())
}
}

View file

@ -1,3 +1,4 @@
mod daily_team_spend;
mod daily_user_spend;
mod end_user_spend;
mod key_spend;
@ -7,6 +8,7 @@ mod team_member_spend;
mod team_spend;
mod user_spend;
pub use daily_team_spend::upsert_daily_team_spend;
pub use daily_user_spend::{DailySpend, upsert_daily_user_spend};
pub use end_user_spend::add_end_user_spend;
pub use key_spend::{EntitySpend, add_key_spend};