Skip to content

Commit 133fce7

Browse files
pietroalbinipietroalbini
committed
add the action column to rate limit tables
1 parent 418185f commit 133fce7

File tree

5 files changed

+68
-21
lines changed

5 files changed

+68
-21
lines changed

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

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
DELETE FROM publish_limit_buckets WHERE action != 0;
2+
ALTER TABLE publish_limit_buckets DROP CONSTRAINT publish_limit_buckets_pkey;
3+
ALTER TABLE publish_limit_buckets ADD CONSTRAINT publish_limit_buckets_pkey PRIMARY KEY (user_id);
4+
ALTER TABLE publish_limit_buckets DROP COLUMN action;
5+
6+
DELETE FROM publish_rate_overrides WHERE action != 0;
7+
ALTER TABLE publish_rate_overrides DROP CONSTRAINT publish_rate_overrides_pkey;
8+
ALTER TABLE publish_rate_overrides ADD CONSTRAINT publish_rate_overrides_pkey PRIMARY KEY (user_id);
9+
ALTER TABLE publish_rate_overrides DROP COLUMN action;

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

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
ALTER TABLE publish_limit_buckets ADD COLUMN action INTEGER NOT NULL DEFAULT 0;
2+
ALTER TABLE publish_limit_buckets DROP CONSTRAINT publish_limit_buckets_pkey;
3+
ALTER TABLE publish_limit_buckets ADD CONSTRAINT publish_limit_buckets_pkey PRIMARY KEY (user_id, action);
4+
5+
ALTER TABLE publish_rate_overrides ADD COLUMN action INTEGER NOT NULL DEFAULT 0;
6+
ALTER TABLE publish_rate_overrides DROP CONSTRAINT publish_rate_overrides_pkey;
7+
ALTER TABLE publish_rate_overrides ADD CONSTRAINT publish_rate_overrides_pkey PRIMARY KEY (user_id, action);

‎src/publish_rate_limit.rs

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,12 @@ use std::time::Duration;
66
use crate::schema::{publish_limit_buckets, publish_rate_overrides};
77
use crate::util::errors::{AppResult, TooManyRequests};
88

