From 04819d8cfdae92e69162f051c5571234b14c5387 Mon Sep 17 00:00:00 2001 From: Sergio Benitez Date: Tue, 24 May 2022 16:47:09 -0700 Subject: [PATCH] Add pool retrieval to sync_db_pools. Generates a new method on attributed types, `pool()`, which returns an opaque reference to a type that can be used to get pooled connections. Also adds a code-generated example to the crate docs which includes real, proper function signatures and fully checked examples. Resolves #1884. Closes #1972. --- contrib/sync_db_pools/codegen/src/database.rs | 52 ++-- .../ui-fail-nightly/database-types.stderr | 28 ++ contrib/sync_db_pools/lib/Cargo.toml | 3 + contrib/sync_db_pools/lib/build.rs | 5 + contrib/sync_db_pools/lib/src/connection.rs | 19 +- contrib/sync_db_pools/lib/src/lib.rs | 262 ++++++++++++++++-- contrib/sync_db_pools/lib/tests/databases.rs | 4 +- core/lib/src/fs/temp_file.rs | 3 +- 8 files changed, 308 insertions(+), 68 deletions(-) create mode 100644 contrib/sync_db_pools/lib/build.rs diff --git a/contrib/sync_db_pools/codegen/src/database.rs b/contrib/sync_db_pools/codegen/src/database.rs index 8ec5e4e9..b5e4bc97 100644 --- a/contrib/sync_db_pools/codegen/src/database.rs +++ b/contrib/sync_db_pools/codegen/src/database.rs @@ -1,18 +1,20 @@ use proc_macro::TokenStream; use devise::{Spanned, Result, ext::SpanDiagnosticExt}; -use crate::syn::{Fields, Data, Type, LitStr, DeriveInput, Ident, Visibility}; +use crate::syn; #[derive(Debug)] struct DatabaseInvocation { + /// The attributes on the attributed structure. + attrs: Vec, /// The name of the structure on which `#[database(..)] struct This(..)` was invoked. - type_name: Ident, + type_name: syn::Ident, /// The visibility of the structure on which `#[database(..)] struct This(..)` was invoked. - visibility: Visibility, + visibility: syn::Visibility, /// The database name as passed in via #[database('database name')]. db_name: String, /// The type inside the structure: struct MyDb(ThisType). - connection_type: Type, + connection_type: syn::Type, } const EXAMPLE: &str = "example: `struct MyDatabase(diesel::SqliteConnection);`"; @@ -24,20 +26,20 @@ const NO_GENERIC_STRUCTS: &str = "`database` attribute cannot be applied to stru fn parse_invocation(attr: TokenStream, input: TokenStream) -> Result { let attr_stream2 = crate::proc_macro2::TokenStream::from(attr); - let string_lit = crate::syn::parse2::(attr_stream2)?; + let string_lit = crate::syn::parse2::(attr_stream2)?; - let input = crate::syn::parse::(input).unwrap(); + let input = crate::syn::parse::(input).unwrap(); if !input.generics.params.is_empty() { return Err(input.generics.span().error(NO_GENERIC_STRUCTS)); } let structure = match input.data { - Data::Struct(s) => s, + syn::Data::Struct(s) => s, _ => return Err(input.span().error(ONLY_ON_STRUCTS_MSG)) }; let inner_type = match structure.fields { - Fields::Unnamed(ref fields) if fields.unnamed.len() == 1 => { + syn::Fields::Unnamed(ref fields) if fields.unnamed.len() == 1 => { let first = fields.unnamed.first().expect("checked length"); first.ty.clone() } @@ -45,6 +47,7 @@ fn parse_invocation(attr: TokenStream, input: TokenStream) -> Result Result Result - #vis struct #guard_type(#root::Connection); + #(#attrs)* #vis struct #guard_type(#root::Connection); }; let pool = quote_spanned!(span => #root::ConnectionPool); @@ -79,32 +83,30 @@ pub fn database_attr(attr: TokenStream, input: TokenStream) -> Result impl #rocket::fairing::Fairing { <#pool>::fairing(#fairing_name, #name) } - /// Retrieves a connection of type `Self` from the `rocket` - /// instance. Returns `Some` as long as `Self::fairing()` has been - /// attached. - pub async fn get_one

