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 175

feat(auth): move auth policy and effect reads into Rust
Some checks failed
Cost map pins recheck / rerun-cost-map-guard (push) Has been cancelled
Details
LiteLLM Rust DB / db-queries (push) Has been cancelled
Details
LiteLLM Rust / rust-lint (push) Has been cancelled
Details
LiteLLM Rust / rust-test (push) Has been cancelled
Details
LiteLLM Rust / rust-wheel (push) Has been cancelled
Details
Terraform Provider / gofmt, vet, build, test (push) Has been cancelled
Details
Terraform Provider / Provider endpoints vs proxy OpenAPI schema (push) Has been cancelled
Details

Browse source
This commit is contained in:
Devin AI 2026-11-09 03:06:20 +00:00
parent eae60c0bc9
commit 012c8b95f9
406 changed files with 50782 additions and 7667 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
.github/workflows/test-unit.yml vendored
View file

@ -475,6 +475,7 @@ jobs:
- shard: rust-bridge-harness
test-path: >-
tests/unit/_auth
tests/test_litellm_rust/test_dispatch.py
tests/test_litellm_rust/test_public_entries.py
tests/test_litellm_rust/chat_completions/test_dispatch.py

20
litellm-rust/Cargo.lock generated
View file

@ -3768,6 +3768,8 @@ name = "litellm-db"
version = "0.1.0"
dependencies = [
"chrono",
"futures-util",
"litellm-secret-value",
"macro_rules_attribute",
"rstest",
"serde",
@ -3776,6 +3778,7 @@ dependencies = [
"strum",
"thiserror 2.0.19",
"tokio",
"veil",
]
[[package]]
@ -4389,6 +4392,7 @@ dependencies = [
"litellm-callbacks-legacy-python",
"litellm-circuit-breaker",
"litellm-core-utils",
"litellm-db",
"litellm-error",
"litellm-host",
"litellm-host-python",
@ -4408,6 +4412,7 @@ dependencies = [
"litellm-redis-resolve",
"litellm-redis-store",
"litellm-redis-types",
"litellm-request-auth-core",
"litellm-router-core",
"litellm-router-types",
"litellm-secret-value",
@ -4534,6 +4539,7 @@ dependencies = [
"litellm-redis-testing",
"litellm-secret-value",
"litellm-tracing",
"moka",
"rstest",
"strum",
"thiserror 2.0.19",
@ -4562,6 +4568,20 @@ dependencies = [
"url",
]
[[package]]
name = "litellm-request-auth-core"
version = "0.1.0"
dependencies = [
"fancy-regex 0.19.2",
"litellm-secret-value",
"rstest",
"serde",
"sha2 0.10.9",
"strum",
"thiserror 2.0.19",
"veil",
]
[[package]]
name = "litellm-router"
version = "0.1.0"

1
litellm-rust/Cargo.toml
View file

@ -9,6 +9,7 @@ license = "MIT"
repository = "https://github.com/BerriAI/litellm"
[workspace.dependencies]
litellm-request-auth-core = { path = "crates/request-auth-core" }
litellm-config = { path = "crates/config" }
litellm-error = { path = "crates/error" }
litellm-router = { path = "crates/router" }

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

@ -0,0 +1,23 @@
{
"db_name": "PostgreSQL",
"query": "SELECT COALESCE(spend, 0.0)::float8 AS \"spend!\" FROM \"LiteLLM_TeamTable\" WHERE team_id=$1\n",
"describe": {
"columns": [
{
"ordinal": 0,
"name": "spend!",
"type_info": "Float8",
"origin": "Expression"
}
],
"parameters": {
"Left": [
"Text"
]
},
"nullable": [
null
]
},
"hash": "00ddffb253f55705f0b742827f9de6d938c46a3108c2385c52d80e86c4473b96"
}

24
litellm-rust/crates/db/.sqlx/query-0136e6b22d40ef0181f3f122c5abc5f270be24cd01b8c39bceea2de70d44608b.json generated Normal file
View file

@ -0,0 +1,24 @@
{
"db_name": "PostgreSQL",
"query": "SELECT jsonb_build_object('spend', COALESCE(SUM(spend), 0.0)) AS \"data!: AuthData\" FROM \"LiteLLM_SpendLogs\" WHERE \"team_id\"=$1 AND \"startTime\" >= $2\n",
"describe": {
"columns": [
{
"ordinal": 0,
"name": "data!: AuthData",
"type_info": "Jsonb",
"origin": "Expression"
}
],
"parameters": {
"Left": [
"Text",
"Timestamp"
]
},
"nullable": [
null
]
},
"hash": "0136e6b22d40ef0181f3f122c5abc5f270be24cd01b8c39bceea2de70d44608b"
}

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

@ -0,0 +1,23 @@
{
"db_name": "PostgreSQL",
"query": "SELECT to_jsonb(a) AS \"data!: AuthData\"\nFROM \"LiteLLM_RetiredAgent\" a\nWHERE a.original_agent_id = ANY($1::text[])\nORDER BY a.original_agent_id\n",
"describe": {
"columns": [
{
"ordinal": 0,
"name": "data!: AuthData",
"type_info": "Jsonb",
"origin": "Expression"
}
],
"parameters": {
"Left": [
"TextArray"
]
},
"nullable": [
null
]
},
"hash": "16fec3e3896126e8bd14726b346a2a45bd52705356fbf94128776db52032fd05"
}

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

@ -0,0 +1,23 @@
{
"db_name": "PostgreSQL",
"query": "SELECT COALESCE(spend, 0.0)::float8 AS \"spend!\" FROM \"LiteLLM_OrganizationTable\" WHERE organization_id=$1\n",
"describe": {
"columns": [
{
"ordinal": 0,
"name": "spend!",
"type_info": "Float8",
"origin": "Expression"
}
],
"parameters": {
"Left": [
"Text"
]
},
"nullable": [
null
]
},
"hash": "1d689acb624e95e13cdadde37817398561f205596a12e7f6eb71dd5ecbf7d49a"
}

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

@ -0,0 +1,23 @@
{
"db_name": "PostgreSQL",
"query": "SELECT COALESCE(spend, 0.0)::float8 AS \"spend!\" FROM \"LiteLLM_ModelAccessGroupBudgetTable\" WHERE access_group_name=$1\n",
"describe": {
"columns": [
{
"ordinal": 0,
"name": "spend!",
"type_info": "Float8",
"origin": "Expression"
}
],
"parameters": {
"Left": [
"Text"
]
},
"nullable": [
null
]
},
"hash": "1f9b9a6524e417b18865fda7b49c78c7c6a44ada9e601f6c74af1e9debd06f21"
}

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

@ -0,0 +1,23 @@
{
"db_name": "PostgreSQL",
"query": "SELECT to_jsonb(g) || jsonb_build_object('litellm_budget_table', (SELECT to_jsonb(b) || jsonb_build_object('allowed_models', COALESCE(to_jsonb(b.allowed_models), '[]'::jsonb)) FROM \"LiteLLM_BudgetTable\" b WHERE b.budget_id = g.budget_id)) AS \"data!: AuthData\"\nFROM \"LiteLLM_ModelAccessGroupBudgetTable\" g\nWHERE g.access_group_name = ANY($1::text[])\nORDER BY g.access_group_name\n",
"describe": {
"columns": [
{
"ordinal": 0,
"name": "data!: AuthData",
"type_info": "Jsonb",
"origin": "Expression"
}
],
"parameters": {
"Left": [
"TextArray"
]
},
"nullable": [
null
]
},
"hash": "21e8660e197f5f9132584f22da32c8ef734194248e2a078af73776398a256090"
}

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

@ -0,0 +1,23 @@
{
"db_name": "PostgreSQL",
"query": "SELECT to_jsonb(u) || jsonb_build_object('teams', COALESCE(to_jsonb(u.teams), '[]'::jsonb),\n 'models', COALESCE(to_jsonb(u.models), '[]'::jsonb),\n 'allowed_cache_controls', COALESCE(to_jsonb(u.allowed_cache_controls), '[]'::jsonb),\n 'policies', COALESCE(to_jsonb(u.policies), '[]'::jsonb), 'organization_memberships', COALESCE((SELECT jsonb_agg(to_jsonb(om)) FROM \"LiteLLM_OrganizationMembership\" om WHERE om.user_id = u.user_id), '[]'::jsonb)) AS \"data!: AuthData\"\nFROM \"LiteLLM_UserTable\" u\nWHERE u.user_id = ANY($1::text[])\nORDER BY u.user_id\n",
"describe": {
"columns": [
{
"ordinal": 0,
"name": "data!: AuthData",
"type_info": "Jsonb",
"origin": "Expression"
}
],
"parameters": {
"Left": [
"TextArray"
]
},
"nullable": [
null
]
},
"hash": "26a5fa1687c580afc0e1c1f7ca71b53e073a2b8f1e5b9007793b1c89df1ae24f"
}

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

24
litellm-rust/crates/db/.sqlx/query-2c4e1f73d4170dd3c6a2b077fbd837fa0adb227764bcaff1c2267f9a4624f99c.json generated Normal file
View file

@ -0,0 +1,24 @@
{
"db_name": "PostgreSQL",
"query": "SELECT COALESCE(spend, 0.0)::float8 AS \"spend!\" FROM \"LiteLLM_TeamMembership\" WHERE user_id=$1 AND team_id=$2\n",
"describe": {
"columns": [
{
"ordinal": 0,
"name": "spend!",
"type_info": "Float8",
"origin": "Expression"
}
],
"parameters": {
"Left": [
"Text",
"Text"
]
},
"nullable": [
null
]
},
"hash": "2c4e1f73d4170dd3c6a2b077fbd837fa0adb227764bcaff1c2267f9a4624f99c"
}

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

@ -0,0 +1,23 @@
{
"db_name": "PostgreSQL",
"query": "SELECT to_jsonb(o) || jsonb_build_object('models', COALESCE(to_jsonb(o.models), '[]'::jsonb), 'litellm_budget_table', (SELECT to_jsonb(b) || jsonb_build_object('allowed_models', COALESCE(to_jsonb(b.allowed_models), '[]'::jsonb)) FROM \"LiteLLM_BudgetTable\" b WHERE b.budget_id = o.budget_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 = o.object_permission_id)) AS \"data!: AuthData\" FROM \"LiteLLM_OrganizationTable\" o WHERE o.organization_alias = $1 ORDER BY o.organization_id\n",
"describe": {
"columns": [
{
"ordinal": 0,
"name": "data!: AuthData",
"type_info": "Jsonb",
"origin": "Expression"
}
],
"parameters": {
"Left": [
"Text"
]
},
"nullable": [
null
]
},
"hash": "31061f5f7b219176565eb4b3ebe723f7346cf3192e7e5eb6d54f0e4b444b20e7"
}

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

@ -0,0 +1,23 @@
{
"db_name": "PostgreSQL",
"query": "SELECT to_jsonb(o) || jsonb_build_object('models', COALESCE(to_jsonb(o.models), '[]'::jsonb), 'litellm_budget_table', (SELECT to_jsonb(b) || jsonb_build_object('allowed_models', COALESCE(to_jsonb(b.allowed_models), '[]'::jsonb)) FROM \"LiteLLM_BudgetTable\" b WHERE b.budget_id = o.budget_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 = o.object_permission_id)) AS \"data!: AuthData\"\nFROM \"LiteLLM_OrganizationTable\" o\nWHERE o.organization_id = ANY($1::text[])\nORDER BY o.organization_id\n",
"describe": {
"columns": [
{
"ordinal": 0,
"name": "data!: AuthData",
"type_info": "Jsonb",
"origin": "Expression"
}
],
"parameters": {
"Left": [
"TextArray"
]
},
"nullable": [
null
]
},
"hash": "3210aa55f379884ca4232cf2a0b20aac76ac927183a15738d776ff9e104375af"
}

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

@ -0,0 +1,23 @@
{
"db_name": "PostgreSQL",
"query": "SELECT to_jsonb(e) || jsonb_build_object('models', COALESCE(to_jsonb(e.models), '[]'::jsonb), 'litellm_budget_table', (SELECT to_jsonb(b) || jsonb_build_object('allowed_models', COALESCE(to_jsonb(b.allowed_models), '[]'::jsonb)) FROM \"LiteLLM_BudgetTable\" b WHERE b.budget_id = e.budget_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 = e.object_permission_id)) AS \"data!: AuthData\"\nFROM \"LiteLLM_EndUserTable\" e\nWHERE e.user_id = ANY($1::text[])\nORDER BY e.user_id\n",
"describe": {
"columns": [
{
"ordinal": 0,
"name": "data!: AuthData",
"type_info": "Jsonb",
"origin": "Expression"
}
],
"parameters": {
"Left": [
"TextArray"
]
},
"nullable": [
null
]
},
"hash": "34107587736bee7d817b940717c26d0ffba418b154ae4e511287256db5168177"
}

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

@ -0,0 +1,23 @@
{
"db_name": "PostgreSQL",
"query": "SELECT to_jsonb(u) || jsonb_build_object('teams', COALESCE(to_jsonb(u.teams), '[]'::jsonb),\n 'models', COALESCE(to_jsonb(u.models), '[]'::jsonb),\n 'allowed_cache_controls', COALESCE(to_jsonb(u.allowed_cache_controls), '[]'::jsonb),\n 'policies', COALESCE(to_jsonb(u.policies), '[]'::jsonb), 'organization_memberships', COALESCE((SELECT jsonb_agg(to_jsonb(om)) FROM \"LiteLLM_OrganizationMembership\" om WHERE om.user_id = u.user_id), '[]'::jsonb)) AS \"data!: AuthData\" FROM \"LiteLLM_UserTable\" u WHERE u.sso_user_id = $1 ORDER BY u.user_id\n",
"describe": {
"columns": [
{
"ordinal": 0,
"name": "data!: AuthData",
"type_info": "Jsonb",
"origin": "Expression"
}
],
"parameters": {
"Left": [
"Text"
]
},
"nullable": [
null
]
},
"hash": "36c30969226b767cf7d4eec854ccf1ff4c7fb7a158a0af6694c4cc3da0d66dd1"
}

25
litellm-rust/crates/db/.sqlx/query-3b2157705b66ebcf6ee0c9805f15926a45b69c60677f527f227a2f8a58020b5b.json generated Normal file
View file

@ -0,0 +1,25 @@
{
"db_name": "PostgreSQL",
"query": "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)) AS \"data!: AuthData\" FROM \"LiteLLM_AgentIdentity\" i WHERE i.provider=$1 AND i.tenant_id=$2 AND i.client_id=$3\n",
"describe": {
"columns": [
{
"ordinal": 0,
"name": "data!: AuthData",
"type_info": "Jsonb",
"origin": "Expression"
}
],
"parameters": {
"Left": [
"Text",
"Text",
"Text"
]
},
"nullable": [
null
]
},
"hash": "3b2157705b66ebcf6ee0c9805f15926a45b69c60677f527f227a2f8a58020b5b"
}

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

@ -0,0 +1,23 @@
{
"db_name": "PostgreSQL",
"query": "SELECT COALESCE(spend, 0.0)::float8 AS \"spend!\" FROM \"LiteLLM_TagTable\" WHERE tag_name=$1\n",
"describe": {
"columns": [
{
"ordinal": 0,
"name": "spend!",
"type_info": "Float8",
"origin": "Expression"
}
],
"parameters": {
"Left": [
"Text"
]
},
"nullable": [
null
]
},
"hash": "498268b2d6edf244dfbe6d2c33327eae2c4d796d2891f6f3816289f8f71d7ea1"
}

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

@ -0,0 +1,23 @@
{
"db_name": "PostgreSQL",
"query": "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)) AS \"data!: AuthData\"\nFROM \"LiteLLM_ObjectPermissionTable\" op\nWHERE op.object_permission_id = ANY($1::text[])\nORDER BY op.object_permission_id\n",
"describe": {
"columns": [
{
"ordinal": 0,
"name": "data!: AuthData",
"type_info": "Jsonb",
"origin": "Expression"
}
],
"parameters": {
"Left": [
"TextArray"
]
},
"nullable": [
null
]
},
"hash": "4ddd384812d4cc34d601873e73294c7e1bafada2968e53f4ecd16cd49d82cd5b"
}

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

@ -0,0 +1,23 @@
{
"db_name": "PostgreSQL",
"query": "SELECT COALESCE(spend, 0.0)::float8 AS \"spend!\" FROM \"LiteLLM_ProjectTable\" WHERE project_id=$1\n",
"describe": {
"columns": [
{
"ordinal": 0,
"name": "spend!",
"type_info": "Float8",
"origin": "Expression"
}
],
"parameters": {
"Left": [
"Text"
]
},
"nullable": [
null
]
},
"hash": "56a5d90aee114ff352a31fbec84ee9dd8dcb36772073095414df0d4493a3e363"
}

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

@ -0,0 +1,23 @@
{
"db_name": "PostgreSQL",
"query": "SELECT COALESCE(spend, 0.0)::float8 AS \"spend!\" FROM \"LiteLLM_VerificationToken\" WHERE token=$1\n",
"describe": {
"columns": [
{
"ordinal": 0,
"name": "spend!",
"type_info": "Float8",
"origin": "Expression"
}
],
"parameters": {
"Left": [
"Text"
]
},
"nullable": [
null
]
},
"hash": "57d69d9d2ad5803d123717bd788223c5ce05000eaf8cb6888a66033a94fe7a09"
}

24
litellm-rust/crates/db/.sqlx/query-65e9b4dc09101e63092902802978e012ff43436560342357afabbc0d5b961c44.json generated Normal file
View file

@ -0,0 +1,24 @@
{
"db_name": "PostgreSQL",
"query": "SELECT to_jsonb(tm) || jsonb_build_object('litellm_budget_table', (SELECT to_jsonb(b) || jsonb_build_object('allowed_models', COALESCE(to_jsonb(b.allowed_models), '[]'::jsonb)) FROM \"LiteLLM_BudgetTable\" b WHERE b.budget_id = tm.budget_id)) AS \"data!: AuthData\"\nFROM \"LiteLLM_TeamMembership\" tm WHERE tm.user_id=$1 AND tm.team_id=$2\n",
"describe": {
"columns": [
{
"ordinal": 0,
"name": "data!: AuthData",
"type_info": "Jsonb",
"origin": "Expression"
}
],
"parameters": {
"Left": [
"Text",
"Text"
]
},
"nullable": [
null
]
},
"hash": "65e9b4dc09101e63092902802978e012ff43436560342357afabbc0d5b961c44"
}

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

@ -0,0 +1,23 @@
{
"db_name": "PostgreSQL",
"query": "SELECT to_jsonb(p) || jsonb_build_object('models', COALESCE(to_jsonb(p.models), '[]'::jsonb), 'litellm_budget_table', (SELECT to_jsonb(b) || jsonb_build_object('allowed_models', COALESCE(to_jsonb(b.allowed_models), '[]'::jsonb)) FROM \"LiteLLM_BudgetTable\" b WHERE b.budget_id = p.budget_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 = p.object_permission_id)) AS \"data!: AuthData\"\nFROM \"LiteLLM_ProjectTable\" p\nWHERE p.project_id = ANY($1::text[])\nORDER BY p.project_id\n",
"describe": {
"columns": [
{
"ordinal": 0,
"name": "data!: AuthData",
"type_info": "Jsonb",
"origin": "Expression"
}
],
"parameters": {
"Left": [
"TextArray"
]
},
"nullable": [
null
]
},
"hash": "667556bca77708a5c027641bb3be14a6dcec24dbe6c5063ef9ef75cb196d886c"
}

25
litellm-rust/crates/db/.sqlx/query-683dd0c1cdf39792c5b0c7ca482fee4c56e187d9f9e3d26d625fc3432e9d5652.json generated Normal file
View file

@ -0,0 +1,25 @@
{
"db_name": "PostgreSQL",
"query": "SELECT jsonb_build_object('window_start', window_start, 'spend', spend) AS \"data!: AuthData\" FROM \"LiteLLM_BudgetWindowSpend\" WHERE entity_type=$1 AND entity_id=$2 AND window_duration=$3\n",
"describe": {
"columns": [
{
"ordinal": 0,
"name": "data!: AuthData",
"type_info": "Jsonb",
"origin": "Expression"
}
],
"parameters": {
"Left": [
"Text",
"Text",
"Text"
]
},
"nullable": [
null
]
},
"hash": "683dd0c1cdf39792c5b0c7ca482fee4c56e187d9f9e3d26d625fc3432e9d5652"
}

25
litellm-rust/crates/db/.sqlx/query-6af954d0490c3cea037cd7bd00993f983adaeaf5ab4166a90ab4a89ae39e2b0b.json generated Normal file
View file

@ -0,0 +1,25 @@
{
"db_name": "PostgreSQL",
"query": "SELECT to_jsonb(i) AS \"data!: AuthData\" FROM \"LiteLLM_RetiredAgentIdentity\" i WHERE i.provider=$1 AND i.tenant_id=$2 AND i.client_id=$3\n",
"describe": {
"columns": [
{
"ordinal": 0,
"name": "data!: AuthData",
"type_info": "Jsonb",
"origin": "Expression"
}
],
"parameters": {
"Left": [
"Text",
"Text",
"Text"
]
},
"nullable": [
null
]
},
"hash": "6af954d0490c3cea037cd7bd00993f983adaeaf5ab4166a90ab4a89ae39e2b0b"
}

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

@ -0,0 +1,23 @@
{
"db_name": "PostgreSQL",
"query": "SELECT to_jsonb(t) || jsonb_build_object('admins', COALESCE(to_jsonb(t.admins), '[]'::jsonb),\n 'members', COALESCE(to_jsonb(t.members), '[]'::jsonb),\n 'models', COALESCE(to_jsonb(t.models), '[]'::jsonb),\n 'team_member_permissions', COALESCE(to_jsonb(t.team_member_permissions), '[]'::jsonb),\n 'access_group_ids', COALESCE(to_jsonb(t.access_group_ids), '[]'::jsonb),\n 'policies', COALESCE(to_jsonb(t.policies), '[]'::jsonb),\n 'default_team_member_models', COALESCE(to_jsonb(t.default_team_member_models), '[]'::jsonb), 'litellm_model_table', (SELECT (to_jsonb(m) - 'aliases') || jsonb_build_object('model_aliases', m.aliases) FROM \"LiteLLM_ModelTable\" m WHERE m.id=t.model_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 = t.object_permission_id)) AS \"data!: AuthData\"\nFROM \"LiteLLM_TeamTable\" t\nWHERE t.team_id = ANY($1::text[])\nORDER BY t.team_id\n",
"describe": {
"columns": [
{
"ordinal": 0,
"name": "data!: AuthData",
"type_info": "Jsonb",
"origin": "Expression"
}
],
"parameters": {
"Left": [
"TextArray"
]
},
"nullable": [
null
]
},
"hash": "786700d34873fa39e57ca9b0c1bec5df311bb3953f9756a5181821af9339cdba"
}

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

@ -0,0 +1,23 @@
{
"db_name": "PostgreSQL",
"query": "WITH resolved AS (\n SELECT COALESCE(v.token, d.active_token_id) AS token\n FROM UNNEST($1::text[]) AS requested(token)\n LEFT JOIN \"LiteLLM_VerificationToken\" v ON v.token=requested.token\n LEFT JOIN \"LiteLLM_DeprecatedVerificationToken\" d ON d.token=requested.token AND d.revoke_at > (clock_timestamp() AT TIME ZONE 'UTC')\n)\nSELECT to_jsonb(view) AS \"data!: AuthData\" FROM (\nSELECT v.*,\nt.spend AS team_spend,\nt.max_budget AS team_max_budget,\nt.soft_budget AS team_soft_budget,\nt.model_max_budget AS team_model_max_budget,\nt.tpm_limit AS team_tpm_limit,\nt.rpm_limit AS team_rpm_limit,\nt.tpd_limit AS team_tpd_limit,\nt.models AS team_models,\nt.metadata AS team_metadata,\nt.blocked AS team_blocked,\nt.team_alias AS team_alias,\nt.members_with_roles AS team_members_with_roles,\nt.object_permission_id AS team_object_permission_id,\nt.organization_id AS org_id,\np.project_alias AS project_alias,\ntm.spend AS team_member_spend,\nb_tm.tpm_limit AS team_member_tpm_limit,\nb_tm.rpm_limit AS team_member_rpm_limit,\nm.aliases AS team_model_aliases,\nb.max_budget AS litellm_budget_table_max_budget,\nb.tpm_limit AS litellm_budget_table_tpm_limit,\nb.rpm_limit AS litellm_budget_table_rpm_limit,\nb.tpd_limit AS litellm_budget_table_tpd_limit,\nb.model_max_budget AS litellm_budget_table_model_max_budget,\nb.soft_budget AS litellm_budget_table_soft_budget,\no.metadata AS organization_metadata,\no.organization_alias AS organization_alias,\nb2.max_budget AS organization_max_budget,\nb2.tpm_limit AS organization_tpm_limit,\nb2.rpm_limit AS organization_rpm_limit\nFROM \"LiteLLM_VerificationToken\" v\nLEFT JOIN \"LiteLLM_TeamTable\" t ON v.team_id=t.team_id\nLEFT JOIN \"LiteLLM_TeamMembership\" tm ON v.team_id=tm.team_id AND tm.user_id=v.user_id\nLEFT JOIN \"LiteLLM_BudgetTable\" b_tm ON tm.budget_id=b_tm.budget_id\nLEFT JOIN \"LiteLLM_ModelTable\" m ON t.model_id=m.id\nLEFT JOIN \"LiteLLM_BudgetTable\" b ON v.budget_id=b.budget_id\nLEFT JOIN \"LiteLLM_ProjectTable\" p ON v.project_id=p.project_id\nLEFT JOIN \"LiteLLM_OrganizationTable\" o ON v.organization_id=o.organization_id\nLEFT JOIN \"LiteLLM_BudgetTable\" b2 ON o.budget_id=b2.budget_id\nWHERE v.token IN (SELECT token FROM resolved)\nORDER BY v.token\n) view\n",
"describe": {
"columns": [
{
"ordinal": 0,
"name": "data!: AuthData",
"type_info": "Jsonb",
"origin": "Expression"
}
],
"parameters": {
"Left": [
"TextArray"
]
},
"nullable": [
null
]
},
"hash": "847628007ea76e788caf20d35219860f7a2da7d3d6d40c209322e0893b99ac69"
}

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

@ -0,0 +1,23 @@
{
"db_name": "PostgreSQL",
"query": "SELECT jsonb_build_object('tag_name', tag_name) AS \"data!: AuthData\" FROM \"LiteLLM_TagTable\" ORDER BY tag_name LIMIT $1\n",
"describe": {
"columns": [
{
"ordinal": 0,
"name": "data!: AuthData",
"type_info": "Jsonb",
"origin": "Expression"
}
],
"parameters": {
"Left": [
"Int8"
]
},
"nullable": [
null
]
},
"hash": "8637bd849446309e631a4eaa47f7b51a9910fb27160356d9cc5dd03eb16985b2"
}

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

@ -0,0 +1,23 @@
{
"db_name": "PostgreSQL",
"query": "SELECT to_jsonb(a) || jsonb_build_object('access_model_names', COALESCE(to_jsonb(a.access_model_names), '[]'::jsonb),\n 'access_mcp_server_ids', COALESCE(to_jsonb(a.access_mcp_server_ids), '[]'::jsonb),\n 'access_agent_ids', COALESCE(to_jsonb(a.access_agent_ids), '[]'::jsonb),\n 'access_skill_names', COALESCE(to_jsonb(a.access_skill_names), '[]'::jsonb),\n 'assigned_team_ids', COALESCE(to_jsonb(a.assigned_team_ids), '[]'::jsonb),\n 'assigned_key_ids', COALESCE(to_jsonb(a.assigned_key_ids), '[]'::jsonb)) AS \"data!: AuthData\"\nFROM \"LiteLLM_AccessGroupTable\" a\nWHERE a.access_group_id = ANY($1::text[])\nORDER BY a.access_group_id\n",
"describe": {
"columns": [
{
"ordinal": 0,
"name": "data!: AuthData",
"type_info": "Jsonb",
"origin": "Expression"
}
],
"parameters": {
"Left": [
"TextArray"
]
},
"nullable": [
null
]
},
"hash": "8d8ac844d36122a8e652c494519d5b7f5ea9ad7a158c5824ab8fc42e1e67dbdb"
}

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-ab4b895f3a1ee6ab9e86f13b4f02b5719d6fcbc961b5234cd5ad3100674b293b.json generated Normal file
View file

@ -0,0 +1,23 @@
{
"db_name": "PostgreSQL",
"query": "SELECT to_jsonb(t) || jsonb_build_object('models', COALESCE(to_jsonb(t.models), '[]'::jsonb), 'litellm_budget_table', (SELECT to_jsonb(b) || jsonb_build_object('allowed_models', COALESCE(to_jsonb(b.allowed_models), '[]'::jsonb)) FROM \"LiteLLM_BudgetTable\" b WHERE b.budget_id = t.budget_id)) AS \"data!: AuthData\"\nFROM \"LiteLLM_TagTable\" t\nWHERE t.tag_name = ANY($1::text[])\nORDER BY t.tag_name\n",
"describe": {
"columns": [
{
"ordinal": 0,
"name": "data!: AuthData",
"type_info": "Jsonb",
"origin": "Expression"
}
],
"parameters": {
"Left": [
"TextArray"
]
},
"nullable": [
null
]
},
"hash": "ab4b895f3a1ee6ab9e86f13b4f02b5719d6fcbc961b5234cd5ad3100674b293b"
}

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

@ -0,0 +1,23 @@
{
"db_name": "PostgreSQL",
"query": "SELECT to_jsonb(u) || jsonb_build_object('teams', COALESCE(to_jsonb(u.teams), '[]'::jsonb),\n 'models', COALESCE(to_jsonb(u.models), '[]'::jsonb),\n 'allowed_cache_controls', COALESCE(to_jsonb(u.allowed_cache_controls), '[]'::jsonb),\n 'policies', COALESCE(to_jsonb(u.policies), '[]'::jsonb), 'organization_memberships', COALESCE((SELECT jsonb_agg(to_jsonb(om)) FROM \"LiteLLM_OrganizationMembership\" om WHERE om.user_id = u.user_id), '[]'::jsonb)) AS \"data!: AuthData\" FROM \"LiteLLM_UserTable\" u WHERE LOWER(u.user_email) = LOWER($1) ORDER BY u.user_id LIMIT 1\n",
"describe": {
"columns": [
{
"ordinal": 0,
"name": "data!: AuthData",
"type_info": "Jsonb",
"origin": "Expression"
}
],
"parameters": {
"Left": [
"Text"
]
},
"nullable": [
null
]
},
"hash": "b0d27eac417a4044836ce7ba4450b4bd94727d6f882992ff4b8419be945175ef"
}

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

@ -0,0 +1,23 @@
{
"db_name": "PostgreSQL",
"query": "SELECT to_jsonb(b) || jsonb_build_object('allowed_models', COALESCE(to_jsonb(b.allowed_models), '[]'::jsonb)) AS \"data!: AuthData\"\nFROM \"LiteLLM_BudgetTable\" b\nWHERE b.budget_id = ANY($1::text[])\nORDER BY b.budget_id\n",
"describe": {
"columns": [
{
"ordinal": 0,
"name": "data!: AuthData",
"type_info": "Jsonb",
"origin": "Expression"
}
],
"parameters": {
"Left": [
"TextArray"
]
},
"nullable": [
null
]
},
"hash": "b23d9fe187171500ebd8593571e9ba9059a84dbd20d9e148c72f23a62f52a981"
}

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

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

@ -0,0 +1,23 @@
{
"db_name": "PostgreSQL",
"query": "SELECT COALESCE(spend, 0.0)::float8 AS \"spend!\" FROM \"LiteLLM_UserTable\" WHERE user_id=$1\n",
"describe": {
"columns": [
{
"ordinal": 0,
"name": "spend!",
"type_info": "Float8",
"origin": "Expression"
}
],
"parameters": {
"Left": [
"Text"
]
},
"nullable": [
null
]
},
"hash": "b840de2b7f4c4dfc7c2909062a4494d92ea3bf7d09bfeb6f6aa30a7bff7da7bf"
}

24
litellm-rust/crates/db/.sqlx/query-c265904d0f03ab912de3cfd5adae562f532616ee7d81288ffdcc9ba34d4762c0.json generated Normal file
View file

@ -0,0 +1,24 @@
{
"db_name": "PostgreSQL",
"query": "SELECT jsonb_build_object('spend', COALESCE(SUM(spend), 0.0)) AS \"data!: AuthData\" FROM \"LiteLLM_SpendLogs\" WHERE \"api_key\"=$1 AND \"startTime\" >= $2\n",
"describe": {
"columns": [
{
"ordinal": 0,
"name": "data!: AuthData",
"type_info": "Jsonb",
"origin": "Expression"
}
],
"parameters": {
"Left": [
"Text",
"Timestamp"
]
},
"nullable": [
null
]
},
"hash": "c265904d0f03ab912de3cfd5adae562f532616ee7d81288ffdcc9ba34d4762c0"
}

23
litellm-rust/crates/db/.sqlx/query-c720cb3a555cd750cd929003591e8b674266007e4b11e388f722c9418f8ccea0.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_id = 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": "c720cb3a555cd750cd929003591e8b674266007e4b11e388f722c9418f8ccea0"
}

25
litellm-rust/crates/db/.sqlx/query-c879e82114e82685c3613fa4bc82f48ade878aae01105b6d7181f95cd5a16864.json generated Normal file
View file

@ -0,0 +1,25 @@
{
"db_name": "PostgreSQL",
"query": "SELECT to_jsonb(i) AS \"data!: AuthData\" FROM \"LiteLLM_VerifiedSubject\" i WHERE i.issuer=$1 AND i.tenant_id=$2 AND i.oid=$3\n",
"describe": {
"columns": [
{
"ordinal": 0,
"name": "data!: AuthData",
"type_info": "Jsonb",
"origin": "Expression"
}
],
"parameters": {
"Left": [
"Text",
"Text",
"Text"
]
},
"nullable": [
null
]
},
"hash": "c879e82114e82685c3613fa4bc82f48ade878aae01105b6d7181f95cd5a16864"
}

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

@ -0,0 +1,23 @@
{
"db_name": "PostgreSQL",
"query": "SELECT jsonb_build_object('access_group_name', access_group_name) AS \"data!: AuthData\" FROM \"LiteLLM_ModelAccessGroupBudgetTable\" ORDER BY access_group_name LIMIT $1\n",
"describe": {
"columns": [
{
"ordinal": 0,
"name": "data!: AuthData",
"type_info": "Jsonb",
"origin": "Expression"
}
],
"parameters": {
"Left": [
"Int8"
]
},
"nullable": [
null
]
},
"hash": "cc5bac434a523f8deb0416b806412422927cedd9981b6b3617aa61b1086ab8f9"
}

27
litellm-rust/crates/db/.sqlx/query-d804605eae54eaeaa01c7f511cb33c8443365b9262f2827e1c0b09f715b0227b.json generated Normal file
View file

File diff suppressed because one or more lines are too long

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

@ -0,0 +1,23 @@
{
"db_name": "PostgreSQL",
"query": "SELECT COALESCE(spend, 0.0)::float8 AS \"spend!\" FROM \"LiteLLM_EndUserTable\" WHERE user_id=$1\n",
"describe": {
"columns": [
{
"ordinal": 0,
"name": "spend!",
"type_info": "Float8",
"origin": "Expression"
}
],
"parameters": {
"Left": [
"Text"
]
},
"nullable": [
null
]
},
"hash": "dbe6bf2326dcb74823f6bfb90149e5d79069036df1c90481aa17b5ee7c6e4293"
}

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

@ -0,0 +1,23 @@
{
"db_name": "PostgreSQL",
"query": "SELECT to_jsonb(t) || jsonb_build_object('admins', COALESCE(to_jsonb(t.admins), '[]'::jsonb),\n 'members', COALESCE(to_jsonb(t.members), '[]'::jsonb),\n 'models', COALESCE(to_jsonb(t.models), '[]'::jsonb),\n 'team_member_permissions', COALESCE(to_jsonb(t.team_member_permissions), '[]'::jsonb),\n 'access_group_ids', COALESCE(to_jsonb(t.access_group_ids), '[]'::jsonb),\n 'policies', COALESCE(to_jsonb(t.policies), '[]'::jsonb),\n 'default_team_member_models', COALESCE(to_jsonb(t.default_team_member_models), '[]'::jsonb), 'litellm_model_table', (SELECT (to_jsonb(m) - 'aliases') || jsonb_build_object('model_aliases', m.aliases) FROM \"LiteLLM_ModelTable\" m WHERE m.id=t.model_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 = t.object_permission_id)) AS \"data!: AuthData\" FROM \"LiteLLM_TeamTable\" t WHERE t.team_alias = $1 ORDER BY t.team_id\n",
"describe": {
"columns": [
{
"ordinal": 0,
"name": "data!: AuthData",
"type_info": "Jsonb",
"origin": "Expression"
}
],
"parameters": {
"Left": [
"Text"
]
},
"nullable": [
null
]
},
"hash": "e6b5061b738fd353df83eaef9dc02f70cecce44f8e5d6673d2521a57de3513d8"
}

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

@ -0,0 +1,23 @@
{
"db_name": "PostgreSQL",
"query": "SELECT jsonb_build_object('user_id', user_id) AS \"data!: AuthData\" FROM \"LiteLLM_EndUserTable\"\nWHERE blocked OR budget_id IS NOT NULL OR allowed_model_region IS NOT NULL OR default_model IS NOT NULL OR object_permission_id IS NOT NULL OR cardinality(models) > 0\nORDER BY user_id LIMIT $1\n",
"describe": {
"columns": [
{
"ordinal": 0,
"name": "data!: AuthData",
"type_info": "Jsonb",
"origin": "Expression"
}
],
"parameters": {
"Left": [
"Int8"
]
},
"nullable": [
null
]
},
"hash": "f6ed52b7b550c41a8d085f1a66bd1dcafc1f0cf69e33cb58ebe8642b75ee8d36"
}

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

@ -6,7 +6,11 @@ license.workspace = true
repository.workspace = true
[dependencies]
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
@ -16,4 +20,3 @@ thiserror.workspace = true
[dev-dependencies]
rstest.workspace = true
tokio.workspace = true

9
litellm-rust/crates/db/queries/auth/access_group.sql Normal file
View file

@ -0,0 +1,9 @@
SELECT to_jsonb(a) || jsonb_build_object('access_model_names', COALESCE(to_jsonb(a.access_model_names), '[]'::jsonb),
'access_mcp_server_ids', COALESCE(to_jsonb(a.access_mcp_server_ids), '[]'::jsonb),
'access_agent_ids', COALESCE(to_jsonb(a.access_agent_ids), '[]'::jsonb),
'access_skill_names', COALESCE(to_jsonb(a.access_skill_names), '[]'::jsonb),
'assigned_team_ids', COALESCE(to_jsonb(a.assigned_team_ids), '[]'::jsonb),
'assigned_key_ids', COALESCE(to_jsonb(a.assigned_key_ids), '[]'::jsonb)) AS "data!: AuthData"
FROM "LiteLLM_AccessGroupTable" a
WHERE a.access_group_id = ANY($1::text[])
ORDER BY a.access_group_id

16
litellm-rust/crates/db/queries/auth/agent.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_id = ANY($1::text[])
ORDER BY a.agent_id

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

2
litellm-rust/crates/db/queries/auth/agent_identity.sql Normal file
View file

@ -0,0 +1,2 @@
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)) AS "data!: AuthData" FROM "LiteLLM_AgentIdentity" i WHERE i.provider=$1 AND i.tenant_id=$2 AND i.client_id=$3

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

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

@ -0,0 +1,4 @@
SELECT to_jsonb(b) || jsonb_build_object('allowed_models', COALESCE(to_jsonb(b.allowed_models), '[]'::jsonb)) AS "data!: AuthData"
FROM "LiteLLM_BudgetTable" b
WHERE b.budget_id = ANY($1::text[])
ORDER BY b.budget_id

1
litellm-rust/crates/db/queries/auth/budget_window.sql Normal file
View file

@ -0,0 +1 @@
SELECT jsonb_build_object('window_start', window_start, 'spend', spend) AS "data!: AuthData" FROM "LiteLLM_BudgetWindowSpend" WHERE entity_type=$1 AND entity_id=$2 AND window_duration=$3

13
litellm-rust/crates/db/queries/auth/end_user.sql Normal file
View file

@ -0,0 +1,13 @@
SELECT to_jsonb(e) || jsonb_build_object('models', COALESCE(to_jsonb(e.models), '[]'::jsonb), 'litellm_budget_table', (SELECT to_jsonb(b) || jsonb_build_object('allowed_models', COALESCE(to_jsonb(b.allowed_models), '[]'::jsonb)) FROM "LiteLLM_BudgetTable" b WHERE b.budget_id = e.budget_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 = e.object_permission_id)) AS "data!: AuthData"
FROM "LiteLLM_EndUserTable" e
WHERE e.user_id = ANY($1::text[])
ORDER BY e.user_id

3
litellm-rust/crates/db/queries/auth/end_user_restricted_registry.sql Normal file
View file

@ -0,0 +1,3 @@
SELECT jsonb_build_object('user_id', user_id) AS "data!: AuthData" FROM "LiteLLM_EndUserTable"
WHERE blocked OR budget_id IS NOT NULL OR allowed_model_region IS NOT NULL OR default_model IS NOT NULL OR object_permission_id IS NOT NULL OR cardinality(models) > 0
ORDER BY user_id LIMIT $1

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

@ -0,0 +1,58 @@
INSERT INTO "LiteLLM_BudgetTable" (budget_id, max_budget, tpm_limit, rpm_limit, allowed_models, created_by, updated_by)
VALUES ('auth-b', 0, 22, 33, NULL, 'auth-test', 'auth-test'),
('auth-b-key', 4.5, 44, 55, ARRAY['key-model'], 'auth-test', 'auth-test'),
('auth-b-org', 13, 66, 77, ARRAY['org-model'], 'auth-test', 'auth-test');
INSERT INTO "LiteLLM_ObjectPermissionTable" (object_permission_id, models, mcp_servers, agents)
VALUES ('auth-op', ARRAY['allowed-model'], NULL, ARRAY['auth-agent']);
INSERT INTO "LiteLLM_OrganizationTable" (organization_id, organization_alias, budget_id, object_permission_id, models, spend, created_by, updated_by)
VALUES ('auth-org', 'Auth Org', 'auth-b-org', 'auth-op', NULL, 13.5, 'auth-test', 'auth-test');
INSERT INTO "LiteLLM_ModelTable" (id, aliases, created_by, updated_by)
VALUES (-922, '{"alias":"actual"}', 'auth-test', 'auth-test');
INSERT INTO "LiteLLM_TeamTable" (team_id, team_alias, organization_id, model_id, object_permission_id, models, admins, members, spend, max_budget, blocked, members_with_roles)
VALUES ('auth-team', 'Auth Team', 'auth-org', -922, 'auth-op', ARRAY['team-model'], NULL, NULL, 2.5, 25, true, '[{"user_id":"auth-u","role":"admin"}]'),
('auth-team-empty', 'Auth Empty', NULL, NULL, NULL, NULL, NULL, NULL, 0, NULL, false, '[]');
INSERT INTO "LiteLLM_UserTable" (user_id, user_email, sso_user_id, models, teams, spend)
VALUES ('auth-u', 'Test_%@Example.com', 'auth-sso', NULL, ARRAY['auth-team'], 3.5),
('auth-other', 'test-xx@example.com', NULL, ARRAY['other-model'], NULL, 0);
INSERT INTO "LiteLLM_OrganizationMembership" (user_id, organization_id, user_role, budget_id)
VALUES ('auth-u', 'auth-org', 'org_admin', 'auth-b');
INSERT INTO "LiteLLM_TeamMembership" (user_id, team_id, spend, budget_id)
VALUES ('auth-u', 'auth-team', 6.5, 'auth-b');
INSERT INTO "LiteLLM_ProjectTable" (project_id, project_alias, team_id, budget_id, object_permission_id, models, spend, created_by, updated_by)
VALUES ('auth-project', 'Auth Project', 'auth-team', 'auth-b', 'auth-op', NULL, 8.5, 'auth-test', 'auth-test');
INSERT INTO "LiteLLM_VerificationToken" (token, user_id, team_id, organization_id, project_id, budget_id, spend, models)
VALUES (repeat('a', 64), 'auth-u', 'auth-team', 'auth-org', 'auth-project', 'auth-b-key', 4.5, ARRAY['key-model']),
(repeat('e', 64), NULL, NULL, NULL, NULL, NULL, 0, ARRAY[]::text[]);
INSERT INTO "LiteLLM_DeprecatedVerificationToken" (id, token, active_token_id, revoke_at)
VALUES ('auth-deprecated', repeat('b', 64), repeat('a', 64), now() + INTERVAL '1 hour'),
('auth-expired', repeat('c', 64), repeat('a', 64), now() - INTERVAL '1 hour'),
('auth-chained', repeat('d', 64), repeat('b', 64), now() + INTERVAL '1 hour'),
('auth-shadowed', repeat('a', 64), repeat('e', 64), now() + INTERVAL '1 hour');
INSERT INTO "LiteLLM_EndUserTable" (user_id, models, budget_id, object_permission_id, spend, blocked)
VALUES ('auth-end', NULL, 'auth-b', 'auth-op', 10.5, false),
('auth-unrestricted', NULL, NULL, NULL, 0, false);
INSERT INTO "LiteLLM_TagTable" (tag_name, models, budget_id, spend)
VALUES ('auth-tag', NULL, 'auth-b', 11.5);
INSERT INTO "LiteLLM_ModelAccessGroupBudgetTable" (access_group_name, budget_id, spend)
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, 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)
VALUES ('auth-retired-binding', 'microsoft_entra', 'auth-issuer', 'auth-tenant', 'auth-retired-client');
INSERT INTO "LiteLLM_RetiredAgent" (original_agent_id) VALUES ('auth-retired-agent');
INSERT INTO "LiteLLM_VerifiedSubject" (subject_id, issuer, tenant_id, oid, user_id)
VALUES ('auth-subject', 'auth-issuer', 'auth-tenant', 'auth-oid', 'auth-u');
INSERT INTO "LiteLLM_BudgetWindowSpend" (entity_type, entity_id, window_duration, window_start, spend, updated_at)
VALUES ('key', repeat('a',64), '1h', '2025-01-01 01:00:00', 2.5, now()),
('team', 'auth-team', '1h', '2025-01-01 01:00:00', 2.5, now());
INSERT INTO "LiteLLM_SpendLogs" (request_id, call_type, api_key, team_id, "startTime", "endTime", spend)
VALUES ('auth-log-before', 'completion', repeat('a',64), 'auth-team', '2025-01-01 00:59:59', '2025-01-01 01:00:01', 100),
('auth-log-boundary', 'completion', repeat('a',64), 'auth-team', '2025-01-01 01:00:00', '2025-01-01 01:00:01', 2.5),
('auth-log-after', 'completion', repeat('a',64), 'auth-team', '2025-01-01 01:00:01', '2025-01-01 01:00:02', 3.5),
('auth-log-other', 'completion', repeat('e',64), 'auth-team-empty', '2025-01-01 01:00:00', '2025-01-01 01:00:01', 500);

