systemprompt_users/repository/
api_key.rs1use 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}