litellm_rs/storage/database/entities/
api_key.rs1use sea_orm::Set;
2use sea_orm::entity::prelude::*;
3use serde::de::DeserializeOwned;
4use serde::{Deserialize, Serialize};
5use std::collections::HashMap;
6
7use crate::utils::error::gateway_error::{GatewayError, Result as GatewayResult};
8
9fn parse_json_field<T: DeserializeOwned>(raw: &str, field: &str) -> GatewayResult<T> {
10 serde_json::from_str(raw).map_err(|_| {
11 GatewayError::Serialization(format!("Invalid persisted API key JSON field '{field}'"))
12 })
13}
14
15fn serialize_json_field<T: Serialize>(value: &T, field: &str) -> GatewayResult<String> {
16 serde_json::to_string(value).map_err(|_| {
17 GatewayError::Serialization(format!("Failed to serialize API key field '{field}'"))
18 })
19}
20
21#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Eq, Serialize, Deserialize)]
23#[sea_orm(table_name = "api_keys")]
24pub struct Model {
25 #[sea_orm(primary_key, auto_increment = false)]
27 pub id: Uuid,
28
29 pub name: String,
31
32 #[sea_orm(unique)]
34 pub key_hash: String,
35
36 pub key_prefix: String,
38
39 pub user_id: Option<Uuid>,
41
42 pub team_id: Option<Uuid>,
44
45 pub permissions: String,
47
48 pub rate_limits: Option<String>,
50
51 pub expires_at: Option<DateTimeWithTimeZone>,
53
54 pub is_active: bool,
56
57 pub last_used_at: Option<DateTimeWithTimeZone>,
59
60 pub usage_stats: String,
62
63 pub extra: Option<String>,
65
66 pub created_at: DateTimeWithTimeZone,
68
69 pub updated_at: DateTimeWithTimeZone,
71
72 pub version: i32,
74}
75
76#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
78pub enum Relation {
79 #[sea_orm(
81 belongs_to = "super::user::Entity",
82 from = "Column::UserId",
83 to = "super::user::Column::Id",
84 on_update = "NoAction",
85 on_delete = "Restrict"
86 )]
87 User,
88}
89
90impl Related<super::user::Entity> for Entity {
91 fn to() -> RelationDef {
92 Relation::User.def()
93 }
94}
95
96impl ActiveModelBehavior for ActiveModel {}
97
98impl Model {
99 pub fn to_domain_api_key(&self) -> GatewayResult<crate::core::models::ApiKey> {
101 use crate::core::models::{Metadata, RateLimits, UsageStats};
102
103 let permissions = parse_json_field::<Vec<String>>(&self.permissions, "permissions")?;
104 let rate_limits = self
105 .rate_limits
106 .as_deref()
107 .map(|raw| parse_json_field::<RateLimits>(raw, "rate_limits"))
108 .transpose()?;
109 let usage_stats = parse_json_field::<UsageStats>(&self.usage_stats, "usage_stats")?;
110 let extra = match self.extra.as_deref() {
111 Some(raw) => parse_json_field::<HashMap<String, serde_json::Value>>(raw, "extra")?,
112 None => HashMap::new(),
113 };
114
115 let metadata = Metadata {
116 id: self.id,
117 created_at: self.created_at.naive_utc().and_utc(),
118 updated_at: self.updated_at.naive_utc().and_utc(),
119 version: self.version as i64,
120 extra,
121 };
122
123 Ok(crate::core::models::ApiKey {
124 metadata,
125 name: self.name.clone(),
126 key_hash: self.key_hash.clone(),
127 key_prefix: self.key_prefix.clone(),
128 user_id: self.user_id,
129 team_id: self.team_id,
130 permissions,
131 rate_limits,
132 expires_at: self.expires_at.map(|dt| dt.naive_utc().and_utc()),
133 is_active: self.is_active,
134 last_used_at: self.last_used_at.map(|dt| dt.naive_utc().and_utc()),
135 usage_stats,
136 })
137 }
138
139 pub fn from_domain_api_key(
141 api_key: &crate::core::models::ApiKey,
142 ) -> GatewayResult<ActiveModel> {
143 let permissions = serialize_json_field(&api_key.permissions, "permissions")?;
144 let rate_limits = api_key
145 .rate_limits
146 .as_ref()
147 .map(|limits| serialize_json_field(limits, "rate_limits"))
148 .transpose()?;
149 let usage_stats = serialize_json_field(&api_key.usage_stats, "usage_stats")?;
150 let extra = if api_key.metadata.extra.is_empty() {
151 None
152 } else {
153 Some(serialize_json_field(&api_key.metadata.extra, "extra")?)
154 };
155
156 Ok(ActiveModel {
157 id: Set(api_key.metadata.id),
158 name: Set(api_key.name.clone()),
159 key_hash: Set(api_key.key_hash.clone()),
160 key_prefix: Set(api_key.key_prefix.clone()),
161 user_id: Set(api_key.user_id),
162 team_id: Set(api_key.team_id),
163 permissions: Set(permissions),
164 rate_limits: Set(rate_limits),
165 expires_at: Set(api_key.expires_at.map(Into::into)),
166 is_active: Set(api_key.is_active),
167 last_used_at: Set(api_key.last_used_at.map(Into::into)),
168 usage_stats: Set(usage_stats),
169 extra: Set(extra),
170 created_at: Set(api_key.metadata.created_at.into()),
171 updated_at: Set(api_key.metadata.updated_at.into()),
172 version: Set(api_key.metadata.version as i32),
173 })
174 }
175}