From c94f20e06efc9993f6e394e691a7809497b6ced6 Mon Sep 17 00:00:00 2001 From: Denis Cornehl Date: Fri, 10 Nov 2023 06:32:08 +0100 Subject: [PATCH] use axum extractor to retrieve a new database connection from the pool --- src/web/error.rs | 11 ++++++-- src/web/extractors.rs | 54 +++++++++++++++++++++++++++++++++++++ src/web/mod.rs | 1 + src/web/releases.rs | 62 +++++++++++++------------------------------ 4 files changed, 83 insertions(+), 45 deletions(-) create mode 100644 src/web/extractors.rs diff --git a/src/web/error.rs b/src/web/error.rs index 47f08ee23..d5660f10b 100644 --- a/src/web/error.rs +++ b/src/web/error.rs @@ -1,13 +1,14 @@ -use std::borrow::Cow; - use crate::{ + db::PoolError, storage::PathNotFoundError, web::{releases::Search, AxumErrorPage}, }; +use anyhow::anyhow; use axum::{ http::StatusCode, response::{IntoResponse, Response as AxumResponse}, }; +use std::borrow::Cow; #[derive(Debug, thiserror::Error)] #[allow(dead_code)] // FIXME: remove after iron is gone @@ -131,6 +132,12 @@ impl From for AxumNope { } } +impl From for AxumNope { + fn from(err: PoolError) -> Self { + AxumNope::InternalError(anyhow!(err)) + } +} + pub(crate) type AxumResult = Result; #[cfg(test)] diff --git a/src/web/extractors.rs b/src/web/extractors.rs new file mode 100644 index 000000000..360f24079 --- /dev/null +++ b/src/web/extractors.rs @@ -0,0 +1,54 @@ +use crate::db::{AsyncPoolClient, Pool}; +use anyhow::Context as _; +use axum::{ + async_trait, + extract::{Extension, FromRequestParts}, + http::request::Parts, + RequestPartsExt, +}; +use std::ops::{Deref, DerefMut}; + +use super::error::AxumNope; + +/// Extractor for a async sqlx database connection. +/// Can be used in normal axum handlers, middleware, or other extractors. +/// +/// For now, we will retrieve a new connection each time the extractor is used. +/// +/// This could be optimized in the future by caching the connection as a request +/// extension, so one request only uses on connection. +#[derive(Debug)] +pub(crate) struct DbConnection(AsyncPoolClient); + +#[async_trait] +impl FromRequestParts for DbConnection +where + S: Send + Sync, +{ + type Rejection = AxumNope; + + async fn from_request_parts(parts: &mut Parts, _state: &S) -> Result { + let Extension(pool) = parts + .extract::>() + .await + .context("could not extract pool extension")?; + + Ok(Self(pool.get_async().await?)) + } +} + +impl Deref for DbConnection { + type Target = sqlx::PgConnection; + + fn deref(&self) -> &Self::Target { + &self.0 + } +} + +impl DerefMut for DbConnection { + fn deref_mut(&mut self) -> &mut Self::Target { + &mut self.0 + } +} + +// TODO: we will write tests for this when async db tests are working diff --git a/src/web/mod.rs b/src/web/mod.rs index 6448fd99a..f98100a79 100644 --- a/src/web/mod.rs +++ b/src/web/mod.rs @@ -15,6 +15,7 @@ pub(crate) mod cache; pub(crate) mod crate_details; mod csp; pub(crate) mod error; +mod extractors; mod features; mod file; mod headers; diff --git a/src/web/releases.rs b/src/web/releases.rs index 093012aec..b0cfa0130 100644 --- a/src/web/releases.rs +++ b/src/web/releases.rs @@ -9,6 +9,7 @@ use crate::{ web::{ axum_parse_uri_with_params, axum_redirect, encode_url_path, error::{AxumNope, AxumResult}, + extractors::DbConnection, match_version_axum, }, BuildQueue, Config, InstanceMetrics, @@ -133,7 +134,7 @@ struct SearchResult { /// /// This delegates to the crates.io search API. async fn get_search_results( - pool: Pool, + conn: &mut sqlx::PgConnection, config: &Config, query_params: &str, ) -> Result { @@ -212,10 +213,6 @@ async fn get_search_results( // So for now we are using the version with the youngest release_time. // This is different from all other release-list views where we show // our latest build. - let mut conn = pool - .get_async() - .await - .context("can't get pool connection")?; let crates: HashMap = sqlx::query!( r#"SELECT crates.name, @@ -276,12 +273,7 @@ impl_axum_webpage! { HomePage = "core/home.html", } -pub(crate) async fn home_page(Extension(pool): Extension) -> AxumResult { - let mut conn = pool - .get_async() - .await - .context("can't get pool connection")?; - +pub(crate) async fn home_page(mut conn: DbConnection) -> AxumResult { let recent_releases = get_releases(&mut conn, 1, RELEASES_IN_HOME, Order::ReleaseTime, true).await?; @@ -298,14 +290,7 @@ impl_axum_webpage! { content_type = "application/xml", } -pub(crate) async fn releases_feed_handler( - Extension(pool): Extension, -) -> AxumResult { - let mut conn = pool - .get_async() - .await - .context("can't get pool connection")?; - +pub(crate) async fn releases_feed_handler(mut conn: DbConnection) -> AxumResult { let recent_releases = get_releases(&mut conn, 1, RELEASES_IN_FEED, Order::ReleaseTime, true).await?; Ok(ReleaseFeed { recent_releases }) @@ -337,15 +322,10 @@ pub(crate) enum ReleaseType { } pub(crate) async fn releases_handler( - pool: Pool, + conn: &mut sqlx::PgConnection, page: Option, release_type: ReleaseType, ) -> AxumResult { - let mut conn = pool - .get_async() - .await - .context("can't get pool connection")?; - let page_number = page.unwrap_or(1); let (description, release_order, latest_only) = match release_type { @@ -368,7 +348,7 @@ pub(crate) async fn releases_handler( }; let releases = get_releases( - &mut conn, + &mut *conn, page_number, RELEASES_IN_RELEASES, release_order, @@ -395,30 +375,30 @@ pub(crate) async fn releases_handler( pub(crate) async fn recent_releases_handler( page: Option>, - Extension(pool): Extension, + mut conn: DbConnection, ) -> AxumResult { - releases_handler(pool, page.map(|p| p.0), ReleaseType::Recent).await + releases_handler(&mut conn, page.map(|p| p.0), ReleaseType::Recent).await } pub(crate) async fn releases_by_stars_handler( page: Option>, - Extension(pool): Extension, + mut conn: DbConnection, ) -> AxumResult { - releases_handler(pool, page.map(|p| p.0), ReleaseType::Stars).await + releases_handler(&mut conn, page.map(|p| p.0), ReleaseType::Stars).await } pub(crate) async fn releases_recent_failures_handler( page: Option>, - Extension(pool): Extension, + mut conn: DbConnection, ) -> AxumResult { - releases_handler(pool, page.map(|p| p.0), ReleaseType::RecentFailures).await + releases_handler(&mut conn, page.map(|p| p.0), ReleaseType::RecentFailures).await } pub(crate) async fn releases_failures_by_stars_handler( page: Option>, - Extension(pool): Extension, + mut conn: DbConnection, ) -> AxumResult { - releases_handler(pool, page.map(|p| p.0), ReleaseType::Failures).await + releases_handler(&mut conn, page.map(|p| p.0), ReleaseType::Failures).await } pub(crate) async fn owner_handler(Path(owner): Path) -> AxumResult { @@ -460,7 +440,7 @@ impl Default for Search { async fn redirect_to_random_crate( config: Arc, metrics: Arc, - pool: Pool, + conn: &mut sqlx::PgConnection, ) -> AxumResult { // We try to find a random crate and redirect to it. // @@ -468,11 +448,6 @@ async fn redirect_to_random_crate( // on the amount of crates with > 100 GH stars over the amount of all crates. // // If random-crate-searches end up being empty, increase that value. - let mut conn = pool - .get_async() - .await - .context("can't get pool connection")?; - let row = sqlx::query!( "WITH params AS ( -- get maximum possible id-value in crates-table @@ -519,6 +494,7 @@ impl_axum_webpage! { } pub(crate) async fn search_handler( + mut conn: DbConnection, Extension(pool): Extension, Extension(config): Extension>, Extension(metrics): Extension>, @@ -534,7 +510,7 @@ pub(crate) async fn search_handler( if params.remove("i-am-feeling-lucky").is_some() || query.contains("::") { // redirect to a random crate if query is empty if query.is_empty() { - return Ok(redirect_to_random_crate(config, metrics, pool) + return Ok(redirect_to_random_crate(config, metrics, &mut conn) .await? .into_response()); } @@ -595,14 +571,14 @@ pub(crate) async fn search_handler( return Err(AxumNope::NoResults); } - get_search_results(pool, &config, &query_params).await? + get_search_results(&mut conn, &config, &query_params).await? } else if !query.is_empty() { let query_params: String = form_urlencoded::Serializer::new(String::new()) .append_pair("q", &query) .append_pair("per_page", &RELEASES_IN_RELEASES.to_string()) .finish(); - get_search_results(pool, &config, &format!("?{}", &query_params)).await? + get_search_results(&mut conn, &config, &format!("?{}", &query_params)).await? } else { return Err(AxumNope::NoResults); };