add the action column to rate limit tables

Author: pietroalbini

migrations/2021-07-18-125813_add_rate_limit_action/down.sql

@@ -0,0 +1,9 @@
+DELETE FROM publish_limit_buckets WHERE action != 0;
+ALTER TABLE publish_limit_buckets DROP CONSTRAINT publish_limit_buckets_pkey;
+ALTER TABLE publish_limit_buckets ADD CONSTRAINT publish_limit_buckets_pkey PRIMARY KEY (user_id);
+ALTER TABLE publish_limit_buckets DROP COLUMN action;
+
+DELETE FROM publish_rate_overrides WHERE action != 0;
+ALTER TABLE publish_rate_overrides DROP CONSTRAINT publish_rate_overrides_pkey;
+ALTER TABLE publish_rate_overrides ADD CONSTRAINT publish_rate_overrides_pkey PRIMARY KEY (user_id);
+ALTER TABLE publish_rate_overrides DROP COLUMN action;

migrations/2021-07-18-125813_add_rate_limit_action/up.sql

@@ -0,0 +1,7 @@
+ALTER TABLE publish_limit_buckets ADD COLUMN action INTEGER NOT NULL DEFAULT 0;
+ALTER TABLE publish_limit_buckets DROP CONSTRAINT publish_limit_buckets_pkey;
+ALTER TABLE publish_limit_buckets ADD CONSTRAINT publish_limit_buckets_pkey PRIMARY KEY (user_id, action);
+
+ALTER TABLE publish_rate_overrides ADD COLUMN action INTEGER NOT NULL DEFAULT 0;
+ALTER TABLE publish_rate_overrides DROP CONSTRAINT publish_rate_overrides_pkey;
+ALTER TABLE publish_rate_overrides ADD CONSTRAINT publish_rate_overrides_pkey PRIMARY KEY (user_id, action);

src/publish_rate_limit.rs

@@ -8,6 +8,12 @@ use crate::schema::{publish_limit_buckets, publish_rate_overrides};
 use crate::sql::{date_part, floor, greatest, interval_part, least};
 use crate::util::errors::{AppResult, TooManyRequests};
 
+crate::pg_enum! {
+    pub enum LimitedAction {
+        PublishNew = 0,
+    }
+}
+
 #[derive(Debug, Clone, Copy)]
 pub struct PublishRateLimit {
     pub rate: Duration,
@@ -40,6 +46,7 @@ struct Bucket {
     user_id: i32,
     tokens: i32,
     last_refill: NaiveDateTime,
+    action: LimitedAction,
 }
 
 impl PublishRateLimit {
@@ -71,7 +78,7 @@ impl PublishRateLimit {
         use self::publish_limit_buckets::dsl::*;
 
         let burst: i32 = publish_rate_overrides::table
-            .find(uploader)
+            .find((uploader, LimitedAction::PublishNew))
             .filter(
                 publish_rate_overrides::expires_at
                     .is_null()
@@ -128,6 +135,7 @@ mod tests {
             user_id: bucket.user_id,
             tokens: 10,
             last_refill: now,
+            action: LimitedAction::PublishNew,
         };
         assert_eq!(expected, bucket);
 
@@ -140,6 +148,7 @@ mod tests {
             user_id: bucket.user_id,
             tokens: 20,
             last_refill: now,
+            action: LimitedAction::PublishNew,
         };
         assert_eq!(expected, bucket);
         Ok(())
@@ -160,6 +169,7 @@ mod tests {
             user_id,
             tokens: 4,
             last_refill: now,
+            action: LimitedAction::PublishNew,
         };
         assert_eq!(expected, bucket);
         Ok(())
@@ -181,6 +191,7 @@ mod tests {
             user_id,
             tokens: 6,
             last_refill: refill_time,
+            action: LimitedAction::PublishNew,
         };
         assert_eq!(expected, bucket);
         Ok(())
@@ -206,6 +217,7 @@ mod tests {
             user_id,
             tokens: 7,
             last_refill: refill_time,
+            action: LimitedAction::PublishNew,
         };
         assert_eq!(expected, bucket);
         Ok(())
@@ -227,6 +239,7 @@ mod tests {
             user_id,
             tokens: 6,
             last_refill: expected_refill_time,
+            action: LimitedAction::PublishNew,
         };
         assert_eq!(expected, bucket);
         Ok(())
@@ -247,6 +260,7 @@ mod tests {
             user_id,
             tokens: 0,
             last_refill: now,
+            action: LimitedAction::PublishNew,
         };
         assert_eq!(expected, bucket);
 
@@ -271,6 +285,7 @@ mod tests {
             user_id,
             tokens: 1,
             last_refill: refill_time,
+            action: LimitedAction::PublishNew,
         };
         assert_eq!(expected, bucket);
 
@@ -293,6 +308,7 @@ mod tests {
             user_id,
             tokens: 10,
             last_refill: refill_time,
+            action: LimitedAction::PublishNew,
         };
         assert_eq!(expected, bucket);
 
@@ -391,6 +407,7 @@ mod tests {
                 user_id: new_user(conn, "new_user")?,
                 tokens,
                 last_refill: now,
+                action: LimitedAction::PublishNew,
             })
             .get_result(conn)
     }

src/schema.rs

@@ -605,7 +605,7 @@ diesel::table! {
     /// Representation of the `publish_limit_buckets` table.
     ///
     /// (Automatically generated by Diesel.)
-    publish_limit_buckets (user_id) {
+    publish_limit_buckets (user_id, action) {
         /// The `user_id` column of the `publish_limit_buckets` table.
         ///
         /// Its SQL type is `Int4`.
@@ -624,14 +624,20 @@ diesel::table! {
         ///
         /// (Automatically generated by Diesel.)
         last_refill -> Timestamp,
+        /// The `action` column of the `publish_limit_buckets` table.
+        ///
+        /// Its SQL type is `Int4`.
+        ///
+        /// (Automatically generated by Diesel.)
+        action -> Int4,
     }
 }
 
 diesel::table! {
     /// Representation of the `publish_rate_overrides` table.
     ///
     /// (Automatically generated by Diesel.)
-    publish_rate_overrides (user_id) {
+    publish_rate_overrides (user_id, action) {
         /// The `user_id` column of the `publish_rate_overrides` table.
         ///
         /// Its SQL type is `Int4`.
@@ -650,6 +656,12 @@ diesel::table! {
         ///
         /// (Automatically generated by Diesel.)
         expires_at -> Nullable<Timestamp>,
+        /// The `action` column of the `publish_rate_overrides` table.
+        ///
+        /// Its SQL type is `Int4`.
+        ///
+        /// (Automatically generated by Diesel.)
+        action -> Int4,
     }
 }
 

src/worker/dump_db/dump-db.toml

@@ -143,11 +143,13 @@ total_downloads = "public"
 
 [publish_limit_buckets.columns]
 user_id = "private"
+action = "private"
 tokens = "private"
 last_refill = "private"
 
 [publish_rate_overrides.columns]
 user_id = "private"
+action = "private"
 burst = "private"
 expires_at = "private"