50
litellm-rust/crates/db/queries/auth/key.sql Normal file
View file

@ -0,0 +1,50 @@
WITH resolved AS (
SELECT COALESCE(v.token, d.active_token_id) AS token
FROM UNNEST($1::text[]) AS requested(token)
LEFT JOIN "LiteLLM_VerificationToken" v ON v.token=requested.token
LEFT JOIN "LiteLLM_DeprecatedVerificationToken" d ON d.token=requested.token AND d.revoke_at > (clock_timestamp() AT TIME ZONE 'UTC')
)
SELECT to_jsonb(view) AS "data!: AuthData" FROM (
SELECT v.*,
t.spend AS team_spend,
t.max_budget AS team_max_budget,
t.soft_budget AS team_soft_budget,
t.model_max_budget AS team_model_max_budget,
t.tpm_limit AS team_tpm_limit,
t.rpm_limit AS team_rpm_limit,
t.tpd_limit AS team_tpd_limit,
t.models AS team_models,
t.metadata AS team_metadata,
t.blocked AS team_blocked,
t.team_alias AS team_alias,
t.members_with_roles AS team_members_with_roles,
t.object_permission_id AS team_object_permission_id,
t.organization_id AS org_id,
p.project_alias AS project_alias,
tm.spend AS team_member_spend,
b_tm.tpm_limit AS team_member_tpm_limit,
b_tm.rpm_limit AS team_member_rpm_limit,
m.aliases AS team_model_aliases,
b.max_budget AS litellm_budget_table_max_budget,
b.tpm_limit AS litellm_budget_table_tpm_limit,
b.rpm_limit AS litellm_budget_table_rpm_limit,
b.tpd_limit AS litellm_budget_table_tpd_limit,
b.model_max_budget AS litellm_budget_table_model_max_budget,
b.soft_budget AS litellm_budget_table_soft_budget,
o.metadata AS organization_metadata,
o.organization_alias AS organization_alias,
b2.max_budget AS organization_max_budget,
b2.tpm_limit AS organization_tpm_limit,
b2.rpm_limit AS organization_rpm_limit
FROM "LiteLLM_VerificationToken" v
LEFT JOIN "LiteLLM_TeamTable" t ON v.team_id=t.team_id
LEFT JOIN "LiteLLM_TeamMembership" tm ON v.team_id=tm.team_id AND tm.user_id=v.user_id
LEFT JOIN "LiteLLM_BudgetTable" b_tm ON tm.budget_id=b_tm.budget_id
LEFT JOIN "LiteLLM_ModelTable" m ON t.model_id=m.id
LEFT JOIN "LiteLLM_BudgetTable" b ON v.budget_id=b.budget_id
LEFT JOIN "LiteLLM_ProjectTable" p ON v.project_id=p.project_id
LEFT JOIN "LiteLLM_OrganizationTable" o ON v.organization_id=o.organization_id
LEFT JOIN "LiteLLM_BudgetTable" b2 ON o.budget_id=b2.budget_id
WHERE v.token IN (SELECT token FROM resolved)
ORDER BY v.token
) view

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

