Skip to content

Commit c0caaa1

Browse files
KobzolKobzol
committed
Store rotation mode status in the DB
1 parent 21730e8 commit c0caaa1

File tree

3 files changed

+116
-28
lines changed

3 files changed

+116
-28
lines changed

‎src/db.rs

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -345,5 +345,8 @@ CREATE UNIQUE INDEX IF NOT EXISTS review_prefs_user_id ON review_prefs(user_id);
345345
",
346346
"
347347
ALTER TABLE review_prefs ADD COLUMN IF NOT EXISTS max_assigned_prs INTEGER DEFAULT NULL;
348+
",
349+
"
350+
ALTER TABLE review_prefs ADD COLUMN IF NOT EXISTS rotation_mode TEXT NOT NULL DEFAULT 'on-rotation';
348351
",
349352
];

‎src/db/review_prefs.rs

Lines changed: 88 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,65 @@
11
use crate::db::users::record_username;
22
use crate::github::{User, UserId};
33
use anyhow::Context;
4+
use bytes::BytesMut;
5+
use postgres_types::{to_sql_checked, FromSql, IsNull, ToSql, Type};
46
use std::collections::HashMap;
7+
use std::error::Error;
8+
9+
#[derive(Debug, Default, Copy, Clone)]
10+
pub enum RotationMode {
11+
/// The reviewer can be automatically assigned by triagebot,
12+
/// and they can be assigned through teams and assign groups.
13+
#[default]
14+
OnRotation,
15+
/// The user is off rotation (e.g. on a vacation) and cannot be assigned automatically
16+
/// nor through teams and assign groups.
17+
OffRotation,
18+
}
19+
20+
impl<'a> FromSql<'a> for RotationMode {
21+
fn from_sql(ty: &Type, raw: &'a [u8]) -> Result<Self, Box<dyn Error + Sync + Send>> {
22+
let value = <&str as FromSql>::from_sql(ty, raw)?;
23+
match value {
24+
"on-rotation" => Ok(Self::OnRotation),
25+
"off-rotation" => Ok(Self::OffRotation),
26+
_ => Err(format!("Unknown value for RotationMode: {value}").into()),
27+
}
28+
}
29+
30+
fn accepts(ty: &Type) -> bool {
31+
<&str as FromSql>::accepts(ty)
32+
}
33+
}
34+
35+
impl ToSql for RotationMode {
36+
fn to_sql(&self, ty: &Type, out: &mut BytesMut) -> Result<IsNull, Box<dyn Error + Sync + Send>>
37+
where
38+
Self: Sized,
39+
{
40+
let value = match self {
41+
RotationMode::OnRotation => "on-rotation",
42+
RotationMode::OffRotation => "off-rotation",
43+
};
44+
<&str as ToSql>::to_sql(&value, ty, out)
45+
}
46+
47+
fn accepts(ty: &Type) -> bool
48+
where
49+
Self: Sized,
50+
{
51+
<&str as FromSql>::accepts(ty)
52+
}
53+
54+
to_sql_checked!();
55+
}
556

657
#[derive(Debug)]
758
pub struct ReviewPrefs {
859
pub id: uuid::Uuid,
960
pub user_id: i64,
1061
pub max_assigned_prs: Option<i32>,
62+
pub rotation_mode: RotationMode,
1163
}
1264

