From b446d589090e73f118462f4d1d4d57823cd03bcf Mon Sep 17 00:00:00 2001 From: yujonglee Date: Mon, 26 Oct 2026 17:05:06 -0700 Subject: [PATCH] feat(db): add typed spend-log insert query (#48234) Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...acf6d06eb22facdf478f88107c9dccdac326b.json | 40 ++ .../db/queries/spend/insert_spend_logs.sql | 26 + litellm-rust/crates/db/src/lib.rs | 4 +- litellm-rust/crates/db/src/spend/mod.rs | 2 + .../crates/db/src/spend/spend_logs.rs | 517 ++++++++++++++++++ 5 files changed, 587 insertions(+), 2 deletions(-) create mode 100644 litellm-rust/crates/db/.sqlx/query-32f3c07ced96324b535bbc703c1acf6d06eb22facdf478f88107c9dccdac326b.json create mode 100644 litellm-rust/crates/db/queries/spend/insert_spend_logs.sql create mode 100644 litellm-rust/crates/db/src/spend/spend_logs.rs diff --git a/litellm-rust/crates/db/.sqlx/query-32f3c07ced96324b535bbc703c1acf6d06eb22facdf478f88107c9dccdac326b.json b/litellm-rust/crates/db/.sqlx/query-32f3c07ced96324b535bbc703c1acf6d06eb22facdf478f88107c9dccdac326b.json new file mode 100644 index 00000000000..c4eb1737d85 --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-32f3c07ced96324b535bbc703c1acf6d06eb22facdf478f88107c9dccdac326b.json @@ -0,0 +1,40 @@ +{ + "db_name": "PostgreSQL", + "query": "INSERT INTO \"LiteLLM_SpendLogs\" (\n request_id, call_type, api_key, spend, total_tokens, prompt_tokens,\n completion_tokens, \"startTime\", \"endTime\", \"completionStartTime\",\n request_duration_ms, model, model_id, model_group, custom_llm_provider, api_base,\n \"user\", metadata, cache_hit, cache_key, request_tags, team_id, organization_id,\n end_user, requester_ip_address, messages, response, proxy_server_request,\n session_id, status, litellm_call_id\n)\nSELECT l.request_id, l.call_type, l.api_key, l.spend, l.total_tokens, l.prompt_tokens,\n l.completion_tokens, l.start_time::timestamp(3), l.end_time::timestamp(3),\n l.completion_start_time::timestamp(3), l.request_duration_ms, l.model,\n l.model_id, l.model_group, l.custom_llm_provider, l.api_base, l.user_id,\n l.metadata, l.cache_hit, l.cache_key, '[]'::jsonb, l.team_id,\n l.organization_id, l.end_user, l.requester_ip_address, '{}'::jsonb,\n '{}'::jsonb, '{}'::jsonb, l.session_id, l.status, l.litellm_call_id\nFROM UNNEST($1::text[], $2::text[], $3::text[], $4::float8[], $5::int4[], $6::int4[],\n $7::int4[], $8::timestamp[], $9::timestamp[], $10::timestamp[], $11::int4[], $12::text[],\n $13::text[], $14::text[], $15::text[], $16::text[], $17::text[],\n $18::jsonb[], $19::text[], $20::text[], $21::text[], $22::text[],\n $23::text[], $24::text[], $25::text[], $26::text[], $27::text[])\n AS l(request_id, call_type, api_key, spend, total_tokens, prompt_tokens,\n completion_tokens, start_time, end_time, completion_start_time,\n request_duration_ms, model, model_id, model_group, custom_llm_provider,\n api_base, user_id, metadata, cache_hit, cache_key, team_id, organization_id,\n end_user, requester_ip_address, session_id, status, litellm_call_id)\nON CONFLICT DO NOTHING\n", + "describe": { + "columns": [], + "parameters": { + "Left": [ + "TextArray", + "TextArray", + "TextArray", + "Float8Array", + "Int4Array", + "Int4Array", + "Int4Array", + "TimestampArray", + "TimestampArray", + "TimestampArray", + "Int4Array", + "TextArray", + "TextArray", + "TextArray", + "TextArray", + "TextArray", + "TextArray", + "JsonbArray", + "TextArray", + "TextArray", + "TextArray", + "TextArray", + "TextArray", + "TextArray", + "TextArray", + "TextArray", + "TextArray" + ] + }, + "nullable": [] + }, + "hash": "32f3c07ced96324b535bbc703c1acf6d06eb22facdf478f88107c9dccdac326b" +} diff --git a/litellm-rust/crates/db/queries/spend/insert_spend_logs.sql b/litellm-rust/crates/db/queries/spend/insert_spend_logs.sql new file mode 100644 index 00000000000..fcffee181a1 --- /dev/null +++ b/litellm-rust/crates/db/queries/spend/insert_spend_logs.sql @@ -0,0 +1,26 @@ +INSERT INTO "LiteLLM_SpendLogs" ( + request_id, call_type, api_key, spend, total_tokens, prompt_tokens, + completion_tokens, "startTime", "endTime", "completionStartTime", + request_duration_ms, model, model_id, model_group, custom_llm_provider, api_base, + "user", metadata, cache_hit, cache_key, request_tags, team_id, organization_id, + end_user, requester_ip_address, messages, response, proxy_server_request, + session_id, status, litellm_call_id +) +SELECT l.request_id, l.call_type, l.api_key, l.spend, l.total_tokens, l.prompt_tokens, + l.completion_tokens, l.start_time::timestamp(3), l.end_time::timestamp(3), + l.completion_start_time::timestamp(3), l.request_duration_ms, l.model, + l.model_id, l.model_group, l.custom_llm_provider, l.api_base, l.user_id, + l.metadata, l.cache_hit, l.cache_key, '[]'::jsonb, l.team_id, + l.organization_id, l.end_user, l.requester_ip_address, '{}'::jsonb, + '{}'::jsonb, '{}'::jsonb, l.session_id, l.status, l.litellm_call_id +FROM UNNEST($1::text[], $2::text[], $3::text[], $4::float8[], $5::int4[], $6::int4[], + $7::int4[], $8::timestamp[], $9::timestamp[], $10::timestamp[], $11::int4[], $12::text[], + $13::text[], $14::text[], $15::text[], $16::text[], $17::text[], + $18::jsonb[], $19::text[], $20::text[], $21::text[], $22::text[], + $23::text[], $24::text[], $25::text[], $26::text[], $27::text[]) + AS l(request_id, call_type, api_key, spend, total_tokens, prompt_tokens, + completion_tokens, start_time, end_time, completion_start_time, + request_duration_ms, model, model_id, model_group, custom_llm_provider, + api_base, user_id, metadata, cache_hit, cache_key, team_id, organization_id, + end_user, requester_ip_address, session_id, status, litellm_call_id) +ON CONFLICT DO NOTHING diff --git a/litellm-rust/crates/db/src/lib.rs b/litellm-rust/crates/db/src/lib.rs index afa962cff0e..511a4bd387a 100644 --- a/litellm-rust/crates/db/src/lib.rs +++ b/litellm-rust/crates/db/src/lib.rs @@ -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, TeamMemberSpend, add_end_user_spend, add_key_spend, add_team_member_spend, - add_team_spend, add_user_spend, begin, lock_team, + 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, }; diff --git a/litellm-rust/crates/db/src/spend/mod.rs b/litellm-rust/crates/db/src/spend/mod.rs index 1acf629af4e..dcfcee41f83 100644 --- a/litellm-rust/crates/db/src/spend/mod.rs +++ b/litellm-rust/crates/db/src/spend/mod.rs @@ -1,5 +1,6 @@ mod end_user_spend; mod key_spend; +mod spend_logs; mod team_lock; mod team_member_spend; mod team_spend; @@ -7,6 +8,7 @@ mod user_spend; pub use end_user_spend::add_end_user_spend; pub use key_spend::{EntitySpend, add_key_spend}; +pub use spend_logs::{SpendLog, SpendLogRow, insert_spend_logs}; pub use team_lock::{begin, lock_team}; pub use team_member_spend::{TeamMemberSpend, add_team_member_spend}; pub use team_spend::add_team_spend; diff --git a/litellm-rust/crates/db/src/spend/spend_logs.rs b/litellm-rust/crates/db/src/spend/spend_logs.rs new file mode 100644 index 00000000000..57a9ff85301 --- /dev/null +++ b/litellm-rust/crates/db/src/spend/spend_logs.rs @@ -0,0 +1,517 @@ +use chrono::{DateTime, NaiveDateTime}; +use sqlx::PgExecutor; + +use crate::Error; + +#[macro_rules_attribute::apply(crate::db_input)] +#[derive(Default)] +pub struct SpendLog { + pub request_id: String, + pub call_type: String, + pub api_key: String, + pub spend: f64, + pub total_tokens: i32, + pub prompt_tokens: i32, + pub completion_tokens: i32, + pub start_time: String, + pub end_time: String, + pub completion_start_time: Option, + pub request_duration_ms: Option, + pub model: String, + pub model_id: String, + pub model_group: String, + pub custom_llm_provider: String, + pub api_base: String, + pub user: String, + pub metadata: serde_json::Value, + pub cache_hit: String, + pub cache_key: String, + pub team_id: String, + pub organization_id: String, + pub end_user: String, + pub requester_ip_address: Option, + pub session_id: Option, + pub status: String, + pub litellm_call_id: Option, +} + +#[macro_rules_attribute::apply(crate::db_row)] +pub struct SpendLogRow { + pub request_id: String, + pub call_type: String, + pub api_key: String, + pub spend: f64, + pub total_tokens: i32, + pub prompt_tokens: i32, + pub completion_tokens: i32, + #[sqlx(rename = "startTime")] + pub start_time: NaiveDateTime, + #[sqlx(rename = "endTime")] + pub end_time: NaiveDateTime, + #[sqlx(rename = "completionStartTime")] + pub completion_start_time: Option, + pub request_duration_ms: Option, + pub model: String, + pub model_id: Option, + pub model_group: Option, + pub custom_llm_provider: Option, + pub api_base: Option, + pub user: Option, + pub metadata: Option, + pub cache_hit: Option, + pub cache_key: Option, + pub request_tags: Option, + pub team_id: Option, + pub organization_id: Option, + pub end_user: Option, + pub requester_ip_address: Option, + pub messages: Option, + pub response: Option, + pub proxy_server_request: Option, + pub session_id: Option, + pub status: Option, + pub litellm_call_id: Option, +} + +fn parse_timestamp(value: &str) -> Result { + DateTime::parse_from_rfc3339(value) + .map(|timestamp| timestamp.naive_utc()) + .or_else(|_| NaiveDateTime::parse_from_str(value, "%Y-%m-%dT%H:%M:%S%.f")) + .or_else(|_| NaiveDateTime::parse_from_str(value, "%Y-%m-%d %H:%M:%S%.f")) +} + +pub async fn insert_spend_logs( + executor: impl PgExecutor<'_>, + logs: &[SpendLog], +) -> Result<(), Error> { + if logs.is_empty() { + return Ok(()); + } + let start_times = logs + .iter() + .map(|log| parse_timestamp(&log.start_time)) + .collect::, _>>()?; + let end_times = logs + .iter() + .map(|log| parse_timestamp(&log.end_time)) + .collect::, _>>()?; + let completion_start_times = logs + .iter() + .map(|log| { + log.completion_start_time + .as_deref() + .map(parse_timestamp) + .transpose() + }) + .collect::, _>>()?; + macro_rules! column { + ($field:ident) => { + logs.iter() + .map(|log| log.$field.clone()) + .collect::>() + }; + } + sqlx::query_file!( + "queries/spend/insert_spend_logs.sql", + &column!(request_id), + &column!(call_type), + &column!(api_key), + &column!(spend), + &column!(total_tokens), + &column!(prompt_tokens), + &column!(completion_tokens), + &start_times, + &end_times, + &completion_start_times as &[Option], + &column!(request_duration_ms) as &[Option], + &column!(model), + &column!(model_id), + &column!(model_group), + &column!(custom_llm_provider), + &column!(api_base), + &column!(user), + &column!(metadata), + &column!(cache_hit), + &column!(cache_key), + &column!(team_id), + &column!(organization_id), + &column!(end_user), + &column!(requester_ip_address) as &[Option], + &column!(session_id) as &[Option], + &column!(status), + &column!(litellm_call_id) as &[Option], + ) + .execute(executor) + .await?; + Ok(()) +} + +#[cfg(test)] +mod tests { + use std::error::Error as StdError; + + use super::{SpendLog, SpendLogRow, insert_spend_logs}; + use crate::test_support::database_pool; + use rstest::{fixture, rstest}; + use serde_json::json; + use sqlx::{Acquire, Connection, PgExecutor, PgPool, types::chrono::NaiveDateTime}; + + #[fixture] + fn logs() -> [SpendLog; 2] { + [ + SpendLog { + request_id: "spend-logs-a".into(), + call_type: "acompletion".into(), + api_key: "key-a".into(), + spend: 0.125, + total_tokens: 17, + prompt_tokens: 11, + completion_tokens: 6, + start_time: "2026-01-02 03:04:05.123".into(), + end_time: "2026-01-02 03:04:06.789".into(), + completion_start_time: Some("2026-01-02 03:04:05.456".into()), + request_duration_ms: Some(1666), + model: "model-a".into(), + model_id: "deployment-a".into(), + model_group: "group-a".into(), + custom_llm_provider: "provider-a".into(), + api_base: "https://a.example.test".into(), + user: "user-a".into(), + metadata: json!({"tenant": "a", "nested": {"tokens": 11}, "flag": false}), + cache_hit: "False".into(), + cache_key: "cache-a".into(), + team_id: "team-a".into(), + organization_id: "org-a".into(), + end_user: "end-user-a".into(), + requester_ip_address: Some("192.0.2.1".into()), + session_id: Some("session-a".into()), + status: "success".into(), + litellm_call_id: Some("call-a".into()), + }, + SpendLog { + request_id: "spend-logs-b".into(), + call_type: "aembedding".into(), + api_key: "key-b".into(), + spend: 2.75, + total_tokens: 43, + prompt_tokens: 29, + completion_tokens: 14, + start_time: "2026-02-03 14:15:16.234".into(), + end_time: "2026-02-03 14:15:18.567".into(), + completion_start_time: None, + request_duration_ms: None, + model: "model-b".into(), + model_id: "deployment-b".into(), + model_group: "group-b".into(), + custom_llm_provider: "provider-b".into(), + api_base: "https://b.example.test".into(), + user: "user-b".into(), + metadata: json!({"tenant": "b", "tags": ["distinct", 7], "nullable": null}), + cache_hit: "True".into(), + cache_key: "cache-b".into(), + team_id: "team-b".into(), + organization_id: "org-b".into(), + end_user: "end-user-b".into(), + requester_ip_address: None, + session_id: None, + status: "failure".into(), + litellm_call_id: None, + }, + ] + } + + async fn read_log( + executor: impl PgExecutor<'_>, + request_id: &str, + ) -> Result { + #[expect( + clippy::disallowed_methods, + reason = "Test assertions are outside the compile-checked query set" + )] + sqlx::query_as::<_, SpendLogRow>( + r#"SELECT * FROM "LiteLLM_SpendLogs" WHERE request_id = $1"#, + ) + .bind(request_id) + .fetch_one(executor) + .await + } + + fn expected_log(log: &SpendLog) -> Result { + Ok(SpendLogRow { + request_id: log.request_id.clone(), + call_type: log.call_type.clone(), + api_key: log.api_key.clone(), + spend: log.spend, + total_tokens: log.total_tokens, + prompt_tokens: log.prompt_tokens, + completion_tokens: log.completion_tokens, + start_time: log.start_time.replace(' ', "T").parse()?, + end_time: log.end_time.replace(' ', "T").parse()?, + completion_start_time: log + .completion_start_time + .as_ref() + .map(|value| value.replace(' ', "T").parse()) + .transpose()?, + request_duration_ms: log.request_duration_ms, + model: log.model.clone(), + model_id: Some(log.model_id.clone()), + model_group: Some(log.model_group.clone()), + custom_llm_provider: Some(log.custom_llm_provider.clone()), + api_base: Some(log.api_base.clone()), + user: Some(log.user.clone()), + metadata: Some(log.metadata.clone()), + cache_hit: Some(log.cache_hit.clone()), + cache_key: Some(log.cache_key.clone()), + request_tags: Some(json!([])), + team_id: Some(log.team_id.clone()), + organization_id: Some(log.organization_id.clone()), + end_user: Some(log.end_user.clone()), + requester_ip_address: log.requester_ip_address.clone(), + messages: Some(json!({})), + response: Some(json!({})), + proxy_server_request: Some(json!({})), + session_id: log.session_id.clone(), + status: Some(log.status.clone()), + litellm_call_id: log.litellm_call_id.clone(), + }) + } + + #[rstest] + #[case::all_columns_and_duplicate()] + #[tokio::test] + #[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"] + async fn inserts_distinct_rows_without_replacing_duplicates( + logs: [SpendLog; 2], + #[future(awt)] database_pool: Result>, + ) -> Result<(), Box> { + let pool = database_pool?; + let mut tx = pool.begin().await?; + insert_spend_logs(&mut *tx, &logs).await?; + let replacement = SpendLog { + request_id: logs[0].request_id.clone(), + ..logs[1].clone() + }; + insert_spend_logs(&mut *tx, &[replacement]).await?; + assert_eq!( + read_log(&mut *tx, &logs[0].request_id).await?, + expected_log(&logs[0])? + ); + assert_eq!( + read_log(&mut *tx, &logs[1].request_id).await?, + expected_log(&logs[1])? + ); + tx.rollback().await?; + Ok(()) + } + + #[rstest] + #[case::partitioned_duplicate()] + #[tokio::test] + #[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"] + async fn inserts_into_partitioned_logs_without_replacing_duplicates( + logs: [SpendLog; 2], + #[future(awt)] database_pool: Result>, + ) -> Result<(), Box> { + let pool = database_pool?; + let mut tx = pool.begin().await?; + #[expect( + clippy::disallowed_methods, + reason = "The test creates a temporary partitioned table at runtime" + )] + sqlx::raw_sql( + r#"CREATE TEMP TABLE "LiteLLM_SpendLogs" ( + LIKE public."LiteLLM_SpendLogs" INCLUDING DEFAULTS INCLUDING GENERATED, + PRIMARY KEY (request_id, "startTime") + ) PARTITION BY RANGE ("startTime"); + CREATE TEMP TABLE spend_logs_default + PARTITION OF pg_temp."LiteLLM_SpendLogs" DEFAULT; + SET LOCAL search_path TO pg_temp, public;"#, + ) + .execute(&mut *tx) + .await?; + tx.clear_cached_statements().await?; + insert_spend_logs(&mut *tx, &logs).await?; + let replacement = SpendLog { + spend: 99.0, + ..logs[0].clone() + }; + insert_spend_logs(&mut *tx, &[replacement]).await?; + #[expect( + clippy::disallowed_methods, + reason = "The partitioned test fixture is created at runtime" + )] + let count: (i64,) = sqlx::query_as(r#"SELECT count(*) FROM pg_temp."LiteLLM_SpendLogs""#) + .fetch_one(&mut *tx) + .await?; + assert_eq!(count.0, 2); + assert_eq!( + read_log(&mut *tx, &logs[0].request_id).await?, + expected_log(&logs[0])? + ); + assert_eq!( + read_log(&mut *tx, &logs[1].request_id).await?, + expected_log(&logs[1])? + ); + let distinct_partition_key = SpendLog { + request_id: logs[0].request_id.clone(), + ..logs[1].clone() + }; + insert_spend_logs(&mut *tx, &[distinct_partition_key]).await?; + #[expect( + clippy::disallowed_methods, + reason = "The partitioned test fixture is created at runtime" + )] + let count: (i64,) = sqlx::query_as(r#"SELECT count(*) FROM pg_temp."LiteLLM_SpendLogs""#) + .fetch_one(&mut *tx) + .await?; + assert_eq!(count.0, 3); + tx.rollback().await?; + Ok(()) + } + + #[rstest] + #[case::positive_offset_across_midnight( + ["2026-01-02T00:04:05.1236+02:00", "2026-01-02T00:04:06.7894+02:00", "2026-01-02T00:04:05.4567+02:00"], + ["2026-01-01T22:04:05.124", "2026-01-01T22:04:06.789", "2026-01-01T22:04:05.457"] + )] + #[case::negative_offset_across_midnight( + ["2026-01-01T23:04:05-02:00", "2026-01-01T23:04:06-02:00", "2026-01-01T23:04:05.456-02:00"], + ["2026-01-02T01:04:05", "2026-01-02T01:04:06", "2026-01-02T01:04:05.456"] + )] + #[case::utc( + ["2026-01-02T03:04:05Z", "2026-01-02T03:04:06Z", "2026-01-02T03:04:05.456Z"], + ["2026-01-02T03:04:05", "2026-01-02T03:04:06", "2026-01-02T03:04:05.456"] + )] + #[case::naive_iso( + ["2026-01-02T03:04:05", "2026-01-02T03:04:06", "2026-01-02T03:04:05.456"], + ["2026-01-02T03:04:05", "2026-01-02T03:04:06", "2026-01-02T03:04:05.456"] + )] + #[case::naive_space( + ["2026-01-02 03:04:05", "2026-01-02 03:04:06", "2026-01-02 03:04:05.456"], + ["2026-01-02T03:04:05", "2026-01-02T03:04:06", "2026-01-02T03:04:05.456"] + )] + #[tokio::test] + #[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"] + async fn stores_utc_timestamps( + #[case] input: [&str; 3], + #[case] expected: [&str; 3], + logs: [SpendLog; 2], + #[future(awt)] database_pool: Result>, + ) -> Result<(), Box> { + let pool = database_pool?; + let mut tx = pool.begin().await?; + #[expect( + clippy::disallowed_methods, + reason = "The test sets a transaction-local timezone to verify timestamp handling" + )] + sqlx::raw_sql("SET LOCAL TIME ZONE 'America/Los_Angeles'") + .execute(&mut *tx) + .await?; + let log = SpendLog { + start_time: input[0].into(), + end_time: input[1].into(), + completion_start_time: Some(input[2].into()), + ..logs[0].clone() + }; + insert_spend_logs(&mut *tx, std::slice::from_ref(&log)).await?; + let row = read_log(&mut *tx, &log.request_id).await?; + assert_eq!(row.start_time, expected[0].parse::()?); + assert_eq!(row.end_time, expected[1].parse::()?); + assert_eq!( + row.completion_start_time, + Some(expected[2].parse::()?) + ); + tx.rollback().await?; + Ok(()) + } + + #[rstest] + #[case::empty_batch()] + #[tokio::test] + #[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"] + async fn empty_input_leaves_rows_unchanged( + #[future(awt)] database_pool: Result>, + ) -> Result<(), Box> { + let pool = database_pool?; + let mut tx = pool.begin().await?; + #[expect( + clippy::disallowed_methods, + reason = "Test assertions are outside the compile-checked query set" + )] + let before: (i64,) = sqlx::query_as(r#"SELECT count(*) FROM "LiteLLM_SpendLogs""#) + .fetch_one(&mut *tx) + .await?; + insert_spend_logs(&mut *tx, &[]).await?; + #[expect( + clippy::disallowed_methods, + reason = "Test assertions are outside the compile-checked query set" + )] + let after: (i64,) = sqlx::query_as(r#"SELECT count(*) FROM "LiteLLM_SpendLogs""#) + .fetch_one(&mut *tx) + .await?; + assert_eq!(before, after); + tx.rollback().await?; + Ok(()) + } + + #[rstest] + #[case::invalid_timestamp("not-a-timestamp", "2026-01-02T03:04:06", None)] + #[case::invalid_end("2026-01-02T03:04:05", "not-a-timestamp", None)] + #[case::invalid_completion( + "2026-01-02T03:04:05", + "2026-01-02T03:04:06", + Some("not-a-timestamp") + )] + #[tokio::test] + #[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"] + async fn invalid_timestamp_allows_caller_rollback( + #[case] start_time: &str, + #[case] end_time: &str, + #[case] completion_start_time: Option<&str>, + logs: [SpendLog; 2], + #[future(awt)] database_pool: Result>, + ) -> Result<(), Box> { + let pool = database_pool?; + let mut outer = pool.begin().await?; + let mut tx = outer.begin().await?; + let valid = SpendLog { + request_id: "spend-logs-rollback-valid".into(), + ..logs[0].clone() + }; + let invalid = SpendLog { + request_id: "spend-logs-rollback-invalid".into(), + start_time: start_time.into(), + end_time: end_time.into(), + completion_start_time: completion_start_time.map(str::to_owned), + ..logs[1].clone() + }; + insert_spend_logs(&mut *tx, std::slice::from_ref(&valid)).await?; + let result = insert_spend_logs(&mut *tx, &[logs[0].clone(), invalid]).await; + assert!(matches!(result, Err(crate::Error::InvalidTimestamp(_)))); + #[expect( + clippy::disallowed_methods, + reason = "Test assertions are outside the compile-checked query set" + )] + let partial_batch: (i64,) = sqlx::query_as( + r#"SELECT count(*) FROM "LiteLLM_SpendLogs" WHERE request_id = ANY($1)"#, + ) + .bind([logs[0].request_id.as_str(), "spend-logs-rollback-invalid"].as_slice()) + .fetch_one(&mut *tx) + .await?; + assert_eq!(partial_batch.0, 0); + tx.rollback().await?; + #[expect( + clippy::disallowed_methods, + reason = "Test assertions are outside the compile-checked query set" + )] + let remaining: (i64,) = sqlx::query_as( + r#"SELECT count(*) FROM "LiteLLM_SpendLogs" WHERE request_id = ANY($1)"#, + ) + .bind([valid.request_id.as_str(), "spend-logs-rollback-invalid"].as_slice()) + .fetch_one(&mut *outer) + .await?; + assert_eq!(remaining.0, 0); + outer.rollback().await?; + Ok(()) + } +}