9+
crate::pg_enum! {
10+
pub enum LimitedAction {
11+
PublishNew = 0,
12+
}
13+
}
14+
915
#[derive(Debug, Clone, Copy)]
1016
pub struct PublishRateLimit {
1117
pub rate: Duration,
@@ -38,6 +44,7 @@ struct Bucket {
3844
user_id: i32,
3945
tokens: i32,
4046
last_refill: NaiveDateTime,
47+
action: LimitedAction,
4148
}
4249

4350
impl PublishRateLimit {
@@ -79,7 +86,7 @@ impl PublishRateLimit {
7986
sql_function!(fn least<T>(x: T, y: T) -> T);
8087

8188
let burst: i32 = publish_rate_overrides::table
82-
.find(uploader)
89+
.find((uploader, LimitedAction::PublishNew))
8390
.filter(
8491
publish_rate_overrides::expires_at
8592
.is_null()
@@ -136,6 +143,7 @@ mod tests {
136143
user_id: bucket.user_id,
137144
tokens: 10,
138145
last_refill: now,
146+
action: LimitedAction::PublishNew,
139147
};
140148
assert_eq!(expected, bucket);
141149

@@ -148,6 +156,7 @@ mod tests {
148156
user_id: bucket.user_id,
149157
tokens: 20,
150158
last_refill: now,
159+
action: LimitedAction::PublishNew,
151160
};
152161
assert_eq!(expected, bucket);
153162
Ok(())
@@ -168,6 +177,7 @@ mod tests {
168177
user_id,
169178
tokens: 4,
170179
last_refill: now,
180+
action: LimitedAction::PublishNew,
171181
};
172182
assert_eq!(expected, bucket);
173183
Ok(())
@@ -189,6 +199,7 @@ mod tests {
189199
user_id,
190200
tokens: 6,
191201
last_refill: refill_time,
202+
action: LimitedAction::PublishNew,
192203
};
193204
assert_eq!(expected, bucket);
194205
Ok(())
@@ -214,6 +225,7 @@ mod tests {
214225
user_id,
215226
tokens: 7,
216227
last_refill: refill_time,
228+
action: LimitedAction::PublishNew,
217229
};
218230
assert_eq!(expected, bucket);
219231
Ok(())
@@ -235,6 +247,7 @@ mod tests {
235247
user_id,
236248
tokens: 6,
237249
last_refill: expected_refill_time,
250+
action: LimitedAction::PublishNew,
238251
};
239252
assert_eq!(expected, bucket);
240253
Ok(())
@@ -255,6 +268,7 @@ mod tests {
255268
user_id,
256269
tokens: 0,
257270
last_refill: now,
271+
action: LimitedAction::PublishNew,
258272
};
259273
assert_eq!(expected, bucket);
260274

@@ -279,6 +293,7 @@ mod tests {
279293
user_id,
280294
tokens: 1,
281295
last_refill: refill_time,
296+
action: LimitedAction::PublishNew,
282297
};
283298
assert_eq!(expected, bucket);
284299

@@ -301,6 +316,7 @@ mod tests {
301316
user_id,
302317
tokens: 10,
303318
last_refill: refill_time,
319+
action: LimitedAction::PublishNew,
304320
};
305321
assert_eq!(expected, bucket);
306322

@@ -399,6 +415,7 @@ mod tests {
399415
user_id: new_user(conn, "new_user")?,
400416
tokens,
401417
last_refill: now,
418+
action: LimitedAction::PublishNew,
402419
})
403420
.get_result(conn)
404421
}

‎src/schema.rs

Lines changed: 32 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -605,7 +605,7 @@ table! {
605605
/// Representation of the `publish_limit_buckets` table.
606606
///
607607
/// (Automatically generated by Diesel.)
608-
publish_limit_buckets (user_id) {
608+
publish_limit_buckets (user_id, action) {
609609
/// The `user_id` column of the `publish_limit_buckets` table.
610610
///
611611
/// Its SQL type is `Int4`.
@@ -624,6 +624,12 @@ table! {
624624
///
625625
/// (Automatically generated by Diesel.)
626626
last_refill -> Timestamp,
627+
/// The `action` column of the `publish_limit_buckets` table.
628+
///
629+
/// Its SQL type is `Int4`.
630+
///
631+
/// (Automatically generated by Diesel.)
632+
action -> Int4,
627633
}
628634
}
629635

@@ -634,7 +640,7 @@ table! {
634640
/// Representation of the `publish_rate_overrides` table.
635641
///
636642
/// (Automatically generated by Diesel.)
637-
publish_rate_overrides (user_id) {
643+
publish_rate_overrides (user_id, action) {
638644
/// The `user_id` column of the `publish_rate_overrides` table.
639645
///
640646
/// Its SQL type is `Int4`.
@@ -653,6 +659,30 @@ table! {
653659
///
654660
/// (Automatically generated by Diesel.)
655661
expires_at -> Nullable<Timestamp>,
662+
/// The `action` column of the `publish_rate_overrides` table.
663+
///
664+
/// Its SQL type is `Int4`.
665+
///
666+
/// (Automatically generated by Diesel.)
667+
action -> Int4,
668+
}
669+
}
670+
671+
table! {
672+
/// Representation of the `recent_crate_downloads` view.
673+
///
674+
/// This data represents the downloads in the last 90 days.
675+
/// This view does not contain realtime data.
676+
/// It is refreshed by the `update-downloads` script.
677+
recent_crate_downloads (crate_id) {
678+
/// The `crate_id` column of the `recent_crate_downloads` view.
679+
///
680+
/// Its SQL type is `Integer`.
681+
crate_id -> Integer,
682+
/// The `downloads` column of the `recent_crate_downloads` table.
683+
///
684+
/// Its SQL type is `BigInt`.
685+
downloads -> BigInt,
656686
}
657687
}
658688

@@ -679,24 +709,6 @@ table! {
679709
}
680710
}
681711

682-
table! {
683-
/// Representation of the `recent_crate_downloads` view.
684-
///
685-
/// This data represents the downloads in the last 90 days.
686-
/// This view does not contain realtime data.
687-
/// It is refreshed by the `update-downloads` script.
688-
recent_crate_downloads (crate_id) {
689-
/// The `crate_id` column of the `recent_crate_downloads` view.
690-
///
691-
/// Its SQL type is `Integer`.
692-
crate_id -> Integer,
693-
/// The `downloads` column of the `recent_crate_downloads` table.
694-
///
695-
/// Its SQL type is `BigInt`.
696-
downloads -> BigInt,
697-
}
698-
}
699-
700712
table! {
701713
use diesel::sql_types::*;
702714
use diesel_full_text_search::{TsVector as Tsvector};

‎src/tasks/dump_db/dump-db.toml

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -138,11 +138,13 @@ total_downloads = "public"
138138

139139
[publish_limit_buckets.columns]
140140
user_id = "private"
141+
action = "private"
141142
tokens = "private"
142143
last_refill = "private"
143144

144145
[publish_rate_overrides.columns]
145146
user_id = "private"
147+
action = "private"
146148
burst = "private"
147149
expires_at = "private"
148150

0 commit comments

Comments
 (0)