model

Author: epompeii

lib/bencher_schema/src/macros/rate_limit.rs

@@ -6,7 +6,7 @@ use crate::{
     error::BencherResource,
     model::{
         organization::QueryOrganization,
-        project::{branch::QueryBranch, QueryProject},
+        project::{branch::QueryBranch, threshold::QueryThreshold, QueryProject},
         user::QueryUser,
     },
 };
@@ -17,6 +17,11 @@ pub const CLAIMED_RATE_LIMIT: u32 = u16::MAX as u32;
 
 #[derive(Debug, thiserror::Error)]
 pub enum RateLimitError {
+    #[error("User ({uuid}) has exceeded the daily rate limit ({UNCLAIMED_RATE_LIMIT}) for {resource} creation. Please, reduce your daily usage.", uuid = user.uuid)]
+    User {
+        user: QueryUser,
+        resource: BencherResource,
+    },
     #[error("Organization ({uuid}) has exceeded the daily rate limit ({UNCLAIMED_RATE_LIMIT}) for {resource} creation. Please, reduce your daily usage.", uuid = organization.uuid)]
     Organization {
         organization: QueryOrganization,
@@ -32,14 +37,14 @@ pub enum RateLimitError {
         project: QueryProject,
         resource: BencherResource,
     },
-    #[error("Branch ({uuid}) has exceeded the daily rate limit ({UNCLAIMED_RATE_LIMIT}) for {resource} creation. Please, reduce your daily usage.", uuid = branch.uuid)]
+    #[error("Branch ({uuid}) has exceeded the daily rate limit ({CLAIMED_RATE_LIMIT}) for {resource} creation. Please, reduce your daily usage.", uuid = branch.uuid)]
     Branch {
         branch: QueryBranch,
         resource: BencherResource,
     },
-    #[error("User ({uuid}) has exceeded the daily rate limit ({UNCLAIMED_RATE_LIMIT}) for {resource} creation. Please, reduce your daily usage.", uuid = user.uuid)]
-    User {
-        user: QueryUser,
+    #[error("Threshold ({uuid}) has exceeded the daily rate limit ({CLAIMED_RATE_LIMIT}) for {resource} creation. Please, reduce your daily usage.", uuid = threshold.uuid)]
+    Threshold {
+        threshold: QueryThreshold,
         resource: BencherResource,
     },
 }

lib/bencher_schema/src/model/organization/mod.rs

@@ -316,7 +316,7 @@ impl InsertOrganization {
                 .filter(schema::organization::created.le(end_time))
                 .count()
                 .get_result::<i64>(conn_lock!(context))
-                .map_err(resource_not_found_err!(Token, (query_user, start_time, end_time)))?
+                .map_err(resource_not_found_err!(Organization, (query_user, start_time, end_time)))?
                 .try_into()
                 .map_err(|e| {
                     issue_error(

lib/bencher_schema/src/model/project/branch/head.rs

@@ -219,7 +219,7 @@ impl InsertHead {
                 .filter(schema::head::created.le(end_time))
                 .count()
                 .get_result::<i64>(conn_lock!(context))
-                .map_err(resource_not_found_err!(Token, (query_branch, start_time, end_time)))?
+                .map_err(resource_not_found_err!(Head, (query_branch, start_time, end_time)))?
                 .try_into()
                 .map_err(|e| {
                     issue_error(

lib/bencher_schema/src/model/project/mod.rs

@@ -513,7 +513,7 @@ impl InsertProject {
                 .filter(schema::project::created.le(end_time))
                 .count()
                 .get_result::<i64>(conn_lock!(context))
-                .map_err(resource_not_found_err!(Token, (query_organization, start_time, end_time)))?
+                .map_err(resource_not_found_err!(Project, (query_organization, start_time, end_time)))?
                 .try_into()
                 .map_err(|e| {
                     issue_error(

lib/bencher_schema/src/model/project/threshold/mod.rs

@@ -88,7 +88,7 @@ impl QueryThreshold {
             (None, None) => Ok(()),
             // No current model but a new model,
             // insert the new model.
-            (None, Some(model)) => self.update_from_model(conn_lock!(context), model),
+            (None, Some(model)) => self.update_from_model(context, model).await,
             // Current model but no new model,
             // remove the current model.
             (Some(_), None) => self.remove_current_model(conn_lock!(context)),
@@ -101,13 +101,19 @@ impl QueryThreshold {
                 if current_model == model {
                     Ok(())
                 } else {
-                    self.update_from_model(conn_lock!(context), model)
+                    self.update_from_model(context, model).await
                 }
             },
         }
     }
 
-    pub fn update_from_model(
+    async fn update_from_model(&self, context: &ApiContext, model: Model) -> Result<(), HttpError> {
+        #[cfg(feature = "plus")]
+        InsertModel::rate_limit(context, self).await?;
+        self.update_from_model_inner(conn_lock!(context), model)
+    }
+
+    fn update_from_model_inner(
         &self,
         conn: &mut DbConnection,
         model: Model,

lib/bencher_schema/src/model/project/threshold/model.rs

@@ -123,6 +123,49 @@ pub struct InsertModel {
 }
 
 impl InsertModel {
+    #[cfg(feature = "plus")]
+    pub async fn rate_limit(
+        context: &crate::ApiContext,
+        query_threshold: &QueryThreshold,
+    ) -> Result<(), HttpError> {
+        use crate::{
+            conn_lock,
+            error::issue_error,
+            macros::rate_limit::{one_day, RateLimitError, CLAIMED_RATE_LIMIT},
+        };
+
+        let resource = BencherResource::Model;
+        let (start_time, end_time) = one_day();
+        let creation_count: u32 = schema::model::table
+                .filter(schema::model::threshold_id.eq(query_threshold.id))
+                .filter(schema::model::created.ge(start_time))
+                .filter(schema::model::created.le(end_time))
+                .count()
+                .get_result::<i64>(conn_lock!(context))
+                .map_err(resource_not_found_err!(Model, (query_threshold, start_time, end_time)))?
+                .try_into()
+                .map_err(|e| {
+                    issue_error(
+                        "Failed to count creation",
+                        &format!("Failed to count {resource} creation for threshold ({uuid}) between {start_time} and {end_time}.", uuid = query_threshold.uuid),
+                    e
+                    )}
+                )?;
+
+        // The only way that new Model can be crated is either through running a Report
+        // or by updating an existing threshold using the API.
+        // The running of a Report will be rate limited already for unclaimed projects,
+        // and the API endpoint to update an existing threshold would require authentication and would therefore be a claimed project.
+        if creation_count >= CLAIMED_RATE_LIMIT {
+            Err(crate::error::too_many_requests(RateLimitError::Threshold {
+                threshold: query_threshold.clone(),
+                resource,
+            }))
+        } else {
+            Ok(())
+        }
+    }
+
     pub fn new(threshold_id: ThresholdId, model: Model) -> Self {
         let Model {
             test,