1365
impl From<tokio_postgres::row::Row> for ReviewPrefs {
@@ -16,6 +68,7 @@ impl From<tokio_postgres::row::Row> for ReviewPrefs {
1668
id: row.get("id"),
1769
user_id: row.get("user_id"),
1870
max_assigned_prs: row.get("max_assigned_prs"),
71+
rotation_mode: row.get("rotation_mode"),
1972
}
2073
}
2174
}
@@ -27,7 +80,7 @@ pub async fn get_review_prefs(
2780
user_id: UserId,
2881
) -> anyhow::Result<Option<ReviewPrefs>> {
2982
let query = "
30-
SELECT id, user_id, max_assigned_prs
83+
SELECT id, user_id, max_assigned_prs, rotation_mode
3184
FROM review_prefs
3285
WHERE review_prefs.user_id = $1;";
3386
let row = db
@@ -56,10 +109,15 @@ pub async fn get_review_prefs_batch<'a>(
56109
.collect();
57110
let lowercase_users: Vec<&str> = lowercase_map.keys().map(|s| s.as_str()).collect();
58111

59-
// The id/user_id/max_assigned_prs columns have to match the names used in
112+
// The id/user_id/max_assigned_prs/rotation_mode columns have to match the names used in
60113
// `From<tokio_postgres::row::Row> for ReviewPrefs`.
61114
let query = "
62-
SELECT lower(u.username) AS username, r.id AS id, r.user_id AS user_id, r.max_assigned_prs AS max_assigned_prs
115+
SELECT
116+
lower(u.username) AS username,
117+
r.id AS id,
118+
r.user_id AS user_id,
119+
r.max_assigned_prs AS max_assigned_prs,
120+
r.rotation_mode AS rotation_mode
63121
FROM review_prefs AS r
64122
JOIN users AS u ON u.user_id = r.user_id
65123
WHERE lower(u.username) = ANY($1);";
@@ -87,28 +145,33 @@ pub async fn upsert_review_prefs(
87145
db: &tokio_postgres::Client,
88146
user: User,
89147
max_assigned_prs: Option<u32>,
148+
rotation_mode: RotationMode,
90149
) -> anyhow::Result<u64, anyhow::Error> {
91150
// We need to have the user stored in the DB to have a valid FK link in review_prefs
92151
record_username(db, user.id, &user.login).await?;
93152

94153
let max_assigned_prs = max_assigned_prs.map(|v| v as i32);
95154
let query = "
96-
INSERT INTO review_prefs(user_id, max_assigned_prs)
97-
VALUES ($1, $2)
155+
INSERT INTO review_prefs(user_id, max_assigned_prs, rotation_mode)
156+
VALUES ($1, $2, $3)
98157
ON CONFLICT (user_id)
99158
DO UPDATE
100-
SET max_assigned_prs = excluded.max_assigned_prs";
159+
SET max_assigned_prs = excluded.max_assigned_prs,
160+
rotation_mode = excluded.rotation_mode";
101161

102162
let res = db
103-
.execute(query, &[&(user.id as i64), &max_assigned_prs])
163+
.execute(
164+
query,
165+
&[&(user.id as i64), &max_assigned_prs, &rotation_mode],
166+
)
104167
.await
105168
.context("Error upserting review preferences")?;
106169
Ok(res)
107170
}
108171

109172
#[cfg(test)]
110173
mod tests {
111-
use crate::db::review_prefs::{get_review_prefs, upsert_review_prefs};
174+
use crate::db::review_prefs::{get_review_prefs, upsert_review_prefs, RotationMode};
112175
use crate::db::users::get_user;
113176
use crate::tests::github::user;
114177
use crate::tests::run_db_test;
@@ -117,7 +180,13 @@ mod tests {
117180
async fn insert_prefs_create_user() {
118181
run_db_test(|ctx| async {
119182
let user = user("Martin", 1);
120-
upsert_review_prefs(&ctx.db_client(), user.clone(), Some(1)).await?;
183+
upsert_review_prefs(
184+
&ctx.db_client(),
185+
user.clone(),
186+
Some(1),
187+
RotationMode::OnRotation,
188+
)
189+
.await?;
121190
assert_eq!(get_user(&ctx.db_client(), user.id).await?.unwrap(), user);
122191

123192
Ok(ctx)
@@ -128,7 +197,13 @@ mod tests {
128197
#[tokio::test]
129198
async fn insert_max_assigned_prs() {
130199
run_db_test(|ctx| async {
131-
upsert_review_prefs(&ctx.db_client(), user("Martin", 1), Some(5)).await?;
200+
upsert_review_prefs(
201+
&ctx.db_client(),
202+
user("Martin", 1),
203+
Some(5),
204+
RotationMode::OnRotation,
205+
)
206+
.await?;
132207
assert_eq!(
133208
get_review_prefs(&ctx.db_client(), 1)
134209
.await?
@@ -147,18 +222,18 @@ mod tests {
147222
run_db_test(|ctx| async {
148223
let db = ctx.db_client();
149224

150-
upsert_review_prefs(&db, user("Martin", 1), Some(5)).await?;
225+
upsert_review_prefs(&db, user("Martin", 1), Some(5), RotationMode::OnRotation).await?;
151226
assert_eq!(
152227
get_review_prefs(&db, 1).await?.unwrap().max_assigned_prs,
153228
Some(5)
154229
);
155-
upsert_review_prefs(&db, user("Martin", 1), Some(10)).await?;
230+
upsert_review_prefs(&db, user("Martin", 1), Some(10), RotationMode::OnRotation).await?;
156231
assert_eq!(
157232
get_review_prefs(&db, 1).await?.unwrap().max_assigned_prs,
158233
Some(10)
159234
);
160235

161-
upsert_review_prefs(&db, user("Martin", 1), None).await?;
236+
upsert_review_prefs(&db, user("Martin", 1), None, RotationMode::OnRotation).await?;
162237
assert_eq!(
163238
get_review_prefs(&db, 1).await?.unwrap().max_assigned_prs,
164239
None

‎src/handlers/assign/tests/tests_candidates.rs

Lines changed: 25 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
//! Tests for `candidate_reviewers_from_names`
22
33
use super::super::*;
4-
use crate::db::review_prefs::upsert_review_prefs;
4+
use crate::db::review_prefs::{upsert_review_prefs, RotationMode};
55
use crate::github::{PullRequestNumber, User};
66
use crate::tests::github::{issue, user};
77
use crate::tests::{run_db_test, TestContext};
@@ -54,10 +54,20 @@ impl AssignCtx {
5454
self
5555
}
5656

57-
async fn set_max_capacity(self, user: &User, capacity: Option<u32>) -> Self {
58-
upsert_review_prefs(self.test_ctx.db_client(), user.clone(), capacity)
59-
.await
60-
.unwrap();
57+
async fn set_review_prefs(
58+
self,
59+
user: &User,
60+
capacity: Option<u32>,
61+
rotation_mode: RotationMode,
62+
) -> Self {
63+
upsert_review_prefs(
64+
self.test_ctx.db_client(),
65+
user.clone(),
66+
capacity,
67+
rotation_mode,
68+
)
69+
.await
70+
.unwrap();
6171
self
6272
}
6373

@@ -109,7 +119,7 @@ async fn no_assigned_prs() {
109119
run_db_test(|ctx| async move {
110120
let user = user("martin", 1);
111121
review_prefs_test(ctx)
112-
.set_max_capacity(&user, Some(3))
122+
.set_review_prefs(&user, Some(3), RotationMode::OnRotation)
113123
.await
114124
.check(&["martin"], Ok(&["martin"]))
115125
.await
@@ -134,7 +144,7 @@ async fn at_max_capacity() {
134144
run_db_test(|ctx| async move {
135145
let user = user("martin", 1);
136146
review_prefs_test(ctx)
137-
.set_max_capacity(&user, Some(3))
147+
.set_review_prefs(&user, Some(3), RotationMode::OnRotation)
138148
.await
139149
.assign_prs(user.id, 3)
140150
.check(
@@ -153,7 +163,7 @@ async fn below_max_capacity() {
153163
run_db_test(|ctx| async move {
154164
let user = user("martin", 1);
155165
review_prefs_test(ctx)
156-
.set_max_capacity(&user, Some(3))
166+
.set_review_prefs(&user, Some(3), RotationMode::OnRotation)
157167
.await
158168
.assign_prs(user.id, 2)
159169
.check(&["martin"], Ok(&["martin"]))
@@ -167,7 +177,7 @@ async fn above_max_capacity() {
167177
run_db_test(|ctx| async move {
168178
let user = user("martin", 1);
169179
review_prefs_test(ctx)
170-
.set_max_capacity(&user, Some(3))
180+
.set_review_prefs(&user, Some(3), RotationMode::OnRotation)
171181
.await
172182
.assign_prs(user.id, 10)
173183
.check(
@@ -186,7 +196,7 @@ async fn max_capacity_zero() {
186196
run_db_test(|ctx| async move {
187197
let user = user("martin", 1);
188198
review_prefs_test(ctx)
189-
.set_max_capacity(&user, Some(0))
199+
.set_review_prefs(&user, Some(0), RotationMode::OnRotation)
190200
.await
191201
.assign_prs(user.id, 0)
192202
.check(
@@ -205,7 +215,7 @@ async fn ignore_username_case() {
205215
run_db_test(|ctx| async move {
206216
let user = user("MARtin", 1);
207217
review_prefs_test(ctx)
208-
.set_max_capacity(&user, Some(3))
218+
.set_review_prefs(&user, Some(3), RotationMode::OnRotation)
209219
.await
210220
.assign_prs(user.id, 3)
211221
.check(
@@ -224,7 +234,7 @@ async fn unlimited_capacity() {
224234
run_db_test(|ctx| async move {
225235
let user = user("martin", 1);
226236
review_prefs_test(ctx)
227-
.set_max_capacity(&user, None)
237+
.set_review_prefs(&user, None, RotationMode::OnRotation)
228238
.await
229239
.assign_prs(user.id, 10)
230240
.check(&["martin"], Ok(&["martin"]))
@@ -240,11 +250,11 @@ async fn multiple_reviewers() {
240250
let teams = toml::toml!(team = ["martin", "jana", "mark", "diana"]);
241251
review_prefs_test(ctx)
242252
.teams(&teams)
243-
.set_max_capacity(&users[0], Some(3))
253+
.set_review_prefs(&users[0], Some(3), RotationMode::OnRotation)
244254
.await
245-
.set_max_capacity(&users[1], Some(4))
255+
.set_review_prefs(&users[1], Some(4), RotationMode::OnRotation)
246256
.await
247-
.set_max_capacity(&users[2], Some(2))
257+
.set_review_prefs(&users[2], Some(2), RotationMode::OnRotation)
248258
.await
249259
.assign_prs(users[0].id, 4)
250260
.assign_prs(users[1].id, 2)

0 commit comments

Comments
 (0)