scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-13 08:52:02 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 174

feat(db): complete auth reads and validate pool lifecycle

Browse source
This commit is contained in:
Devin AI 2026-11-09 02:28:49 +00:00
parent 977237ada1
commit 087ffb6218
15 changed files with 627 additions and 31 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
litellm-rust/Cargo.lock generated
View file

@ -3768,6 +3768,7 @@ name = "litellm-db"
version = "0.1.0"
dependencies = [
"chrono",
"futures-util",
"litellm-secret-value",
"macro_rules_attribute",
"rstest",

23
litellm-rust/crates/db/.sqlx/query-26f616a74c671fb1fddd79b167f4f571c6fae8cd6d09a7f67fd3f41025e260be.json generated Normal file
View file

@ -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"
}

23
litellm-rust/crates/db/.sqlx/query-a22cd61e78998c35ddbe1e514911a0cb3af72f78c4c6f90890131faa1a15d3db.json generated Normal file
View file

@ -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"
}

23
litellm-rust/crates/db/.sqlx/query-b7d3bcef6d8d9fdd8eb3a51c1b6c97afbec425ba252db103d956f15c43c23c28.json generated Normal file
View file

@ -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"
}

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

@ -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

4
litellm-rust/crates/db/queries/auth/agent_access_group.sql Normal file
View file

@ -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

16
litellm-rust/crates/db/queries/auth/agent_name.sql Normal file
View file

@ -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

6
litellm-rust/crates/db/queries/auth/fixture.sql
View file

@ -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)

4
litellm-rust/crates/db/queries/auth/managed_vector_store.sql Normal file
View file

@ -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

115
litellm-rust/crates/db/src/auth.rs
View file

@ -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<Vec<AuthRow>, 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<Option<f64>, 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<dyn StdError>> {
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<dyn StdError>> {
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<dyn StdError>> {
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<dyn StdError>> {
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<dyn StdError>> {
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))
));
}

41
litellm-rust/crates/db/src/error.rs
View file

@ -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(())
}

84
litellm-rust/crates/db/src/executor.rs Normal file
View file

@ -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<E> {
pub executor: E,
pub persistent: bool,
}
struct QueryPersistence<Q> {
query: Q,
persistent: bool,
}
impl<'q, Q: Execute<'q, Postgres>> Execute<'q, Postgres> for QueryPersistence<Q> {
fn sql(self) -> SqlStr {
self.query.sql()
}
fn statement(&self) -> Option<&PgStatement> {
self.query.statement()
}
fn take_arguments(&mut self) -> Result<Option<PgArguments>, sqlx::error::BoxDynError> {
self.query.take_arguments()
}
fn persistent(&self) -> bool {
self.persistent
}
}
impl<'c, E: Executor<'c, Database = Postgres>> Executor<'c> for StatementPersistence<E> {
type Database = Postgres;
fn fetch_many<'e, 'q: 'e, Q>(
self,
query: Q,
) -> BoxStream<'e, Result<Either<PgQueryResult, PgRow>, 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<Option<PgRow>, 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<PgStatement, sqlx::Error>>
where
'c: 'e,
{
self.executor.prepare_with(sql, parameters)
}
fn describe<'e>(self, sql: SqlStr) -> BoxFuture<'e, Result<Describe<Postgres>, sqlx::Error>>
where
'c: 'e,
{
self.executor.describe(sql)
}
}

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

@ -10,6 +10,7 @@ mod auth;
mod daily_spend;
mod end_user;
mod error;
mod executor;
mod failure;
mod pool;
mod spend;

302
litellm-rust/crates/db/src/pool.rs
View file

@ -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<DatabaseSettings, Box<dyn StdError>> {
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<DatabaseSettings, Box<dyn StdError>>,
) -> Result<(), Box<dyn StdError>> {
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<DatabaseSettings, Box<dyn StdError>>,
) -> Result<(), Box<dyn StdError>> {
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<DatabaseSettings, Box<dyn StdError>>,
) -> Result<(), Box<dyn StdError>> {
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<PgPool, Box<dyn StdError>>,
) -> Result<(), Box<dyn StdError>> {
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<DatabaseSettings, Box<dyn StdError>>,
) -> Result<(), Box<dyn StdError>> {
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<DatabaseSettings, Box<dyn StdError>>,
) -> Result<(), Box<dyn StdError>> {
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(())
}
}

14
litellm-rust/crates/python-bridge/src/db.rs
View file

@ -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::<Vec<_>>(),
))
@ -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)
})