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