Skip to main content

systemprompt_users/repository/
api_key.rs

1//! API-key persistence on the user repository.
2//!
3//! Copyright (c) systemprompt.io — Business Source License 1.1.
4//! See <https://systemprompt.io> for licensing details.
5
6use chrono::{DateTime, Utc};
7use systemprompt_identifiers::{ApiKeyId, UserId};
8use systemprompt_models::attribution::ScopeBinding;
9
10use crate::error::Result;
11use crate::models::{ApiKeyLimits, UserApiKey, UserApiKeyRow};
12use crate::repository::UserRepository;
13
14#[derive(Debug)]
15pub struct CreateApiKeyParams<'a> {
16    pub id: &'a ApiKeyId,
17    pub user_id: &'a UserId,
18    pub name: &'a str,
19    pub key_prefix: &'a str,
20    pub key_hash: &'a str,
21    pub expires_at: Option<DateTime<Utc>>,
22    pub limits: &'a ApiKeyLimits,
23    pub scopes: &'a [ScopeBinding],
24    pub purpose: Option<&'a str>,
25}
26
27impl UserRepository {
28    pub async fn create_api_key(&self, params: CreateApiKeyParams<'_>) -> Result<UserApiKey> {
29        let mut tx = self.write_pool.begin().await?;
30        let row = sqlx::query_as!(
31            UserApiKeyRow,
32            r#"
33            INSERT INTO user_api_keys
34                (id, user_id, name, key_prefix, key_hash, expires_at,
35                 model_allowlist, budget_microdollars, max_requests, request_window_seconds,
36                 purpose)
37            VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11)
38            RETURNING id, user_id, name, key_prefix, key_hash,
39                      created_at, last_used_at, expires_at, revoked_at,
40                      model_allowlist, budget_microdollars, max_requests, request_window_seconds,
41                      purpose,
42                      ARRAY[]::TEXT[] AS "scope_dimensions!", ARRAY[]::TEXT[] AS "scope_values!"
43            "#,
44            params.id.as_str(),
45            params.user_id.as_str(),
46            params.name,
47            params.key_prefix,
48            params.key_hash,
49            params.expires_at,
50            params.limits.model_allowlist.as_deref(),
51            params.limits.budget_microdollars,
52            params.limits.max_requests,
53            params.limits.request_window_seconds,
54            params.purpose,
55        )
56        .fetch_one(&mut *tx)
57        .await?;
58        if !params.scopes.is_empty() {
59            let dimensions: Vec<String> = params
60                .scopes
61                .iter()
62                .map(|s| s.dimension.as_str().to_owned())
63                .collect();
64            let values: Vec<String> = params.scopes.iter().map(|s| s.value.clone()).collect();
65            sqlx::query!(
66                r#"
67                INSERT INTO user_api_key_scopes (key_id, dimension, value)
68                SELECT $1, t.dimension, t.value
69                FROM UNNEST($2::text[], $3::text[]) AS t(dimension, value)
70                "#,
71                params.id.as_str(),
72                &dimensions,
73                &values,
74            )
75            .execute(&mut *tx)
76            .await?;
77        }
78        tx.commit().await?;
79        let mut key = UserApiKey::from(row);
80        key.scopes = params.scopes.to_vec();
81        Ok(key)
82    }
83
84    pub async fn find_active_api_key_by_prefix(
85        &self,
86        key_prefix: &str,
87    ) -> Result<Option<UserApiKey>> {
88        let row = sqlx::query_as!(
89            UserApiKeyRow,
90            r#"
91            SELECT k.id, k.user_id, k.name, k.key_prefix, k.key_hash,
92                   k.created_at, k.last_used_at, k.expires_at, k.revoked_at,
93                   k.model_allowlist, k.budget_microdollars, k.max_requests,
94                   k.request_window_seconds, k.purpose,
95                   ARRAY(SELECT s.dimension FROM user_api_key_scopes s
96                         WHERE s.key_id = k.id ORDER BY s.dimension) AS "scope_dimensions!",
97                   ARRAY(SELECT s.value FROM user_api_key_scopes s
98                         WHERE s.key_id = k.id ORDER BY s.dimension) AS "scope_values!"
99            FROM user_api_keys k
100            WHERE k.key_prefix = $1
101              AND k.revoked_at IS NULL
102            "#,
103            key_prefix,
104        )
105        .fetch_optional(&*self.write_pool)
106        .await
107        .map(|row| row.map(UserApiKey::from))?;
108        Ok(row)
109    }
110
111    pub async fn list_api_keys_for_user(&self, user_id: &UserId) -> Result<Vec<UserApiKey>> {
112        let rows = sqlx::query_as!(
113            UserApiKeyRow,
114            r#"
115            SELECT k.id, k.user_id, k.name, k.key_prefix, k.key_hash,
116                   k.created_at, k.last_used_at, k.expires_at, k.revoked_at,
117                   k.model_allowlist, k.budget_microdollars, k.max_requests,
118                   k.request_window_seconds, k.purpose,
119                   ARRAY(SELECT s.dimension FROM user_api_key_scopes s
120                         WHERE s.key_id = k.id ORDER BY s.dimension) AS "scope_dimensions!",
121                   ARRAY(SELECT s.value FROM user_api_key_scopes s
122                         WHERE s.key_id = k.id ORDER BY s.dimension) AS "scope_values!"
123            FROM user_api_keys k
124            WHERE k.user_id = $1
125            ORDER BY k.created_at DESC
126            "#,
127            user_id.as_str(),
128        )
129        .fetch_all(&*self.pool)
130        .await
131        .map(|rows| rows.into_iter().map(UserApiKey::from).collect::<Vec<_>>())?;
132        Ok(rows)
133    }
134
135    pub async fn revoke_api_key(&self, id: &ApiKeyId, user_id: &UserId) -> Result<bool> {
136        let result = sqlx::query!(
137            r#"
138            UPDATE user_api_keys
139            SET revoked_at = CURRENT_TIMESTAMP
140            WHERE id = $1 AND user_id = $2 AND revoked_at IS NULL
141            "#,
142            id.as_str(),
143            user_id.as_str(),
144        )
145        .execute(&*self.write_pool)
146        .await?;
147        Ok(result.rows_affected() > 0)
148    }
149
150    pub async fn list_revoked_api_key_ids_for_user(&self, user_id: &UserId) -> Result<Vec<String>> {
151        let rows = sqlx::query_scalar!(
152            r#"
153            SELECT id
154            FROM user_api_keys
155            WHERE user_id = $1 AND revoked_at IS NOT NULL
156            ORDER BY revoked_at DESC
157            "#,
158            user_id.as_str(),
159        )
160        .fetch_all(&*self.write_pool)
161        .await?;
162        Ok(rows)
163    }
164
165    pub async fn touch_api_key_usage(&self, id: &ApiKeyId) -> Result<()> {
166        sqlx::query!(
167            r#"
168            UPDATE user_api_keys
169            SET last_used_at = CURRENT_TIMESTAMP
170            WHERE id = $1
171            "#,
172            id.as_str(),
173        )
174        .execute(&*self.write_pool)
175        .await?;
176        Ok(())
177    }
178}