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("id", Uuid::new_v4().to_string());
43        active.set("name", input.name);
44        active.set("start", start.into_sql_value());
45        active.set("prefix", input.prefix);
46        active.set("key", input.key_hash);
47        active.set("reference_id", input.reference_id);
48        active.set("config_id", input.config_id);
49        active.set("refill_interval", input.refill_interval);
50        active.set("refill_amount", input.refill_amount);
51        active.set("last_refill_at", None::<chrono::DateTime<Utc>>);
52        active.set("enabled", input.enabled);
53        active.set("rate_limit_enabled", input.rate_limit_enabled);
54        active.set("rate_limit_time_window", input.rate_limit_time_window);
55        active.set("rate_limit_max", input.rate_limit_max);
56        active.set("request_count", Some(0.0_f64));
57        active.set("remaining", input.remaining);
58        active.set("last_request", None::<chrono::DateTime<Utc>>);
59        active.set(
60            "expires_at",
61            parse_optional_rfc3339(input.expires_at.as_deref(), "expires_at")?,
62        );
63        active.set("created_at", now);
64        active.set("updated_at", now);
65        active.set("permissions", input.permissions);
66        active.set("metadata", input.metadata);
67        let inserted = if sqlite_cast {
68            // Keep one bound INSERT with every ordinary model value. Only the
69            // invalid-surrogate SQLite start needs a bytes-to-TEXT expression;
70            // other engines and valid strings retain the ordinary insert path.
71            let mut sql = Sql::with(self.exec().engine(), "INSERT INTO ");
72            sql.ident(Model::TABLE);
73            sql.push(" (");
74            let present = active.present().collect::<Vec<_>>();
75            sql.column_list(&present.iter().map(|(name, _)| *name).collect::<Vec<_>>());
76            sql.push(") SELECT ");
77            for (index, (name, value)) in present.into_iter().enumerate() {
78                if index > 0 {
79                    sql.push(", ");
80                }
81                if name == "start" {
82                    sql.push("CAST(");
83                    sql.bind(value.clone());
84                    sql.push(" AS text)");
85                } else {
86                    sql.bind(value.clone());
87                }
88            }
89            model::returning::<Model>(&mut sql);
90            self.exec()
91                .fetch_optional::<Model>(sql)
92                .await?
93                .ok_or_else(record_not_inserted)?
94        } else {
95            model::insert::<Model>(self.exec(), &active).await?
96        };
97        Ok(ApiKey::from(&inserted))
98    }
99
100    async fn get_api_key_by_id(&self, id: &str) -> AuthResult<Option<ApiKey>> {
101        let mut sql = model::by_id::<Model>(self.exec(), id);
102        model::limit_one(&mut sql);
103        Ok(self
104            .exec()
105            .fetch_optional::<Model>(sql)
106            .await?
107            .map(|model| ApiKey::from(&model)))
108    }
109
110    async fn get_api_key_by_hash(&self, hash: &str) -> AuthResult<Option<ApiKey>> {
111        let mut sql = model::select_model::<Model>(self.exec());
112        sql.push(" WHERE ");
113        sql.compare(Model::TABLE, "key", " = ", hash);
114        model::limit_one(&mut sql);
115        Ok(self
116            .exec()
117            .fetch_optional::<Model>(sql)
118            .await?
119            .map(|model| ApiKey::from(&model)))
120    }
121
122    async fn list_api_keys_by_reference(&self, reference_id: &str) -> AuthResult<Vec<ApiKey>> {
123        // Explicit ASC order matches TS insertion-order behavior and avoids
124        // nondeterministic results across database backends.
125        let mut sql = model::select_model::<Model>(self.exec());
126        sql.push(" WHERE ");
127        sql.compare(Model::TABLE, "reference_id", " = ", reference_id);
128        sql.push(" ORDER BY ");
129        sql.column(Model::TABLE, "created_at");
130        sql.push(" ASC");
131        Ok(self
132            .exec()
133            .fetch_all::<Model>(sql)
134            .await?
135            .iter()
136            .map(ApiKey::from)
137            .collect())
138    }
139
140    async fn update_api_key(&self, id: &str, update: UpdateApiKey) -> AuthResult<ApiKey> {
141        let mut sql = model::by_id::<Model>(self.exec(), id);
142        model::limit_one(&mut sql);
143        let Some(model) = self.exec().fetch_optional::<Model>(sql).await? else {
144            return Err(AuthError::not_found("API Key not found"));
145        };
146
147        let active = apply_update_fields(model.into_active(), update)?;
148        model::update::<Model>(self.exec(), &active)
149            .await?
150            .map(|model_2| ApiKey::from(&model_2))
151            .ok_or_else(record_not_updated)
152    }
153
154    #[expect(
155        clippy::as_conversions,
156        clippy::cast_precision_loss,
157        reason = "JavaScript-compatible numbers deliberately retain IEEE754 rounding and guarded integer coercion at the wire or adapter boundary"
158    )]
159    async fn consume_api_key_usage(
160        &self,
161        id: &str,
162        global_rate_limit_enabled: bool,
163    ) -> AuthResult<ConsumeApiKeyResult> {
164        // SQLite must acquire its write reservation before reading usage counters.
165        let transaction = self.pool().begin(true).await?;
166        let exec = Exec::Tx(&transaction);
167        let id = id.to_owned();
168        let result = async {
169            let mut select = model::by_id::<Model>(exec, id.as_str());
170            model::limit_one(&mut select);
171            lock_exclusive(&mut select);
172            let Some(model) = exec.fetch_optional::<Model>(select).await? else {
173                return Err(AuthError::not_found("API Key not found"));
174            };
175
176            let now = Utc::now();
177            let mut update = UpdateApiKey::default();
178
179            if let Some(remaining) = model.remaining {
180                if let (Some(interval), Some(amount)) = (model.refill_interval, model.refill_amount)
181                    && interval != 0.0
182                    && amount != 0.0
183                    && (now.timestamp_millis()
184                        - model
185                            .last_refill_at
186                            .unwrap_or(model.created_at)
187                            .timestamp_millis()) as f64
188                        > interval
189                {
190                    update.remaining = Some(amount - 1.0);
191                    update.last_refill_at = Some(Some(now.to_rfc3339()));
192                } else if remaining > 0.0 {
193                    update.remaining = Some(remaining - 1.0);
194                } else {
195                    return Ok(ConsumeApiKeyResult::UsageExhausted);
196                }
197            }
198
199            if global_rate_limit_enabled && model.rate_limit_enabled {
200                if let (Some(window), Some(max)) =
201                    (model.rate_limit_time_window, model.rate_limit_max)
202                {
203                    let elapsed = model
204                        .last_request
205                        .map(|last| (now.timestamp_millis() - last.timestamp_millis()) as f64);
206                    if let Some(elapsed) = elapsed
207                        && elapsed <= window
208                        && model.request_count.unwrap_or(0.0) >= max
209                    {
210                        // TS consumes quota before rejecting a rate-limited request.
211                        // A rejection preserves the rate-limit window and updated_at.
212                        if update.remaining.is_some() {
213                            let updated_at = model.updated_at;
214                            let mut active = apply_update_fields(model.into_active(), update)?;
215                            active.set("updated_at", updated_at);
216                            drop(
217                                model::update::<Model>(exec, &active)
218                                    .await?
219                                    .ok_or_else(record_not_updated)?,
220                            );
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(drop)
261    }
262
263    async fn delete_expired_api_keys(&self) -> AuthResult<usize> {
264        // Single query: DELETE FROM api_keys WHERE expires_at IS NOT NULL AND expires_at < NOW()
265        // Matches TS: adapter.deleteMany({ where: [{ field: "expiresAt", operator: "lt", value: new Date() }, ...] })
266        let mut sql = Sql::with(self.exec().engine(), "DELETE FROM ");
267        sql.ident(Model::TABLE);
268        sql.push(" WHERE ");
269        sql.column(Model::TABLE, "expires_at");
270        sql.push(" IS NOT NULL AND ");
271        sql.compare(Model::TABLE, "expires_at", " < ", Utc::now());
272        let deleted = self.exec().execute(sql).await?;
273        usize::try_from(deleted)
274            .map_err(|_error| AuthError::internal("Affected row count exceeds usize"))
275    }
276}
277
278/// Apply `UpdateApiKey` fields to an API key row.
279fn apply_update_fields(mut active: ActiveRow, update: UpdateApiKey) -> AuthResult<ActiveRow> {
280    if let Some(name) = update.name {
281        active.set("name", Some(name));
282    }
283    if let Some(enabled) = update.enabled {
284        active.set("enabled", enabled);
285    }
286    if let Some(remaining) = update.remaining {
287        active.set("remaining", Some(remaining));
288    }
289    if let Some(rate_limit_enabled) = update.rate_limit_enabled {
290        active.set("rate_limit_enabled", rate_limit_enabled);
291    }
292    if let Some(rate_limit_time_window) = update.rate_limit_time_window {
293        active.set("rate_limit_time_window", Some(rate_limit_time_window));
294    }
295    if let Some(rate_limit_max) = update.rate_limit_max {
296        active.set("rate_limit_max", Some(rate_limit_max));
297    }
298    if let Some(refill_interval) = update.refill_interval {
299        active.set("refill_interval", Some(refill_interval));
300    }
301    if let Some(refill_amount) = update.refill_amount {
302        active.set("refill_amount", Some(refill_amount));
303    }
304    if let Some(permissions) = update.permissions {
305        active.set("permissions", Some(permissions));
306    }
307    if let Some(metadata) = update.metadata {
308        active.set("metadata", Some(metadata));
309    }
310    if let Some(expires_at) = update.expires_at {
311        active.set(
312            "expires_at",
313            parse_optional_rfc3339(expires_at.as_deref(), "expires_at")?,
314        );
315    }
316    if let Some(last_request) = update.last_request {
317        active.set(
318            "last_request",
319            parse_optional_rfc3339(last_request.as_deref(), "last_request")?,
320        );
321    }
322    if let Some(request_count) = update.request_count {
323        active.set("request_count", Some(request_count));
324    }
325    if let Some(last_refill_at) = update.last_refill_at {
326        active.set(
327            "last_refill_at",
328            parse_optional_rfc3339(last_refill_at.as_deref(), "last_refill_at")?,
329        );
330    }
331    active.set("updated_at", Utc::now());
332    Ok(active)
333}
alibi_sqlx/store/api_keys.rs

alibi_sqlx/store/
api_keys.rs
