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 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}