diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 0e72af1bb3b..7565f7911fa 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -3768,6 +3768,7 @@ name = "litellm-db" version = "0.1.0" dependencies = [ "chrono", + "futures-util", "litellm-secret-value", "macro_rules_attribute", "rstest", diff --git a/litellm-rust/crates/db/.sqlx/query-26f616a74c671fb1fddd79b167f4f571c6fae8cd6d09a7f67fd3f41025e260be.json b/litellm-rust/crates/db/.sqlx/query-26f616a74c671fb1fddd79b167f4f571c6fae8cd6d09a7f67fd3f41025e260be.json new file mode 100644 index 00000000000..3c8e9b5c55f --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-26f616a74c671fb1fddd79b167f4f571c6fae8cd6d09a7f67fd3f41025e260be.json @@ -0,0 +1,23 @@ +{ + "db_name": "PostgreSQL", + "query": "SELECT to_jsonb(v) AS \"data!: AuthData\"\nFROM \"LiteLLM_ManagedVectorStoresTable\" v\nWHERE v.vector_store_id = ANY($1::text[])\nORDER BY v.vector_store_id\n", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "data!: AuthData", + "type_info": "Jsonb", + "origin": "Expression" + } + ], + "parameters": { + "Left": [ + "TextArray" + ] + }, + "nullable": [ + null + ] + }, + "hash": "26f616a74c671fb1fddd79b167f4f571c6fae8cd6d09a7f67fd3f41025e260be" +} diff --git a/litellm-rust/crates/db/.sqlx/query-a22cd61e78998c35ddbe1e514911a0cb3af72f78c4c6f90890131faa1a15d3db.json b/litellm-rust/crates/db/.sqlx/query-a22cd61e78998c35ddbe1e514911a0cb3af72f78c4c6f90890131faa1a15d3db.json new file mode 100644 index 00000000000..a27f0fd3ea7 --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-a22cd61e78998c35ddbe1e514911a0cb3af72f78c4c6f90890131faa1a15d3db.json @@ -0,0 +1,23 @@ +{ + "db_name": "PostgreSQL", + "query": "SELECT to_jsonb(a) || jsonb_build_object('extra_headers', COALESCE(to_jsonb(a.extra_headers), '[]'::jsonb),\n 'agent_access_groups', COALESCE(to_jsonb(a.agent_access_groups), '[]'::jsonb),\n 'access_group_ids', COALESCE(to_jsonb(a.access_group_ids), '[]'::jsonb), 'identity', (SELECT to_jsonb(i) || jsonb_build_object('required_roles', COALESCE(to_jsonb(i.required_roles), '[]'::jsonb),\n 'required_scopes', COALESCE(to_jsonb(i.required_scopes), '[]'::jsonb)) FROM \"LiteLLM_AgentIdentity\" i WHERE i.agent_id=a.agent_id), 'object_permission', (SELECT to_jsonb(op) || jsonb_build_object('mcp_servers', COALESCE(to_jsonb(op.mcp_servers), '[]'::jsonb),\n 'mcp_access_groups', COALESCE(to_jsonb(op.mcp_access_groups), '[]'::jsonb),\n 'mcp_toolsets', COALESCE(to_jsonb(op.mcp_toolsets), '[]'::jsonb),\n 'blocked_tools', COALESCE(to_jsonb(op.blocked_tools), '[]'::jsonb),\n 'vector_stores', COALESCE(to_jsonb(op.vector_stores), '[]'::jsonb),\n 'agents', COALESCE(to_jsonb(op.agents), '[]'::jsonb),\n 'agent_access_groups', COALESCE(to_jsonb(op.agent_access_groups), '[]'::jsonb),\n 'models', COALESCE(to_jsonb(op.models), '[]'::jsonb),\n 'search_tools', COALESCE(to_jsonb(op.search_tools), '[]'::jsonb),\n 'skills', COALESCE(to_jsonb(op.skills), '[]'::jsonb)) FROM \"LiteLLM_ObjectPermissionTable\" op WHERE op.object_permission_id = a.object_permission_id)) AS \"data!: AuthData\"\nFROM \"LiteLLM_AgentsTable\" a\nWHERE a.agent_name = ANY($1::text[])\nORDER BY a.agent_id\n", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "data!: AuthData", + "type_info": "Jsonb", + "origin": "Expression" + } + ], + "parameters": { + "Left": [ + "TextArray" + ] + }, + "nullable": [ + null + ] + }, + "hash": "a22cd61e78998c35ddbe1e514911a0cb3af72f78c4c6f90890131faa1a15d3db" +} diff --git a/litellm-rust/crates/db/.sqlx/query-b7d3bcef6d8d9fdd8eb3a51c1b6c97afbec425ba252db103d956f15c43c23c28.json b/litellm-rust/crates/db/.sqlx/query-b7d3bcef6d8d9fdd8eb3a51c1b6c97afbec425ba252db103d956f15c43c23c28.json new file mode 100644 index 00000000000..4e904c41350 --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-b7d3bcef6d8d9fdd8eb3a51c1b6c97afbec425ba252db103d956f15c43c23c28.json @@ -0,0 +1,23 @@ +{ + "db_name": "PostgreSQL", + "query": "SELECT jsonb_build_object('agent_id', a.agent_id) AS \"data!: AuthData\"\nFROM \"LiteLLM_AgentsTable\" a\nWHERE a.agent_access_groups && $1::text[]\nORDER BY a.agent_id\n", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "data!: AuthData", + "type_info": "Jsonb", + "origin": "Expression" + } + ], + "parameters": { + "Left": [ + "TextArray" + ] + }, + "nullable": [ + null + ] + }, + "hash": "b7d3bcef6d8d9fdd8eb3a51c1b6c97afbec425ba252db103d956f15c43c23c28" +} diff --git a/litellm-rust/crates/db/Cargo.toml b/litellm-rust/crates/db/Cargo.toml index 6fc1949abc4..d377eca2086 100644 --- a/litellm-rust/crates/db/Cargo.toml +++ b/litellm-rust/crates/db/Cargo.toml @@ -10,6 +10,7 @@ veil.workspace = true litellm-secret-value.workspace = true tokio.workspace = true chrono.workspace = true +futures-util.workspace = true macro_rules_attribute.workspace = true serde.workspace = true serde_json.workspace = true diff --git a/litellm-rust/crates/db/queries/auth/agent_access_group.sql b/litellm-rust/crates/db/queries/auth/agent_access_group.sql new file mode 100644 index 00000000000..750669b5b4c --- /dev/null +++ b/litellm-rust/crates/db/queries/auth/agent_access_group.sql @@ -0,0 +1,4 @@ +SELECT jsonb_build_object('agent_id', a.agent_id) AS "data!: AuthData" +FROM "LiteLLM_AgentsTable" a +WHERE a.agent_access_groups && $1::text[] +ORDER BY a.agent_id diff --git a/litellm-rust/crates/db/queries/auth/agent_name.sql b/litellm-rust/crates/db/queries/auth/agent_name.sql new file mode 100644 index 00000000000..01776e8d95e --- /dev/null +++ b/litellm-rust/crates/db/queries/auth/agent_name.sql @@ -0,0 +1,16 @@ +SELECT to_jsonb(a) || jsonb_build_object('extra_headers', COALESCE(to_jsonb(a.extra_headers), '[]'::jsonb), + 'agent_access_groups', COALESCE(to_jsonb(a.agent_access_groups), '[]'::jsonb), + 'access_group_ids', COALESCE(to_jsonb(a.access_group_ids), '[]'::jsonb), 'identity', (SELECT to_jsonb(i) || jsonb_build_object('required_roles', COALESCE(to_jsonb(i.required_roles), '[]'::jsonb), + 'required_scopes', COALESCE(to_jsonb(i.required_scopes), '[]'::jsonb)) FROM "LiteLLM_AgentIdentity" i WHERE i.agent_id=a.agent_id), 'object_permission', (SELECT to_jsonb(op) || jsonb_build_object('mcp_servers', COALESCE(to_jsonb(op.mcp_servers), '[]'::jsonb), + 'mcp_access_groups', COALESCE(to_jsonb(op.mcp_access_groups), '[]'::jsonb), + 'mcp_toolsets', COALESCE(to_jsonb(op.mcp_toolsets), '[]'::jsonb), + 'blocked_tools', COALESCE(to_jsonb(op.blocked_tools), '[]'::jsonb), + 'vector_stores', COALESCE(to_jsonb(op.vector_stores), '[]'::jsonb), + 'agents', COALESCE(to_jsonb(op.agents), '[]'::jsonb), + 'agent_access_groups', COALESCE(to_jsonb(op.agent_access_groups), '[]'::jsonb), + 'models', COALESCE(to_jsonb(op.models), '[]'::jsonb), + 'search_tools', COALESCE(to_jsonb(op.search_tools), '[]'::jsonb), + 'skills', COALESCE(to_jsonb(op.skills), '[]'::jsonb)) FROM "LiteLLM_ObjectPermissionTable" op WHERE op.object_permission_id = a.object_permission_id)) AS "data!: AuthData" +FROM "LiteLLM_AgentsTable" a +WHERE a.agent_name = ANY($1::text[]) +ORDER BY a.agent_id diff --git a/litellm-rust/crates/db/queries/auth/fixture.sql b/litellm-rust/crates/db/queries/auth/fixture.sql index 4031cc03b72..f50348acc1b 100644 --- a/litellm-rust/crates/db/queries/auth/fixture.sql +++ b/litellm-rust/crates/db/queries/auth/fixture.sql @@ -37,8 +37,10 @@ INSERT INTO "LiteLLM_ModelAccessGroupBudgetTable" (access_group_name, budget_id, VALUES ('auth-model-group', 'auth-b', 12.5); INSERT INTO "LiteLLM_AccessGroupTable" (access_group_id, access_group_name, access_model_names, access_agent_ids, access_mcp_server_ids) VALUES ('auth-access', 'Auth Access', ARRAY['group-model'], ARRAY['auth-agent'], NULL); -INSERT INTO "LiteLLM_AgentsTable" (agent_id, agent_name, agent_card_params, object_permission_id, identity_managed, enabled, extra_headers, created_by, updated_by) -VALUES ('auth-agent', 'Auth Agent', '{}', 'auth-op', true, false, NULL, 'auth-test', 'auth-test'); +INSERT INTO "LiteLLM_AgentsTable" (agent_id, agent_name, agent_card_params, object_permission_id, identity_managed, enabled, extra_headers, agent_access_groups, created_by, updated_by) +VALUES ('auth-agent', 'Auth Agent', '{}', 'auth-op', true, false, NULL, ARRAY['auth-group'], 'auth-test', 'auth-test'); +INSERT INTO "LiteLLM_ManagedVectorStoresTable" (vector_store_id, custom_llm_provider, litellm_params, team_id, updated_at) +VALUES ('auth-vector', 'openai', '{"api_base":"https://example.com"}', 'auth-team', now()); INSERT INTO "LiteLLM_AgentIdentity" (agent_id, provider, issuer, tenant_id, client_id, revision, active, required_roles, required_scopes) VALUES ('auth-agent', 'microsoft_entra', 'auth-issuer', 'auth-tenant', 'auth-client', 'auth-revision', false, NULL, ARRAY['read']); INSERT INTO "LiteLLM_RetiredAgentIdentity" (binding_id, provider, issuer, tenant_id, client_id) diff --git a/litellm-rust/crates/db/queries/auth/managed_vector_store.sql b/litellm-rust/crates/db/queries/auth/managed_vector_store.sql new file mode 100644 index 00000000000..06765a45ac2 --- /dev/null +++ b/litellm-rust/crates/db/queries/auth/managed_vector_store.sql @@ -0,0 +1,4 @@ +SELECT to_jsonb(v) AS "data!: AuthData" +FROM "LiteLLM_ManagedVectorStoresTable" v +WHERE v.vector_store_id = ANY($1::text[]) +ORDER BY v.vector_store_id diff --git a/litellm-rust/crates/db/src/auth.rs b/litellm-rust/crates/db/src/auth.rs index 5f1a0fbe5d4..840b517120d 100644 --- a/litellm-rust/crates/db/src/auth.rs +++ b/litellm-rust/crates/db/src/auth.rs @@ -1,4 +1,4 @@ -use crate::Error; +use crate::{Error, executor::StatementPersistence}; use sqlx::{Executor, Postgres}; #[derive(Clone, PartialEq, serde::Serialize, sqlx::Type, veil::Redact)] @@ -29,12 +29,18 @@ pub enum AuthRead { ObjectPermission, #[strum(serialize = "access_group")] AccessGroup, + #[strum(serialize = "managed_vector_store")] + ManagedVectorStore, #[strum(serialize = "tag")] Tag, #[strum(serialize = "model_access_group")] ModelAccessGroup, #[strum(serialize = "agent")] Agent, + #[strum(serialize = "agent_name")] + AgentName, + #[strum(serialize = "agent_access_group")] + AgentAccessGroup, #[strum(serialize = "retired_agent")] RetiredAgent, #[strum(serialize = "key")] @@ -119,7 +125,12 @@ pub async fn auth_read<'e>( executor: impl Executor<'e, Database = Postgres>, kind: AuthRead, ids: &[String], + persistent: bool, ) -> Result, Error> { + let executor = StatementPersistence { + executor, + persistent, + }; if ids.len() > 1000 || ids.iter().any(|id| id.len() > 65536) { return Err(Error::InvalidAuthArguments); } @@ -164,6 +175,11 @@ pub async fn auth_read<'e>( .fetch_all(executor) .await? } + AuthRead::ManagedVectorStore => { + sqlx::query_file_as!(AuthRow, "queries/auth/managed_vector_store.sql", ids) + .fetch_all(executor) + .await? + } AuthRead::Tag => { sqlx::query_file_as!(AuthRow, "queries/auth/tag.sql", ids) .fetch_all(executor) @@ -179,6 +195,16 @@ pub async fn auth_read<'e>( .fetch_all(executor) .await? } + AuthRead::AgentName => { + sqlx::query_file_as!(AuthRow, "queries/auth/agent_name.sql", ids) + .fetch_all(executor) + .await? + } + AuthRead::AgentAccessGroup => { + sqlx::query_file_as!(AuthRow, "queries/auth/agent_access_group.sql", ids) + .fetch_all(executor) + .await? + } AuthRead::RetiredAgent => { sqlx::query_file_as!(AuthRow, "queries/auth/retired_agent.sql", ids) .fetch_all(executor) @@ -333,9 +359,8 @@ pub async fn auth_read<'e>( exact(ids, 3)?; let start = chrono::DateTime::parse_from_rfc3339(&ids[2]) .map(|value| value.naive_utc()) - .or_else(|_| { - chrono::NaiveDateTime::parse_from_str(&ids[2], "%Y-%m-%dT%H:%M:%S%.f") - })?; + .or_else(|_| chrono::NaiveDateTime::parse_from_str(&ids[2], "%Y-%m-%dT%H:%M:%S%.f")) + .map_err(|_| Error::InvalidAuthArguments)?; match entity(&ids[0])? { "key" => { hashed_key(&ids[1])?; @@ -368,7 +393,12 @@ pub async fn auth_spend<'e>( executor: impl Executor<'e, Database = Postgres>, kind: AuthSpend, id: &str, + persistent: bool, ) -> Result, Error> { + let executor = StatementPersistence { + executor, + persistent, + }; if id.len() > 65536 { return Err(Error::InvalidAuthArguments); } @@ -447,6 +477,8 @@ mod tests { } #[rstest] + #[case::vector_store(AuthRead::ManagedVectorStore, &["auth-vector"], "/litellm_params", json!({"api_base":"https://example.com"}))] + #[case::vector_store_nullable(AuthRead::ManagedVectorStore, &["auth-vector"], "/vector_store_metadata", json!(null))] #[case::user(AuthRead::User, &["auth-u"], "/organization_memberships/0/user_role", json!("org_admin"))] #[case::team(AuthRead::Team, &["auth-team"], "/litellm_model_table/model_aliases/alias", json!("actual"))] #[case::team_null_lists(AuthRead::Team, &["auth-team-empty"], "/models", json!([]))] @@ -461,6 +493,8 @@ mod tests { #[case::model_access_group(AuthRead::ModelAccessGroup, &["auth-model-group"], "/litellm_budget_table/max_budget", json!(0.0))] #[case::member(AuthRead::Member, &["auth-u","auth-team"], "/litellm_budget_table/rpm_limit", json!(33))] #[case::agent(AuthRead::Agent, &["auth-agent"], "/identity/active", json!(false))] + #[case::agent_name(AuthRead::AgentName, &["Auth Agent"], "/identity/active", json!(false))] + #[case::agent_group(AuthRead::AgentAccessGroup, &["other-group", "auth-group"], "", json!({"agent_id":"auth-agent"}))] #[case::retired_agent(AuthRead::RetiredAgent, &["auth-retired-agent"], "/original_agent_id", json!("auth-retired-agent"))] #[case::agent_identity(AuthRead::AgentIdentity, &["microsoft_entra","auth-tenant","auth-client"], "/required_roles", json!([]))] #[case::retired_agent_identity(AuthRead::RetiredAgentIdentity, &["microsoft_entra","auth-tenant","auth-retired-client"], "/binding_id", json!("auth-retired-binding"))] @@ -487,7 +521,7 @@ mod tests { ) -> Result<(), Box> { let pool = database_pool?; let mut tx = seeded(&pool).await?; - let rows = auth_read(&mut *tx, kind, &strings(ids)).await?; + let rows = auth_read(&mut *tx, kind, &strings(ids), false).await?; assert_eq!(rows.len(), 1); let actual = rows[0].data.0.pointer(pointer).expect("projected field"); if expected.is_number() { @@ -513,7 +547,13 @@ mod tests { ) -> Result<(), Box> { let pool = database_pool?; let mut tx = seeded(&pool).await?; - let rows = auth_read(&mut *tx, AuthRead::Key, &[token.to_string().repeat(64)]).await?; + let rows = auth_read( + &mut *tx, + AuthRead::Key, + &[token.to_string().repeat(64)], + false, + ) + .await?; assert_eq!(rows.len(), usize::from(found)); if !found { return Ok(()); @@ -541,7 +581,7 @@ mod tests { ) -> Result<(), Box> { let pool = database_pool?; let mut tx = seeded(&pool).await?; - let rows = auth_read(&mut *tx, AuthRead::Key, &["e".repeat(64)]).await?; + let rows = auth_read(&mut *tx, AuthRead::Key, &["e".repeat(64)], false).await?; assert_eq!(rows.len(), 1); assert_eq!(rows[0].data.0["team_models"], Value::Null); assert_eq!(rows[0].data.0["team_blocked"], Value::Null); @@ -559,17 +599,24 @@ mod tests { let pool = database_pool?; let mut tx = seeded(&pool).await?; let refs = json!(["auth-u", "auth-team", "auth-u", "auth-org", "auth-project"]).to_string(); - let rows = auth_read(&mut *tx, AuthRead::Prefetch, &[refs]).await?; - let user = auth_read(&mut *tx, AuthRead::User, &strings(&["auth-u"])).await?; - let team = auth_read(&mut *tx, AuthRead::Team, &strings(&["auth-team"])).await?; + let rows = auth_read(&mut *tx, AuthRead::Prefetch, &[refs], false).await?; + let user = auth_read(&mut *tx, AuthRead::User, &strings(&["auth-u"]), false).await?; + let team = auth_read(&mut *tx, AuthRead::Team, &strings(&["auth-team"]), false).await?; let member = auth_read( &mut *tx, AuthRead::Member, &strings(&["auth-u", "auth-team"]), + false, + ) + .await?; + let org = auth_read(&mut *tx, AuthRead::Org, &strings(&["auth-org"]), false).await?; + let project = auth_read( + &mut *tx, + AuthRead::Project, + &strings(&["auth-project"]), + false, ) .await?; - let org = auth_read(&mut *tx, AuthRead::Org, &strings(&["auth-org"])).await?; - let project = auth_read(&mut *tx, AuthRead::Project, &strings(&["auth-project"])).await?; assert_eq!(rows.len(), 1); assert_eq!( rows[0].data.0, @@ -583,6 +630,7 @@ mod tests { &mut *tx, AuthRead::Prefetch, &[json!([null, null, null, null, null]).to_string()], + false, ) .await?; assert_eq!( @@ -604,16 +652,22 @@ mod tests { &mut *tx, AuthRead::User, &strings(&["auth-u", "missing", "auth-u"]), + false, ) .await?; assert_eq!(rows.len(), 1); assert_eq!(rows[0].data.0["user_id"], "auth-u"); - assert!(auth_read(&mut *tx, AuthRead::User, &[]).await?.is_empty()); + assert!( + auth_read(&mut *tx, AuthRead::User, &[], false) + .await? + .is_empty() + ); assert!( auth_read( &mut *tx, AuthRead::Member, - &strings(&["auth-other", "auth-team"]) + &strings(&["auth-other", "auth-team"]), + false ) .await? .is_empty() @@ -622,7 +676,8 @@ mod tests { auth_read( &mut *tx, AuthRead::AgentIdentity, - &strings(&["microsoft_entra", "other-tenant", "auth-client"]) + &strings(&["microsoft_entra", "other-tenant", "auth-client"]), + false ) .await? .is_empty() @@ -631,16 +686,27 @@ mod tests { auth_read( &mut *tx, AuthRead::UserEmail, - &strings(&["test%@example.com"]) + &strings(&["test%@example.com"]), + false ) .await? .is_empty() ); assert!( - auth_read(&mut *tx, AuthRead::User, &strings(&["' OR 1=1 --"])) + auth_read(&mut *tx, AuthRead::User, &strings(&["' OR 1=1 --"]), false) .await? .is_empty() ); + assert!( + auth_read( + &mut *tx, + AuthRead::AgentAccessGroup, + &strings(&["other-group"]), + false + ) + .await? + .is_empty() + ); Ok(()) } @@ -665,13 +731,16 @@ mod tests { ) -> Result<(), Box> { let pool = database_pool?; let mut tx = seeded(&pool).await?; - assert_eq!(auth_spend(&mut *tx, kind, &id).await?, Some(expected)); + assert_eq!( + auth_spend(&mut *tx, kind, &id, false).await?, + Some(expected) + ); let absent = match kind { AuthSpend::Key => "f".repeat(64), AuthSpend::Member => "missing:missing".into(), _ => "missing".into(), }; - assert_eq!(auth_spend(&mut *tx, kind, &absent).await?, None); + assert_eq!(auth_spend(&mut *tx, kind, &absent, false).await?, None); Ok(()) } @@ -688,7 +757,7 @@ mod tests { ) -> Result<(), Box> { let pool = database_pool?; let mut tx = seeded(&pool).await?; - let rows = auth_read(&mut *tx, kind, &ids).await?; + let rows = auth_read(&mut *tx, kind, &ids, false).await?; assert_eq!(rows.len(), 1); assert_eq!(rows[0].data.0["spend"], expected); Ok(()) @@ -711,7 +780,7 @@ mod tests { .unwrap(); pool.close().await; assert!(matches!( - auth_read(&pool, kind, &ids).await, + auth_read(&pool, kind, &ids, false).await, Err(crate::Error::InvalidAuthArguments) )); } @@ -724,11 +793,11 @@ mod tests { .unwrap(); pool.close().await; assert!(matches!( - auth_read(&pool, AuthRead::User, &strings(&["id"])).await, + auth_read(&pool, AuthRead::User, &strings(&["id"]), false).await, Err(crate::Error::Query(sqlx::Error::PoolClosed)) )); assert!(matches!( - auth_spend(&pool, AuthSpend::User, "id").await, + auth_spend(&pool, AuthSpend::User, "id", false).await, Err(crate::Error::Query(sqlx::Error::PoolClosed)) )); } diff --git a/litellm-rust/crates/db/src/error.rs b/litellm-rust/crates/db/src/error.rs index 52024d4ef25..5c5bdee7382 100644 --- a/litellm-rust/crates/db/src/error.rs +++ b/litellm-rust/crates/db/src/error.rs @@ -17,6 +17,25 @@ pub enum Error { } impl Error { + pub fn is_database_failure(&self) -> bool { + match self { + Self::ConnectTimeout(_) => true, + Self::Query(error) => matches!( + error, + sqlx::Error::Database(_) + | sqlx::Error::Io(_) + | sqlx::Error::Tls(_) + | sqlx::Error::PoolTimedOut + | sqlx::Error::PoolClosed + ), + Self::InvalidSetting(_) + | Self::InvalidAuthArguments + | Self::NotConfigured + | Self::InvalidSpendBatch { .. } + | Self::InvalidTimestamp(_) => false, + } + } + pub fn is_retryable(&self) -> bool { match self { Self::Query(sqlx::Error::Database(error)) => { @@ -50,6 +69,24 @@ mod tests { assert_eq!(Error::from(error).is_retryable(), retryable); } + #[rstest] + #[case::closed_pool(sqlx::Error::PoolClosed, true)] + #[case::pool_timeout(sqlx::Error::PoolTimedOut, true)] + #[case::io(sqlx::Error::Io(std::io::ErrorKind::ConnectionReset.into()), true)] + #[case::tls( + sqlx::Error::Tls(Box::new(std::io::Error::from(std::io::ErrorKind::ConnectionReset))), + true + )] + #[case::decode( + sqlx::Error::Decode(Box::new(std::io::Error::from(std::io::ErrorKind::InvalidData))), + false + )] + #[case::missing_column(sqlx::Error::ColumnNotFound("spend".into()), false)] + #[case::missing_row(sqlx::Error::RowNotFound, false)] + fn distinguishes_io_from_native_invariants(#[case] error: sqlx::Error, #[case] io: bool) { + assert_eq!(Error::from(error).is_database_failure(), io); + } + #[rstest] #[case::connection_capacity("53300", true)] #[case::deadlock("40P01", true)] @@ -80,7 +117,9 @@ mod tests { error.as_database_error().and_then(|error| error.code()), Some(code.into()) ); - assert_eq!(Error::from(error).is_retryable(), retryable); + let error = Error::from(error); + assert!(error.is_database_failure()); + assert_eq!(error.is_retryable(), retryable); tx.rollback().await?; Ok(()) } diff --git a/litellm-rust/crates/db/src/executor.rs b/litellm-rust/crates/db/src/executor.rs new file mode 100644 index 00000000000..a656c53e408 --- /dev/null +++ b/litellm-rust/crates/db/src/executor.rs @@ -0,0 +1,84 @@ +use futures_util::{future::BoxFuture, stream::BoxStream}; +use sqlx::{ + Describe, Either, Execute, Executor, Postgres, SqlStr, + postgres::{PgArguments, PgQueryResult, PgRow, PgStatement, PgTypeInfo}, +}; + +#[derive(Debug)] +pub(crate) struct StatementPersistence { + pub executor: E, + pub persistent: bool, +} + +struct QueryPersistence { + query: Q, + persistent: bool, +} + +impl<'q, Q: Execute<'q, Postgres>> Execute<'q, Postgres> for QueryPersistence { + fn sql(self) -> SqlStr { + self.query.sql() + } + + fn statement(&self) -> Option<&PgStatement> { + self.query.statement() + } + + fn take_arguments(&mut self) -> Result, sqlx::error::BoxDynError> { + self.query.take_arguments() + } + + fn persistent(&self) -> bool { + self.persistent + } +} + +impl<'c, E: Executor<'c, Database = Postgres>> Executor<'c> for StatementPersistence { + type Database = Postgres; + + fn fetch_many<'e, 'q: 'e, Q>( + self, + query: Q, + ) -> BoxStream<'e, Result, sqlx::Error>> + where + 'c: 'e, + Q: 'q + Execute<'q, Postgres>, + { + self.executor.fetch_many(QueryPersistence { + query, + persistent: self.persistent, + }) + } + + fn fetch_optional<'e, 'q: 'e, Q>( + self, + query: Q, + ) -> BoxFuture<'e, Result, sqlx::Error>> + where + 'c: 'e, + Q: 'q + Execute<'q, Postgres>, + { + self.executor.fetch_optional(QueryPersistence { + query, + persistent: self.persistent, + }) + } + + fn prepare_with<'e>( + self, + sql: SqlStr, + parameters: &'e [PgTypeInfo], + ) -> BoxFuture<'e, Result> + where + 'c: 'e, + { + self.executor.prepare_with(sql, parameters) + } + + fn describe<'e>(self, sql: SqlStr) -> BoxFuture<'e, Result, sqlx::Error>> + where + 'c: 'e, + { + self.executor.describe(sql) + } +} diff --git a/litellm-rust/crates/db/src/lib.rs b/litellm-rust/crates/db/src/lib.rs index 471918ed8c7..3264e3a8592 100644 --- a/litellm-rust/crates/db/src/lib.rs +++ b/litellm-rust/crates/db/src/lib.rs @@ -10,6 +10,7 @@ mod auth; mod daily_spend; mod end_user; mod error; +mod executor; mod failure; mod pool; mod spend; diff --git a/litellm-rust/crates/db/src/pool.rs b/litellm-rust/crates/db/src/pool.rs index 63d6df7ca19..74a48e2cd5d 100644 --- a/litellm-rust/crates/db/src/pool.rs +++ b/litellm-rust/crates/db/src/pool.rs @@ -74,6 +74,8 @@ pub struct DatabasePools { reader_available: bool, writer_unbounded: bool, reader_unbounded: bool, + writer_persistent: bool, + reader_persistent: bool, } async fn acquire( @@ -129,6 +131,10 @@ impl DatabasePools { reader_unbounded: reader .as_ref() .is_some_and(|settings| settings.acquire_timeout.is_zero()), + writer_persistent: writer.statement_cache_capacity > 0, + reader_persistent: reader + .as_ref() + .is_some_and(|settings| settings.statement_cache_capacity > 0), }) } @@ -136,6 +142,14 @@ impl DatabasePools { self.reader_available } + pub fn persistent(&self, writer: bool) -> bool { + if !writer && self.reader_available { + self.reader_persistent + } else { + self.writer_persistent + } + } + pub fn pool(&self, writer: bool) -> &PgPool { match (&self.reader, writer, self.reader_available) { (Some(pool), false, true) => pool, @@ -191,3 +205,291 @@ impl DatabasePools { } } } + +#[cfg(test)] +mod tests { + use std::{error::Error as StdError, time::Duration}; + + use litellm_secret_value::SecretValue; + use rstest::{fixture, rstest}; + use sqlx::{ConnectOptions, PgPool, postgres::PgPoolOptions}; + + use super::{DatabasePools, DatabaseSettings, acquire}; + use crate::test_support::database_pool; + + #[fixture] + async fn settings() -> Result> { + let pool = database_pool().await?; + let options = pool.connect_options(); + let settings = DatabaseSettings { + host: options.get_socket().map_or_else( + || options.get_host().to_owned(), + |socket| socket.to_string_lossy().into_owned(), + ), + port: options.get_port(), + username: options.get_username().into(), + password: options.to_url_lossy().password().map(SecretValue::from), + database: options + .get_database() + .ok_or("database must be explicit")? + .into(), + schema: "public".into(), + ssl_mode: options.get_ssl_mode(), + ssl_root_cert: None, + max_connections: 1, + acquire_timeout: Duration::from_secs(2), + connect_timeout: Duration::from_secs(2), + idle_timeout: None, + statement_cache_capacity: 0, + }; + pool.close().await; + Ok(settings) + } + + #[rstest] + #[case::writer_only(false, false)] + #[case::reader_healthy(true, false)] + #[case::reader_unavailable(true, true)] + #[tokio::test] + #[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"] + async fn reader_selection_and_shutdown( + #[case] configured: bool, + #[case] unavailable: bool, + #[future(awt)] settings: Result>, + ) -> Result<(), Box> { + let settings = DatabaseSettings { + statement_cache_capacity: 1, + ..settings? + }; + let reader = configured.then(|| DatabaseSettings { + database: if unavailable { + "native_auth_missing_database".into() + } else { + settings.database.clone() + }, + connect_timeout: Duration::from_millis(100), + statement_cache_capacity: 0, + ..settings.clone() + }); + let pools = DatabasePools::connect(settings, reader).await?; + assert_eq!(pools.reader_available(), configured && !unavailable); + assert!(pools.persistent(true)); + assert_eq!(pools.persistent(false), !pools.reader_available()); + let mut writer = pools.acquire(true).await?; + #[expect( + clippy::disallowed_methods, + reason = "Test verifies connection provenance rather than a production query" + )] + let (writer_pid, app, schema): (i32, String, String) = sqlx::query_as( + "SELECT pg_backend_pid(), current_setting('application_name'), current_schema()", + ) + .fetch_one(&mut *writer) + .await?; + assert_eq!(app, "native-auth"); + assert_eq!(schema, "public"); + drop(writer); + let mut reader = pools.acquire(false).await?; + #[expect( + clippy::disallowed_methods, + reason = "Test verifies pool selection by backend identity" + )] + let reader_pid: i32 = sqlx::query_scalar("SELECT pg_backend_pid()") + .fetch_one(&mut *reader) + .await?; + assert_eq!(writer_pid == reader_pid, !configured || unavailable); + drop(reader); + let clone = pools.clone(); + pools.close().await; + clone.close().await; + assert!(matches!( + clone.acquire(false).await, + Err(crate::Error::Query(sqlx::Error::PoolClosed)) + )); + assert!( + clone + .set_passwords(Some(SecretValue::from("next")), None) + .is_err() + ); + Ok(()) + } + + #[rstest] + #[case::uncached(0, 0)] + #[case::cached(4, 2)] + #[tokio::test] + #[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"] + async fn respects_statement_persistence( + #[case] capacity: usize, + #[case] expected: i64, + #[future(awt)] settings: Result>, + ) -> Result<(), Box> { + let settings = DatabaseSettings { + statement_cache_capacity: capacity, + ..settings? + }; + let pools = DatabasePools::connect(settings, None).await?; + let mut connection = pools.acquire(false).await?; + crate::auth_read( + &mut *connection, + crate::AuthRead::User, + &[], + pools.persistent(false), + ) + .await?; + crate::auth_spend( + &mut *connection, + crate::AuthSpend::User, + "missing", + pools.persistent(false), + ) + .await?; + #[expect( + clippy::disallowed_methods, + reason = "Test observes server-side prepared statements on the selected connection" + )] + let prepared: i64 = sqlx::query_scalar("SELECT count(*) FROM pg_prepared_statements") + .persistent(false) + .fetch_one(&mut *connection) + .await?; + assert_eq!(prepared, expected); + drop(connection); + pools.close().await; + Ok(()) + } + + #[rstest] + #[case::reader_survives(true)] + #[case::both_unavailable(false)] + #[tokio::test] + #[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"] + async fn unavailable_writer( + #[case] healthy_reader: bool, + #[future(awt)] settings: Result>, + ) -> Result<(), Box> { + let settings = settings?; + let writer = DatabaseSettings { + database: "native_auth_missing_database".into(), + connect_timeout: Duration::from_millis(100), + acquire_timeout: Duration::from_millis(100), + ..settings.clone() + }; + let result = DatabasePools::connect(writer, healthy_reader.then_some(settings)).await; + if !healthy_reader { + assert!(result.is_err()); + return Ok(()); + } + let pools = result?; + assert!(pools.reader_available()); + assert!(pools.acquire(false).await.is_ok()); + assert!(pools.acquire(true).await.is_err()); + pools.close().await; + Ok(()) + } + + #[rstest] + #[case::bounded(false)] + #[case::unbounded(true)] + #[tokio::test] + #[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"] + async fn acquisition_deadline_and_cancellation( + #[case] unbounded: bool, + #[future(awt)] database_pool: Result>, + ) -> Result<(), Box> { + let source = database_pool?; + let pool = PgPoolOptions::new() + .max_connections(1) + .acquire_timeout(Duration::from_millis(20)) + .connect_with((*source.connect_options()).clone()) + .await?; + let held = pool.acquire().await?; + let result = + tokio::time::timeout(Duration::from_millis(150), acquire(&pool, unbounded)).await; + if unbounded { + assert!(result.is_err()); + } else { + assert!(matches!(result, Ok(Err(sqlx::Error::PoolTimedOut)))); + } + drop(held); + let recovered = + tokio::time::timeout(Duration::from_secs(2), acquire(&pool, unbounded)).await??; + drop(recovered); + pool.close().await; + assert!(matches!( + acquire(&pool, unbounded).await, + Err(sqlx::Error::PoolClosed) + )); + Ok(()) + } + + #[rstest] + #[tokio::test] + #[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"] + async fn query_error_does_not_replay_on_writer( + #[future(awt)] settings: Result>, + ) -> Result<(), Box> { + let settings = settings?; + let reader = DatabaseSettings { + schema: "native_auth_missing_schema".into(), + ..settings.clone() + }; + let pools = DatabasePools::connect(settings, Some(reader)).await?; + let mut read = pools.acquire(false).await?; + assert!( + crate::auth_read( + &mut *read, + crate::AuthRead::User, + &[], + pools.persistent(false) + ) + .await + .is_err() + ); + drop(read); + let mut write = pools.acquire(true).await?; + assert!( + crate::auth_read( + &mut *write, + crate::AuthRead::User, + &[], + pools.persistent(true) + ) + .await? + .is_empty() + ); + drop(write); + pools.close().await; + Ok(()) + } + + #[rstest] + #[tokio::test] + #[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"] + async fn zero_timeouts_and_secret_redaction( + #[future(awt)] settings: Result>, + ) -> Result<(), Box> { + let settings = DatabaseSettings { + acquire_timeout: Duration::ZERO, + connect_timeout: Duration::ZERO, + ..settings? + }; + let debug = format!( + "{:?}", + DatabaseSettings { + password: Some(SecretValue::from("sensitive-test-password")), + ..settings.clone() + } + ); + assert!(!debug.contains("sensitive-test-password")); + assert!(debug.contains("[REDACTED]")); + let pools = DatabasePools::connect(settings, None).await?; + assert!(!pools.reader_available()); + assert!(pools.acquire(false).await.is_ok()); + assert!( + pools + .set_passwords(None, Some(SecretValue::from("reader"))) + .is_err() + ); + pools.close().await; + Ok(()) + } +} diff --git a/litellm-rust/crates/python-bridge/src/db.rs b/litellm-rust/crates/python-bridge/src/db.rs index 77534b641fc..4f9ed6eec16 100644 --- a/litellm-rust/crates/python-bridge/src/db.rs +++ b/litellm-rust/crates/python-bridge/src/db.rs @@ -4,7 +4,7 @@ use litellm_db::{AuthRead, AuthSpend, DatabasePools, DatabaseSettings}; use litellm_host_python::{Pythonized, ready_future, run_async_value}; use litellm_secret_value::SecretValue; use pyo3::{ - exceptions::{PyRuntimeError, PyValueError}, + exceptions::{PyOSError, PyRuntimeError, PyValueError}, prelude::*, types::PyNone, }; @@ -15,6 +15,9 @@ pub(crate) struct NativeDatabase { } fn db_error(error: litellm_db::Error) -> PyErr { + if error.is_database_failure() { + return PyOSError::new_err(litellm_error::Report::new(&error).to_string()); + } match error { litellm_db::Error::InvalidSetting(_) | litellm_db::Error::InvalidAuthArguments => { PyValueError::new_err(litellm_error::Report::new(&error).to_string()) @@ -138,9 +141,10 @@ impl NativeDatabase { .ok_or_else(|| db_error(litellm_db::Error::NotConfigured))?; run_async_value(py, async move { let mut connection = pools.acquire(writer).await.map_err(db_error)?; - let rows = litellm_db::auth_read(&mut *connection, kind, &ids) - .await - .map_err(db_error)?; + let rows = + litellm_db::auth_read(&mut *connection, kind, &ids, pools.persistent(writer)) + .await + .map_err(db_error)?; Ok(Pythonized( rows.into_iter().map(|row| row.data).collect::>(), )) @@ -162,7 +166,7 @@ impl NativeDatabase { .ok_or_else(|| db_error(litellm_db::Error::NotConfigured))?; run_async_value(py, async move { let mut connection = pools.acquire(false).await.map_err(db_error)?; - litellm_db::auth_spend(&mut *connection, kind, &id) + litellm_db::auth_spend(&mut *connection, kind, &id, pools.persistent(false)) .await .map_err(db_error) })