use mnm_core::types::RateLimitOverride;
use sqlx::PgPool;
use time::OffsetDateTime;
use uuid::Uuid;
use crate::error::{Result, StoreError};
const COLS: &str = "id, cidr::text AS cidr, limit_rps, expires_at, note, created_by, created_at";
pub async fn insert(
pool: &PgPool,
cidr: &str,
limit_rps: i32,
expires_at: OffsetDateTime,
note: Option<&str>,
created_by: &str,
) -> Result<RateLimitOverride> {
let row = sqlx::query_as::<_, OverrideRow>(
"INSERT INTO rate_limit_override (cidr, limit_rps, expires_at, note, created_by) \
VALUES (network($1::inet), $2, $3, $4, $5) \
RETURNING id, cidr::text AS cidr, limit_rps, expires_at, note, created_by, created_at",
)
.bind(cidr)
.bind(limit_rps)
.bind(expires_at)
.bind(note)
.bind(created_by)
.fetch_one(pool)
.await?;
Ok(row.into())
}
pub async fn list_active(pool: &PgPool) -> Result<Vec<RateLimitOverride>> {
let rows = sqlx::query_as::<_, OverrideRow>(&format!(
"SELECT {COLS} FROM rate_limit_override \
WHERE expires_at > now() ORDER BY created_at DESC"
))
.fetch_all(pool)
.await?;
Ok(rows.into_iter().map(Into::into).collect())
}
pub async fn get_by_id(pool: &PgPool, id: Uuid) -> Result<RateLimitOverride> {
let row = sqlx::query_as::<_, OverrideRow>(&format!(
"SELECT {COLS} FROM rate_limit_override WHERE id = $1"
))
.bind(id)
.fetch_one(pool)
.await?;
Ok(row.into())
}
#[derive(Debug, Default, Clone)]
pub struct RateLimitPatch {
pub expires_at: Option<OffsetDateTime>,
pub limit_rps: Option<i32>,
pub note: Option<String>,
}
pub async fn update(pool: &PgPool, id: Uuid, patch: RateLimitPatch) -> Result<RateLimitOverride> {
let row = sqlx::query_as::<_, OverrideRow>(
"UPDATE rate_limit_override SET \
expires_at = COALESCE($2, expires_at), \
limit_rps = COALESCE($3, limit_rps), \
note = COALESCE($4, note) \
WHERE id = $1 \
RETURNING id, cidr::text AS cidr, limit_rps, expires_at, note, created_by, created_at",
)
.bind(id)
.bind(patch.expires_at)
.bind(patch.limit_rps)
.bind(patch.note)
.fetch_optional(pool)
.await?;
row.map_or(Err(StoreError::NotFound), |r| Ok(r.into()))
}
pub async fn delete(pool: &PgPool, id: Uuid) -> Result<RateLimitOverride> {
let row = sqlx::query_as::<_, OverrideRow>(
"DELETE FROM rate_limit_override WHERE id = $1 \
RETURNING id, cidr::text AS cidr, limit_rps, expires_at, note, created_by, created_at",
)
.bind(id)
.fetch_optional(pool)
.await?;
row.map_or(Err(StoreError::NotFound), |r| Ok(r.into()))
}
#[derive(sqlx::FromRow)]
struct OverrideRow {
id: Uuid,
cidr: String,
limit_rps: i32,
expires_at: OffsetDateTime,
note: Option<String>,
created_by: String,
created_at: OffsetDateTime,
}
impl From<OverrideRow> for RateLimitOverride {
fn from(r: OverrideRow) -> Self {
Self {
id: r.id,
cidr: r.cidr,
limit_rps: r.limit_rps,
expires_at: r.expires_at,
note: r.note,
created_by: r.created_by,
created_at: r.created_at,
}
}
}