Refactor storage of TreeState in the DB

Author: Kobzol

.sqlx/query-6d68067c7121438630d77590ed644d4bc654978836fe5ceda2f0041aeba68e64.json .sqlx/query-cc7019b45037e155c4ce8f0bad9df0dcbe339fb59f93aec8bb6d34719ee042d3.json

@@ -4,14 +4,14 @@ use std::{
     str::FromStr,
 };
 
-use chrono::{DateTime, Utc};
-pub use client::PgDbClient;
-use sqlx::error::BoxDynError;
-
 use crate::{
     bors::RollupMode,
     github::{GithubRepoName, PullRequestNumber},
 };
+use chrono::{DateTime, Utc};
+pub use client::PgDbClient;
+use sqlx::error::BoxDynError;
+use sqlx::{Database, Postgres};
 
 mod client;
 pub(crate) mod operations;
@@ -211,6 +211,30 @@ pub enum TreeState {
     },
 }
 
+impl sqlx::Type<sqlx::Postgres> for TreeState {
+    fn type_info() -> sqlx::postgres::PgTypeInfo {
+        <(Option<i32>, Option<String>) as sqlx::Type<sqlx::Postgres>>::type_info()
+    }
+}
+
+impl sqlx::Decode<'_, sqlx::Postgres> for TreeState {
+    fn decode(value: <Postgres as Database>::ValueRef<'_>) -> Result<Self, BoxDynError> {
+        let data = <(Option<i32>, Option<String>) as sqlx::Decode<sqlx::Postgres>>::decode(value)?;
+        match data {
+            (Some(priority), Some(source)) => Ok(TreeState::Closed {
+                priority: priority as u32,
+                source,
+            }),
+            (None, None) => Ok(TreeState::Open),
+            _ => Err(
+                "Cannot deserialize TreeState, priority is non-NULL, but source is NULL"
+                    .to_string()
+                    .into(),
+            ),
+        }
+    }
+}
+
 /// Represents a repository configuration.
 pub struct RepoModel {
     pub id: PrimaryKey,

src/database/mod.rs

@@ -456,9 +456,17 @@ pub(crate) async fn get_repository(
     executor: impl PgExecutor<'_>,
     repo: &GithubRepoName,
 ) -> anyhow::Result<Option<RepoModel>> {
-    let repo = sqlx::query!(
+    let repo = sqlx::query_as!(
+        RepoModel,
         r#"
-        SELECT id, name as "name: GithubRepoName", tree_state, treeclosed_src, created_at
+        SELECT
+            id,
+            name as "name: GithubRepoName",
+            (
+                tree_state,
+                treeclosed_src
+            ) AS "tree_state!: TreeState",
+            created_at
         FROM repository
         WHERE name = $1
         "#,
@@ -467,23 +475,7 @@ pub(crate) async fn get_repository(
     .fetch_optional(executor)
     .await?;
 
-    Ok(repo.map(|repo| {
-        let tree_state = match repo.tree_state {
-            None => TreeState::Open,
-            Some(priority) => TreeState::Closed {
-                priority: priority as u32,
-                source: repo
-                    .treeclosed_src
-                    .expect("treeclosed_src is NULL even though tree_state is non-NULL"),
-            },
-        };
-        RepoModel {
-            id: repo.id,
-            name: repo.name,
-            tree_state,
-            created_at: repo.created_at,
-        }
-    }))
+    Ok(repo)
 }
 
 /// Updates the tree state of a repository.