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

feat(db): add bounded SQLx auth reads and native pools

Browse source
This commit is contained in:
Devin AI 2026-11-09 02:13:48 +00:00
parent dd55f2c292
commit 684d945bcf
83 changed files with 2307 additions and 9 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
litellm-rust/Cargo.lock generated
View file

@ -3768,6 +3768,7 @@ name = "litellm-db"
version = "0.1.0"
dependencies = [
"chrono",
"litellm-secret-value",
"macro_rules_attribute",
"rstest",
"serde",
@ -3776,6 +3777,7 @@ dependencies = [
"strum",
"thiserror 2.0.19",
"tokio",
"veil",
]
[[package]]
@ -4389,6 +4391,7 @@ dependencies = [
"litellm-callbacks-legacy-python",
"litellm-circuit-breaker",
"litellm-core-utils",
"litellm-db",
"litellm-error",
"litellm-host",
"litellm-host-python",

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

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

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

@ -6,6 +6,9 @@ license.workspace = true
repository.workspace = true
[dependencies]
veil.workspace = true
litellm-secret-value.workspace = true
tokio.workspace = true
chrono.workspace = true
macro_rules_attribute.workspace = true
serde.workspace = true
@ -16,4 +19,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

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

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

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

@ -0,0 +1,56 @@
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, created_by, updated_by)
VALUES ('auth-agent', 'Auth Agent', '{}', 'auth-op', true, false, NULL, 'auth-test', 'auth-test');
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

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

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

@ -0,0 +1,745 @@
use crate::Error;
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 = "tag")]
Tag,
#[strum(serialize = "model_access_group")]
ModelAccessGroup,
#[strum(serialize = "agent")]
Agent,
#[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],
) -> Result<Vec<AuthRow>, Error> {
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::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::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")
})?;
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,
) -> Result<Option<f64>, Error> {
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::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::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)).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)]).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)]).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]).await?;
let user = auth_read(&mut *tx, AuthRead::User, &strings(&["auth-u"])).await?;
let team = auth_read(&mut *tx, AuthRead::Team, &strings(&["auth-team"])).await?;
let member = auth_read(
&mut *tx,
AuthRead::Member,
&strings(&["auth-u", "auth-team"]),
)
.await?;
let org = auth_read(&mut *tx, AuthRead::Org, &strings(&["auth-org"])).await?;
let project = auth_read(&mut *tx, AuthRead::Project, &strings(&["auth-project"])).await?;
assert_eq!(rows.len(), 1);
assert_eq!(
rows[0].data.0,
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()],
)
.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"]),
)
.await?;
assert_eq!(rows.len(), 1);
assert_eq!(rows[0].data.0["user_id"], "auth-u");
assert!(auth_read(&mut *tx, AuthRead::User, &[]).await?.is_empty());
assert!(
auth_read(
&mut *tx,
AuthRead::Member,
&strings(&["auth-other", "auth-team"])
)
.await?
.is_empty()
);
assert!(
auth_read(
&mut *tx,
AuthRead::AgentIdentity,
&strings(&["microsoft_entra", "other-tenant", "auth-client"])
)
.await?
.is_empty()
);
assert!(
auth_read(
&mut *tx,
AuthRead::UserEmail,
&strings(&["test%@example.com"])
)
.await?
.is_empty()
);
assert!(
auth_read(&mut *tx, AuthRead::User, &strings(&["' OR 1=1 --"]))
.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).await?, Some(expected));
let absent = match kind {
AuthSpend::Key => "f".repeat(64),
AuthSpend::Member => "missing:missing".into(),
_ => "missing".into(),
};
assert_eq!(auth_spend(&mut *tx, kind, &absent).await?, None);
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).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).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"])).await,
Err(crate::Error::Query(sqlx::Error::PoolClosed))
));
assert!(matches!(
auth_spend(&pool, AuthSpend::User, "id").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]"));
}
}

8
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")]

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

@ -6,18 +6,22 @@ macro_rules_attribute::attribute_alias! {
#[derive(Clone, Debug, PartialEq, sqlx::FromRow)];
}
mod auth;
mod daily_spend;
mod end_user;
mod error;
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,

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

@ -0,0 +1,193 @@
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,
}
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()),
})
}
pub fn reader_available(&self) -> bool {
self.reader_available
}
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;
}
}
}

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

@ -19,6 +19,7 @@ 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

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

@ -1,19 +1,170 @@
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::{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 {
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)
.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)
.await
.map_err(db_error)
})
}
}