controllers/user/session: Migrate to diesel-async database queries

Author: Turbo87

src/controllers/user/session.rs

@@ -3,11 +3,10 @@ use axum::Json;
 use axum_extra::json;
 use axum_extra::response::ErasedJson;
 use diesel::QueryResult;
-use diesel_async::async_connection_wrapper::AsyncConnectionWrapper;
+use diesel_async::AsyncPgConnection;
 use http::request::Parts;
 use oauth2::reqwest::http_client;
 use oauth2::{AuthorizationCode, CsrfToken, Scope, TokenResponse};
-use tokio::runtime::Handle;
 
 use crate::app::AppState;
 use crate::email::Emails;
@@ -16,7 +15,7 @@ use crate::middleware::session::SessionExtension;
 use crate::models::{NewUser, User};
 use crate::schema::users;
 use crate::tasks::spawn_blocking;
-use crate::util::diesel::{is_read_only_error, Conn};
+use crate::util::diesel::is_read_only_error;
 use crate::util::errors::{bad_request, server_error, AppResult};
 use crate::views::EncodableMe;
 use crates_io_github::GithubUser;
@@ -90,54 +89,48 @@ pub async fn authorize(
     session: SessionExtension,
     req: Parts,
 ) -> AppResult<Json<EncodableMe>> {
+    // Make sure that the state we just got matches the session state that we
+    // should have issued earlier.
+    let session_state = session.remove("github_oauth_state").map(CsrfToken::new);
+    if !session_state.is_some_and(|state| query.state.secret() == state.secret()) {
+        return Err(bad_request("invalid state parameter"));
+    }
+
+    // Fetch the access token from GitHub using the code we just got
     let app_clone = app.clone();
     let request_log = req.request_log().clone();
-
-    let conn = app.db_write().await?;
-    spawn_blocking(move || {
-        let conn: &mut AsyncConnectionWrapper<_> = &mut conn.into();
-
-        // Make sure that the state we just got matches the session state that we
-        // should have issued earlier.
-        let session_state = session.remove("github_oauth_state").map(CsrfToken::new);
-        if !session_state.is_some_and(|state| query.state.secret() == state.secret()) {
-            return Err(bad_request("invalid state parameter"));
-        }
-
-        // Fetch the access token from GitHub using the code we just got
-        let token = app
+    let token = spawn_blocking(move || {
+        app_clone
             .github_oauth
             .exchange_code(query.code)
             .request(http_client)
             .map_err(|err| {
                 request_log.add("cause", err);
                 server_error("Error obtaining token")
-            })?;
+            })
+    })
+    .await?;
 
-        let token = token.access_token();
+    let token = token.access_token();
 
-        // Fetch the user info from GitHub using the access token we just got and create a user record
-        let ghuser = Handle::current().block_on(app.github.current_user(token))?;
-        let user = save_user_to_database(&ghuser, token.secret(), &app.emails, conn)?;
+    // Fetch the user info from GitHub using the access token we just got and create a user record
+    let ghuser = app.github.current_user(token).await?;
 
-        // Log in by setting a cookie and the middleware authentication
-        session.insert("user_id".to_string(), user.id.to_string());
+    let mut conn = app.db_write().await?;
+    let user = save_user_to_database(&ghuser, token.secret(), &app.emails, &mut conn).await?;
 
-        Ok(())
-    })
-    .await?;
+    // Log in by setting a cookie and the middleware authentication
+    session.insert("user_id".to_string(), user.id.to_string());
 
-    super::me::me(app_clone, req).await
+    super::me::me(app, req).await
 }
 