(__rocket: &#rocket::Rocket

) -> Option - where P: #rocket::Phase, - { - <#pool>::get_one(&__rocket).await.map(Self) + /// Returns an opaque type that represents the connection pool + /// backing connections of type `Self`. + pub fn pool(__rocket: &#rocket::Rocket

) -> Option<&#pool> { + <#pool>::pool(&__rocket) } - /// Runs the provided closure on a thread from a threadpool. The - /// closure will be passed an `&mut r2d2::PooledConnection`. - /// `.await`ing the return value of this function yields the value - /// returned by the closure. + /// Runs the provided function `__f` in an async-safe blocking + /// thread. pub async fn run(&self, __f: F) -> R - where - F: FnOnce(&mut #conn_type) -> R + Send + 'static, - R: Send + 'static, + where F: FnOnce(&mut #conn_type) -> R + Send + 'static, + R: Send + 'static, { self.0.run(__f).await } + + /// Retrieves a connection of type `Self` from the `rocket` instance. + pub async fn get_one(__rocket: &#rocket::Rocket

) -> Option { + <#pool>::get_one(&__rocket).await.map(Self) + } } #[#rocket::async_trait] diff --git a/contrib/sync_db_pools/codegen/tests/ui-fail-nightly/database-types.stderr b/contrib/sync_db_pools/codegen/tests/ui-fail-nightly/database-types.stderr index 8c5c073b..b3f92220 100644 --- a/contrib/sync_db_pools/codegen/tests/ui-fail-nightly/database-types.stderr +++ b/contrib/sync_db_pools/codegen/tests/ui-fail-nightly/database-types.stderr @@ -11,6 +11,20 @@ note: required by a bound in `rocket_sync_db_pools::Connection` | pub struct Connection { | ^^^^^^^^ required by this bound in `rocket_sync_db_pools::Connection` +error[E0277]: the trait bound `Unknown: Poolable` is not satisfied + --> tests/ui-fail-nightly/database-types.rs:5:1 + | +5 | #[database("foo")] + | ^^^^^^^^^^^^^^^^^^ the trait `Poolable` is not implemented for `Unknown` + | + = help: the trait `Poolable` is implemented for `SqliteConnection` +note: required by a bound in `ConnectionPool` + --> $WORKSPACE/contrib/sync_db_pools/lib/src/connection.rs + | + | pub struct ConnectionPool { + | ^^^^^^^^ required by this bound in `ConnectionPool` + = note: this error originates in the attribute macro `database` (in Nightly builds, run with -Z macro-backtrace for more info) + error[E0277]: the trait bound `Vec: Poolable` is not satisfied --> tests/ui-fail-nightly/database-types.rs:9:10 | @@ -23,3 +37,17 @@ note: required by a bound in `rocket_sync_db_pools::Connection` | | pub struct Connection { | ^^^^^^^^ required by this bound in `rocket_sync_db_pools::Connection` + +error[E0277]: the trait bound `Vec: Poolable` is not satisfied + --> tests/ui-fail-nightly/database-types.rs:8:1 + | +8 | #[database("foo")] + | ^^^^^^^^^^^^^^^^^^ the trait `Poolable` is not implemented for `Vec` + | + = help: the trait `Poolable` is implemented for `SqliteConnection` +note: required by a bound in `ConnectionPool` + --> $WORKSPACE/contrib/sync_db_pools/lib/src/connection.rs + | + | pub struct ConnectionPool { + | ^^^^^^^^ required by this bound in `ConnectionPool` + = note: this error originates in the attribute macro `database` (in Nightly builds, run with -Z macro-backtrace for more info) diff --git a/contrib/sync_db_pools/lib/Cargo.toml b/contrib/sync_db_pools/lib/Cargo.toml index 1d8c2277..adc9d11d 100644 --- a/contrib/sync_db_pools/lib/Cargo.toml +++ b/contrib/sync_db_pools/lib/Cargo.toml @@ -43,5 +43,8 @@ version = "0.5.0-rc.2" path = "../../../core/lib" default-features = false +[build-dependencies] +version_check = "0.9.1" + [package.metadata.docs.rs] all-features = true diff --git a/contrib/sync_db_pools/lib/build.rs b/contrib/sync_db_pools/lib/build.rs new file mode 100644 index 00000000..99369685 --- /dev/null +++ b/contrib/sync_db_pools/lib/build.rs @@ -0,0 +1,5 @@ +fn main() { + if let Some(true) = version_check::is_feature_flaggable() { + println!("cargo:rustc-cfg=nightly"); + } +} diff --git a/contrib/sync_db_pools/lib/src/connection.rs b/contrib/sync_db_pools/lib/src/connection.rs index 47e1b23e..e3a0fd0a 100644 --- a/contrib/sync_db_pools/lib/src/connection.rs +++ b/contrib/sync_db_pools/lib/src/connection.rs @@ -93,13 +93,13 @@ impl ConnectionPool { }) } - async fn get(&self) -> Result, ()> { + pub async fn get(&self) -> Option> { let duration = std::time::Duration::from_secs(self.config.timeout as u64); let permit = match timeout(duration, self.semaphore.clone().acquire_owned()).await { Ok(p) => p.expect("internal invariant broken: semaphore should not be closed"), Err(_) => { error_!("database connection retrieval timed out"); - return Err(()); + return None; } }; @@ -107,22 +107,22 @@ impl ConnectionPool { .expect("internal invariant broken: self.pool is Some"); match run_blocking(move || pool.get_timeout(duration)).await { - Ok(c) => Ok(Connection { + Ok(c) => Some(Connection { connection: Arc::new(Mutex::new(Some(c))), permit: Some(permit), _marker: PhantomData, }), Err(e) => { error_!("failed to get a database connection: {}", e); - Err(()) + None } } } #[inline] pub async fn get_one(rocket: &Rocket

) -> Option> { - match rocket.state::() { - Some(pool) => match pool.get().await.ok() { + match Self::pool(rocket) { + Some(pool) => match pool.get().await { Some(conn) => Some(conn), None => { error_!("no connections available for `{}`", std::any::type_name::()); @@ -137,13 +137,12 @@ impl ConnectionPool { } #[inline] - pub async fn get_pool(rocket: &Rocket

) -> Option { - rocket.state::().cloned() + pub fn pool(rocket: &Rocket

) -> Option<&Self> { + rocket.state::() } } impl Connection { - #[inline] pub async fn run(&self, f: F) -> R where F: FnOnce(&mut C) -> R + Send + 'static, R: Send + 'static, @@ -207,7 +206,7 @@ impl<'r, K: 'static, C: Poolable> FromRequest<'r> for Connection { #[inline] async fn from_request(request: &'r Request<'_>) -> Outcome { match request.rocket().state::>() { - Some(c) => c.get().await.into_outcome(Status::ServiceUnavailable), + Some(c) => c.get().await.into_outcome((Status::ServiceUnavailable, ())), None => { error_!("Missing database fairing for `{}`", std::any::type_name::()); Outcome::Failure((Status::InternalServerError, ())) diff --git a/contrib/sync_db_pools/lib/src/lib.rs b/contrib/sync_db_pools/lib/src/lib.rs index ef05ac10..b3876742 100644 --- a/contrib/sync_db_pools/lib/src/lib.rs +++ b/contrib/sync_db_pools/lib/src/lib.rs @@ -100,10 +100,10 @@ //! //! ### `Rocket.toml` //! -//! To configure a database via `Rocket.toml`, add a table for each database -//! to the `databases` table where the key is a name of your choice. The table -//! should have a `url` key and, optionally, a `pool_size` key. This looks as -//! follows: +//! To configure a database via `Rocket.toml`, add a table for each database to +//! the `databases` table where the key is a name of your choice. The table +//! should have a `url` key and, optionally, `pool_size` and `timeout` keys. +//! This looks as follows: //! //! ```toml //! # Option 1: @@ -114,9 +114,11 @@ //! [global.databases.my_db] //! url = "postgres://root:root@localhost/my_db" //! -//! # With a `pool_size` key: -//! [global.databases] -//! sqlite_db = { url = "db.sqlite", pool_size = 20 } +//! # With `pool_size` and `timeout` keys: +//! [global.databases.sqlite_db] +//! url = "db.sqlite" +//! pool_size = 20 +//! timeout = 5 //! ``` //! //! The table _requires_ one key: @@ -127,6 +129,8 @@ //! //! * `pool_size` - the size of the pool, i.e., the number of connections to //! pool (defaults to the configured number of workers * 4) +//! * `timeout` - max number of seconds to wait for a connection to become +//! available (defaults to `5`) //! //! Additional options may be required or supported by other adapters. //! @@ -144,7 +148,8 @@ //! fn rocket() -> _ { //! let db: Map<_, Value> = map! { //! "url" => "db.sqlite".into(), -//! "pool_size" => 10.into() +//! "pool_size" => 10.into(), +//! "timeout" => 5.into(), //! }; //! //! let figment = rocket::Config::figment() @@ -182,30 +187,19 @@ //! database. This corresponds to the database name set as the database's //! configuration key. //! -//! The macro generates a [`FromRequest`] implementation for the decorated type, -//! allowing the type to be used as a request guard. This implementation -//! retrieves a connection from the database pool or fails with a -//! `Status::ServiceUnavailable` if connecting to the database times out. +//! See [`ExampleDb`](example::ExampleDb) for everything that the macro +//! generates. Specifically, it generates: //! -//! The macro also generates three inherent methods on the decorated type: +//! * A [`FromRequest`] implementation for the decorated type. +//! * A [`Sentinel`](rocket::Sentinel) implementation for the decorated type. +//! * A [`fairing()`](example::ExampleDb::fairing()) method to initialize the +//! database. +//! * A [`run()`](example::ExampleDb::run()) method to execute blocking +//! database operations in an `async`-safe manner. +//! * A [`pool()`](example::ExampleDb::pool()) method to retrieve the +//! backing connection pool. //! -//! * `fn fairing() -> impl Fairing` -//! -//! Returns a fairing that initializes the associated database connection -//! pool. -//! -//! * `async fn get_one(&Rocket

) -> Option` -//! -//! Retrieves a connection wrapper from the configured pool. Returns `Some` -//! as long as `Self::fairing()` has been attached. -//! -//! * `async fn run(&self, impl FnOnce(&mut Db) -> R + Send + 'static) -> R` -//! -//! Runs the specified function or closure, providing it access to the -//! underlying database connection (`&mut Db`). Returns the value returned -//! by the function or closure. -//! -//! The attribute can only be applied to unit-like structs with one type. The +//! The attribute can only be applied to tuple structs with one field. The //! internal type of the structure must implement [`Poolable`]. //! //! ```rust @@ -358,6 +352,7 @@ #![doc(html_root_url = "https://api.rocket.rs/v0.5-rc/rocket_sync_db_pools")] #![doc(html_favicon_url = "https://rocket.rs/images/favicon.ico")] #![doc(html_logo_url = "https://rocket.rs/images/logo-boxed.png")] +#![cfg_attr(nightly, feature(doc_cfg))] #[doc(hidden)] #[macro_use] @@ -392,3 +387,210 @@ pub use self::error::Error; pub use rocket_sync_db_pools_codegen::*; pub use self::connection::*; + +/// Example of code generated by the `#[database]` attribute. +#[cfg(all(nightly, doc, feature = "diesel_sqlite_pool"))] +pub mod example { + use crate::diesel; + + /// Example of code generated by the `#[database]` attribute. + /// + /// This implementation of `ExampleDb` was generated by: + /// + /// ```rust + /// use rocket_sync_db_pools::{database, diesel}; + /// + /// #[database("example")] + /// pub struct ExampleDb(diesel::SqliteConnection); + /// ``` + pub struct ExampleDb(crate::Connection); + + impl ExampleDb { + /// Returns a fairing that initializes the database connection pool + /// associated with `Self`. + /// + /// The fairing _must_ be attached before `Self` can be used as a + /// request guard. + /// + /// # Example + /// + /// ```rust + /// # #[macro_use] extern crate rocket; + /// # #[macro_use] extern crate rocket_sync_db_pools; + /// # + /// # #[cfg(feature = "diesel_sqlite_pool")] { + /// use rocket_sync_db_pools::diesel; + /// + /// #[database("my_db")] + /// struct MyConn(diesel::SqliteConnection); + /// + /// #[launch] + /// fn rocket() -> _ { + /// rocket::build().attach(MyConn::fairing()) + /// } + /// # } + /// ``` + pub fn fairing() -> impl crate::rocket::fairing::Fairing { + >::fairing( + "'example' Database Pool", + "example", + ) + } + + /// Returns an opaque type that represents the connection pool backing + /// connections of type `Self` _as long as_ the fairing returned by + /// [`Self::fairing()`] is attached and has run on `__rocket`. + /// + /// The returned pool is `Clone`. Values of type `Self` can be retrieved + /// from the pool by calling `pool.get().await` which has the same + /// signature and semantics as [`Self::get_one()`]. + /// + /// # Example + /// + /// ```rust + /// # #[macro_use] extern crate rocket; + /// # #[macro_use] extern crate rocket_sync_db_pools; + /// # + /// # #[cfg(feature = "diesel_sqlite_pool")] { + /// use rocket::tokio::{task, time}; + /// use rocket::fairing::AdHoc; + /// use rocket_sync_db_pools::diesel; + /// + /// #[database("my_db")] + /// struct MyConn(diesel::SqliteConnection); + /// + /// #[launch] + /// fn rocket() -> _ { + /// rocket::build() + /// .attach(MyConn::fairing()) + /// .attach(AdHoc::try_on_ignite("Background DB", |rocket| async { + /// let pool = match MyConn::pool(&rocket) { + /// Some(pool) => pool.clone(), + /// None => return Err(rocket) + /// }; + /// + /// // Start a background task that runs some database + /// // operation every 10 seconds. If a connection isn't + /// // available, retries 10 + timeout seconds later. + /// tokio::task::spawn(async move { + /// loop { + /// time::sleep(time::Duration::from_secs(10)).await; + /// if let Some(conn) = pool.get().await { + /// conn.run(|c| { /* perform db ops */ }).await; + /// } + /// } + /// }); + /// + /// Ok(rocket) + /// })) + /// } + /// # } + /// ``` + pub fn pool( + __rocket: &crate::rocket::Rocket

, + ) -> Option<&crate::ConnectionPool> + { + >::pool( + &__rocket, + ) + } + + /// Runs the provided function `__f` in an async-safe blocking thread. + /// The function is supplied with a mutable reference to the raw + /// connection (a value of type `&mut Self.0`). `.await`ing the return + /// value of this function yields the value returned by `__f`. + /// + /// # Example + /// + /// ```rust + /// # #[macro_use] extern crate rocket; + /// # #[macro_use] extern crate rocket_sync_db_pools; + /// # + /// # #[cfg(feature = "diesel_sqlite_pool")] { + /// use rocket_sync_db_pools::diesel; + /// + /// #[database("my_db")] + /// struct MyConn(diesel::SqliteConnection); + /// + /// #[get("/")] + /// async fn f(conn: MyConn) { + /// // The type annotation is illustrative and isn't required. + /// let result = conn.run(|c: &mut diesel::SqliteConnection| { + /// // Use `c`. + /// }).await; + /// } + /// # } + /// ``` + pub async fn run(&self, __f: F) -> R + where + F: FnOnce(&mut diesel::SqliteConnection) -> R + Send + 'static, + R: Send + 'static, + { + self.0.run(__f).await + } + + /// Retrieves a connection of type `Self` from the `rocket` instance. + /// Returns `Some` as long as `Self::fairing()` has been attached and + /// there is a connection available within at most `timeout` seconds. + pub async fn get_one( + __rocket: &crate::rocket::Rocket

, + ) -> Option { + >::get_one( + &__rocket, + ) + .await + .map(Self) + } + } + + /// Retrieves a connection from the database pool or fails with a + /// `Status::ServiceUnavailable` if doing so times out. + impl<'r> crate::rocket::request::FromRequest<'r> for ExampleDb { + type Error = (); + #[allow( + clippy::let_unit_value, + clippy::no_effect_underscore_binding, + clippy::shadow_same, + clippy::type_complexity, + clippy::type_repetition_in_bounds, + clippy::used_underscore_binding + )] + fn from_request<'life0, 'async_trait>( + __r: &'r crate::rocket::request::Request<'life0>, + ) -> ::core::pin::Pin< + Box< + dyn ::core::future::Future< + Output = crate::rocket::request::Outcome, + > + ::core::marker::Send + + 'async_trait, + >, + > + where + 'r: 'async_trait, + 'life0: 'async_trait, + Self: 'async_trait, + { + Box::pin(async move { + if let ::core::option::Option::Some(__ret) = ::core::option::Option::None::< + crate::rocket::request::Outcome, + > { + return __ret; + } + let __r = __r; + let __ret: crate::rocket::request::Outcome = { + < crate :: Connection < Self , diesel :: SqliteConnection > > + :: from_request (__r) . await . map (Self) + }; + #[allow(unreachable_code)] + __ret + }) + } + } + impl crate::rocket::Sentinel for ExampleDb { + fn abort( + __r: &crate::rocket::Rocket, + ) -> bool { + >::abort(__r) + } + } +} diff --git a/contrib/sync_db_pools/lib/tests/databases.rs b/contrib/sync_db_pools/lib/tests/databases.rs index bf0e5cb8..0bccbc56 100644 --- a/contrib/sync_db_pools/lib/tests/databases.rs +++ b/contrib/sync_db_pools/lib/tests/databases.rs @@ -2,8 +2,8 @@ mod databases_tests { use rocket_sync_db_pools::{database, diesel}; - #[database("foo")] - struct TempStorage(diesel::SqliteConnection); + #[database("example")] + struct ExampleDb(diesel::SqliteConnection); #[database("bar")] struct PrimaryDb(diesel::PgConnection); diff --git a/core/lib/src/fs/temp_file.rs b/core/lib/src/fs/temp_file.rs index 1c6df57e..7cfbab18 100644 --- a/core/lib/src/fs/temp_file.rs +++ b/core/lib/src/fs/temp_file.rs @@ -458,7 +458,8 @@ impl<'v> TempFile<'v> { let mut file = File::from_std(file); let fut = data.open(limit).stream_to(tokio::io::BufWriter::new(&mut file)); - let n = fut.await?; + let n = fut.await; + let n = n?; let temp_file = TempFile::File { content_type, file_name, path: Either::Left(temp_path),