use super::entities::api_key::{ActiveModel, Column, Entity};
use super::entities::api_key_start::ApiKeyStart;
use super::{SeaOrmStore, map_db_err};
use crate::schema::AuthSchema;
use alibi_core::error::{AuthError, AuthResult};
use alibi_core::store::adapter::parse_optional_rfc3339;
use alibi_core::store::{ApiKeyStore, ConsumeApiKeyResult};
use alibi_core::types::{ApiKey, CreateApiKey, UpdateApiKey};
use async_trait::async_trait;
use chrono::Utc;
use sea_orm::{
ActiveModelTrait, ColumnTrait, EntityTrait, IntoActiveModel, Iterable, QueryFilter, QueryOrder,
QuerySelect, QueryTrait, Set,
};
use uuid::Uuid;
#[async_trait]
impl<S> ApiKeyStore for SeaOrmStore<S>
where
S: AuthSchema + Send + Sync,
{
async fn consume_api_key_usage_from_snapshot(
&self,
observed: &ApiKey,
global_rate_limit_enabled: bool,
) -> AuthResult<ConsumeApiKeyResult> {
self.consume_usage_phases(observed, global_rate_limit_enabled)
.await
}
async fn create_api_key(&self, input: CreateApiKey) -> AuthResult<ApiKey> {
let now = Utc::now();
let start = input
.start
.map(|start| {
ApiKeyStart::prepare(&start, self.scoped_connection().get_database_backend())
})
.transpose()?;
let sqlite_cast = start
.as_ref()
.is_some_and(ApiKeyStart::requires_sqlite_cast);
let model = ActiveModel {
id: Set(self
.generated_id(&self.db, "apikey", "apikey", "id")
.await?
.unwrap_or_else(|| Uuid::new_v4().to_string())),
name: Set(input.name),
start: Set(start),
prefix: Set(input.prefix),
key_hash: Set(input.key_hash),
reference_id: Set(input.reference_id),
config_id: Set(input.config_id),
refill_interval: Set(input.refill_interval.map(crate::Number::from)),
refill_amount: Set(input.refill_amount.map(crate::Number::from)),
last_refill_at: Set(None),
enabled: Set(Some(input.enabled)),
rate_limit_enabled: Set(Some(input.rate_limit_enabled)),
rate_limit_time_window: Set(input.rate_limit_time_window.map(crate::Number::from)),
rate_limit_max: Set(input.rate_limit_max.map(crate::Number::from)),
request_count: Set(Some(crate::Number(0.0))),
remaining: Set(input.remaining.map(crate::Number::from)),
last_request: Set(None),
expires_at: Set(parse_optional_rfc3339(
input.expires_at.as_deref(),
"expires_at",
)?),
created_at: Set(now),
updated_at: Set(now),
permissions: Set(input.permissions),
metadata: Set(input.metadata),
};
let inserted = if sqlite_cast {
use sea_orm::sea_query::{Expr, ExprTrait, Query};
let mut insert = Entity::insert(model.clone());
let mut values = Query::select();
for column in Column::iter() {
let value = Expr::val(model.get(column).unwrap());
_ = values.expr(if matches!(column, Column::Start) {
value.cast_as("text")
} else {
value
});
}
_ = insert
.query()
.select_from(values)
.map_err(|error| AuthError::internal(error.to_string()))?;
insert.exec_with_returning(self.scoped_connection()).await
} else {
model.insert(self.scoped_connection()).await
};
inserted
.map(|model| ApiKey::from(&model))
.map_err(map_db_err)
}
async fn get_api_key_by_id(&self, id: &str) -> AuthResult<Option<ApiKey>> {
Entity::find_by_id(id.to_owned())
.one(self.scoped_connection())
.await
.map(|model| model.map(|model| ApiKey::from(&model)))
.map_err(map_db_err)
}
async fn get_api_key_by_hash(&self, hash: &str) -> AuthResult<Option<ApiKey>> {
Entity::find()
.filter(Column::KeyHash.eq(hash))
.one(self.scoped_connection())
.await
.map(|model| model.map(|model| ApiKey::from(&model)))
.map_err(map_db_err)
}
async fn list_api_keys_by_reference(&self, reference_id: &str) -> AuthResult<Vec<ApiKey>> {
Entity::find()
.filter(Column::ReferenceId.eq(reference_id))
.order_by_asc(Column::CreatedAt)
.all(self.scoped_connection())
.await
.map(|models| models.iter().map(ApiKey::from).collect())
.map_err(map_db_err)
}
async fn update_api_key(&self, id: &str, update: UpdateApiKey) -> AuthResult<ApiKey> {
let Some(model) = Entity::find_by_id(id.to_owned())
.one(self.scoped_connection())
.await
.map_err(map_db_err)?
else {
return Err(AuthError::not_found("API Key not found"));
};
let active = apply_update_fields(model.into_active_model(), update)?;
active
.update(self.scoped_connection())
.await
.map(|model| ApiKey::from(&model))
.map_err(map_db_err)
}
#[expect(
clippy::as_conversions,
clippy::cast_precision_loss,
reason = "JavaScript-compatible numbers deliberately retain IEEE754 rounding and guarded integer coercion at the wire or adapter boundary"
)]
async fn consume_api_key_usage(
&self,
id: &str,
global_rate_limit_enabled: bool,
) -> AuthResult<ConsumeApiKeyResult> {
let transaction = self.scoped_connection().begin_immediate().await?;
let txn = &transaction;
let id = id.to_owned();
let result = async {
let Some(model) = Entity::find_by_id(id.clone())
.lock_exclusive()
.one(txn)
.await
.map_err(map_db_err)?
else {
return Err(AuthError::not_found("API Key not found"));
};
let now = Utc::now();
let mut update = UpdateApiKey::default();
if let Some(remaining) = model.remaining.map(|value| value.0) {
if let (Some(interval), Some(amount)) = (
model.refill_interval.map(|value| value.0),
model.refill_amount.map(|value| value.0),
) && interval != 0.0
&& amount != 0.0
&& (now.timestamp_millis()
- model
.last_refill_at
.unwrap_or(model.created_at)
.timestamp_millis()) as f64
> interval
{
update.remaining = Some(amount - 1.0);
update.last_refill_at = Some(Some(now.to_rfc3339()));
} else if remaining > 0.0 {
update.remaining = Some(remaining - 1.0);
} else {
return Ok(ConsumeApiKeyResult::UsageExhausted);
}
}
if global_rate_limit_enabled && model.rate_limit_enabled.unwrap_or(false) {
if let (Some(window), Some(max)) = (
model.rate_limit_time_window.map(|value| value.0),
model.rate_limit_max.map(|value| value.0),
) {
let elapsed = model
.last_request
.map(|last| (now.timestamp_millis() - last.timestamp_millis()) as f64);
if let Some(elapsed) = elapsed
&& elapsed <= window
&& model.request_count.map_or(0.0, |value| value.0) >= max
{
if update.remaining.is_some() {
let updated_at = model.updated_at;
let mut active =
apply_update_fields(model.into_active_model(), update)?;
active.updated_at = Set(updated_at);
_ = active.update(txn).await.map_err(map_db_err)?;
}
return Ok(ConsumeApiKeyResult::RateLimited {
try_again_in: (window - elapsed).ceil(),
});
}
update.request_count =
Some(if elapsed.is_none_or(|elapsed| elapsed > window) {
1.0
} else {
model.request_count.map_or(0.0, |value| value.0) + 1.0
});
update.last_request = Some(Some(now.to_rfc3339()));
}
} else {
update.last_request = Some(Some(now.to_rfc3339()));
}
let active = apply_update_fields(model.into_active_model(), update)?;
let updated = active.update(txn).await.map_err(map_db_err)?;
Ok(ConsumeApiKeyResult::Allowed(Box::new(ApiKey::from(
&updated,
))))
}
.await;
if result.is_ok() {
transaction.commit().await.map_err(map_db_err)?;
} else {
transaction.rollback().await.map_err(map_db_err)?;
}
result
}
async fn delete_api_key(&self, id: &str) -> AuthResult<()> {
Entity::delete_by_id(id.to_owned())
.exec(self.scoped_connection())
.await
.map(|_| ())
.map_err(map_db_err)
}
async fn delete_expired_api_keys(&self) -> AuthResult<usize> {
Entity::delete_many()
.filter(Column::ExpiresAt.is_not_null())
.filter(Column::ExpiresAt.lt(Utc::now()))
.exec(self.scoped_connection())
.await
.map_err(map_db_err)
.and_then(|result| {
usize::try_from(result.rows_affected)
.map_err(|_error| AuthError::internal("Affected row count exceeds usize"))
})
}
}
fn apply_update_fields(mut active: ActiveModel, update: UpdateApiKey) -> AuthResult<ActiveModel> {
if let Some(name) = update.name {
active.name = Set(Some(name));
}
if let Some(enabled) = update.enabled {
active.enabled = Set(Some(enabled));
}
if let Some(remaining) = update.remaining {
active.remaining = Set(Some(crate::Number(remaining)));
}
if let Some(rate_limit_enabled) = update.rate_limit_enabled {
active.rate_limit_enabled = Set(Some(rate_limit_enabled));
}
if let Some(rate_limit_time_window) = update.rate_limit_time_window {
active.rate_limit_time_window = Set(Some(crate::Number(rate_limit_time_window)));
}
if let Some(rate_limit_max) = update.rate_limit_max {
active.rate_limit_max = Set(Some(crate::Number(rate_limit_max)));
}
if let Some(refill_interval) = update.refill_interval {
active.refill_interval = Set(Some(crate::Number(refill_interval)));
}
if let Some(refill_amount) = update.refill_amount {
active.refill_amount = Set(Some(crate::Number(refill_amount)));
}
if let Some(permissions) = update.permissions {
active.permissions = Set(Some(permissions));
}
if let Some(metadata) = update.metadata {
active.metadata = Set(Some(metadata));
}
if let Some(expires_at) = update.expires_at {
active.expires_at = Set(parse_optional_rfc3339(expires_at.as_deref(), "expires_at")?);
}
if let Some(last_request) = update.last_request {
active.last_request = Set(parse_optional_rfc3339(
last_request.as_deref(),
"last_request",
)?);
}
if let Some(request_count) = update.request_count {
active.request_count = Set(Some(crate::Number(request_count)));
}
if let Some(last_refill_at) = update.last_refill_at {
active.last_refill_at = Set(parse_optional_rfc3339(
last_refill_at.as_deref(),
"last_refill_at",
)?);
}
active.updated_at = Set(Utc::now());
Ok(active)
}