2
litellm-rust/crates/db/queries/auth/member.sql Normal file
View file

@ -0,0 +1,2 @@
SELECT to_jsonb(tm) || jsonb_build_object('litellm_budget_table', (SELECT to_jsonb(b) || jsonb_build_object('allowed_models', COALESCE(to_jsonb(b.allowed_models), '[]'::jsonb)) FROM "LiteLLM_BudgetTable" b WHERE b.budget_id = tm.budget_id)) AS "data!: AuthData"
FROM "LiteLLM_TeamMembership" tm WHERE tm.user_id=$1 AND tm.team_id=$2

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

@ -0,0 +1,4 @@
SELECT to_jsonb(g) || jsonb_build_object('litellm_budget_table', (SELECT to_jsonb(b) || jsonb_build_object('allowed_models', COALESCE(to_jsonb(b.allowed_models), '[]'::jsonb)) FROM "LiteLLM_BudgetTable" b WHERE b.budget_id = g.budget_id)) AS "data!: AuthData"
FROM "LiteLLM_ModelAccessGroupBudgetTable" g
WHERE g.access_group_name = ANY($1::text[])
ORDER BY g.access_group_name

1
litellm-rust/crates/db/queries/auth/model_access_group_registry.sql Normal file
View file

@ -0,0 +1 @@
SELECT jsonb_build_object('access_group_name', access_group_name) AS "data!: AuthData" FROM "LiteLLM_ModelAccessGroupBudgetTable" ORDER BY access_group_name LIMIT $1