-fn save_user_to_database(
+async fn save_user_to_database(
     user: &GithubUser,
     access_token: &str,
     emails: &Emails,
-    conn: &mut impl Conn,
+    conn: &mut AsyncPgConnection,
 ) -> AppResult<User> {
-    use diesel::prelude::*;
-
     let new_user = NewUser::new(
         user.id,
         &user.login,
@@ -146,23 +139,30 @@ fn save_user_to_database(
         access_token,
     );
 
-    match new_user.create_or_update(user.email.as_deref(), emails, conn) {
+    match new_user
+        .create_or_update(user.email.as_deref(), emails, conn)
+        .await
+    {
         Ok(user) => Ok(user),
         Err(error) if is_read_only_error(&error) => {
             // If we're in read only mode, we can't update their details
             // just look for an existing user
-            find_user_by_gh_id(conn, user.id)?.ok_or_else(|| error.into())
+            find_user_by_gh_id(conn, user.id)
+                .await?
+                .ok_or_else(|| error.into())
         }
         Err(error) => Err(error.into()),
     }
 }
 
-fn find_user_by_gh_id(conn: &mut impl Conn, gh_id: i32) -> QueryResult<Option<User>> {
+async fn find_user_by_gh_id(conn: &mut AsyncPgConnection, gh_id: i32) -> QueryResult<Option<User>> {
     use diesel::prelude::*;
+    use diesel_async::RunQueryDsl;
 
     users::table
         .filter(users::gh_id.eq(gh_id))
         .first(conn)
+        .await
         .optional()
 }
 
@@ -175,20 +175,24 @@ pub async fn logout(session: SessionExtension) -> Json<bool> {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::test_util::test_db_connection;
+    use crates_io_test_db::TestDatabase;
+    use diesel_async::AsyncConnection;
 
-    #[test]
-    fn gh_user_with_invalid_email_doesnt_fail() {
+    #[tokio::test]
+    async fn gh_user_with_invalid_email_doesnt_fail() {
         let emails = Emails::new_in_memory();
-        let (_test_db, conn) = &mut test_db_connection();
+
+        let test_db = TestDatabase::new();
+        let mut conn = AsyncPgConnection::establish(test_db.url()).await.unwrap();
+
         let gh_user = GithubUser {
             email: Some("String.Format(\"{0}.{1}@live.com\", FirstName, LastName)".into()),
             name: Some("My Name".into()),
             login: "github_user".into(),
             id: -1,
             avatar_url: None,
         };
-        let result = save_user_to_database(&gh_user, "arbitrary_token", &emails, conn);
+        let result = save_user_to_database(&gh_user, "arbitrary_token", &emails, &mut conn).await;
 
         assert!(
             result.is_ok(),

src/models/user.rs

@@ -1,5 +1,6 @@
 use chrono::NaiveDateTime;
-use diesel_async::AsyncPgConnection;
+use diesel_async::scoped_futures::ScopedFutureExt;
+use diesel_async::{AsyncConnection, AsyncPgConnection};
 use secrecy::SecretString;
 
 use crate::app::App;
@@ -171,66 +172,72 @@ impl<'a> NewUser<'a> {
     }
 
     /// Inserts the user into the database, or updates an existing one.
-    pub fn create_or_update(
+    pub async fn create_or_update(
         &self,
         email: Option<&'a str>,
         emails: &Emails,
-        conn: &mut impl Conn,
+        conn: &mut AsyncPgConnection,
     ) -> QueryResult<User> {
         use diesel::dsl::sql;
         use diesel::insert_into;
         use diesel::pg::upsert::excluded;
         use diesel::sql_types::Integer;
-        use diesel::RunQueryDsl;
+        use diesel_async::RunQueryDsl;
 
         conn.transaction(|conn| {
-            let user: User = insert_into(users::table)
-                .values(self)
-                // We need the `WHERE gh_id > 0` condition here because `gh_id` set
-                // to `-1` indicates that we were unable to find a GitHub ID for
-                // the associated GitHub login at the time that we backfilled
-                // GitHub IDs. Therefore, there are multiple records in production
-                // that have a `gh_id` of `-1` so we need to exclude those when
-                // considering uniqueness of `gh_id` values. The `> 0` condition isn't
-                // necessary for most fields in the database to be used as a conflict
-                // target :)
-                .on_conflict(sql::<Integer>("(gh_id) WHERE gh_id > 0"))
-                .do_update()
-                .set((
-                    users::gh_login.eq(excluded(users::gh_login)),
-                    users::name.eq(excluded(users::name)),
-                    users::gh_avatar.eq(excluded(users::gh_avatar)),
-                    users::gh_access_token.eq(excluded(users::gh_access_token)),
-                ))
-                .get_result(conn)?;
-
-            // To send the user an account verification email
-            if let Some(user_email) = email {
-                let new_email = NewEmail {
-                    user_id: user.id,
-                    email: user_email,
-                };
-
-                let token = insert_into(emails::table)
-                    .values(&new_email)
-                    .on_conflict_do_nothing()
-                    .returning(emails::token)
-                    .get_result::<String>(conn)
-                    .optional()?
-                    .map(SecretString::from);
-
-                if let Some(token) = token {
-                    // Swallows any error. Some users might insert an invalid email address here.
-                    let email = UserConfirmEmail {
-                        user_name: &user.gh_login,
-                        domain: &emails.domain,
-                        token,
+            async move {
+                let user: User = insert_into(users::table)
+                    .values(self)
+                    // We need the `WHERE gh_id > 0` condition here because `gh_id` set
+                    // to `-1` indicates that we were unable to find a GitHub ID for
+                    // the associated GitHub login at the time that we backfilled
+                    // GitHub IDs. Therefore, there are multiple records in production
+                    // that have a `gh_id` of `-1` so we need to exclude those when
+                    // considering uniqueness of `gh_id` values. The `> 0` condition isn't
+                    // necessary for most fields in the database to be used as a conflict
+                    // target :)
+                    .on_conflict(sql::<Integer>("(gh_id) WHERE gh_id > 0"))
+                    .do_update()
+                    .set((
+                        users::gh_login.eq(excluded(users::gh_login)),
+                        users::name.eq(excluded(users::name)),
+                        users::gh_avatar.eq(excluded(users::gh_avatar)),
+                        users::gh_access_token.eq(excluded(users::gh_access_token)),
+                    ))
+                    .get_result(conn)
+                    .await?;
+
+                // To send the user an account verification email
+                if let Some(user_email) = email {
+                    let new_email = NewEmail {
+                        user_id: user.id,
+                        email: user_email,
                     };
-                    let _ = emails.send(user_email, email);
+
+                    let token = insert_into(emails::table)
+                        .values(&new_email)
+                        .on_conflict_do_nothing()
+                        .returning(emails::token)
+                        .get_result::<String>(conn)
+                        .await
+                        .optional()?
+                        .map(SecretString::from);
+
+                    if let Some(token) = token {
+                        // Swallows any error. Some users might insert an invalid email address here.
+                        let email = UserConfirmEmail {
+                            user_name: &user.gh_login,
+                            domain: &emails.domain,
+                            token,
+                        };
+                        let _ = emails.async_send(user_email, email).await;
+                    }
                 }
-            }
 
-            Ok(user)
+                Ok(user)
+            }
+            .scope_boxed()
         })
+        .await
     }
 }