litellm_rs/core/keys/
database_repository.rs

1//! Database-backed API key repository implementation.
2
3use super::db_mapping::{from_domain_api_key, to_domain_api_key};
4use super::db_update::apply_update_config;
5use super::repository::KeyRepository;
6use super::types::{KeyStatus, ManagedApiKey, UpdateKeyConfig, UsageRecord};
7use crate::storage::StorageLayer;
8use crate::utils::error::gateway_error::{GatewayError, Result};
9use async_trait::async_trait;
10use std::sync::Arc;
11use uuid::Uuid;
12
13#[derive(Debug, Clone)]
14pub struct DatabaseKeyRepository {
15    storage: Arc<StorageLayer>,
16}
17
18impl DatabaseKeyRepository {
19    pub fn new(storage: Arc<StorageLayer>) -> Self {
20        Self { storage }
21    }
22}
23
24#[async_trait]
25impl KeyRepository for DatabaseKeyRepository {
26    async fn create(&self, key: ManagedApiKey) -> Result<ManagedApiKey> {
27        let domain_key = to_domain_api_key(&key)?;
28        let stored = self.storage.db().create_api_key(&domain_key).await?;
29        Ok(from_domain_api_key(&stored))
30    }
31
32    async fn find_by_id(&self, id: Uuid) -> Result<Option<ManagedApiKey>> {
33        let key = self.storage.db().find_api_key_by_id(id).await?;
34        Ok(key.map(|k| from_domain_api_key(&k)))
35    }
36
37    async fn find_by_hash(&self, key_hash: &str) -> Result<Option<ManagedApiKey>> {
38        let key = self.storage.db().find_api_key_by_hash(key_hash).await?;
39        Ok(key.map(|k| from_domain_api_key(&k)))
40    }
41
42    async fn update(&self, id: Uuid, config: UpdateKeyConfig) -> Result<ManagedApiKey> {
43        let mut domain_key = self
44            .storage
45            .db()
46            .find_api_key_by_id(id)
47            .await?
48            .ok_or_else(|| GatewayError::not_found("API key not found"))?;
49
50        apply_update_config(&mut domain_key, config)?;
51        let updated = self.storage.db().update_api_key(&domain_key).await?;
52        Ok(from_domain_api_key(&updated))
53    }
54
55    async fn update_status(&self, id: Uuid, status: KeyStatus) -> Result<()> {
56        match status {
57            KeyStatus::Active => self.storage.db().set_api_key_active(id, true).await,
58            KeyStatus::Revoked => self.storage.db().set_api_key_active(id, false).await,
59            KeyStatus::Expired => Err(GatewayError::validation(
60                "Cannot set key status to expired manually",
61            )),
62        }
63    }
64
65    async fn update_last_used(&self, id: Uuid) -> Result<()> {
66        self.storage.db().update_api_key_last_used(id).await
67    }
68
69    async fn update_usage(&self, id: Uuid, record: UsageRecord) -> Result<()> {
70        self.storage
71            .db()
72            .update_api_key_usage(
73                id,
74                record.requests,
75                record.tokens,
76                record.cost,
77                record.unpriced,
78                record.pricing_policy.as_deref(),
79            )
80            .await
81    }
82
83    async fn list_by_user(&self, user_id: Uuid) -> Result<Vec<ManagedApiKey>> {
84        let keys = self.storage.db().list_api_keys_by_user(user_id).await?;
85        Ok(keys.into_iter().map(|k| from_domain_api_key(&k)).collect())
86    }
87
88    async fn list_by_team(&self, team_id: Uuid) -> Result<Vec<ManagedApiKey>> {
89        let keys = self.storage.db().list_api_keys_by_team(team_id).await?;
90        Ok(keys.into_iter().map(|k| from_domain_api_key(&k)).collect())
91    }
92
93    async fn list_all(
94        &self,
95        status: Option<KeyStatus>,
96        limit: Option<usize>,
97        offset: Option<usize>,
98    ) -> Result<Vec<ManagedApiKey>> {
99        let keys = self
100            .storage
101            .db()
102            .list_api_keys(status, limit, offset)
103            .await?;
104        Ok(keys.into_iter().map(|k| from_domain_api_key(&k)).collect())
105    }
106
107    async fn delete(&self, id: Uuid) -> Result<()> {
108        self.storage.db().delete_api_key(id).await
109    }
110
111    async fn delete_expired(&self) -> Result<u64> {
112        self.storage.db().delete_expired_api_keys().await
113    }
114
115    async fn count(&self, status: Option<KeyStatus>) -> Result<u64> {
116        self.storage.db().count_api_keys(status).await
117    }
118}
litellm_rs/core/keys/database_repository.rs

litellm_rs/core/keys/
database_repository.rs
