Skip to content

Commit c6ba29d

Browse files
UrgauUrgau
authored
Merge pull request #1970 from Kobzol/on-rotation-through-db
Store vacation status in the database
2 parents ecec1ab + bd7aae3 commit c6ba29d

File tree

4 files changed

+206
-49
lines changed

4 files changed

+206
-49
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: 111 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,65 @@
11
use crate::db::users::record_username;
22
use crate::github::{User, UserId};
33
use anyhow::Context;
4-
use serde::Serialize;
4+
use bytes::BytesMut;
5+
use postgres_types::{to_sql_checked, FromSql, IsNull, ToSql, Type};
56
use std::collections::HashMap;
7+
use std::error::Error;
68

7-
#[derive(Debug, Serialize)]
9+
#[derive(Debug, Default, Copy, Clone, PartialEq)]
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+
}
56+
57+
#[derive(Debug)]
858
pub struct ReviewPrefs {
959
pub id: uuid::Uuid,
1060
pub user_id: i64,
1161
pub max_assigned_prs: Option<i32>,
62+
pub rotation_mode: RotationMode,
1263
}
1364

1465
impl From<tokio_postgres::row::Row> for ReviewPrefs {
@@ -17,6 +68,7 @@ impl From<tokio_postgres::row::Row> for ReviewPrefs {
1768
id: row.get("id"),
1869
user_id: row.get("user_id"),
1970
max_assigned_prs: row.get("max_assigned_prs"),
71+
rotation_mode: row.get("rotation_mode"),
2072
}
2173
}
2274
}
@@ -28,7 +80,7 @@ pub async fn get_review_prefs(
2880
user_id: UserId,
2981
) -> anyhow::Result<Option<ReviewPrefs>> {
3082
let query = "
31-
SELECT id, user_id, max_assigned_prs
83+
SELECT id, user_id, max_assigned_prs, rotation_mode
3284
FROM review_prefs
3385
WHERE review_prefs.user_id = $1;";
3486
let row = db
@@ -57,10 +109,15 @@ pub async fn get_review_prefs_batch<'a>(
57109
.collect();
58110
let lowercase_users: Vec<&str> = lowercase_map.keys().map(|s| s.as_str()).collect();
59111

60-
// 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
61113
// `From<tokio_postgres::row::Row> for ReviewPrefs`.
62114
let query = "
63-
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
64121
FROM review_prefs AS r
65122
JOIN users AS u ON u.user_id = r.user_id
66123
WHERE lower(u.username) = ANY($1);";
@@ -88,28 +145,33 @@ pub async fn upsert_review_prefs(
88145
db: &tokio_postgres::Client,
89146
user: User,
90147
max_assigned_prs: Option<u32>,
148+
rotation_mode: RotationMode,
91149
) -> anyhow::Result<u64, anyhow::Error> {
92150
// We need to have the user stored in the DB to have a valid FK link in review_prefs
93151
record_username(db, user.id, &user.login).await?;
94152

95153
let max_assigned_prs = max_assigned_prs.map(|v| v as i32);
96154
let query = "
97-
INSERT INTO review_prefs(user_id, max_assigned_prs)
98-
VALUES ($1, $2)
155+
INSERT INTO review_prefs(user_id, max_assigned_prs, rotation_mode)
156+
VALUES ($1, $2, $3)
99157
ON CONFLICT (user_id)
100158
DO UPDATE
101-
SET max_assigned_prs = excluded.max_assigned_prs";
159+
SET max_assigned_prs = excluded.max_assigned_prs,
160+
rotation_mode = excluded.rotation_mode";
102161

103162
let res = db
104-
.execute(query, &[&(user.id as i64), &max_assigned_prs])
163+
.execute(
164+
query,
165+
&[&(user.id as i64), &max_assigned_prs, &rotation_mode],
166+
)
105167
.await
106168
.context("Error upserting review preferences")?;
107169
Ok(res)
108170
}
109171

110172
#[cfg(test)]
111173
mod tests {
112-
use crate::db::review_prefs::{get_review_prefs, upsert_review_prefs};
174+
use crate::db::review_prefs::{get_review_prefs, upsert_review_prefs, RotationMode};
113175
use crate::db::users::get_user;
114176
use crate::tests::github::user;
115177
use crate::tests::run_db_test;
@@ -118,7 +180,13 @@ mod tests {
118180
async fn insert_prefs_create_user() {
119181
run_db_test(|ctx| async {
120182
let user = user("Martin", 1);
121-
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?;
122190
assert_eq!(get_user(&ctx.db_client(), user.id).await?.unwrap(), user);
123191

124192
Ok(ctx)
@@ -129,7 +197,13 @@ mod tests {
129197
#[tokio::test]
130198
async fn insert_max_assigned_prs() {
131199
run_db_test(|ctx| async {
132-
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?;
133207
assert_eq!(
134208
get_review_prefs(&ctx.db_client(), 1)
135209
.await?
@@ -148,18 +222,18 @@ mod tests {
148222
run_db_test(|ctx| async {
149223
let db = ctx.db_client();
150224

151-
upsert_review_prefs(&db, user("Martin", 1), Some(5)).await?;
225+
upsert_review_prefs(&db, user("Martin", 1), Some(5), RotationMode::OnRotation).await?;
152226
assert_eq!(
153227
get_review_prefs(&db, 1).await?.unwrap().max_assigned_prs,
154228
Some(5)
155229
);
156-
upsert_review_prefs(&db, user("Martin", 1), Some(10)).await?;
230+
upsert_review_prefs(&db, user("Martin", 1), Some(10), RotationMode::OnRotation).await?;
157231
assert_eq!(
158232
get_review_prefs(&db, 1).await?.unwrap().max_assigned_prs,
159233
Some(10)
160234
);
161235

162-
upsert_review_prefs(&db, user("Martin", 1), None).await?;
236+
upsert_review_prefs(&db, user("Martin", 1), None, RotationMode::OnRotation).await?;
163237
assert_eq!(
164238
get_review_prefs(&db, 1).await?.unwrap().max_assigned_prs,
165239
None
@@ -169,4 +243,26 @@ mod tests {
169243
})
170244
.await;
171245
}
246+
247+
#[tokio::test]
248+
async fn set_rotation_mode() {
249+
run_db_test(|ctx| async {
250+
let db = ctx.db_client();
251+
let user = user("Martin", 1);
252+
253+
upsert_review_prefs(&db, user.clone(), Some(5), RotationMode::OnRotation).await?;
254+
assert_eq!(
255+
get_review_prefs(&db, 1).await?.unwrap().rotation_mode,
256+
RotationMode::OnRotation
257+
);
258+
upsert_review_prefs(&db, user.clone(), Some(10), RotationMode::OffRotation).await?;
259+
assert_eq!(
260+
get_review_prefs(&db, 1).await?.unwrap().rotation_mode,
261+
RotationMode::OffRotation
262+
);
263+
264+
Ok(ctx)
265+
})
266+
.await;
267+
}
172268
}

‎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)