systemprompt_users/repository/
rate_limit_bucket.rs

1//! Persistence for `user_rate_limit_buckets`: the replica-shared counters
2//! behind the global per-user HTTP rate limit.
3//!
4//! Copyright (c) systemprompt.io — Business Source License 1.1.
5//! See <https://systemprompt.io> for licensing details.
6
7use std::sync::Arc;
8
9use chrono::{DateTime, Utc};
10use sqlx::PgPool;
11use systemprompt_database::DbPool;
12use systemprompt_identifiers::UserId;
13
14use crate::error::Result;
15
16#[derive(Clone, Debug)]
17pub struct UserRateLimitBucketRepository {
18    write_pool: Arc<PgPool>,
19}
20
21impl UserRateLimitBucketRepository {
22    pub fn new(db: &DbPool) -> Result<Self> {
23        let write_pool = db.write_pool_arc()?;
24        Ok(Self { write_pool })
25    }
26
27    pub async fn hit(
28        &self,
29        user_id: &UserId,
30        scope: &str,
31        window_start: DateTime<Utc>,
32    ) -> Result<i64> {
33        let row = sqlx::query!(
34            r#"
35            INSERT INTO user_rate_limit_buckets (user_id, scope, window_start, hits)
36            VALUES ($1, $2, $3, 1)
37            ON CONFLICT (user_id, scope, window_start)
38            DO UPDATE SET hits = user_rate_limit_buckets.hits + 1
39            RETURNING hits
40            "#,
41            user_id.as_str(),
42            scope,
43            window_start,
44        )
45        .fetch_one(&*self.write_pool)
46        .await?;
47        Ok(row.hits)
48    }
49
50    pub async fn prune(&self, before: DateTime<Utc>) -> Result<u64> {
51        let result = sqlx::query!(
52            r#"DELETE FROM user_rate_limit_buckets WHERE window_start < $1"#,
53            before,
54        )
55        .execute(&*self.write_pool)
56        .await?;
57        Ok(result.rows_affected())
58    }
59}
systemprompt_users/repository/rate_limit_bucket.rs

systemprompt_users/repository/
rate_limit_bucket.rs