13
litellm-rust/crates/db/queries/auth/object_permission.sql Normal file
View file

@ -0,0 +1,13 @@
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)) AS "data!: AuthData"
FROM "LiteLLM_ObjectPermissionTable" op
WHERE op.object_permission_id = ANY($1::text[])
ORDER BY op.object_permission_id

13
litellm-rust/crates/db/queries/auth/org.sql Normal file
View file

@ -0,0 +1,13 @@
SELECT to_jsonb(o) || jsonb_build_object('models', COALESCE(to_jsonb(o.models), '[]'::jsonb), 'litellm_budget_table', (SELECT to_jsonb(b) || jsonb_build_object('allowed_models', COALESCE(to_jsonb(b.allowed_models), '[]'::jsonb)) FROM "LiteLLM_BudgetTable" b WHERE b.budget_id = o.budget_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 = o.object_permission_id)) AS "data!: AuthData"
FROM "LiteLLM_OrganizationTable" o
WHERE o.organization_id = ANY($1::text[])
ORDER BY o.organization_id

10
litellm-rust/crates/db/queries/auth/org_alias.sql Normal file
View file

@ -0,0 +1,10 @@
SELECT to_jsonb(o) || jsonb_build_object('models', COALESCE(to_jsonb(o.models), '[]'::jsonb), 'litellm_budget_table', (SELECT to_jsonb(b) || jsonb_build_object('allowed_models', COALESCE(to_jsonb(b.allowed_models), '[]'::jsonb)) FROM "LiteLLM_BudgetTable" b WHERE b.budget_id = o.budget_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 = o.object_permission_id)) AS "data!: AuthData" FROM "LiteLLM_OrganizationTable" o WHERE o.organization_alias = $1 ORDER BY o.organization_id

42
litellm-rust/crates/db/queries/auth/prefetch.sql Normal file
View file

@ -0,0 +1,42 @@
SELECT jsonb_build_object(
'user_row', (SELECT to_jsonb(u) || jsonb_build_object('teams', COALESCE(to_jsonb(u.teams), '[]'::jsonb),
'models', COALESCE(to_jsonb(u.models), '[]'::jsonb),
'allowed_cache_controls', COALESCE(to_jsonb(u.allowed_cache_controls), '[]'::jsonb),
'policies', COALESCE(to_jsonb(u.policies), '[]'::jsonb), 'organization_memberships', COALESCE((SELECT jsonb_agg(to_jsonb(om)) FROM "LiteLLM_OrganizationMembership" om WHERE om.user_id = u.user_id), '[]'::jsonb)) FROM "LiteLLM_UserTable" u WHERE u.user_id=$1),
'team_row', (SELECT to_jsonb(t) || jsonb_build_object('admins', COALESCE(to_jsonb(t.admins), '[]'::jsonb),
'members', COALESCE(to_jsonb(t.members), '[]'::jsonb),
'models', COALESCE(to_jsonb(t.models), '[]'::jsonb),
'team_member_permissions', COALESCE(to_jsonb(t.team_member_permissions), '[]'::jsonb),
'access_group_ids', COALESCE(to_jsonb(t.access_group_ids), '[]'::jsonb),
'policies', COALESCE(to_jsonb(t.policies), '[]'::jsonb),
'default_team_member_models', COALESCE(to_jsonb(t.default_team_member_models), '[]'::jsonb), 'litellm_model_table', (SELECT (to_jsonb(m) - 'aliases') || jsonb_build_object('model_aliases', m.aliases) FROM "LiteLLM_ModelTable" m WHERE m.id=t.model_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 = t.object_permission_id)) FROM "LiteLLM_TeamTable" t WHERE t.team_id=$2),
'membership_row', (SELECT to_jsonb(tm) || jsonb_build_object('litellm_budget_table', (SELECT to_jsonb(b) || jsonb_build_object('allowed_models', COALESCE(to_jsonb(b.allowed_models), '[]'::jsonb)) FROM "LiteLLM_BudgetTable" b WHERE b.budget_id = tm.budget_id)) FROM "LiteLLM_TeamMembership" tm WHERE tm.user_id=$3 AND tm.team_id=$2),
'organization_row', (SELECT to_jsonb(o) || jsonb_build_object('models', COALESCE(to_jsonb(o.models), '[]'::jsonb), 'litellm_budget_table', (SELECT to_jsonb(b) || jsonb_build_object('allowed_models', COALESCE(to_jsonb(b.allowed_models), '[]'::jsonb)) FROM "LiteLLM_BudgetTable" b WHERE b.budget_id = o.budget_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 = o.object_permission_id)) FROM "LiteLLM_OrganizationTable" o WHERE o.organization_id=$4),
'project_row', (SELECT to_jsonb(p) || jsonb_build_object('models', COALESCE(to_jsonb(p.models), '[]'::jsonb), 'litellm_budget_table', (SELECT to_jsonb(b) || jsonb_build_object('allowed_models', COALESCE(to_jsonb(b.allowed_models), '[]'::jsonb)) FROM "LiteLLM_BudgetTable" b WHERE b.budget_id = p.budget_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 = p.object_permission_id)) FROM "LiteLLM_ProjectTable" p WHERE p.project_id=$5)) AS "data!: AuthData"

13
litellm-rust/crates/db/queries/auth/project.sql Normal file
View file

@ -0,0 +1,13 @@
SELECT to_jsonb(p) || jsonb_build_object('models', COALESCE(to_jsonb(p.models), '[]'::jsonb), 'litellm_budget_table', (SELECT to_jsonb(b) || jsonb_build_object('allowed_models', COALESCE(to_jsonb(b.allowed_models), '[]'::jsonb)) FROM "LiteLLM_BudgetTable" b WHERE b.budget_id = p.budget_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 = p.object_permission_id)) AS "data!: AuthData"
FROM "LiteLLM_ProjectTable" p
WHERE p.project_id = ANY($1::text[])
ORDER BY p.project_id

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

@ -0,0 +1,4 @@
SELECT to_jsonb(a) AS "data!: AuthData"
FROM "LiteLLM_RetiredAgent" a
WHERE a.original_agent_id = ANY($1::text[])
ORDER BY a.original_agent_id

1
litellm-rust/crates/db/queries/auth/retired_agent_identity.sql Normal file
View file

@ -0,0 +1 @@
SELECT to_jsonb(i) AS "data!: AuthData" FROM "LiteLLM_RetiredAgentIdentity" i WHERE i.provider=$1 AND i.tenant_id=$2 AND i.client_id=$3

1
litellm-rust/crates/db/queries/auth/spend_end_user.sql Normal file
View file

@ -0,0 +1 @@
SELECT COALESCE(spend, 0.0)::float8 AS "spend!" FROM "LiteLLM_EndUserTable" WHERE user_id=$1

1
litellm-rust/crates/db/queries/auth/spend_key.sql Normal file
View file

@ -0,0 +1 @@
SELECT COALESCE(spend, 0.0)::float8 AS "spend!" FROM "LiteLLM_VerificationToken" WHERE token=$1

1
litellm-rust/crates/db/queries/auth/spend_member.sql Normal file
View file

@ -0,0 +1 @@
SELECT COALESCE(spend, 0.0)::float8 AS "spend!" FROM "LiteLLM_TeamMembership" WHERE user_id=$1 AND team_id=$2

1
litellm-rust/crates/db/queries/auth/spend_model_access_group.sql Normal file
View file

@ -0,0 +1 @@
SELECT COALESCE(spend, 0.0)::float8 AS "spend!" FROM "LiteLLM_ModelAccessGroupBudgetTable" WHERE access_group_name=$1

1
litellm-rust/crates/db/queries/auth/spend_org.sql Normal file
View file

@ -0,0 +1 @@
SELECT COALESCE(spend, 0.0)::float8 AS "spend!" FROM "LiteLLM_OrganizationTable" WHERE organization_id=$1

1
litellm-rust/crates/db/queries/auth/spend_project.sql Normal file
View file

@ -0,0 +1 @@
SELECT COALESCE(spend, 0.0)::float8 AS "spend!" FROM "LiteLLM_ProjectTable" WHERE project_id=$1

1
litellm-rust/crates/db/queries/auth/spend_tag.sql Normal file
View file

@ -0,0 +1 @@
SELECT COALESCE(spend, 0.0)::float8 AS "spend!" FROM "LiteLLM_TagTable" WHERE tag_name=$1

1
litellm-rust/crates/db/queries/auth/spend_team.sql Normal file
View file

@ -0,0 +1 @@
SELECT COALESCE(spend, 0.0)::float8 AS "spend!" FROM "LiteLLM_TeamTable" WHERE team_id=$1

