controllers/token: Remove last remaining spawn_blocking() call (#9988)

Turbo87 · web-flow · commit d8888c67e912 · 2024-11-18T14:33:54.000+01:00
diff --git a/src/controllers/token.rs b/src/controllers/token.rs
@@ -6,7 +6,6 @@ use crate::views::EncodableApiTokenWithToken;
 use crate::app::AppState;
 use crate::auth::AuthCheck;
 use crate::models::token::{CrateScope, EndpointScope};
-use crate::tasks::spawn_blocking;
 use crate::util::diesel::prelude::*;
 use crate::util::errors::{bad_request, AppResult};
 use axum::extract::{Path, Query};
@@ -17,7 +16,7 @@ use axum_extra::response::ErasedJson;
 use chrono::NaiveDateTime;
 use diesel::data_types::PgInterval;
 use diesel::dsl::{now, IntervalDsl};
-use diesel_async::async_connection_wrapper::AsyncConnectionWrapper;
+use diesel_async::RunQueryDsl;
 use http::request::Parts;
 use http::StatusCode;
 
@@ -42,8 +41,6 @@ pub async fn list(
     Query(params): Query<GetParams>,
     req: Parts,
 ) -> AppResult<ErasedJson> {
-    use diesel_async::RunQueryDsl;
-
     let mut conn = app.db_read_prefer_primary().await?;
     let auth = AuthCheck::only_cookie().check(&req, &mut conn).await?;
     let user = auth.user();
@@ -91,89 +88,85 @@ pub async fn new(
 
     let mut conn = app.db_write().await?;
     let auth = AuthCheck::default().check(&parts, &mut conn).await?;
-    spawn_blocking(move || {
-        use diesel::RunQueryDsl;
-
-        let conn: &mut AsyncConnectionWrapper<_> = &mut conn.into();
 
-        if auth.api_token_id().is_some() {
-            return Err(bad_request(
-                "cannot use an API token to create a new API token",
-            ));
-        }
+    if auth.api_token_id().is_some() {
+        return Err(bad_request(
+            "cannot use an API token to create a new API token",
+        ));
+    }
 
-        let user = auth.user();
+    let user = auth.user();
 
-        let max_token_per_user = 500;
-        let count: i64 = ApiToken::belonging_to(user).count().get_result(conn)?;
-        if count >= max_token_per_user {
-            return Err(bad_request(format!(
-                "maximum tokens per user is: {max_token_per_user}"
-            )));
-        }
+    let max_token_per_user = 500;
+    let count: i64 = ApiToken::belonging_to(user)
+        .count()
+        .get_result(&mut conn)
+        .await?;
+    if count >= max_token_per_user {
+        return Err(bad_request(format!(
+            "maximum tokens per user is: {max_token_per_user}"
+        )));
+    }
 
-        let crate_scopes = new
-            .api_token
-            .crate_scopes
-            .map(|scopes| {
-                scopes
-                    .into_iter()
-                    .map(CrateScope::try_from)
-                    .collect::<Result<Vec<_>, _>>()
-            })
-            .transpose()
-            .map_err(|_err| bad_request("invalid crate scope"))?;
-
-        let endpoint_scopes = new
-            .api_token
-            .endpoint_scopes
-            .map(|scopes| {
-                scopes
-                    .into_iter()
-                    .map(|scope| EndpointScope::try_from(scope.as_bytes()))
-                    .collect::<Result<Vec<_>, _>>()
-            })
-            .transpose()
-            .map_err(|_err| bad_request("invalid endpoint scope"))?;
-
-        let recipient = user.email(conn)?;
-
-        let api_token = ApiToken::insert_with_scopes(
-            conn,
-            user.id,
-            &new.api_token.name,
-            crate_scopes,
-            endpoint_scopes,
-            new.api_token.expired_at,
-        )?;
-
-        if let Some(recipient) = recipient {
-            let email = NewTokenEmail {
-                token_name: &new.api_token.name,
-                user_name: &user.gh_login,
-                domain: &app.emails.domain,
-            };
-
-            // At this point the token has been created so failing to send the
-            // email should not cause an error response to be returned to the
-            // caller.
-            let email_ret = app.emails.send(&recipient, email);
-            if let Err(e) = email_ret {
-                error!("Failed to send token creation email: {e}")
-            }
+    let crate_scopes = new
+        .api_token
+        .crate_scopes
+        .map(|scopes| {
+            scopes
+                .into_iter()
+                .map(CrateScope::try_from)
+                .collect::<Result<Vec<_>, _>>()
+        })
+        .transpose()
+        .map_err(|_err| bad_request("invalid crate scope"))?;
+
+    let endpoint_scopes = new
+        .api_token
+        .endpoint_scopes
+        .map(|scopes| {
+            scopes
+                .into_iter()
+                .map(|scope| EndpointScope::try_from(scope.as_bytes()))
+                .collect::<Result<Vec<_>, _>>()
+        })
+        .transpose()
+        .map_err(|_err| bad_request("invalid endpoint scope"))?;
+
+    let recipient = user.async_email(&mut conn).await?;
+
+    let api_token = ApiToken::insert_with_scopes(
+        &mut conn,
+        user.id,
+        &new.api_token.name,
+        crate_scopes,
+        endpoint_scopes,
+        new.api_token.expired_at,
+    )
+    .await?;
+
+    if let Some(recipient) = recipient {
+        let email = NewTokenEmail {
+            token_name: &new.api_token.name,
+            user_name: &user.gh_login,
+            domain: &app.emails.domain,
+        };
+
+        // At this point the token has been created so failing to send the
+        // email should not cause an error response to be returned to the
+        // caller.
+        let email_ret = app.emails.async_send(&recipient, email).await;
+        if let Err(e) = email_ret {
+            error!("Failed to send token creation email: {e}")
         }
+    }
 
-        let api_token = EncodableApiTokenWithToken::from(api_token);
+    let api_token = EncodableApiTokenWithToken::from(api_token);
 
-        Ok(json!({ "api_token": api_token }))
-    })
-    .await
+    Ok(json!({ "api_token": api_token }))
 }
 
 /// Handles the `GET /me/tokens/:id` route.
 pub async fn show(app: AppState, Path(id): Path<i32>, req: Parts) -> AppResult<ErasedJson> {
-    use diesel_async::RunQueryDsl;
-
     let mut conn = app.db_write().await?;
     let auth = AuthCheck::default().check(&req, &mut conn).await?;
     let user = auth.user();
@@ -188,8 +181,6 @@ pub async fn show(app: AppState, Path(id): Path<i32>, req: Parts) -> AppResult<E
 
 /// Handles the `DELETE /me/tokens/:id` route.
 pub async fn revoke(app: AppState, Path(id): Path<i32>, req: Parts) -> AppResult<ErasedJson> {
-    use diesel_async::RunQueryDsl;
-
     let mut conn = app.db_write().await?;
     let auth = AuthCheck::default().check(&req, &mut conn).await?;
     let user = auth.user();
@@ -203,8 +194,6 @@ pub async fn revoke(app: AppState, Path(id): Path<i32>, req: Parts) -> AppResult
 
 /// Handles the `DELETE /tokens/current` route.
 pub async fn revoke_current(app: AppState, req: Parts) -> AppResult<Response> {
-    use diesel_async::RunQueryDsl;
-
     let mut conn = app.db_write().await?;
     let auth = AuthCheck::default().check(&req, &mut conn).await?;
     let api_token_id = auth
diff --git a/src/models/token.rs b/src/models/token.rs
@@ -1,13 +1,14 @@
 mod scopes;
 
 use chrono::NaiveDateTime;
-use diesel_async::AsyncPgConnection;
+use diesel::dsl::now;
+use diesel::prelude::*;
+use diesel_async::scoped_futures::ScopedFutureExt;
+use diesel_async::{AsyncConnection, AsyncPgConnection, RunQueryDsl};
 
 pub use self::scopes::{CrateScope, EndpointScope};
 use crate::models::User;
 use crate::schema::api_tokens;
-use crate::util::diesel::prelude::*;
-use crate::util::diesel::Conn;
 use crate::util::rfc3339;
 use crate::util::token::{HashedToken, PlainToken};
 
@@ -35,20 +36,22 @@ pub struct ApiToken {
 
 impl ApiToken {
     /// Generates a new named API token for a user
-    pub fn insert(conn: &mut impl Conn, user_id: i32, name: &str) -> QueryResult<CreatedApiToken> {
-        Self::insert_with_scopes(conn, user_id, name, None, None, None)
+    pub async fn insert(
+        conn: &mut AsyncPgConnection,
+        user_id: i32,
+        name: &str,
+    ) -> QueryResult<CreatedApiToken> {
+        Self::insert_with_scopes(conn, user_id, name, None, None, None).await
     }
 
-    pub fn insert_with_scopes(
-        conn: &mut impl Conn,
+    pub async fn insert_with_scopes(
+        conn: &mut AsyncPgConnection,
         user_id: i32,
         name: &str,
         crate_scopes: Option<Vec<CrateScope>>,
         endpoint_scopes: Option<Vec<EndpointScope>>,
         expired_at: Option<NaiveDateTime>,
     ) -> QueryResult<CreatedApiToken> {
-        use diesel::RunQueryDsl;
-
         let token = PlainToken::generate();
 
         let model: ApiToken = diesel::insert_into(api_tokens::table)
@@ -61,7 +64,8 @@ impl ApiToken {
                 api_tokens::expired_at.eq(expired_at),
             ))
             .returning(ApiToken::as_returning())
-            .get_result(conn)?;
+            .get_result(conn)
+            .await?;
 
         Ok(CreatedApiToken {
             plaintext: token,
@@ -73,10 +77,6 @@ impl ApiToken {
         conn: &mut AsyncPgConnection,
         token: &HashedToken,
     ) -> QueryResult<ApiToken> {
-        use diesel::{dsl::now, update};
-        use diesel_async::scoped_futures::ScopedFutureExt;
-        use diesel_async::{AsyncConnection, RunQueryDsl};
-
         let tokens = api_tokens::table
             .filter(api_tokens::revoked.eq(false))
             .filter(
@@ -91,7 +91,7 @@ impl ApiToken {
         let token = conn
             .transaction(|conn| {
                 async move {
-                    update(tokens)
+                    diesel::update(tokens)
                         .set(api_tokens::last_used_at.eq(now.nullable()))
                         .returning(ApiToken::as_returning())
                         .get_result(conn)
diff --git a/src/tests/routes/me/tokens/create.rs b/src/tests/routes/me/tokens/create.rs
@@ -42,11 +42,11 @@ async fn create_token_no_name() {
 #[tokio::test(flavor = "multi_thread")]
 async fn create_token_exceeded_tokens_per_user() {
     let (app, _, user) = TestApp::init().with_user().await;
-    let mut conn = app.db_conn();
+    let mut conn = app.async_db_conn().await;
     let id = user.as_model().id;
 
     for i in 0..1000 {
-        assert_ok!(ApiToken::insert(&mut conn, id, &format!("token {i}")));
+        assert_ok!(ApiToken::insert(&mut conn, id, &format!("token {i}")).await);
     }
 
     let response = user.put::<()>("/api/v1/me/tokens", NEW_BAR).await;
diff --git a/src/tests/routes/me/tokens/get.rs b/src/tests/routes/me/tokens/get.rs
@@ -27,22 +27,25 @@ async fn show() {
 #[tokio::test(flavor = "multi_thread")]
 async fn show_token_with_scopes() {
     let (app, _, user) = TestApp::init().with_user().await;
-    let mut conn = app.db_conn();
+    let mut conn = app.async_db_conn().await;
     let user_model = user.as_model();
     let id = user_model.id;
 
-    assert_ok!(ApiToken::insert(&mut conn, id, "bar"));
-    let token = assert_ok!(ApiToken::insert_with_scopes(
-        &mut conn,
-        id,
-        "baz",
-        Some(vec![
-            CrateScope::try_from("serde").unwrap(),
-            CrateScope::try_from("serde-*").unwrap()
-        ]),
-        Some(vec![EndpointScope::PublishUpdate]),
-        Some((Utc::now() - Duration::days(31)).naive_utc()),
-    ));
+    assert_ok!(ApiToken::insert(&mut conn, id, "bar").await);
+    let token = assert_ok!(
+        ApiToken::insert_with_scopes(
+            &mut conn,
+            id,
+            "baz",
+            Some(vec![
+                CrateScope::try_from("serde").unwrap(),
+                CrateScope::try_from("serde-*").unwrap()
+            ]),
+            Some(vec![EndpointScope::PublishUpdate]),
+            Some((Utc::now() - Duration::days(31)).naive_utc()),
+        )
+        .await
+    );
 
     let url = format!("/api/v1/me/tokens/{}", token.model.id);
     let response = user.get::<()>(&url).await;
@@ -63,11 +66,11 @@ async fn show_with_anonymous_user() {
 #[tokio::test(flavor = "multi_thread")]
 async fn show_other_user_token() {
     let (app, _, user1) = TestApp::init().with_user().await;
-    let mut conn = app.db_conn();
+    let mut conn = app.async_db_conn().await;
     let user2 = app.db_new_user("baz").await;
     let user2 = user2.as_model();
 
-    let token = assert_ok!(ApiToken::insert(&mut conn, user2.id, "bar"));
+    let token = assert_ok!(ApiToken::insert(&mut conn, user2.id, "bar").await);
 
     let url = format!("/api/v1/me/tokens/{}", token.model.id);
     let response = user1.get::<()>(&url).await;
diff --git a/src/tests/routes/me/tokens/list.rs b/src/tests/routes/me/tokens/list.rs
diff --git a/src/tests/util.rs b/src/tests/util.rs
