Skip to main content

litellm_rs/storage/database/entities/
api_key.rs

1use 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/// API key database model
22#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Eq, Serialize, Deserialize)]
23#[sea_orm(table_name = "api_keys")]
24pub struct Model {
25    /// API key ID (UUID)
26    #[sea_orm(primary_key, auto_increment = false)]
27    pub id: Uuid,
28
29    /// Display name
30    pub name: String,
31
32    /// Hashed secret value
33    #[sea_orm(unique)]
34    pub key_hash: String,
35
36    /// Prefix shown to users
37    pub key_prefix: String,
38
39    /// Optional owner user
40    pub user_id: Option<Uuid>,
41
42    /// Optional owner team
43    pub team_id: Option<Uuid>,
44
45    /// JSON array serialized as text
46    pub permissions: String,
47
48    /// JSON object serialized as text
49    pub rate_limits: Option<String>,
50
51    /// Optional expiration timestamp
52    pub expires_at: Option<DateTimeWithTimeZone>,
53
54    /// Whether key is active
55    pub is_active: bool,
56
57    /// Last usage timestamp
58    pub last_used_at: Option<DateTimeWithTimeZone>,
59
60    /// Usage statistics JSON serialized as text
61    pub usage_stats: String,
62
63    /// Metadata extra fields JSON serialized as text
64    pub extra: Option<String>,
65
66    /// Creation timestamp
67    pub created_at: DateTimeWithTimeZone,
68
69    /// Last update timestamp
70    pub updated_at: DateTimeWithTimeZone,
71
72    /// Optimistic lock version
73    pub version: i32,
74}
75
76/// API key entity relations
77#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
78pub enum Relation {
79    /// API key belongs to a user
80    #[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    /// Convert SeaORM model to domain API key model
100    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    /// Convert domain API key model to SeaORM active model
140    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}