1
litellm-rust/crates/db/queries/auth/spend_user.sql Normal file
View file

@ -0,0 +1 @@
SELECT COALESCE(spend, 0.0)::float8 AS "spend!" FROM "LiteLLM_UserTable" WHERE user_id=$1

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

@ -0,0 +1,4 @@
SELECT to_jsonb(t) || jsonb_build_object('models', COALESCE(to_jsonb(t.models), '[]'::jsonb), 'litellm_budget_table', (SELECT to_jsonb(b) || jsonb_build_object('allowed_models', COALESCE(to_jsonb(b.allowed_models), '[]'::jsonb)) FROM "LiteLLM_BudgetTable" b WHERE b.budget_id = t.budget_id)) AS "data!: AuthData"
FROM "LiteLLM_TagTable" t
WHERE t.tag_name = ANY($1::text[])
ORDER BY t.tag_name

1
litellm-rust/crates/db/queries/auth/tag_registry.sql Normal file
View file

@ -0,0 +1 @@
SELECT jsonb_build_object('tag_name', tag_name) AS "data!: AuthData" FROM "LiteLLM_TagTable" ORDER BY tag_name LIMIT $1

19
litellm-rust/crates/db/queries/auth/team.sql Normal file
View file

@ -0,0 +1,19 @@
SELECT to_jsonb(t) || jsonb_build_object('admins', COALESCE(to_jsonb(t.admins), '[]'::jsonb),
'members', COALESCE(to_jsonb(t.members), '[]'::jsonb),
'models', COALESCE(to_jsonb(t.models), '[]'::jsonb),
'team_member_permissions', COALESCE(to_jsonb(t.team_member_permissions), '[]'::jsonb),
'access_group_ids', COALESCE(to_jsonb(t.access_group_ids), '[]'::jsonb),
'policies', COALESCE(to_jsonb(t.policies), '[]'::jsonb),
'default_team_member_models', COALESCE(to_jsonb(t.default_team_member_models), '[]'::jsonb), 'litellm_model_table', (SELECT (to_jsonb(m) - 'aliases') || jsonb_build_object('model_aliases', m.aliases) FROM "LiteLLM_ModelTable" m WHERE m.id=t.model_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 = t.object_permission_id)) AS "data!: AuthData"
FROM "LiteLLM_TeamTable" t
WHERE t.team_id = ANY($1::text[])
ORDER BY t.team_id

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

@ -0,0 +1,16 @@
SELECT to_jsonb(t) || jsonb_build_object('admins', COALESCE(to_jsonb(t.admins), '[]'::jsonb),
'members', COALESCE(to_jsonb(t.members), '[]'::jsonb),
'models', COALESCE(to_jsonb(t.models), '[]'::jsonb),
'team_member_permissions', COALESCE(to_jsonb(t.team_member_permissions), '[]'::jsonb),
'access_group_ids', COALESCE(to_jsonb(t.access_group_ids), '[]'::jsonb),
'policies', COALESCE(to_jsonb(t.policies), '[]'::jsonb),
'default_team_member_models', COALESCE(to_jsonb(t.default_team_member_models), '[]'::jsonb), 'litellm_model_table', (SELECT (to_jsonb(m) - 'aliases') || jsonb_build_object('model_aliases', m.aliases) FROM "LiteLLM_ModelTable" m WHERE m.id=t.model_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 = t.object_permission_id)) AS "data!: AuthData" FROM "LiteLLM_TeamTable" t WHERE t.team_alias = $1 ORDER BY t.team_id

7
litellm-rust/crates/db/queries/auth/user.sql Normal file
View file

@ -0,0 +1,7 @@
SELECT to_jsonb(u) || jsonb_build_object('teams', COALESCE(to_jsonb(u.teams), '[]'::jsonb),
'models', COALESCE(to_jsonb(u.models), '[]'::jsonb),
'allowed_cache_controls', COALESCE(to_jsonb(u.allowed_cache_controls), '[]'::jsonb),
'policies', COALESCE(to_jsonb(u.policies), '[]'::jsonb), 'organization_memberships', COALESCE((SELECT jsonb_agg(to_jsonb(om)) FROM "LiteLLM_OrganizationMembership" om WHERE om.user_id = u.user_id), '[]'::jsonb)) AS "data!: AuthData"
FROM "LiteLLM_UserTable" u
WHERE u.user_id = ANY($1::text[])
ORDER BY u.user_id

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

@ -0,0 +1,4 @@
SELECT to_jsonb(u) || jsonb_build_object('teams', COALESCE(to_jsonb(u.teams), '[]'::jsonb),
'models', COALESCE(to_jsonb(u.models), '[]'::jsonb),
'allowed_cache_controls', COALESCE(to_jsonb(u.allowed_cache_controls), '[]'::jsonb),
'policies', COALESCE(to_jsonb(u.policies), '[]'::jsonb), 'organization_memberships', COALESCE((SELECT jsonb_agg(to_jsonb(om)) FROM "LiteLLM_OrganizationMembership" om WHERE om.user_id = u.user_id), '[]'::jsonb)) AS "data!: AuthData" FROM "LiteLLM_UserTable" u WHERE LOWER(u.user_email) = LOWER($1) ORDER BY u.user_id LIMIT 1

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

@ -0,0 +1,4 @@
SELECT to_jsonb(u) || jsonb_build_object('teams', COALESCE(to_jsonb(u.teams), '[]'::jsonb),
'models', COALESCE(to_jsonb(u.models), '[]'::jsonb),
'allowed_cache_controls', COALESCE(to_jsonb(u.allowed_cache_controls), '[]'::jsonb),
'policies', COALESCE(to_jsonb(u.policies), '[]'::jsonb), 'organization_memberships', COALESCE((SELECT jsonb_agg(to_jsonb(om)) FROM "LiteLLM_OrganizationMembership" om WHERE om.user_id = u.user_id), '[]'::jsonb)) AS "data!: AuthData" FROM "LiteLLM_UserTable" u WHERE u.sso_user_id = $1 ORDER BY u.user_id

1
litellm-rust/crates/db/queries/auth/verified_subject.sql Normal file
View file

@ -0,0 +1 @@
SELECT to_jsonb(i) AS "data!: AuthData" FROM "LiteLLM_VerifiedSubject" i WHERE i.issuer=$1 AND i.tenant_id=$2 AND i.oid=$3

1
litellm-rust/crates/db/queries/auth/window_key_spend.sql Normal file
View file

@ -0,0 +1 @@
SELECT jsonb_build_object('spend', COALESCE(SUM(spend), 0.0)) AS "data!: AuthData" FROM "LiteLLM_SpendLogs" WHERE "api_key"=$1 AND "startTime" >= $2

1
litellm-rust/crates/db/queries/auth/window_team_spend.sql Normal file
View file

@ -0,0 +1 @@
SELECT jsonb_build_object('spend', COALESCE(SUM(spend), 0.0)) AS "data!: AuthData" FROM "LiteLLM_SpendLogs" WHERE "team_id"=$1 AND "startTime" >= $2

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

