use super::*;use crate::db::tables::rate_buckets;use sea_orm::{ColumnTrait, EntityTrait, QueryFilter}; impl Database { /// Saves the rate limit for the given user and rate limit name if the last_refill is later /// than the currently saved timestamp. pub async fn save_rate_buckets(&self, buckets: &[rate_buckets::Model]) -> Result<()> { if buckets.is_empty() { return Ok(()); } self.transaction(|tx| async move { rate_buckets::Entity::insert_many(buckets.iter().map(|bucket| { rate_buckets::ActiveModel { user_id: ActiveValue::Set(bucket.user_id), rate_limit_name: ActiveValue::Set(bucket.rate_limit_name.clone()), token_count: ActiveValue::Set(bucket.token_count), last_refill: ActiveValue::Set(bucket.last_refill), } })) .on_conflict( OnConflict::columns([ rate_buckets::Column::UserId, rate_buckets::Column::RateLimitName, ]) .update_columns([ rate_buckets::Column::TokenCount, rate_buckets::Column::LastRefill, ]) .to_owned(), ) .exec(&*tx) .await?; Ok(()) }) .await } /// Retrieves the rate limit for the given user and rate limit name. pub async fn get_rate_bucket( &self, user_id: UserId, rate_limit_name: &str, ) -> Result<Option<rate_buckets::Model>> { self.transaction(|tx| async move { let rate_limit = rate_buckets::Entity::find() .filter(rate_buckets::Column::UserId.eq(user_id)) .filter(rate_buckets::Column::RateLimitName.eq(rate_limit_name)) .one(&*tx) .await?; Ok(rate_limit) }) .await }}
use super::*;
use crate::db::tables::rate_buckets;
use sea_orm::{ColumnTrait, EntityTrait, QueryFilter};
impl Database {
/// Saves the rate limit for the given user and rate limit name if the last_refill is later
/// than the currently saved timestamp.
pub async fn save_rate_buckets(&self, buckets: &[rate_buckets::Model]) -> Result<()> {
if buckets.is_empty() {
return Ok(());
}
self.transaction(|tx| async move {
rate_buckets::Entity::insert_many(buckets.iter().map(|bucket| {
rate_buckets::ActiveModel {
user_id: ActiveValue::Set(bucket.user_id),
rate_limit_name: ActiveValue::Set(bucket.rate_limit_name.clone()),
token_count: ActiveValue::Set(bucket.token_count),
last_refill: ActiveValue::Set(bucket.last_refill),
}
}))
.on_conflict(
OnConflict::columns([
rate_buckets::Column::UserId,
rate_buckets::Column::RateLimitName,
])
.update_columns([
rate_buckets::Column::TokenCount,
rate_buckets::Column::LastRefill,
])
.to_owned(),
)
.exec(&*tx)
.await?;
Ok(())
})
.await
}
/// Retrieves the rate limit for the given user and rate limit name.
pub async fn get_rate_bucket(
&self,
user_id: UserId,
rate_limit_name: &str,
) -> Result<Option<rate_buckets::Model>> {
self.transaction(|tx| async move {
let rate_limit = rate_buckets::Entity::find()
.filter(rate_buckets::Column::UserId.eq(user_id))
.filter(rate_buckets::Column::RateLimitName.eq(rate_limit_name))
.one(&*tx)
.await?;
Ok(rate_limit)
})
.await
}
}