alibi_sqlx/store/
api_keys.rs

1use super::entities::api_key::Model;
2use super::entities::api_key_start::ApiKeyStart;
3use super::{SqlxStore, lock_exclusive};
4use crate::error::{record_not_inserted, record_not_updated};
5use crate::model::{self, ActiveRow, SqlxModel};
6use crate::pool::Exec;
7use crate::schema::AuthSchema;
8use crate::sql::Sql;
9use crate::value::SqlxValue;
10use alibi_core::error::{AuthError, AuthResult};
11use alibi_core::store::adapter::parse_optional_rfc3339;
12use alibi_core::store::{ApiKeyStore, ConsumeApiKeyResult};
13use alibi_core::types::{ApiKey, CreateApiKey, UpdateApiKey};
14use async_trait::async_trait;
15use chrono::Utc;
16use uuid::Uuid;
17
18#[async_trait]
19impl<S> ApiKeyStore for SqlxStore<S>
20where
21    S: AuthSchema + Send + Sync,
22{
23    async fn consume_api_key_usage_from_snapshot(
24        &self,
25        observed: &ApiKey,
26        global_rate_limit_enabled: bool,
27    ) -> AuthResult<ConsumeApiKeyResult> {
28        self.consume_usage_phases(observed, global_rate_limit_enabled)
29            .await
30    }
31
32    async fn create_api_key(&self, input: CreateApiKey) -> AuthResult<ApiKey> {
33        let now = Utc::now();
34        let start = input
35            .start
36            .map(|start| ApiKeyStart::prepare(&start, self.exec().engine()))
37            .transpose()?;
38        let sqlite_cast = start
39            .as_ref()
40            .is_some_and(ApiKeyStart::requires_sqlite_cast);
41        let mut active = ActiveRow::new();
42        active.set(
43            "id",
44            self.generated_id(self.exec(), "apikey", "api_keys", "id")
45                .await?
46                .unwrap_or_else(|| Uuid::new_v4().to_string()),
47        );
48        active.set("name", input.name);
49        active.set("start", start.into_sql_value());
50        active.set("prefix", input.prefix);
51        active.set("key", input.key_hash);
52        active.set("reference_id", input.reference_id);
53        active.set("config_id", input.config_id);
54        active.set("refill_interval", input.refill_interval);
55        active.set("refill_amount", input.refill_amount);
56        active.set("last_refill_at", None::<chrono::DateTime<Utc>>);
57        active.set("enabled", input.enabled);
58        active.set("rate_limit_enabled", input.rate_limit_enabled);
59        active.set("rate_limit_time_window", input.rate_limit_time_window);
60        active.set("rate_limit_max", input.rate_limit_max);
61        active.set("request_count", Some(0.0_f64));
62        active.set("remaining", input.remaining);
63        active.set("last_request", None::<chrono::DateTime<Utc>>);
64        active.set(
65            "expires_at",
66            parse_optional_rfc3339(input.expires_at.as_deref(), "expires_at")?,
67        );
68        active.set("created_at", now);
69        active.set("updated_at", now);
70        active.set("permissions", input.permissions);
71        active.set("metadata", input.metadata);
72        let inserted = if sqlite_cast {
73            // Only an invalid-surrogate SQLite start needs a bytes-to-TEXT expression.
74            let mut sql = Sql::with(self.exec().engine(), "INSERT INTO ");
75            sql.ident(Model::TABLE);
76            sql.push(" (");
77            let present = active.present().collect::<Vec<_>>();
78            sql.column_list(&present.iter().map(|(name, _)| *name).collect::<Vec<_>>());
79            sql.push(") SELECT ");
80            for (index, (name, value)) in present.into_iter().enumerate() {
81                if index > 0 {
82                    sql.push(", ");
83                }
84                if name == "start" {
85                    sql.push("CAST(");
86                    sql.bind(value.clone());
87                    sql.push(" AS text)");
88                } else {
89                    sql.bind(value.clone());
90                }
91            }
92            model::returning::<Model>(&mut sql);
93            self.exec()
94                .fetch_optional::<Model>(sql)
95                .await?
96                .ok_or_else(record_not_inserted)?
97        } else {
98            model::insert::<Model>(self.exec(), &active).await?
99        };
100        Ok(ApiKey::from(&inserted))
101    }
102
103    async fn get_api_key_by_id(&self, id: &str) -> AuthResult<Option<ApiKey>> {
104        let mut sql = model::by_id::<Model>(self.exec(), id);
105        model::limit_one(&mut sql);
106        Ok(self
107            .exec()
108            .fetch_optional::<Model>(sql)
109            .await?
110            .map(|model| ApiKey::from(&model)))
111    }
112
113    async fn get_api_key_by_hash(&self, hash: &str) -> AuthResult<Option<ApiKey>> {
114        let mut sql = model::select_model::<Model>(self.exec());
115        sql.push(" WHERE ");
116        sql.compare(Model::TABLE, "key", " = ", hash);
117        model::limit_one(&mut sql);
118        Ok(self
119            .exec()
120            .fetch_optional::<Model>(sql)
121            .await?
122            .map(|model| ApiKey::from(&model)))
123    }
124
125    async fn list_api_keys_by_reference(&self, reference_id: &str) -> AuthResult<Vec<ApiKey>> {
126        // Explicit ASC order matches TS insertion order on every backend.
127        let mut sql = model::select_model::<Model>(self.exec());
128        sql.push(" WHERE ");
129        sql.compare(Model::TABLE, "reference_id", " = ", reference_id);
130        sql.push(" ORDER BY ");
131        sql.column(Model::TABLE, "created_at");
132        sql.push(" ASC");
133        Ok(self
134            .exec()
135            .fetch_all::<Model>(sql)
136            .await?
137            .iter()
138            .map(ApiKey::from)
139            .collect())
140    }
141
142    async fn update_api_key(&self, id: &str, update: UpdateApiKey) -> AuthResult<ApiKey> {
143        let mut sql = model::by_id::<Model>(self.exec(), id);
144        model::limit_one(&mut sql);
145        let Some(model) = self.exec().fetch_optional::<Model>(sql).await? else {
146            return Err(AuthError::not_found("API Key not found"));
147        };
148
149        let active = apply_update_fields(model.into_active(), update)?;
150        model::update::<Model>(self.exec(), &active)
151            .await?
152            .map(|model| ApiKey::from(&model))
153            .ok_or_else(record_not_updated)
154    }
155
156    #[expect(
157        clippy::as_conversions,
158        clippy::cast_precision_loss,
159        reason = "JavaScript-compatible numbers deliberately retain IEEE754 rounding and guarded integer coercion at the wire or adapter boundary"
160    )]
161    async fn consume_api_key_usage(
162        &self,
163        id: &str,
164        global_rate_limit_enabled: bool,
165    ) -> AuthResult<ConsumeApiKeyResult> {
166        // SQLite must acquire its write reservation before reading usage counters.
167        let transaction = self.begin(true).await?;
168        let exec = Exec::tx(&transaction);
169        let id = id.to_owned();
170        let result = async {
171            let mut select = model::by_id::<Model>(exec, id.as_str());
172            model::limit_one(&mut select);
173            lock_exclusive(&mut select);
174            let Some(model) = exec.fetch_optional::<Model>(select).await? else {
175                return Err(AuthError::not_found("API Key not found"));
176            };
177
178            let now = Utc::now();
179            let mut update = UpdateApiKey::default();
180
181            if let Some(remaining) = model.remaining {
182                if let (Some(interval), Some(amount)) = (model.refill_interval, model.refill_amount)
183                    && interval != 0.0
184                    && amount != 0.0
185                    && (now.timestamp_millis()
186                        - model
187                            .last_refill_at
188                            .unwrap_or(model.created_at)
189                            .timestamp_millis()) as f64
190                        > interval
191                {
192                    update.remaining = Some(amount - 1.0);
193                    update.last_refill_at = Some(Some(now.to_rfc3339()));
194                } else if remaining > 0.0 {
195                    update.remaining = Some(remaining - 1.0);
196                } else {
197                    return Ok(ConsumeApiKeyResult::UsageExhausted);
198                }
199            }
200
201            if global_rate_limit_enabled && model.rate_limit_enabled {
202                if let (Some(window), Some(max)) =
203                    (model.rate_limit_time_window, model.rate_limit_max)
204                {
205                    let elapsed = model
206                        .last_request
207                        .map(|last| (now.timestamp_millis() - last.timestamp_millis()) as f64);
208                    if let Some(elapsed) = elapsed
209                        && elapsed <= window
210                        && model.request_count.unwrap_or(0.0) >= max
211                    {
212                        // TS consumes quota before rejecting a rate-limited request.
213                        // A rejection preserves the rate-limit window and updated_at.
214                        if update.remaining.is_some() {
215                            let updated_at = model.updated_at;
216                            let mut active = apply_update_fields(model.into_active(), update)?;
217                            active.set("updated_at", updated_at);
218                            _ = model::update::<Model>(exec, &active)
219                                .await?
220                                .ok_or_else(record_not_updated)?;
221                        }
222                        return Ok(ConsumeApiKeyResult::RateLimited {
223                            try_again_in: (window - elapsed).ceil(),
224                        });
225                    }
226
227                    update.request_count =
228                        Some(if elapsed.is_none_or(|elapsed| elapsed > window) {
229                            1.0
230                        } else {
231                            model.request_count.unwrap_or(0.0) + 1.0
232                        });
233                    update.last_request = Some(Some(now.to_rfc3339()));
234                }
235            } else {
236                update.last_request = Some(Some(now.to_rfc3339()));
237            }
238
239            let active = apply_update_fields(model.into_active(), update)?;
240            let updated = model::update::<Model>(exec, &active)
241                .await?
242                .ok_or_else(record_not_updated)?;
243            Ok(ConsumeApiKeyResult::Allowed(Box::new(ApiKey::from(
244                &updated,
245            ))))
246        }
247        .await;
248        if result.is_ok() {
249            transaction.commit().await?;
250        } else {
251            transaction.rollback().await?;
252        }
253        result
254    }
255
256    async fn delete_api_key(&self, id: &str) -> AuthResult<()> {
257        self.exec()
258            .execute(model::delete_by_id::<Model>(self.exec(), id))
259            .await
260            .map(|_| ())
261    }
262
263    async fn delete_expired_api_keys(&self) -> AuthResult<usize> {
264        let mut sql = Sql::with(self.exec().engine(), "DELETE FROM ");
265        sql.ident(Model::TABLE);
266        sql.push(" WHERE ");
267        sql.column(Model::TABLE, "expires_at");
268        sql.push(" IS NOT NULL AND ");
269        sql.compare(Model::TABLE, "expires_at", " < ", Utc::now());
270        let deleted = self.exec().execute(sql).await?;
271        usize::try_from(deleted)
272            .map_err(|_error| AuthError::internal("Affected row count exceeds usize"))
273    }
274}
275
276/// Apply `UpdateApiKey` fields to an API key row.
277fn apply_update_fields(mut active: ActiveRow, update: UpdateApiKey) -> AuthResult<ActiveRow> {
278    if let Some(name) = update.name {
279        active.set("name", Some(name));
280    }
281    if let Some(enabled) = update.enabled {
282        active.set("enabled", enabled);
283    }
284    if let Some(remaining) = update.remaining {
285        active.set("remaining", Some(remaining));
286    }
287    if let Some(rate_limit_enabled) = update.rate_limit_enabled {
288        active.set("rate_limit_enabled", rate_limit_enabled);
289    }
290    if let Some(rate_limit_time_window) = update.rate_limit_time_window {
291        active.set("rate_limit_time_window", Some(rate_limit_time_window));
292    }
293    if let Some(rate_limit_max) = update.rate_limit_max {
294        active.set("rate_limit_max", Some(rate_limit_max));
295    }
296    if let Some(refill_interval) = update.refill_interval {
297        active.set("refill_interval", Some(refill_interval));
298    }
299    if let Some(refill_amount) = update.refill_amount {
300        active.set("refill_amount", Some(refill_amount));
301    }
302    if let Some(permissions) = update.permissions {
303        active.set("permissions", Some(permissions));
304    }
305    if let Some(metadata) = update.metadata {
306        active.set("metadata", Some(metadata));
307    }
308    if let Some(expires_at) = update.expires_at {
309        active.set(
310            "expires_at",
311            parse_optional_rfc3339(expires_at.as_deref(), "expires_at")?,
312        );
313    }
314    if let Some(last_request) = update.last_request {
315        active.set(
316            "last_request",
317            parse_optional_rfc3339(last_request.as_deref(), "last_request")?,
318        );
319    }
320    if let Some(request_count) = update.request_count {
321        active.set("request_count", Some(request_count));
322    }
323    if let Some(last_refill_at) = update.last_refill_at {
324        active.set(
325            "last_refill_at",
326            parse_optional_rfc3339(last_refill_at.as_deref(), "last_refill_at")?,
327        );
328    }
329    active.set("updated_at", Utc::now());
330    Ok(active)
331}
alibi_sqlx/store/api_keys.rs

alibi_sqlx/store/
api_keys.rs