@ -0,0 +1,814 @@
use crate::{Error, executor::StatementPersistence};
use sqlx::{Executor, Postgres};
#[derive(Clone, PartialEq, serde::Serialize, sqlx::Type, veil::Redact)]
#[serde(transparent)]
#[sqlx(transparent)]
pub struct AuthData(#[redact(with = "[REDACTED]")] pub serde_json::Value);
#[macro_rules_attribute::apply(crate::db_row)]
pub struct AuthRow {
pub data: AuthData,
}
#[derive(Clone, Copy, Debug, strum::EnumString, strum::VariantArray)]
pub enum AuthRead {
#[strum(serialize = "user")]
User,
#[strum(serialize = "team")]
Team,
#[strum(serialize = "org")]
Org,
#[strum(serialize = "project")]
Project,
#[strum(serialize = "end_user")]
EndUser,
#[strum(serialize = "budget")]
Budget,
#[strum(serialize = "object_permission")]
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")]
Key,
#[strum(serialize = "member")]
Member,
#[strum(serialize = "agent_identity")]
AgentIdentity,
#[strum(serialize = "retired_agent_identity")]
RetiredAgentIdentity,
#[strum(serialize = "verified_subject")]
VerifiedSubject,
#[strum(serialize = "team_alias")]
TeamAlias,
#[strum(serialize = "org_alias")]
OrgAlias,
#[strum(serialize = "user_sso")]
UserSso,
#[strum(serialize = "user_email")]
UserEmail,
#[strum(serialize = "tag_registry")]
TagRegistry,
#[strum(serialize = "model_access_group_registry")]
ModelAccessGroupRegistry,
#[strum(serialize = "end_user_restricted_registry")]
EndUserRestrictedRegistry,
#[strum(serialize = "prefetch")]
Prefetch,
#[strum(serialize = "budget_window")]
BudgetWindow,
#[strum(serialize = "window_spend")]
WindowSpend,
}
#[derive(Clone, Copy, Debug, strum::EnumString, strum::VariantArray)]
pub enum AuthSpend {
#[strum(serialize = "key")]
Key,
#[strum(serialize = "user")]
User,
#[strum(serialize = "team")]
Team,
#[strum(serialize = "org")]
Org,
#[strum(serialize = "project")]
Project,
#[strum(serialize = "end_user")]
EndUser,
#[strum(serialize = "tag")]
Tag,
#[strum(serialize = "model_access_group")]
ModelAccessGroup,
#[strum(serialize = "member")]
Member,
}
fn exact(ids: &[String], count: usize) -> Result<(), Error> {
if ids.len() == count {
Ok(())
} else {
Err(Error::InvalidAuthArguments)
}
}
fn hashed_key(id: &str) -> Result<(), Error> {
if id.len() == 64 && id.bytes().all(|byte| byte.is_ascii_hexdigit()) {
Ok(())
} else {
Err(Error::InvalidAuthArguments)
}
}
fn entity(value: &str) -> Result<&'static str, Error> {
match value {
"Key" => Ok("key"),
"Team" => Ok("team"),
_ => Err(Error::InvalidAuthArguments),
}
}
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);
}
let rows = match kind {
AuthRead::User => {
sqlx::query_file_as!(AuthRow, "queries/auth/user.sql", ids)
.fetch_all(executor)
.await?
}
AuthRead::Team => {
sqlx::query_file_as!(AuthRow, "queries/auth/team.sql", ids)
.fetch_all(executor)
.await?
}
AuthRead::Org => {
sqlx::query_file_as!(AuthRow, "queries/auth/org.sql", ids)
.fetch_all(executor)
.await?
}
AuthRead::Project => {
sqlx::query_file_as!(AuthRow, "queries/auth/project.sql", ids)
.fetch_all(executor)
.await?
}
AuthRead::EndUser => {
sqlx::query_file_as!(AuthRow, "queries/auth/end_user.sql", ids)
.fetch_all(executor)
.await?
}
AuthRead::Budget => {
sqlx::query_file_as!(AuthRow, "queries/auth/budget.sql", ids)
.fetch_all(executor)
.await?
}
AuthRead::ObjectPermission => {
sqlx::query_file_as!(AuthRow, "queries/auth/object_permission.sql", ids)
.fetch_all(executor)
.await?
}
AuthRead::AccessGroup => {
sqlx::query_file_as!(AuthRow, "queries/auth/access_group.sql", ids)
.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)
.await?
}
AuthRead::ModelAccessGroup => {
sqlx::query_file_as!(AuthRow, "queries/auth/model_access_group.sql", ids)
.fetch_all(executor)
.await?
}
AuthRead::Agent => {
sqlx::query_file_as!(AuthRow, "queries/auth/agent.sql", ids)
.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)
.await?
}
AuthRead::Key => {
ids.iter().try_for_each(|id| hashed_key(id))?;
sqlx::query_file_as!(AuthRow, "queries/auth/key.sql", ids)
.fetch_all(executor)
.await?
}
AuthRead::Member => {
exact(ids, 2)?;
sqlx::query_file_as!(AuthRow, "queries/auth/member.sql", &ids[0], &ids[1])
.fetch_all(executor)
.await?
}
AuthRead::AgentIdentity => {
exact(ids, 3)?;
sqlx::query_file_as!(
AuthRow,
"queries/auth/agent_identity.sql",
&ids[0],
&ids[1],
&ids[2]
)
.fetch_all(executor)
.await?
}
AuthRead::RetiredAgentIdentity => {
exact(ids, 3)?;
sqlx::query_file_as!(
AuthRow,
"queries/auth/retired_agent_identity.sql",
&ids[0],
&ids[1],
&ids[2]
)
.fetch_all(executor)
.await?
}
AuthRead::VerifiedSubject => {
exact(ids, 3)?;
sqlx::query_file_as!(
AuthRow,
"queries/auth/verified_subject.sql",
&ids[0],
&ids[1],
&ids[2]
)
.fetch_all(executor)
.await?
}
AuthRead::TeamAlias => {
exact(ids, 1)?;
sqlx::query_file_as!(AuthRow, "queries/auth/team_alias.sql", &ids[0])
.fetch_all(executor)
.await?
}
AuthRead::OrgAlias => {
exact(ids, 1)?;
sqlx::query_file_as!(AuthRow, "queries/auth/org_alias.sql", &ids[0])
.fetch_all(executor)
.await?
}
AuthRead::UserSso => {
exact(ids, 1)?;
sqlx::query_file_as!(AuthRow, "queries/auth/user_sso.sql", &ids[0])
.fetch_all(executor)
.await?
}
AuthRead::UserEmail => {
exact(ids, 1)?;
sqlx::query_file_as!(AuthRow, "queries/auth/user_email.sql", &ids[0])
.fetch_all(executor)
.await?
}
AuthRead::TagRegistry => {
exact(ids, 1)?;
let limit: i64 = ids[0]
.parse()
.ok()
.filter(|value| (1..=100_001).contains(value))
.ok_or(Error::InvalidAuthArguments)?;
sqlx::query_file_as!(AuthRow, "queries/auth/tag_registry.sql", limit)
.fetch_all(executor)
.await?
}
AuthRead::ModelAccessGroupRegistry => {
exact(ids, 1)?;
let limit: i64 = ids[0]
.parse()
.ok()
.filter(|value| (1..=100_001).contains(value))
.ok_or(Error::InvalidAuthArguments)?;
sqlx::query_file_as!(
AuthRow,
"queries/auth/model_access_group_registry.sql",
limit
)
.fetch_all(executor)
.await?
}
AuthRead::EndUserRestrictedRegistry => {
exact(ids, 1)?;
let limit: i64 = ids[0]
.parse()
.ok()
.filter(|value| (1..=100_001).contains(value))
.ok_or(Error::InvalidAuthArguments)?;
sqlx::query_file_as!(
AuthRow,
"queries/auth/end_user_restricted_registry.sql",
limit
)
.fetch_all(executor)
.await?
}
AuthRead::Prefetch => {
exact(ids, 1)?;
let refs: [Option<String>; 5] =
serde_json::from_str(&ids[0]).map_err(|_| Error::InvalidAuthArguments)?;
sqlx::query_file_as!(
AuthRow,
"queries/auth/prefetch.sql",
refs[0].as_deref(),
refs[1].as_deref(),
refs[2].as_deref(),
refs[3].as_deref(),
refs[4].as_deref()
)
.fetch_all(executor)
.await?
}
AuthRead::BudgetWindow => {
exact(ids, 3)?;
let entity = entity(&ids[0])?;
if entity == "key" {
hashed_key(&ids[1])?;
}
sqlx::query_file_as!(
AuthRow,
"queries/auth/budget_window.sql",
entity,
&ids[1],
&ids[2]
)
.fetch_all(executor)
.await?
}
AuthRead::WindowSpend => {
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"))
.map_err(|_| Error::InvalidAuthArguments)?;
match entity(&ids[0])? {
"key" => {
hashed_key(&ids[1])?;
sqlx::query_file_as!(
AuthRow,
"queries/auth/window_key_spend.sql",
&ids[1],
start
)
.fetch_all(executor)
.await?
}
_ => {
sqlx::query_file_as!(
AuthRow,
"queries/auth/window_team_spend.sql",
&ids[1],
start
)
.fetch_all(executor)
.await?
}
}
}
};
Ok(rows)
}
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);
}
let row = match kind {
AuthSpend::Key => {
hashed_key(id)?;
sqlx::query_file!("queries/auth/spend_key.sql", id)
.fetch_optional(executor)
.await?
.map(|row| row.spend)
}
AuthSpend::User => sqlx::query_file!("queries/auth/spend_user.sql", id)
.fetch_optional(executor)
.await?
.map(|row| row.spend),
AuthSpend::Team => sqlx::query_file!("queries/auth/spend_team.sql", id)
.fetch_optional(executor)
.await?
.map(|row| row.spend),
AuthSpend::Org => sqlx::query_file!("queries/auth/spend_org.sql", id)
.fetch_optional(executor)
.await?
.map(|row| row.spend),
AuthSpend::Project => sqlx::query_file!("queries/auth/spend_project.sql", id)
.fetch_optional(executor)
.await?
.map(|row| row.spend),
AuthSpend::EndUser => sqlx::query_file!("queries/auth/spend_end_user.sql", id)
.fetch_optional(executor)
.await?
.map(|row| row.spend),
AuthSpend::Tag => sqlx::query_file!("queries/auth/spend_tag.sql", id)
.fetch_optional(executor)
.await?
.map(|row| row.spend),
AuthSpend::ModelAccessGroup => {
sqlx::query_file!("queries/auth/spend_model_access_group.sql", id)
.fetch_optional(executor)
.await?
.map(|row| row.spend)
}
AuthSpend::Member => {
let (user_id, team_id) = id.rsplit_once(':').ok_or(Error::InvalidAuthArguments)?;
sqlx::query_file!("queries/auth/spend_member.sql", user_id, team_id)
.fetch_optional(executor)
.await?
.map(|row| row.spend)
}
};
Ok(row)
}
#[cfg(test)]
mod tests {
use super::{AuthRead, AuthSpend, auth_read, auth_spend};
use crate::test_support::database_pool;
use rstest::rstest;
use serde_json::{Value, json};
use sqlx::{PgPool, Postgres, Transaction};
use std::error::Error as StdError;
async fn seeded(pool: &PgPool) -> Result<Transaction<'_, Postgres>, sqlx::Error> {
let mut tx = pool.begin().await?;
#[expect(
clippy::disallowed_methods,
reason = "Transaction-local test fixture outside the production query set"
)]
sqlx::raw_sql(include_str!("../queries/auth/fixture.sql"))
.execute(&mut *tx)
.await?;
Ok(tx)
}
fn strings(values: &[&str]) -> Vec<String> {
values.iter().map(|value| (*value).to_owned()).collect()
}
#[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!([]))]
#[case::team_missing_relation(AuthRead::Team, &["auth-team-empty"], "/object_permission", json!(null))]
#[case::org(AuthRead::Org, &["auth-org"], "/litellm_budget_table/max_budget", json!(13.0))]
#[case::project(AuthRead::Project, &["auth-project"], "/litellm_budget_table/max_budget", json!(0.0))]
#[case::end_user(AuthRead::EndUser, &["auth-end"], "/models", json!([]))]
#[case::budget(AuthRead::Budget, &["auth-b"], "/allowed_models", json!([]))]
#[case::object_permission(AuthRead::ObjectPermission, &["auth-op"], "/mcp_servers", json!([]))]
#[case::access_group(AuthRead::AccessGroup, &["auth-access"], "/access_model_names", json!(["group-model"]))]
#[case::tag(AuthRead::Tag, &["auth-tag"], "/models", json!([]))]
#[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"))]
#[case::verified_subject(AuthRead::VerifiedSubject, &["auth-issuer","auth-tenant","auth-oid"], "/user_id", json!("auth-u"))]
#[case::team_alias(AuthRead::TeamAlias, &["Auth Team"], "/team_id", json!("auth-team"))]
#[case::org_alias(AuthRead::OrgAlias, &["Auth Org"], "/organization_id", json!("auth-org"))]
#[case::sso(AuthRead::UserSso, &["auth-sso"], "/user_id", json!("auth-u"))]
#[case::email_literal_insensitive(AuthRead::UserEmail, &["test_%@example.COM"], "/user_id", json!("auth-u"))]
#[case::tags_registry(AuthRead::TagRegistry, &["5001"], "/tag_name", json!("auth-tag"))]
#[case::model_registry(AuthRead::ModelAccessGroupRegistry, &["5001"], "/access_group_name", json!("auth-model-group"))]
#[case::restricted_registry(AuthRead::EndUserRestrictedRegistry, &["5001"], "/user_id", json!("auth-end"))]
#[case::budget_window(AuthRead::BudgetWindow, &["Team","auth-team","1h"], "/spend", json!(2.5))]
#[case::window_spend_offset(AuthRead::WindowSpend, &["Team","auth-team","2025-01-01T02:00:00+01:00"], "/spend", json!(6.0))]
#[case::window_spend_naive(AuthRead::WindowSpend, &["Team","auth-team","2025-01-01T01:00:00"], "/spend", json!(6.0))]
#[case::window_spend_no_rows(AuthRead::WindowSpend, &["Team","absent","2025-01-01T01:00:00Z"], "/spend", json!(0.0))]
#[tokio::test]
#[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"]
async fn projects_auth_rows(
#[case] kind: AuthRead,
#[case] ids: &[&str],
#[case] pointer: &str,
#[case] expected: Value,
#[future(awt)] database_pool: Result<PgPool, Box<dyn StdError>>,
) -> Result<(), Box<dyn StdError>> {
let pool = database_pool?;
let mut tx = seeded(&pool).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() {
assert_eq!(actual.as_f64(), expected.as_f64());
} else {
assert_eq!(actual, &expected);
}
Ok(())
}
#[rstest]
#[case::active('a', true)]
#[case::deprecated('b', true)]
#[case::expired('c', false)]
#[case::no_chaining('d', false)]
#[case::missing('f', false)]
#[tokio::test]
#[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"]
async fn combined_key_join_and_deprecation(
#[case] token: char,
#[case] found: bool,
#[future(awt)] database_pool: Result<PgPool, Box<dyn StdError>>,
) -> 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)],
false,
)
.await?;
assert_eq!(rows.len(), usize::from(found));
if !found {
return Ok(());
}
let row = &rows[0].data.0;
assert_eq!(row["token"], "a".repeat(64));
assert_eq!(row["team_spend"], 2.5);
assert_eq!(row["team_blocked"], true);
assert_eq!(row["org_id"], "auth-org");
assert_eq!(row["team_models"], json!(["team-model"]));
assert_eq!(row["team_model_aliases"], json!({"alias":"actual"}));
assert_eq!(row["team_member_spend"], 6.5);
assert_eq!(row["team_member_tpm_limit"], 22);
assert_eq!(row["litellm_budget_table_max_budget"], 4.5);
assert_eq!(row["organization_max_budget"], 13.0);
assert_eq!(row["project_alias"], "Auth Project");
Ok(())
}
#[rstest]
#[tokio::test]
#[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"]
async fn combined_key_keeps_null_join_values(
#[future(awt)] database_pool: Result<PgPool, Box<dyn StdError>>,
) -> 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)], 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);
assert_eq!(rows[0].data.0["organization_max_budget"], Value::Null);
assert_eq!(rows[0].data.0["spend"], 0.0);
Ok(())
}
#[rstest]
#[tokio::test]
#[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"]
async fn prefetch_has_same_nested_projections(
#[future(awt)] database_pool: Result<PgPool, Box<dyn StdError>>,
) -> Result<(), Box<dyn StdError>> {
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], 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?;
assert_eq!(rows.len(), 1);
assert_eq!(
rows[0].data.0,
json!({
"user_row": user[0].data.0, "team_row": team[0].data.0,
"membership_row": member[0].data.0, "organization_row": org[0].data.0,
"project_row": project[0].data.0
})
);
let absent = auth_read(
&mut *tx,
AuthRead::Prefetch,
&[json!([null, null, null, null, null]).to_string()],
false,
)
.await?;
assert_eq!(
absent[0].data.0,
json!({"user_row":null,"team_row":null,"membership_row":null,"organization_row":null,"project_row":null})
);
Ok(())
}
#[rstest]
#[tokio::test]
#[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"]
async fn batch_ids_preserve_tenant_boundaries_and_missing_rows(
#[future(awt)] database_pool: Result<PgPool, Box<dyn StdError>>,
) -> Result<(), Box<dyn StdError>> {
let pool = database_pool?;
let mut tx = seeded(&pool).await?;
let rows = auth_read(
&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, &[], false)
.await?
.is_empty()
);
assert!(
auth_read(
&mut *tx,
AuthRead::Member,
&strings(&["auth-other", "auth-team"]),
false
)
.await?
.is_empty()
);
assert!(
auth_read(
&mut *tx,
AuthRead::AgentIdentity,
&strings(&["microsoft_entra", "other-tenant", "auth-client"]),
false
)
.await?
.is_empty()
);
assert!(
auth_read(
&mut *tx,
AuthRead::UserEmail,
&strings(&["test%@example.com"]),
false
)
.await?
.is_empty()
);
assert!(
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(())
}
#[rstest]
#[case::key(AuthSpend::Key, "a".repeat(64), 4.5)]
#[case::user(AuthSpend::User, "auth-u".into(), 3.5)]
#[case::zero(AuthSpend::User, "auth-other".into(), 0.0)]
#[case::team(AuthSpend::Team, "auth-team".into(), 2.5)]
#[case::org(AuthSpend::Org, "auth-org".into(), 13.5)]
#[case::project(AuthSpend::Project, "auth-project".into(), 8.5)]
#[case::end_user(AuthSpend::EndUser, "auth-end".into(), 10.5)]
#[case::tag(AuthSpend::Tag, "auth-tag".into(), 11.5)]
#[case::model_group(AuthSpend::ModelAccessGroup, "auth-model-group".into(), 12.5)]
#[case::member(AuthSpend::Member, "auth-u:auth-team".into(), 6.5)]
#[tokio::test]
#[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"]
async fn spend_reads(
#[case] kind: AuthSpend,
#[case] id: String,
#[case] expected: f64,
#[future(awt)] database_pool: Result<PgPool, Box<dyn StdError>>,
) -> Result<(), Box<dyn StdError>> {
let pool = database_pool?;
let mut tx = seeded(&pool).await?;
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, false).await?, None);
Ok(())
}
#[rstest]
#[case::key_window(AuthRead::BudgetWindow, vec!["Key".into(), "a".repeat(64), "1h".into()], 2.5)]
#[case::key_logs(AuthRead::WindowSpend, vec!["Key".into(), "a".repeat(64), "2025-01-01T01:00:00Z".into()], 6.0)]
#[tokio::test]
#[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"]
async fn key_window_spend(
#[case] kind: AuthRead,
#[case] ids: Vec<String>,
#[case] expected: f64,
#[future(awt)] database_pool: Result<PgPool, Box<dyn StdError>>,
) -> Result<(), Box<dyn StdError>> {
let pool = database_pool?;
let mut tx = seeded(&pool).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(())
}
#[rstest]
#[case::member(AuthRead::Member, vec!["one".into()])]
#[case::registry(AuthRead::TagRegistry, vec!["0".into()])]
#[case::huge_registry(AuthRead::TagRegistry, vec!["100002".into()])]
#[case::raw_token(AuthRead::Key, vec!["sk-raw-secret".into()])]
#[case::prefetch_shape(AuthRead::Prefetch, vec!["[]".into()])]
#[case::unknown_window(AuthRead::BudgetWindow, strings(&["Unknown","id","1h"]))]
#[tokio::test]
async fn invalid_arguments_fail_before_queries(
#[case] kind: AuthRead,
#[case] ids: Vec<String>,
) {
let pool = sqlx::postgres::PgPoolOptions::new()
.connect_lazy("postgresql://localhost/unused")
.unwrap();
pool.close().await;
assert!(matches!(
auth_read(&pool, kind, &ids, false).await,
Err(crate::Error::InvalidAuthArguments)
));
}
#[rstest]
#[tokio::test]
async fn database_errors_are_not_misses() {
let pool = sqlx::postgres::PgPoolOptions::new()
.connect_lazy("postgresql://localhost/unused")
.unwrap();
pool.close().await;
assert!(matches!(
auth_read(&pool, AuthRead::User, &strings(&["id"]), false).await,
Err(crate::Error::Query(sqlx::Error::PoolClosed))
));
assert!(matches!(
auth_spend(&pool, AuthSpend::User, "id", false).await,
Err(crate::Error::Query(sqlx::Error::PoolClosed))
));
}
#[rstest]
fn projected_credentials_are_redacted() {
let row = super::AuthRow {
data: super::AuthData(json!({"password":"credential-value"})),
};
let debug = format!("{row:?}");
assert!(!debug.contains("credential-value"));
assert!(debug.contains("[REDACTED]"));
}
}

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

@ -1,5 +1,13 @@
#[derive(Debug, thiserror::Error)]
pub enum Error {
#[error("invalid database setting {0}")]
InvalidSetting(&'static str),
#[error("invalid auth read arguments")]
InvalidAuthArguments,
#[error("native database is not configured")]
NotConfigured,
#[error("database connection timed out")]
ConnectTimeout(#[from] tokio::time::error::Elapsed),
#[error("database query failed")]
Query(#[from] sqlx::Error),
#[error("spend batch has {ids} ids and {costs} costs")]
@ -9,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)) => {
@ -42,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)]
@ -72,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)
}
}

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

