rwf2
/
Rocket
mirror of https://github.com/rwf2/Rocket.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

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.
This commit is contained in:
Sergio Benitez 2022-05-24 16:47:09 -07:00
parent 5cb70ec58c
commit 04819d8cfd
8 changed files with 308 additions and 68 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

52
contrib/sync_db_pools/codegen/src/database.rs
View File

@ -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<syn::Attribute>,
/// 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<DatabaseInvocation> {
let attr_stream2 = crate::proc_macro2::TokenStream::from(attr);
let string_lit = crate::syn::parse2::<LitStr>(attr_stream2)?;
let string_lit = crate::syn::parse2::<syn::LitStr>(attr_stream2)?;
let input = crate::syn::parse::<DeriveInput>(input).unwrap();
let input = crate::syn::parse::<syn::DeriveInput>(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<DatabaseInv
};
Ok(DatabaseInvocation {
attrs: input.attrs,
type_name: input.ident,
visibility: input.vis,
db_name: string_lit.value(),
@ -59,6 +62,7 @@ pub fn database_attr(attr: TokenStream, input: TokenStream) -> Result<TokenStrea
// Store everything we're going to need to generate code.
let conn_type = &invocation.connection_type;
let name = &invocation.db_name;
let attrs = &invocation.attrs;
let guard_type = &invocation.type_name;
let vis = &invocation.visibility;
let fairing_name = format!("'{}' Database Pool", name);
@ -69,7 +73,7 @@ pub fn database_attr(attr: TokenStream, input: TokenStream) -> Result<TokenStrea
let rocket = quote!(#root::rocket);
let request_guard_type = quote_spanned! { span =>
#vis struct #guard_type(#root::Connection<Self, #conn_type>);
#(#attrs)* #vis struct #guard_type(#root::Connection<Self, #conn_type>);
};
let pool = quote_spanned!(span => #root::ConnectionPool<Self, #conn_type>);
@ -79,32 +83,30 @@ pub fn database_attr(attr: TokenStream, input: TokenStream) -> Result<TokenStrea
#request_guard_type
impl #guard_type {
/// Returns a fairing that initializes the associated database
/// connection pool.
/// Returns a fairing that initializes the database connection pool.
pub fn fairing() -> 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<P>(__rocket: &#rocket::Rocket<P>) -> Option<Self>
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<P: #rocket::Phase>(__rocket: &#rocket::Rocket<P>) -> 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<F, R>(&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<P: #rocket::Phase>(__rocket: &#rocket::Rocket<P>) -> Option<Self> {
<#pool>::get_one(&__rocket).await.map(Self)
}
}
#[#rocket::async_trait]

28
contrib/sync_db_pools/codegen/tests/ui-fail-nightly/database-types.stderr
View File

@ -11,6 +11,20 @@ note: required by a bound in `rocket_sync_db_pools::Connection`
| pub struct Connection<K, C: Poolable> {
| ^^^^^^^^ 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<K, C: Poolable> {
| ^^^^^^^^ 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<i32>: 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<K, C: Poolable> {
| ^^^^^^^^ required by this bound in `rocket_sync_db_pools::Connection`
error[E0277]: the trait bound `Vec<i32>: Poolable` is not satisfied
--> tests/ui-fail-nightly/database-types.rs:8:1
|
8 | #[database("foo")]
| ^^^^^^^^^^^^^^^^^^ the trait `Poolable` is not implemented for `Vec<i32>`
|
= 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<K, C: Poolable> {
| ^^^^^^^^ 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)

3
contrib/sync_db_pools/lib/Cargo.toml
View File

@ -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

5
contrib/sync_db_pools/lib/build.rs Normal file
View File

@ -0,0 +1,5 @@
fn main() {
if let Some(true) = version_check::is_feature_flaggable() {
println!("cargo:rustc-cfg=nightly");
}
}

19
contrib/sync_db_pools/lib/src/connection.rs
View File

@ -93,13 +93,13 @@ impl<K: 'static, C: Poolable> ConnectionPool<K, C> {
})
}
async fn get(&self) -> Result<Connection<K, C>, ()> {
pub async fn get(&self) -> Option<Connection<K, C>> {
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<K: 'static, C: Poolable> ConnectionPool<K, C> {
.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<P: Phase>(rocket: &Rocket<P>) -> Option<Connection<K, C>> {
match rocket.state::<Self>() {
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::<K>());
@ -137,13 +137,12 @@ impl<K: 'static, C: Poolable> ConnectionPool<K, C> {
}
#[inline]
pub async fn get_pool<P: Phase>(rocket: &Rocket<P>) -> Option<Self> {
rocket.state::<Self>().cloned()
pub fn pool<P: Phase>(rocket: &Rocket<P>) -> Option<&Self> {
rocket.state::<Self>()
}
}
impl<K: 'static, C: Poolable> Connection<K, C> {
#[inline]
pub async fn run<F, R>(&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<K, C> {
#[inline]
async fn from_request(request: &'r Request<'_>) -> Outcome<Self, ()> {
match request.rocket().state::<ConnectionPool<K, C>>() {
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::<K>());
Outcome::Failure((Status::InternalServerError, ()))

262
contrib/sync_db_pools/lib/src/lib.rs
View File

@ -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<P: Phase>(&Rocket<P>) -> Option<Self>`
//!
//! Retrieves a connection wrapper from the configured pool. Returns `Some`
//! as long as `Self::fairing()` has been attached.
//!
//! * `async fn run<R: Send + 'static>(&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<Self, diesel::SqliteConnection>);
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 {
<crate::ConnectionPool<Self, diesel::SqliteConnection>>::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<P: crate::rocket::Phase>(
__rocket: &crate::rocket::Rocket<P>,
) -> Option<&crate::ConnectionPool<Self, diesel::SqliteConnection>>
{
<crate::ConnectionPool<Self, diesel::SqliteConnection>>::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<F, R>(&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<P: crate::rocket::Phase>(
__rocket: &crate::rocket::Rocket<P>,
) -> Option<Self> {
<crate::ConnectionPool<Self, diesel::SqliteConnection>>::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<Self, ()>,
> + ::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<Self, ()>,
> {
return __ret;
}
let __r = __r;
let __ret: crate::rocket::request::Outcome<Self, ()> = {
< 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<crate::rocket::Ignite>,
) -> bool {
<crate::Connection<Self, diesel::SqliteConnection>>::abort(__r)
}
}
}

4
contrib/sync_db_pools/lib/tests/databases.rs
View File

@ -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);

3
core/lib/src/fs/temp_file.rs
View File

@ -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),