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

systemprompt_users/repository/
api_key.rs