@ -6,18 +6,23 @@ macro_rules_attribute::attribute_alias! {
#[derive(Clone, Debug, PartialEq, sqlx::FromRow)];
}
mod auth;
mod daily_spend;
mod end_user;
mod error;
mod executor;
mod failure;
mod pool;
mod spend;
#[cfg(test)]
mod test_support;
pub use auth::{AuthRead, AuthRow, AuthSpend, auth_read, auth_spend};
pub use daily_spend::{DailyEndUserRow, DailyTeamRow, DailyUserRow};
pub use end_user::{EndUser, RateLimits, find_end_user};
pub use error::Error;
pub use failure::DbFailure;
pub use pool::{DatabasePools, DatabaseSettings};
pub use spend::{
DailySpend, DailyTable, EntitySpend, SpendLog, SpendLogRow, TeamMemberSpend,
add_end_user_spend, add_key_spend, add_team_member_spend, add_team_spend, add_user_spend,

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

@ -0,0 +1,495 @@
use std::time::Duration;
use litellm_secret_value::SecretValue;
use sqlx::{
ConnectOptions, PgPool,
postgres::{PgConnectOptions, PgPoolOptions, PgSslMode},
};
use crate::Error;
#[derive(Clone, Debug)]
pub struct DatabaseSettings {
pub host: String,
pub port: u16,
pub username: String,
pub password: Option<SecretValue>,
pub database: String,
pub schema: String,
pub ssl_mode: PgSslMode,
pub ssl_root_cert: Option<String>,
pub max_connections: u32,
pub acquire_timeout: Duration,
pub connect_timeout: Duration,
pub idle_timeout: Option<Duration>,
pub statement_cache_capacity: usize,
}
impl DatabaseSettings {
fn pool(&self) -> Result<PgPool, Error> {
if self.max_connections == 0 {
return Err(Error::InvalidSetting("max_connections"));
}
if self.host.is_empty()
|| self.username.is_empty()
|| self.database.is_empty()
|| self.schema.is_empty()
{
return Err(Error::InvalidSetting("connection identity"));
}
let options = PgConnectOptions::new_without_pgpass()
.host(&self.host)
.port(self.port)
.username(&self.username)
.password(self.password.as_ref().map_or("", SecretValue::expose))
.database(&self.database)
.application_name("native-auth")
.ssl_mode(self.ssl_mode)
.statement_cache_capacity(self.statement_cache_capacity)
.options([(
"search_path",
format!("\"{}\"", self.schema.replace('"', "\"\"")),
)])
.disable_statement_logging();
let options = match &self.ssl_root_cert {
Some(cert) => options.ssl_root_cert(cert),
None => options,
};
Ok(PgPoolOptions::new()
.max_connections(self.max_connections)
.acquire_timeout(if self.acquire_timeout.is_zero() {
Duration::from_secs(30)
} else {
self.acquire_timeout
})
.idle_timeout(self.idle_timeout)
.connect_lazy_with(options))
}
}
#[derive(Clone, Debug)]
pub struct DatabasePools {
writer: PgPool,
reader: Option<PgPool>,
reader_available: bool,
writer_unbounded: bool,
reader_unbounded: bool,
writer_persistent: bool,
reader_persistent: bool,
}
async fn acquire(
pool: &PgPool,
unbounded: bool,
) -> Result<sqlx::pool::PoolConnection<sqlx::Postgres>, sqlx::Error> {
loop {
match pool.acquire().await {
Err(sqlx::Error::PoolTimedOut) if unbounded => continue,
result => return result,
}
}
}
async fn connect_pool(pool: &PgPool, settings: &DatabaseSettings) -> Result<(), Error> {
let connection = acquire(pool, settings.acquire_timeout.is_zero());
if settings.connect_timeout.is_zero() {
connection.await?;
} else {
tokio::time::timeout(settings.connect_timeout, connection).await??;
}
Ok(())
}
impl DatabasePools {
pub async fn connect(
writer: DatabaseSettings,
reader: Option<DatabaseSettings>,
) -> Result<Self, Error> {
let writer_pool = writer.pool()?;
let reader_pool = reader.as_ref().map(DatabaseSettings::pool).transpose()?;
let writer_result = connect_pool(&writer_pool, &writer).await;
let reader_result = match (&reader_pool, &reader) {
(Some(pool), Some(settings)) => Some(connect_pool(pool, settings).await),
_ => None,
};
let reader_available = reader_result.as_ref().is_some_and(Result::is_ok);
drop(reader_result);
if let Err(error) = writer_result
&& !reader_available
{
writer_pool.close().await;
if let Some(pool) = reader_pool {
pool.close().await;
}
return Err(error);
}
Ok(Self {
writer: writer_pool,
reader: reader_pool,
reader_available,
writer_unbounded: writer.acquire_timeout.is_zero(),
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),
})
}
pub fn reader_available(&self) -> bool {
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,
_ => &self.writer,
}
}
pub async fn acquire(
&self,
writer: bool,
) -> Result<sqlx::pool::PoolConnection<sqlx::Postgres>, Error> {
let unbounded = if !writer && self.reader_available {
self.reader_unbounded
} else {
self.writer_unbounded
};
Ok(acquire(self.pool(writer), unbounded).await?)
}
pub fn set_passwords(
&self,
writer: Option<SecretValue>,
reader: Option<SecretValue>,
) -> Result<(), Error> {
if self.writer.is_closed() {
return Err(sqlx::Error::PoolClosed.into());
}
if reader.is_some() && self.reader.is_none() {
return Err(Error::InvalidSetting("reader_password"));
}
if let Some(password) = writer {
self.writer.set_connect_options(
(*self.writer.connect_options())
.clone()
.password(password.expose()),
);
}
if let (Some(pool), Some(password)) = (&self.reader, reader) {
pool.set_connect_options(
(*pool.connect_options())
.clone()
.password(password.expose()),
);
}
Ok(())
}
pub async fn close(&self) {
if let Some(reader) = &self.reader {
tokio::join!(self.writer.close(), reader.close());
} else {
self.writer.close().await;
}
}
}
#[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(())
}
}

2
litellm-rust/crates/python-bridge/Cargo.toml
View file

@ -19,6 +19,8 @@ huggingface = ["litellm-token-counter/huggingface"]
tiktoken = ["litellm-token-counter/tiktoken"]
[dependencies]
litellm-db.workspace = true
litellm-request-auth-core.workspace = true
litellm-secret-value.workspace = true
fancy-regex.workspace = true
litellm-tracing.workspace = true

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

@ -1,19 +1,174 @@
use litellm_host_python::ready_future;
use pyo3::{prelude::*, types::PyNone};
use std::time::Duration;
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::{PyOSError, PyRuntimeError, PyValueError},
prelude::*,
types::PyNone,
};
#[pyclass(frozen)]
pub(crate) struct NativeDatabase {}
pub(crate) struct NativeDatabase {
pools: Option<DatabasePools>,
}
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())
}
_ => PyRuntimeError::new_err(litellm_error::Report::new(&error).to_string()),
}
}
fn duration(seconds: f64, name: &'static str) -> PyResult<Duration> {
Duration::try_from_secs_f64(seconds)
.map_err(|_| db_error(litellm_db::Error::InvalidSetting(name)))
}
fn settings(value: &Bound<'_, PyAny>) -> PyResult<DatabaseSettings> {
let ssl_mode: String = value.getattr("ssl_mode")?.extract()?;
let idle_timeout: Option<f64> = value.getattr("idle_timeout")?.extract()?;
Ok(DatabaseSettings {
host: value.getattr("host")?.extract()?,
port: value.getattr("port")?.extract()?,
username: value.getattr("username")?.extract()?,
password: value
.getattr("password")?
.extract::<Option<String>>()?
.map(SecretValue::from),
database: value.getattr("database")?.extract()?,
schema: value.getattr("schema")?.extract()?,
ssl_mode: ssl_mode
.parse()
.map_err(|_| db_error(litellm_db::Error::InvalidSetting("ssl_mode")))?,
ssl_root_cert: value.getattr("ssl_root_cert")?.extract()?,
max_connections: value.getattr("max_connections")?.extract()?,
acquire_timeout: duration(
value.getattr("acquire_timeout")?.extract()?,
"acquire_timeout",
)?,
connect_timeout: duration(
value.getattr("connect_timeout")?.extract()?,
"connect_timeout",
)?,
idle_timeout: idle_timeout
.map(|seconds| duration(seconds, "idle_timeout"))
.transpose()?,
statement_cache_capacity: value.getattr("statement_cache_capacity")?.extract()?,
})
}
// TODO: take typed writer/reader settings and a token provider, then open a sqlx PgPool, once the first
// repository moves to native. The Prisma DATABASE_URL carries Prisma-only params and rotating IAM tokens
#[pymethods]
impl NativeDatabase {
#[staticmethod]
fn connect(py: Python<'_>) -> PyResult<Bound<'_, PyAny>> {
ready_future(py, Bound::new(py, NativeDatabase {})?.as_any())
#[pyo3(signature = (writer_settings=None, reader_settings=None))]
fn connect<'py>(
py: Python<'py>,
writer_settings: Option<&Bound<'py, PyAny>>,
reader_settings: Option<&Bound<'py, PyAny>>,
) -> PyResult<Bound<'py, PyAny>> {
let writer = writer_settings.map(settings).transpose()?;
let reader = reader_settings.map(settings).transpose()?;
let Some(writer) = writer else {
if reader.is_some() {
return Err(db_error(litellm_db::Error::InvalidSetting(
"writer_settings",
)));
}
return ready_future(py, Bound::new(py, Self { pools: None })?.as_any());
};
run_async_value(py, async move {
let pools = DatabasePools::connect(writer, reader)
.await
.map_err(db_error)?;
Ok(Self { pools: Some(pools) })
})
}
fn close<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
ready_future(py, PyNone::get(py).as_any())
let Some(pools) = self.pools.clone() else {
return ready_future(py, PyNone::get(py).as_any());
};
run_async_value(py, async move {
pools.close().await;
Ok(())
})
}
#[getter]
fn reader_available(&self) -> bool {
self.pools
.as_ref()
.is_some_and(DatabasePools::reader_available)
}
#[pyo3(signature = (writer_password=None, reader_password=None))]
fn set_passwords(
&self,
writer_password: Option<String>,
reader_password: Option<String>,
) -> PyResult<()> {
self.pools
.as_ref()
.ok_or_else(|| db_error(litellm_db::Error::NotConfigured))?
.set_passwords(
writer_password.map(SecretValue::from),
reader_password.map(SecretValue::from),
)
.map_err(db_error)
}
#[pyo3(signature = (kind, ids, writer=false))]
fn auth_read<'py>(
&self,
py: Python<'py>,
kind: &str,
ids: Vec<String>,
writer: bool,
) -> PyResult<Bound<'py, PyAny>> {
let kind: AuthRead = kind
.parse()
.map_err(|_| db_error(litellm_db::Error::InvalidAuthArguments))?;
let pools = self
.pools
.clone()
.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, pools.persistent(writer))
.await
.map_err(db_error)?;
Ok(Pythonized(
rows.into_iter().map(|row| row.data).collect::<Vec<_>>(),
))
})
}
fn auth_spend<'py>(
&self,
py: Python<'py>,
kind: &str,
id: String,
) -> PyResult<Bound<'py, PyAny>> {
let kind: AuthSpend = kind
.parse()
.map_err(|_| db_error(litellm_db::Error::InvalidAuthArguments))?;
let pools = self
.pools
.clone()
.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, pools.persistent(false))
.await
.map_err(db_error)
})
}
}

135
litellm-rust/crates/python-bridge/src/lib.rs
View file

@ -13,6 +13,7 @@ mod model_registry;
mod preflight;
mod python_settings;
mod redis;
mod request_auth;
mod router;
mod routes;
mod secrets;
@ -42,7 +43,7 @@ mod _native {
#[pymodule_export]
use crate::redis::client::NativeRedisClient;
#[pymodule_export]
use crate::redis::store::auth_cache::{NativeAuthCache, PyLocalSizes};
use crate::redis::store::auth_cache::{NativeAuthCache, NativeAuthSubscription, PyLocalSizes};
#[pymodule_export]
use crate::redis::store::rate_limits::{
NativeRateLimits, PyAtomicApplied, PyAtomicOverLimit, PyGaugeOverLimit, PyGaugesAcquired,
@ -61,6 +62,48 @@ mod _native {
NativePending, NativeRedisStoreError, router_cooldowns::NativeRouterCooldowns,
};
#[pymodule_export]
use crate::request_auth::{
NativeAuthContinuation, request_auth_access, request_auth_admission,
request_auth_agent_access, request_auth_agent_ceiling, request_auth_agent_ceiling_source,
request_auth_agent_grants, request_auth_agent_group_load, request_auth_agent_invocation,
request_auth_budget_alert, request_auth_budget_alert_plan, request_auth_budget_exemption,
request_auth_budget_price, request_auth_budget_route, request_auth_budget_window_start,
request_auth_budgets, request_auth_claude_membership, request_auth_claude_normalization,
request_auth_codex_model, request_auth_codex_normalization, request_auth_codex_turn,
request_auth_common_budget, request_auth_common_budget_scopes, request_auth_common_flow,
request_auth_common_read, request_auth_context, request_auth_context_end_user,
request_auth_context_member, request_auth_context_models, request_auth_context_percentage,
request_auth_context_temporary, request_auth_context_throttle, request_auth_context_token,
request_auth_custom_flow, request_auth_end_user_apply, request_auth_end_user_assignments,
request_auth_end_user_default_id, request_auth_end_user_lookup,
request_auth_end_user_normalize_id, request_auth_end_user_registry,
request_auth_end_user_validation, request_auth_expired, request_auth_fallback_access,
request_auth_fallback_candidate, request_auth_key_model_access,
request_auth_litellm_header_token, request_auth_llm_route, request_auth_managed,
request_auth_managed_model_grants, request_auth_managed_models, request_auth_member_limit,
request_auth_member_models, request_auth_model_access, request_auth_model_attribution,
request_auth_model_budget_aliases, request_auth_model_budget_check,
request_auth_model_budget_query, request_auth_model_budget_resolve,
request_auth_model_budget_scope, request_auth_model_budget_team_exclusions,
request_auth_model_groups, request_auth_model_helpers, request_auth_model_projection,
request_auth_org_admin, request_auth_org_admin_raw, request_auth_org_route,
request_auth_owner_teams, request_auth_personal_budget,
request_auth_personal_model_listing, request_auth_personal_model_request,
request_auth_personal_models, request_auth_personal_owner, request_auth_reservation,
request_auth_resolved_models, request_auth_route_access, request_auth_route_enable,
request_auth_route_role, request_auth_route_role_api, request_auth_route_role_expand,
request_auth_route_role_inside, request_auth_route_role_jwt,
request_auth_route_role_listed, request_auth_route_scope, request_auth_scheme_credential,
request_auth_select_api_key, request_auth_select_websocket_api_key,
request_auth_spend_counter_policy, request_auth_start, request_auth_team_aliases,
request_auth_team_override, request_auth_unpriced_model,
};
#[pymodule_export]
use crate::request_auth::{
request_auth_budget_lookup_emails, request_auth_budget_lookup_member,
request_auth_budget_lookup_organization, request_auth_budget_lookup_soft,
};
#[pymodule_export]
use crate::router::{RouterCall, RouterCoreHandle};
#[pymodule_export]
use crate::routes::audio_transcription::{atranscription, transcription};
@ -115,6 +158,95 @@ mod tests {
Python::initialize();
Python::attach(|py| {
let mut expected = vec![
"NativeAuthContinuation",
"request_auth_access",
"request_auth_admission",
"request_auth_agent_access",
"request_auth_agent_ceiling",
"request_auth_agent_ceiling_source",
"request_auth_agent_grants",
"request_auth_agent_group_load",
"request_auth_agent_invocation",
"request_auth_budget_alert",
"request_auth_budget_alert_plan",
"request_auth_budget_exemption",
"request_auth_budget_lookup_emails",
"request_auth_budget_lookup_member",
"request_auth_budget_lookup_organization",
"request_auth_budget_lookup_soft",
"request_auth_budget_price",
"request_auth_budget_route",
"request_auth_budget_window_start",
"request_auth_budgets",
"request_auth_common_budget",
"request_auth_common_budget_scopes",
"request_auth_common_flow",
"request_auth_common_read",
"request_auth_context",
"request_auth_context_end_user",
"request_auth_context_member",
"request_auth_context_models",
"request_auth_context_percentage",
"request_auth_context_temporary",
"request_auth_context_throttle",
"request_auth_context_token",
"request_auth_custom_flow",
"request_auth_end_user_apply",
"request_auth_end_user_assignments",
"request_auth_end_user_default_id",
"request_auth_end_user_lookup",
"request_auth_end_user_normalize_id",
"request_auth_end_user_registry",
"request_auth_end_user_validation",
"request_auth_expired",
"request_auth_fallback_access",
"request_auth_fallback_candidate",
"request_auth_key_model_access",
"request_auth_litellm_header_token",
"request_auth_llm_route",
"request_auth_managed",
"request_auth_managed_model_grants",
"request_auth_managed_models",
"request_auth_member_limit",
"request_auth_member_models",
"request_auth_model_access",
"request_auth_model_attribution",
"request_auth_model_budget_check",
"request_auth_model_budget_aliases",
"request_auth_model_budget_scope",
"request_auth_model_budget_team_exclusions",
"request_auth_model_budget_query",
"request_auth_model_budget_resolve",
"request_auth_model_groups",
"request_auth_model_helpers",
"request_auth_model_projection",
"request_auth_org_admin",
"request_auth_org_admin_raw",
"request_auth_org_route",
"request_auth_owner_teams",
"request_auth_personal_budget",
"request_auth_personal_model_listing",
"request_auth_personal_model_request",
"request_auth_personal_models",
"request_auth_personal_owner",
"request_auth_reservation",
"request_auth_resolved_models",
"request_auth_route_access",
"request_auth_route_enable",
"request_auth_route_role",
"request_auth_route_role_api",
"request_auth_route_role_expand",
"request_auth_route_role_inside",
"request_auth_route_role_jwt",
"request_auth_route_role_listed",
"request_auth_scheme_credential",
"request_auth_select_api_key",
"request_auth_select_websocket_api_key",
"request_auth_spend_counter_policy",
"request_auth_start",
"request_auth_team_aliases",
"request_auth_team_override",
"request_auth_unpriced_model",
"RustBridgeDeclined",
"RustUpstreamError",
"ForkedAfterNativeRuntimeStarted",
@ -141,6 +273,7 @@ mod tests {
"NativeRequestRedis",
"NativeSend",
"NativeAuthCache",
"NativeAuthSubscription",
"NativeRateLimits",
"WindowCount",
"AtomicApplied",

92
litellm-rust/crates/python-bridge/src/redis/store/auth_cache.rs
View file

@ -1,5 +1,9 @@
use std::sync::Arc;
use futures_util::{StreamExt, stream::BoxStream};
use litellm_redis_client::SubscriptionEvent;
use tokio::sync::Mutex;
use litellm_python_compat::Value;
use litellm_redis_store::auth_cache::{AuthCache, Settings};
use pyo3::prelude::*;
@ -44,6 +48,35 @@ impl NativeAuthCache {
self.0.has_redis()
}
fn attach(&self, client: &NativeRedisClient) -> bool {
self.0.attach(remote(client))
}
fn current_generation(&self) -> u64 {
self.0.current_generation()
}
#[pyo3(signature = (key, value, generation, ttl_seconds = None))]
fn set_if_generation(
&self,
key: &str,
value: PyValue,
generation: u64,
ttl_seconds: Option<f64>,
) -> PyResult<bool> {
Ok(self
.0
.set_if_generation(key, value.0, generation, ttl(ttl_seconds)?))
}
fn subscribe(&self, client: &NativeRedisClient, channel: &str) -> NativeAuthSubscription {
self.0.subscriber_connected(false);
NativeAuthSubscription {
cache: self.0.clone(),
stream: Arc::new(Mutex::new(Some(client.client.subscribe_events(channel)))),
}
}
/// `request` answers a key its pre-call reads already found absent without another `GET`.
#[pyo3(signature = (key, request = None))]
fn get<'py>(
@ -245,6 +278,65 @@ impl NativeAuthCache {
}
}
#[pyclass(frozen)]
pub(crate) struct NativeAuthSubscription {
cache: Arc<AuthCache>,
stream: Arc<Mutex<Option<BoxStream<'static, SubscriptionEvent>>>>,
}
#[pymethods]
impl NativeAuthSubscription {
fn is_connected(&self) -> bool {
self.cache.local_ready()
}
fn next<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
let cache = self.cache.clone();
let stream = self.stream.clone();
litellm_host_python::run_async_value(py, async move {
let mut guard = stream.lock().await;
let Some(stream) = guard.as_mut() else {
return Ok(None::<Py<pyo3::types::PyBytes>>);
};
match stream.next().await {
Some(SubscriptionEvent::Connected) => cache.subscriber_connected(true),
Some(SubscriptionEvent::Message(payload)) => {
if let Ok(message) = serde_json::from_slice::<serde_json::Value>(&payload)
&& let Some(key) =
message.get("cache_key").and_then(serde_json::Value::as_str)
&& message
.get("new_value")
.and_then(serde_json::Value::as_f64)
.is_none()
{
cache.delete_local(key).map_err(kind_error)?;
}
return Ok(Some(Python::attach(|py| {
pyo3::types::PyBytes::new(py, &payload).unbind()
})));
}
_ => cache.subscriber_connected(false),
}
Ok(None)
})
}
fn close<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
self.cache.subscriber_connected(false);
let stream = self.stream.clone();
litellm_host_python::run_async_value(py, async move {
stream.lock().await.take();
Ok(())
})
}
}
impl Drop for NativeAuthSubscription {
fn drop(&mut self) {
self.cache.subscriber_connected(false);
}
}
/// Entries the management-object and the user-key partitions hold in memory.
#[pyclass(frozen, eq, get_all, name = "LocalSizes")]
#[derive(PartialEq)]

1086
litellm-rust/crates/python-bridge/src/request_auth.rs Normal file
View file

File diff suppressed because it is too large Load diff

4
litellm-rust/crates/redis-client/src/client.rs
View file

@ -155,6 +155,10 @@ impl RedisClient {
self.0.subscriber.clone().messages(channel.to_owned())
}
pub fn subscribe_events(&self, channel: &str) -> BoxStream<'static, crate::SubscriptionEvent> {
self.0.subscriber.clone().events(channel.to_owned())
}
pub fn pipeline(&self) -> Pipeline {
Pipeline::new(self.0.namespace.clone())
}

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

@ -23,6 +23,7 @@ pub use litellm_redis_types::{
};
pub use namespace::Namespace;
pub use pipeline::Pipeline;
pub use pubsub::SubscriptionEvent;
pub use redis::{
BasicAuth, ErrorKind as RedisErrorKind, RedisError, StreamingCredentialsProvider, Value,
cluster_routing::Slot,

89
litellm-rust/crates/redis-client/src/pubsub.rs
View file

@ -1,7 +1,10 @@
use std::{sync::Arc, time::Duration};
use futures_util::{StreamExt, stream::BoxStream};
use redis::{ConnectionInfo, RedisResult, aio::PubSubStream};
use redis::{
ConnectionInfo, RedisResult,
aio::{PubSubSink, PubSubStream},
};
use crate::{
ConnectError, Credentials, Standalone, Target, Timeouts, config,
@ -11,6 +14,18 @@ use crate::{
// litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py: _BACKOFF_INITIAL_SECONDS, _BACKOFF_MAX_SECONDS
const BACKOFF_INITIAL: Duration = Duration::from_secs(5);
const BACKOFF_MAX: Duration = Duration::from_secs(60);
const IDLE_PROBE: Duration = Duration::from_secs(15);
pub enum SubscriptionEvent {
Connected,
Disconnected,
Message(Vec<u8>),
}
struct Live {
sink: PubSubSink,
stream: PubSubStream,
}
/// A subscription the way redis-py's `PubSub` holds one: a dedicated RESP2 connection that sends
/// `SUBSCRIBE`, authenticated once when it connects.
@ -66,43 +81,76 @@ impl Subscriber {
}
pub(crate) fn messages(self, channel: String) -> BoxStream<'static, Vec<u8>> {
self.events(channel)
.filter_map(|event| async move {
match event {
SubscriptionEvent::Message(payload) => Some(payload),
_ => None,
}
})
.boxed()
}
pub(crate) fn events(self, channel: String) -> BoxStream<'static, SubscriptionEvent> {
futures_util::stream::unfold(
(self, channel, None::<PubSubStream>, BACKOFF_INITIAL),
|(subscriber, channel, live, backoff)| async move {
(self, channel, None::<Live>, BACKOFF_INITIAL, false),
|(subscriber, channel, live, backoff, retry)| async move {
let mut live = live;
let mut backoff = backoff;
if retry {
tokio::time::sleep(backoff).await;
}
loop {
let Some(current) = live.as_mut() else {
match subscriber.subscribe(&channel).await {
let event = match subscriber.subscribe(&channel).await {
Ok(subscribed) => {
live = Some(subscribed);
backoff = BACKOFF_INITIAL;
SubscriptionEvent::Connected
}
Err(error) => {
tracing::warn!(%error, ?backoff, "Redis subscribe failed, retrying");
tokio::time::sleep(backoff).await;
backoff = (backoff * 2).min(BACKOFF_MAX);
Err(_) => {
if retry {
backoff = (backoff * 2).min(BACKOFF_MAX);
}
SubscriptionEvent::Disconnected
}
}
continue;
};
let retry = live.is_none();
return Some((event, (subscriber, channel, live, backoff, retry)));
};
match current.next().await {
Some(message) => {
match tokio::time::timeout(IDLE_PROBE, current.stream.next()).await {
Ok(Some(message)) => {
let payload = message.get_payload_bytes().to_vec();
return Some((payload, (subscriber, channel, live, backoff)));
return Some((
SubscriptionEvent::Message(payload),
(subscriber, channel, live, backoff, false),
));
}
None => live = None,
Err(_) => {
if matches!(
tokio::time::timeout(
subscriber.connect_timeout,
current.sink.ping::<String>()
)
.await,
Ok(Ok(_))
) {
continue;
}
}
Ok(None) => {}
}
tracing::warn!(?backoff, "Redis subscription connection lost, resubscribing");
tokio::time::sleep(backoff).await;
backoff = (backoff * 2).min(BACKOFF_MAX);
return Some((
SubscriptionEvent::Disconnected,
(subscriber, channel, None, backoff, true),
));
}
},
)
.boxed()
}
async fn subscribe(&self, channel: &str) -> RedisResult<PubSubStream> {
async fn subscribe(&self, channel: &str) -> RedisResult<Live> {
match &self.source {
Source::Sentinel(sentinel) => {
let master = sentinel.master().await?;
@ -122,7 +170,7 @@ impl Subscriber {
}
}
async fn subscribe_on(&self, info: ConnectionInfo, channel: &str) -> RedisResult<PubSubStream> {
async fn subscribe_on(&self, info: ConnectionInfo, channel: &str) -> RedisResult<Live> {
let settings =
config::authenticated(info.redis_settings().clone(), &self.credentials).await?;
let client = redis::Client::open(info.set_redis_settings(settings))?;
@ -135,7 +183,8 @@ impl Subscriber {
))
})??;
subscribe_connected(&mut pubsub, channel, self.connect_timeout).await?;
Ok(pubsub.into_on_message())
let (sink, stream) = pubsub.split();
Ok(Live { sink, stream })
}
}

81
litellm-rust/crates/redis-client/tests/pubsub.rs
View file

@ -7,7 +7,7 @@ use litellm_circuit_breaker::{Breaker, Config};
use litellm_redis_client::commands;
use litellm_redis_client::{
Credentials, DEFAULT_TIMEOUT, Error, Namespace, Node, RedisClient, RedisConfig, Sentinel,
Target, Timeouts,
SubscriptionEvent, Target, Timeouts,
};
use litellm_redis_testing::{RedisDeployment, RedisSentinel, SERVICE_NAME, Topology, eventually};
use rstest::rstest;
@ -16,6 +16,85 @@ use tokio::sync::mpsc::{UnboundedReceiver, unbounded_channel};
const CHANNEL: &str = "litellm_proxy.auth_cache_invalidation";
#[rstest]
#[tokio::test]
async fn a_silent_dead_subscription_is_detected_and_reconnected(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let server = RedisDeployment::start(topology).await.unwrap();
let proxy = support::proxied(&server).await;
let client = RedisClient::new(RedisConfig {
target: proxy.target.clone(),
credentials: Credentials::NONE,
namespace: Namespace::default(),
timeouts: Timeouts::new(Duration::from_millis(500), None),
breaker: None,
events: None,
health_check: None,
keepalive: true,
})
.unwrap();
let mut events = client.subscribe_events(CHANNEL);
assert!(matches!(
tokio::time::timeout(Duration::from_secs(5), events.next())
.await
.unwrap(),
Some(SubscriptionEvent::Connected)
));
proxy.blackhole_open_connections();
assert!(matches!(
tokio::time::timeout(Duration::from_secs(20), events.next())
.await
.unwrap(),
Some(SubscriptionEvent::Disconnected)
));
assert!(matches!(
tokio::time::timeout(Duration::from_secs(10), events.next())
.await
.unwrap(),
Some(SubscriptionEvent::Connected)
));
}
#[rstest]
#[tokio::test]
async fn lifecycle_reports_disconnect_before_reconnect(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let server = RedisDeployment::start(topology).await.unwrap();
let client = connect(&server, "", Duration::from_millis(500));
let mut events = client.subscribe_events(CHANNEL);
assert!(matches!(
tokio::time::timeout(Duration::from_secs(5), events.next())
.await
.unwrap(),
Some(SubscriptionEvent::Connected)
));
server.restart().await.unwrap();
assert!(matches!(
tokio::time::timeout(Duration::from_secs(5), events.next())
.await
.unwrap(),
Some(SubscriptionEvent::Disconnected)
));
assert!(matches!(
tokio::time::timeout(Duration::from_secs(15), events.next())
.await
.unwrap(),
Some(SubscriptionEvent::Connected)
));
eventually(async || {
client
.send(commands::publish(CHANNEL, b"after".as_slice()))
.await
.ok()
})
.await;
assert!(
matches!(tokio::time::timeout(Duration::from_secs(5), events.next()).await.unwrap(), Some(SubscriptionEvent::Message(bytes)) if bytes == b"after")
);
}
fn listen(client: &RedisClient, channel: &str) -> UnboundedReceiver<Vec<u8>> {
let (sender, receiver) = unbounded_channel();
let mut messages = client.subscribe(channel);

1
litellm-rust/crates/redis-store/Cargo.toml
View file

@ -11,6 +11,7 @@ litellm-cache.workspace = true
litellm-cache-memory.workspace = true
litellm-python-compat.workspace = true
litellm-redis-client.workspace = true
moka = { workspace = true, features = ["sync"] }
strum.workspace = true
thiserror.workspace = true
tokio = { workspace = true, features = ["macros", "rt", "sync", "time"] }

Some files were not shown because too many files have changed in this diff Show more