From f7f068be11efb33844d308dec7bcfcd770e2bffa Mon Sep 17 00:00:00 2001 From: Jeb Rosen Date: Mon, 31 May 2021 09:15:00 -0700 Subject: [PATCH] Initial implementation of async DB pooling. --- Cargo.toml | 2 + contrib/db_pools/README.md | 60 ++++ contrib/db_pools/codegen/Cargo.toml | 17 + contrib/db_pools/codegen/src/database.rs | 72 +++++ contrib/db_pools/codegen/src/lib.rs | 41 +++ contrib/db_pools/lib/Cargo.toml | 40 +++ contrib/db_pools/lib/src/config.rs | 64 ++++ contrib/db_pools/lib/src/database.rs | 119 +++++++ contrib/db_pools/lib/src/error.rs | 41 +++ contrib/db_pools/lib/src/lib.rs | 382 +++++++++++++++++++++++ contrib/db_pools/lib/src/pool.rs | 287 +++++++++++++++++ examples/README.md | 3 +- examples/databases/Cargo.toml | 6 +- examples/databases/src/sqlx.rs | 67 ++-- scripts/test.sh | 15 + 15 files changed, 1174 insertions(+), 42 deletions(-) create mode 100644 contrib/db_pools/README.md create mode 100644 contrib/db_pools/codegen/Cargo.toml create mode 100644 contrib/db_pools/codegen/src/database.rs create mode 100644 contrib/db_pools/codegen/src/lib.rs create mode 100644 contrib/db_pools/lib/Cargo.toml create mode 100644 contrib/db_pools/lib/src/config.rs create mode 100644 contrib/db_pools/lib/src/database.rs create mode 100644 contrib/db_pools/lib/src/error.rs create mode 100644 contrib/db_pools/lib/src/lib.rs create mode 100644 contrib/db_pools/lib/src/pool.rs diff --git a/Cargo.toml b/Cargo.toml index 52fc6a68..8ec081ef 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -3,6 +3,8 @@ members = [ "core/lib/", "core/codegen/", "core/http/", + "contrib/db_pools/codegen/", + "contrib/db_pools/lib/", "contrib/sync_db_pools/codegen/", "contrib/sync_db_pools/lib/", "contrib/dyn_templates/", diff --git a/contrib/db_pools/README.md b/contrib/db_pools/README.md new file mode 100644 index 00000000..755974a9 --- /dev/null +++ b/contrib/db_pools/README.md @@ -0,0 +1,60 @@ +# `db_pools` [![ci.svg]][ci] [![crates.io]][crate] [![docs.svg]][crate docs] + +[crates.io]: https://img.shields.io/crates/v/rocket_db_pools.svg +[crate]: https://crates.io/crates/rocket_db_pools +[docs.svg]: https://img.shields.io/badge/web-master-red.svg?style=flat&label=docs&colorB=d33847 +[crate docs]: https://api.rocket.rs/master/rocket_db_pools +[ci.svg]: https://github.com/SergioBenitez/Rocket/workflows/CI/badge.svg +[ci]: https://github.com/SergioBenitez/Rocket/actions + +This crate provides traits, utilities, and a procedural macro for configuring +and accessing database connection pools in Rocket. + +## Usage + +First, enable the feature corresponding to your database type: + +```toml +[dependencies.rocket_db_pools] +version = "0.1.0-dev" +features = ["sqlx_sqlite"] +``` + +A full list of supported databases and their associated feature names is +available in the [crate docs]. In whichever configuration source you choose, +configure a `databases` dictionary with a key for each database, here +`sqlite_logs` in a TOML source: + +```toml +[default.databases] +sqlite_logs = { url = "/path/to/database.sqlite" } +``` + +In your application's source code: + +```rust +#[macro_use] extern crate rocket; +use rocket::serde::json::Json; + +use rocket_db_pools::{Database, sqlx}; + +#[derive(Database)] +#[database("sqlite_logs")] +struct LogsDb(sqlx::SqlitePool); + +type LogsDbConn = ::Connection; + +#[get("/logs/")] +async fn get_logs(mut db: LogsDbConn, id: usize) -> Result>> { + let logs = sqlx::query!("SELECT text FROM logs;").execute(&mut *db).await?; + + Ok(Json(logs)) +} + +#[launch] +fn rocket() -> _ { + rocket::build().attach(LogsDb::fairing()) +} +``` + +See the [crate docs] for full details. diff --git a/contrib/db_pools/codegen/Cargo.toml b/contrib/db_pools/codegen/Cargo.toml new file mode 100644 index 00000000..be90e892 --- /dev/null +++ b/contrib/db_pools/codegen/Cargo.toml @@ -0,0 +1,17 @@ +[package] +name = "rocket_db_pools_codegen" +version = "0.1.0-dev" +authors = ["Sergio Benitez ", "Jeb Rosen "] +description = "Procedural macros for rocket_db_pools." +repository = "https://github.com/SergioBenitez/Rocket/contrib/db_pools" +readme = "../README.md" +keywords = ["rocket", "framework", "database", "pools"] +license = "MIT OR Apache-2.0" +edition = "2018" + +[lib] +proc-macro = true + +[dependencies] +devise = "0.3" +quote = "1" diff --git a/contrib/db_pools/codegen/src/database.rs b/contrib/db_pools/codegen/src/database.rs new file mode 100644 index 00000000..f3d168a8 --- /dev/null +++ b/contrib/db_pools/codegen/src/database.rs @@ -0,0 +1,72 @@ +use proc_macro::TokenStream; + +use devise::{DeriveGenerator, FromMeta, MapperBuild, Support, ValidatorBuild}; +use devise::proc_macro2_diagnostics::SpanDiagnosticExt; +use devise::syn::{Fields, spanned::Spanned}; + +#[derive(Debug, FromMeta)] +struct DatabaseAttribute { + #[meta(naked)] + name: String, +} + +const ONE_DATABASE_ATTR: &str = "`Database` derive requires exactly one \ + `#[database(\"\")] attribute`"; +const ONE_UNNAMED_FIELD: &str = "`Database` derive can only be applied to \ + structs with exactly one unnamed field"; + +pub fn derive_database(input: TokenStream) -> TokenStream { + DeriveGenerator::build_for(input, quote!(impl rocket_db_pools::Database)) + .support(Support::TupleStruct) + .validator(ValidatorBuild::new() + .struct_validate(|_, struct_| { + if struct_.fields.len() == 1 { + Ok(()) + } else { + return Err(struct_.fields.span().error(ONE_UNNAMED_FIELD)) + } + }) + ) + .inner_mapper(MapperBuild::new() + .try_struct_map(|_, struct_| { + let krate = quote_spanned!(struct_.span() => ::rocket_db_pools); + let db_name = match DatabaseAttribute::one_from_attrs("database", &struct_.attrs)? { + Some(attr) => attr.name, + None => return Err(struct_.span().error(ONE_DATABASE_ATTR)), + }; + let fairing_name = format!("'{}' Database Pool", db_name); + + let pool_type = match &struct_.fields { + Fields::Unnamed(f) => &f.unnamed[0].ty, + _ => unreachable!("Support::TupleStruct"), + }; + + Ok(quote_spanned! { struct_.span() => + const NAME: &'static str = #db_name; + type Pool = #pool_type; + fn fairing() -> #krate::Fairing { + #krate::Fairing::new(#fairing_name) + } + fn pool(&self) -> &Self::Pool { &self.0 } + }) + }) + ) + .outer_mapper(MapperBuild::new() + .try_struct_map(|_, struct_| { + let decorated_type = &struct_.ident; + let pool_type = match &struct_.fields { + Fields::Unnamed(f) => &f.unnamed[0].ty, + _ => unreachable!("Support::TupleStruct"), + }; + + Ok(quote_spanned! { struct_.span() => + impl From<#pool_type> for #decorated_type { + fn from(pool: #pool_type) -> Self { + Self(pool) + } + } + }) + }) + ) + .to_tokens() +} diff --git a/contrib/db_pools/codegen/src/lib.rs b/contrib/db_pools/codegen/src/lib.rs new file mode 100644 index 00000000..9eb765f8 --- /dev/null +++ b/contrib/db_pools/codegen/src/lib.rs @@ -0,0 +1,41 @@ +#![recursion_limit="256"] + +#![warn(rust_2018_idioms)] + +//! # `rocket_databases` - Code Generation +//! +//! This crate implements the code generation portion of the `rocket_databases` +//! crate. + +#[macro_use] extern crate quote; + +mod database; + +use proc_macro::TokenStream; + +/// Defines a database type and implements [`Database`] on it. +/// +/// ```ignore +/// #[derive(Database)] +/// #[database("database_name")] +/// struct Db(PoolType); +/// ``` +/// +/// `PoolType` must implement [`Pool`]. +/// +/// This macro generates the following code, implementing the [`Database`] trait +/// on the struct. Custom implementations of `Database` should usually also +/// start with roughly this code: +/// +/// ```ignore +/// impl Database for Db { +/// const NAME: &'static str = "config_name"; +/// type Pool = PoolType; +/// fn fairing() -> Fairing { Fairing::new(|p| Self(p)) } +/// fn pool(&self) -> &Self::Pool { &self.0 } +/// } +/// ``` +#[proc_macro_derive(Database, attributes(database))] +pub fn derive_database(input: TokenStream) -> TokenStream { + crate::database::derive_database(input) +} diff --git a/contrib/db_pools/lib/Cargo.toml b/contrib/db_pools/lib/Cargo.toml new file mode 100644 index 00000000..4b6cd23b --- /dev/null +++ b/contrib/db_pools/lib/Cargo.toml @@ -0,0 +1,40 @@ +[package] +name = "rocket_db_pools" +version = "0.1.0-dev" +authors = ["Sergio Benitez ", "Jeb Rosen "] +description = "Rocket async database pooling support" +repository = "https://github.com/SergioBenitez/Rocket/contrib/db_pools" +readme = "../README.md" +keywords = ["rocket", "framework", "database", "pools"] +license = "MIT OR Apache-2.0" +edition = "2018" + +[features] +deadpool_postgres = ["deadpool-postgres"] +deadpool_redis = ["deadpool-redis"] +sqlx_mysql = ["sqlx", "sqlx/mysql"] +sqlx_postgres = ["sqlx", "sqlx/postgres"] +sqlx_sqlite = ["sqlx", "sqlx/sqlite"] + +[dependencies] +rocket_db_pools_codegen = { path = "../codegen" } + +# integration-specific +deadpool-postgres = { version = "0.9", default-features = false, optional = true } +deadpool-redis = { version = "0.8", default-features = false, optional = true } +mongodb = { version = "1", default-features = false, features = ["tokio-runtime"], optional = true } +mysql_async = { version = "0.27", default-features = false, optional = true } +redis = { version = "0.20", default-features = false, features = ["aio", "tokio-comp"] } +sqlx = { version = "0.5", default-features = false, features = ["runtime-tokio-rustls"], optional = true } + +[dependencies.rocket] +path = "../../../core/lib" +default-features = false + +[package.metadata.docs.rs] +all-features = true + +[dev-dependencies.rocket] +path = "../../../core/lib" +default-features = false +features = ["json"] diff --git a/contrib/db_pools/lib/src/config.rs b/contrib/db_pools/lib/src/config.rs new file mode 100644 index 00000000..c4909cb1 --- /dev/null +++ b/contrib/db_pools/lib/src/config.rs @@ -0,0 +1,64 @@ +use rocket::figment::{self, Figment, providers::Serialized}; +use rocket::serde::{Deserialize, Serialize}; +use rocket::{Build, Rocket}; + +/// A base `Config` for any `Pool` type. +/// +/// For the following configuration: +/// +/// ```toml +/// [global.databases.my_database] +/// url = "postgres://root:root@localhost/my_database" +/// pool_size = 10 +/// ``` +/// +/// ...the following struct would be passed to [`Pool::initialize()`]: +/// +/// ```rust +/// # use rocket_db_pools::Config; +/// Config { +/// url: "postgres://root:root@localhost/my_database".into(), +/// pool_size: 10, +/// timeout: 5, +/// }; +/// ``` +/// +/// If you want to implement your own custom database adapter and need some more +/// configuration options, you may need to define a custom `Config` struct. +/// +/// [`Pool::initialize()`]: crate::Pool::initialize +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)] +#[serde(crate = "rocket::serde")] +pub struct Config { + /// Connection URL specified in the Rocket configuration. + pub url: String, + /// Initial pool size. Defaults to the number of Rocket workers * 4. + pub pool_size: u32, + /// How long to wait, in seconds, for a new connection before timing out. + /// Defaults to `5`. + // FIXME: Use `time`. + pub timeout: u8, +} + +impl Config { + pub fn from(db_name: &str, rocket: &Rocket) -> Result { + Self::figment(db_name, rocket).extract::() + } + + pub fn figment(db_name: &str, rocket: &Rocket) -> Figment { + let db_key = format!("databases.{}", db_name); + let default_pool_size = rocket.figment() + .extract_inner::(rocket::Config::WORKERS) + .map(|workers| workers * 4) + .ok(); + + let figment = Figment::from(rocket.figment()) + .focus(&db_key) + .join(Serialized::default("timeout", 5)); + + match default_pool_size { + Some(pool_size) => figment.join(Serialized::default("pool_size", pool_size)), + None => figment + } + } +} diff --git a/contrib/db_pools/lib/src/database.rs b/contrib/db_pools/lib/src/database.rs new file mode 100644 index 00000000..2aa461b1 --- /dev/null +++ b/contrib/db_pools/lib/src/database.rs @@ -0,0 +1,119 @@ +use rocket::fairing::{Info, Kind}; +use rocket::futures::future::BoxFuture; +use rocket::http::Status; +use rocket::request::{FromRequest, Outcome, Request}; +use rocket::yansi::Paint; +use rocket::{Build, Ignite, Rocket, Sentinel}; + +use crate::{Error, Pool}; + +/// Trait implemented to define a database connection pool. +pub trait Database: Sized + Send + Sync + 'static { + /// The name of this connection pool in the configuration. + const NAME: &'static str; + + /// The underlying connection type returned by this pool. + /// Must implement [`Pool`]. + type Pool: Pool; + + /// Returns a fairing that attaches this connection pool to the server. + fn fairing() -> Fairing; + + /// Direct shared access to the underlying database pool + fn pool(&self) -> &Self::Pool; + + /// get().await returns a connection from the pool (or an error) + fn get(&self) -> BoxFuture<'_, Result, ::GetError>> { + Box::pin(async move { self.pool().get().await.map(Connection)} ) + } +} + +/// A connection. The underlying connection type is determined by `D`, which +/// must implement [`Database`]. +pub struct Connection(::Connection); + +impl std::ops::Deref for Connection { + type Target = ::Connection; + fn deref(&self) -> &Self::Target { + &self.0 + } +} + +impl std::ops::DerefMut for Connection { + fn deref_mut(&mut self) -> &mut Self::Target { + &mut self.0 + } +} + +#[rocket::async_trait] +impl<'r, D: Database> FromRequest<'r> for Connection { + type Error = Error<::GetError>; + + async fn from_request(req: &'r Request<'_>) -> Outcome { + let db: &D = match req.rocket().state() { + Some(p) => p, + _ => { + let dbtype = Paint::default(std::any::type_name::()).bold(); + let fairing = Paint::default(format!("{}::fairing()", dbtype)).wrap().bold(); + error!("requesting `{}` DB connection without attaching `{}`.", dbtype, fairing); + info_!("Attach `{}` to use database connection pooling.", fairing); + return Outcome::Failure((Status::InternalServerError, Error::UnattachedFairing)); + } + }; + + match db.pool().get().await { + Ok(conn) => Outcome::Success(Connection(conn)), + Err(e) => Outcome::Failure((Status::ServiceUnavailable, Error::Db(e))), + } + } +} + +impl Sentinel for Connection { + fn abort(rocket: &Rocket) -> bool { + if rocket.state::().is_none() { + let dbtype = Paint::default(std::any::type_name::()).bold(); + let fairing = Paint::default(format!("{}::fairing()", dbtype)).wrap().bold(); + error!("requesting `{}` DB connection without attaching `{}`.", dbtype, fairing); + info_!("Attach `{}` to use database connection pooling.", fairing); + return true; + } + + false + } +} + +/// The database fairing for pool types created with the `pool!` macro. +pub struct Fairing(&'static str, std::marker::PhantomData); + +impl> Fairing { + /// Create a new database fairing with the given constructor. This + /// constructor will be called to create an instance of `D` after the pool + /// is initialized and before it is placed into managed state. + pub fn new(fairing_name: &'static str) -> Self { + Self(fairing_name, std::marker::PhantomData) + } +} + +#[rocket::async_trait] +impl> rocket::fairing::Fairing for Fairing { + fn info(&self) -> Info { + Info { + name: self.0, + kind: Kind::Ignite, + } + } + + async fn on_ignite(&self, rocket: Rocket) -> Result, Rocket> { + let pool = match ::initialize(D::NAME, &rocket).await { + Ok(p) => p, + Err(e) => { + error!("error initializing database connection pool: {}", e); + return Err(rocket); + } + }; + + let db: D = pool.into(); + + Ok(rocket.manage(db)) + } +} diff --git a/contrib/db_pools/lib/src/error.rs b/contrib/db_pools/lib/src/error.rs new file mode 100644 index 00000000..b8b3a77b --- /dev/null +++ b/contrib/db_pools/lib/src/error.rs @@ -0,0 +1,41 @@ +use std::fmt; + +use rocket::figment; + +/// A general error type designed for the `Poolable` trait. +/// +/// [`Pool::initialize`] can return an error for any of several reasons: +/// +/// * Missing or incorrect configuration, including some syntax errors +/// * An error connecting to the database. +/// +/// [`Pool::initialize`]: crate::Pool::initialize +#[derive(Debug)] +pub enum Error { + /// A database-specific error occurred + Db(E), + + /// An error occurred in the configuration + Figment(figment::Error), + + /// Required fairing was not attached + UnattachedFairing, +} + +impl fmt::Display for Error { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Error::Db(e) => e.fmt(f), + Error::Figment(e) => write!(f, "bad configuration: {}", e), + Error::UnattachedFairing => write!(f, "required database fairing was not attached"), + } + } +} + +impl std::error::Error for Error {} + +impl From for Error { + fn from(e: figment::Error) -> Self { + Self::Figment(e) + } +} diff --git a/contrib/db_pools/lib/src/lib.rs b/contrib/db_pools/lib/src/lib.rs new file mode 100644 index 00000000..d8f8d67a --- /dev/null +++ b/contrib/db_pools/lib/src/lib.rs @@ -0,0 +1,382 @@ +//! Traits, utilities, and a macro for easy database connection pooling. +//! +//! # Overview +//! +//! This crate provides traits, utilities, and a procedural macro for +//! configuring and accessing database connection pools in Rocket. A _database +//! connection pool_ is a data structure that maintains active database +//! connections for later use in the application. +//! +//! Databases are individually configured through Rocket's regular configuration +//! mechanisms. Connecting a Rocket application to a database using this library +//! occurs in three simple steps: +//! +//! 1. Configure your databases in `Rocket.toml`. +//! (see [Configuration](#configuration)) +//! 2. Associate a Database type and fairing with each database. +//! (see [Guard Types](#guard-types)) +//! 3. Use the request guard to retrieve a connection in a handler. +//! (see [Handlers](#handlers)) +//! +//! For a list of supported databases, see [Provided Databases](#provided). This +//! support can be easily extended by implementing the [`Pool`] trait. See +//! [Extending](#extending) for more. +//! +//! ## Example +//! +//! Before using this library, the feature corresponding to your database type +//! in `rocket_db_pools` must be enabled: +//! +//! ```toml +//! [dependencies.rocket_db_pools] +//! version = "0.1.0-dev" +//! features = ["sqlx_sqlite"] +//! ``` +//! +//! See [Provided](#provided) for a list of supported database and their +//! associated feature name. +//! +//! In whichever configuration source you choose, configure a `databases` +//! dictionary with an internal dictionary for each database, here `sqlite_logs` +//! in a TOML source: +//! +//! ```toml +//! [default.databases] +//! sqlite_logs = { url = "/path/to/database.sqlite" } +//! ``` +//! +//! In your application's source code, one-time: +//! +//! ```rust +//! # #[macro_use] extern crate rocket; +//! # #[cfg(feature = "sqlx_sqlite")] +//! # mod test { +//! use rocket_db_pools::{Database, Connection, sqlx}; +//! +//! #[derive(Database)] +//! #[database("sqlite_logs")] +//! struct LogsDb(sqlx::SqlitePool); +//! +//! type LogsDbConn = Connection; +//! +//! #[launch] +//! fn rocket() -> _ { +//! rocket::build().attach(LogsDb::fairing()) +//! } +//! # } fn main() {} +//! ``` +//! +//! These steps can be repeated as many times as necessary to configure +//! multiple databases. +//! +//! Whenever a connection to the database is needed: +//! +//! ```rust +//! # #[macro_use] extern crate rocket; +//! # #[macro_use] extern crate rocket_db_pools; +//! # +//! # #[cfg(feature = "sqlx_sqlite")] +//! # mod test { +//! # use rocket::serde::json::Json; +//! # use rocket_db_pools::{Database, Connection, sqlx}; +//! # +//! # #[derive(Database)] +//! # #[database("sqlite_logs")] +//! # struct LogsDb(sqlx::SqlitePool); +//! # type LogsDbConn = Connection; +//! # +//! # type Result = std::result::Result; +//! # +//! #[get("/logs/")] +//! async fn get_logs(conn: LogsDbConn, id: usize) -> Result>> { +//! # /* +//! let logs = sqlx::query!().await?; +//! Ok(Json(logs)) +//! # */ +//! # Ok(Json(vec![])) +//! } +//! # } fn main() {} +//! ``` +//! +//! # Usage +//! +//! ## Configuration +//! +//! Databases can be configured as any other values. Using the default +//! configuration provider, either via `Rocket.toml` or environment variables. +//! You can also use a custom provider. +//! +//! ### `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: +//! +//! ```toml +//! # Option 1: +//! [global.databases] +//! sqlite_db = { url = "db.sqlite" } +//! +//! # Option 2: +//! [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 } +//! ``` +//! +//! Most databases use the default [`Config`] type, for which one key is required: +//! +//! * `url` - the URl to the database +//! +//! And one optional key is accepted: +//! +//! * `pool_size` - the size of the pool, i.e., the number of connections to +//! pool (defaults to the configured number of workers * 4) +//! TODO: currently ignored by most `Pool` implementations. +//! +//! Different options may be required or supported by other adapters, according +//! to the type specified by [`Pool::Config`]. +//! +//! ### Procedurally +//! +//! Databases can also be configured procedurally via `rocket::custom()`. +//! The example below does just this: +//! +//! ```rust +//! # #[cfg(feature = "sqlx_sqlite")] { +//! # use rocket::launch; +//! use rocket::figment::{value::{Map, Value}, util::map}; +//! +//! #[launch] +//! fn rocket() -> _ { +//! let db: Map<_, Value> = map! { +//! "url" => "db.sqlite".into(), +//! "pool_size" => 10.into() +//! }; +//! +//! let figment = rocket::Config::figment() +//! .merge(("databases", map!["my_db" => db])); +//! +//! rocket::custom(figment) +//! } +//! # rocket(); +//! # } +//! ``` +//! +//! ### Environment Variables +//! +//! Lastly, databases can be configured via environment variables by specifying +//! the `databases` table as detailed in the [Environment Variables +//! configuration +//! guide](https://rocket.rs/master/guide/configuration/#environment-variables): +//! +//! ```bash +//! ROCKET_DATABASES='{my_db={url="db.sqlite"}}' +//! ``` +//! +//! Multiple databases can be specified in the `ROCKET_DATABASES` environment variable +//! as well by comma separating them: +//! +//! ```bash +//! ROCKET_DATABASES='{my_db={url="db.sqlite"},my_pg_db={url="postgres://root:root@localhost/my_pg_db"}}' +//! ``` +//! +//! ## Database Types +//! +//! Once a database has been configured, the `#[derive(Database)]` macro can be +//! used to tie a type in your application to a configured database. The derive +//! accepts a single attribute, `#[database("name")]` that indicates the +//! name of the database. This corresponds to the database name set as the +//! database's configuration key. +//! +//! The [`Database`] trait provides a method, `fairing()`, which places an +//! instance of the decorated type in managed state; thus, the database pool can +//! be accessed with a `&State` request guard. +//! +//! The [`Connection`] type also implements [`FromRequest`], allowing it 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 fails or times out. +//! +//! The derive can only be applied to unit-like structs with one type. The +//! internal type of the structure must implement [`Pool`]. +//! +//! ```rust +//! # #[macro_use] extern crate rocket_db_pools; +//! # #[cfg(feature = "sqlx_sqlite")] +//! # mod test { +//! use rocket_db_pools::{Database, sqlx}; +//! +//! #[derive(Database)] +//! #[database("my_db")] +//! struct MyDatabase(sqlx::SqlitePool); +//! # } +//! ``` +//! +//! Other databases can be used by specifying their respective [`Pool`] type: +//! +//! ```rust +//! # #[macro_use] extern crate rocket_db_pools; +//! # #[cfg(feature = "deadpool_postgres")] +//! # mod test { +//! use rocket_db_pools::{Database, deadpool_postgres}; +//! +//! #[derive(Database)] +//! #[database("my_pg_db")] +//! struct MyPgDatabase(deadpool_postgres::Pool); +//! # } +//! ``` +//! +//! The fairing returned from the `fairing()` method _must_ be attached for the +//! request guards to succeed. Putting the pieces together, a use of +//! `#[derive(Database)]` looks as follows: +//! +//! ```rust +//! # #[macro_use] extern crate rocket; +//! # #[macro_use] extern crate rocket_db_pools; +//! # +//! # #[cfg(feature = "sqlx_sqlite")] { +//! # use rocket::figment::{value::{Map, Value}, util::map}; +//! use rocket_db_pools::{Database, sqlx}; +//! +//! #[derive(Database)] +//! #[database("my_db")] +//! struct MyDatabase(sqlx::SqlitePool); +//! +//! #[launch] +//! fn rocket() -> _ { +//! # let db: Map<_, Value> = map![ +//! # "url" => "db.sqlite".into(), "pool_size" => 10.into() +//! # ]; +//! # let figment = rocket::Config::figment().merge(("databases", map!["my_db" => db])); +//! rocket::custom(figment).attach(MyDatabase::fairing()) +//! } +//! # } +//! ``` +//! +//! ## Handlers +//! +//! Finally, access your type via `State` in a handler to access +//! the database connection pool: +//! +//! ```rust +//! # #[macro_use] extern crate rocket; +//! # #[macro_use] extern crate rocket_db_pools; +//! # +//! # #[cfg(feature = "sqlx_sqlite")] +//! # mod test { +//! # use rocket_db_pools::{Database, Connection, sqlx}; +//! use rocket::State; +//! +//! #[derive(Database)] +//! #[database("my_db")] +//! struct MyDatabase(sqlx::SqlitePool); +//! +//! #[get("/")] +//! fn my_handler(conn: &State) { +//! // ... +//! } +//! # } +//! ``` +//! +//! Alternatively, access a single connection directly via the `Connection` +//! request guard: +//! +//! ```rust +//! # #[macro_use] extern crate rocket; +//! # #[macro_use] extern crate rocket_db_pools; +//! # +//! # #[cfg(feature = "sqlx_sqlite")] +//! # mod test { +//! # use rocket_db_pools::{Database, Connection, sqlx}; +//! # type Data = (); +//! #[derive(Database)] +//! #[database("my_db")] +//! struct MyDatabase(sqlx::SqlitePool); +//! +//! type MyConnection = Connection; +//! +//! async fn load_from_db(conn: &mut sqlx::SqliteConnection) -> Data { +//! // Do something with connection, return some data. +//! # () +//! } +//! +//! #[get("/")] +//! async fn my_handler(mut conn: MyConnection) -> Data { +//! load_from_db(&mut conn).await +//! } +//! # } +//! ``` +//! +//! # Database Support +//! +//! Built-in support is provided for many popular databases and drivers. Support +//! can be easily extended by [`Pool`] implementations. +//! +//! ## Provided +//! +//! The list below includes all presently supported database adapters and their +//! corresponding [`Pool`] type. +//! +// Note: Keep this table in sync with site/guite/6-state.md +//! | Kind | Driver | Version | `Pool` Type | Feature | +//! |----------|-----------------------|-----------|--------------------------------|------------------------| +//! | MySQL | [sqlx] | `0.5` | [`sqlx::MySqlPool`] | `sqlx_mysql` | +//! | Postgres | [sqlx] | `0.5` | [`sqlx::PgPool`] | `sqlx_postgres` | +//! | Sqlite | [sqlx] | `0.5` | [`sqlx::SqlitePool`] | `sqlx_sqlite` | +//! | Mongodb | [mongodb] | `2.0.0-beta` | [`mongodb::Client`] | `mongodb` | +//! | MySQL | [mysql_async] | `0.27` | [`mysql_async::Pool`] | `mysql_async` | +//! | Postgres | [deadpool-postgres] | `0.8` | [`deadpool_postgres::Pool`] | `deadpool_postgres` | +//! | Redis | [deadpool-redis] | `0.8` | [`deadpool_redis::Pool`] | `deadpool_redis` | +//! +//! [sqlx]: https://docs.rs/sqlx/0.5/sqlx/ +//! [deadpool-postgres]: https://docs.rs/deadpool-postgres/0.8/deadpool_postgres/ +//! [deadpool-redis]: https://docs.rs/deadpool-redis/0.8/deadpool_redis/ +//! [mongodb]: https://docs.rs/mongodb/2.0.0-beta/mongodb/index.html +//! [mysql_async]: https://docs.rs/mysql_async/0.27/mysql_async/ +//! +//! The above table lists all the supported database adapters in this library. +//! In order to use particular `Pool` type that's included in this library, +//! you must first enable the feature listed in the "Feature" column. The +//! interior type of your decorated database type should match the type in the +//! "`Pool` Type" column. +//! +//! ## Extending +//! +//! Extending Rocket's support to your own custom database adapter is as easy as +//! implementing the [`Pool`] trait. See the documentation for [`Pool`] +//! for more details on how to implement it. +//! +//! [`FromRequest`]: rocket::request::FromRequest +//! [request guards]: rocket::request::FromRequest +//! [`Database`]: crate::Database +//! [`Pool`]: crate::Pool + +#![doc(html_root_url = "https://api.rocket.rs/master/rocket_db_pools")] +#![doc(html_favicon_url = "https://rocket.rs/images/favicon.ico")] +#![doc(html_logo_url = "https://rocket.rs/images/logo-boxed.png")] + +#[doc(hidden)] +#[macro_use] +pub extern crate rocket; + +#[cfg(feature = "deadpool_postgres")] pub use deadpool_postgres; +#[cfg(feature = "deadpool_redis")] pub use deadpool_redis; +#[cfg(feature = "mysql_async")] pub use mysql_async; +#[cfg(feature = "mongodb")] pub use mongodb; +#[cfg(feature = "sqlx")] pub use sqlx; + +mod config; +mod database; +mod error; +mod pool; + +pub use self::config::Config; +pub use self::database::{Connection, Database, Fairing}; +pub use self::error::Error; +pub use self::pool::Pool; + +pub use rocket_db_pools_codegen::*; diff --git a/contrib/db_pools/lib/src/pool.rs b/contrib/db_pools/lib/src/pool.rs new file mode 100644 index 00000000..9c932533 --- /dev/null +++ b/contrib/db_pools/lib/src/pool.rs @@ -0,0 +1,287 @@ +use rocket::async_trait; +use rocket::{Build, Rocket}; + +use crate::{Config, Error}; + +/// This trait is implemented on connection pool types that can be used with the +/// [`Database`] derive macro. +/// +/// `Pool` determines how the connection pool is initialized from configuration, +/// such as a connection string and optional pool size, along with the returned +/// `Connection` type. +/// +/// Implementations of this trait should use `async_trait`. +/// +/// ## Example +/// +/// ``` +/// use rocket::{Build, Rocket}; +/// +/// #[derive(Debug)] +/// struct Error { /* ... */ } +/// # impl std::fmt::Display for Error { +/// # fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { +/// # unimplemented!("example") +/// # } +/// # } +/// # impl std::error::Error for Error { } +/// +/// struct Pool { /* ... */ } +/// struct Connection { /* .. */ } +/// +/// #[rocket::async_trait] +/// impl rocket_db_pools::Pool for Pool { +/// type Connection = Connection; +/// type InitError = Error; +/// type GetError = Error; +/// +/// async fn initialize(db_name: &str, rocket: &Rocket) +/// -> Result> +/// { +/// unimplemented!("example") +/// } +/// +/// async fn get(&self) -> Result { +/// unimplemented!("example") +/// } +/// } +/// ``` +#[async_trait] +pub trait Pool: Sized + Send + Sync + 'static { + /// The type returned by get(). + type Connection; + + /// The error type returned by `initialize`. + type InitError: std::error::Error; + + /// The error type returned by `get`. + type GetError: std::error::Error; + + /// Constructs a pool from a [Value](rocket::figment::value::Value). + /// + /// It is up to each implementor of `Pool` to define its accepted + /// configuration value(s) via the `Config` associated type. Most + /// integrations provided in `rocket_db_pools` use [`Config`], which + /// accepts a (required) `url` and an (optional) `pool_size`. + /// + /// ## Errors + /// + /// This method returns an error if the configuration is not compatible, or + /// if creating a pool failed due to an unavailable database server, + /// insufficient resources, or another database-specific error. + async fn initialize(db_name: &str, rocket: &Rocket) + -> Result>; + + /// Asynchronously gets a connection from the factory or pool. + /// + /// ## Errors + /// + /// This method returns an error if a connection could not be retrieved, + /// such as a preconfigured timeout elapsing or when the database server is + /// unavailable. + async fn get(&self) -> Result; +} + +#[cfg(feature = "deadpool_postgres")] +#[async_trait] +impl Pool for deadpool_postgres::Pool { + type Connection = deadpool_postgres::Client; + type InitError = deadpool_postgres::tokio_postgres::Error; + type GetError = deadpool_postgres::PoolError; + + async fn initialize(db_name: &str, rocket: &Rocket) + -> std::result::Result> + { + let config = Config::from(db_name, rocket)?; + let manager = deadpool_postgres::Manager::new( + config.url.parse().map_err(Error::Db)?, + // TODO: add TLS support in config + deadpool_postgres::tokio_postgres::NoTls, + ); + let mut pool_config = deadpool_postgres::PoolConfig::new(config.pool_size as usize); + pool_config.timeouts.wait = Some(std::time::Duration::from_secs(config.timeout.into())); + + Ok(deadpool_postgres::Pool::from_config(manager, pool_config)) + } + + async fn get(&self) -> Result { + self.get().await + } +} + +#[cfg(feature = "deadpool_redis")] +#[async_trait] +impl Pool for deadpool_redis::Pool { + type Connection = deadpool_redis::ConnectionWrapper; + type InitError = deadpool_redis::redis::RedisError; + type GetError = deadpool_redis::PoolError; + + async fn initialize(db_name: &str, rocket: &Rocket) + -> std::result::Result> + { + let config = Config::from(db_name, rocket)?; + let manager = deadpool_redis::Manager::new(config.url).map_err(Error::Db)?; + + let mut pool_config = deadpool_redis::PoolConfig::new(config.pool_size as usize); + pool_config.timeouts.wait = Some(std::time::Duration::from_secs(config.timeout.into())); + + Ok(deadpool_redis::Pool::from_config(manager, pool_config)) + } + + async fn get(&self) -> Result { + self.get().await + } +} + +#[cfg(feature = "mongodb")] +#[async_trait] +impl Pool for mongodb::Client { + type Connection = mongodb::Client; + type InitError = mongodb::error::Error; + type GetError = std::convert::Infallible; + + async fn initialize(db_name: &str, rocket: &Rocket) + -> std::result::Result> + { + let config = Config::from(db_name, rocket)?; + let mut options = mongodb::options::ClientOptions::parse(&config.url) + .await + .map_err(Error::Db)?; + options.max_pool_size = Some(config.pool_size); + options.wait_queue_timeout = Some(std::time::Duration::from_secs(config.timeout.into())); + + mongodb::Client::with_options(options).map_err(Error::Db) + } + + async fn get(&self) -> Result { + Ok(self.clone()) + } +} + +#[cfg(feature = "mysql_async")] +#[async_trait] +impl Pool for mysql_async::Pool { + type Connection = mysql_async::Conn; + type InitError = mysql_async::Error; + type GetError = mysql_async::Error; + + async fn initialize(db_name: &str, rocket: &Rocket) + -> std::result::Result> + { + use rocket::figment::{self, error::{Actual, Kind}}; + + let config = Config::from(db_name, rocket)?; + let original_opts = mysql_async::Opts::from_url(&config.url) + .map_err(|_| figment::Error::from(Kind::InvalidValue( + Actual::Str(config.url.to_string()), + "mysql connection string".to_string() + )))?; + + let new_pool_opts = original_opts.pool_opts() + .clone() + .with_constraints( + mysql_async::PoolConstraints::new(0, config.pool_size as usize) + .expect("usize can't be < 0") + ); + + // TODO: timeout + + let opts = mysql_async::OptsBuilder::from_opts(original_opts) + .pool_opts(new_pool_opts); + + Ok(mysql_async::Pool::new(opts)) + } + + async fn get(&self) -> std::result::Result { + self.get_conn().await + } +} + +#[cfg(feature = "sqlx_mysql")] +#[async_trait] +impl Pool for sqlx::MySqlPool { + type Connection = sqlx::pool::PoolConnection; + type InitError = sqlx::Error; + type GetError = sqlx::Error; + + async fn initialize(db_name: &str, rocket: &Rocket) + -> std::result::Result> + { + use sqlx::ConnectOptions; + + let config = Config::from(db_name, rocket)?; + let mut opts = config.url.parse::() + .map_err(Error::Db)?; + opts.disable_statement_logging(); + sqlx::pool::PoolOptions::new() + .max_connections(config.pool_size) + .connect_timeout(std::time::Duration::from_secs(config.timeout.into())) + .connect_with(opts) + .await + .map_err(Error::Db) + } + + async fn get(&self) -> std::result::Result { + self.acquire().await + } +} + +#[cfg(feature = "sqlx_postgres")] +#[async_trait] +impl Pool for sqlx::PgPool { + type Connection = sqlx::pool::PoolConnection; + type InitError = sqlx::Error; + type GetError = sqlx::Error; + + async fn initialize(db_name: &str, rocket: &Rocket) + -> std::result::Result> + { + use sqlx::ConnectOptions; + + let config = Config::from(db_name, rocket)?; + let mut opts = config.url.parse::() + .map_err(Error::Db)?; + opts.disable_statement_logging(); + sqlx::pool::PoolOptions::new() + .max_connections(config.pool_size) + .connect_timeout(std::time::Duration::from_secs(config.timeout.into())) + .connect_with(opts) + .await + .map_err(Error::Db) + } + + async fn get(&self) -> std::result::Result { + self.acquire().await + } +} + +#[cfg(feature = "sqlx_sqlite")] +#[async_trait] +impl Pool for sqlx::SqlitePool { + type Connection = sqlx::pool::PoolConnection; + type InitError = sqlx::Error; + type GetError = sqlx::Error; + + async fn initialize(db_name: &str, rocket: &Rocket) + -> std::result::Result> + { + use sqlx::ConnectOptions; + + let config = Config::from(db_name, rocket)?; + let mut opts = config.url.parse::() + .map_err(Error::Db)? + .create_if_missing(true); + opts.disable_statement_logging(); + + dbg!(sqlx::pool::PoolOptions::new() + .max_connections(config.pool_size) + .connect_timeout(std::time::Duration::from_secs(config.timeout.into()))) + .connect_with(opts) + .await + .map_err(Error::Db) + } + + async fn get(&self) -> std::result::Result { + self.acquire().await + } +} diff --git a/examples/README.md b/examples/README.md index dddea7b1..9a3a0997 100644 --- a/examples/README.md +++ b/examples/README.md @@ -35,7 +35,8 @@ This directory contains projects showcasing Rocket's features. * **[`databases`](./databases)** - Implements a CRUD-like "blog" JSON API backed by a SQLite database driven by each of `sqlx`, `diesel`, and `rusqlite`. Runs migrations automatically for the former two drivers. Uses - `contrib` database support for the latter two drivers. + `contrib` database support for all drivers (`rocket_db_pools` for the first; + `rocket_sync_db_pools` for the other latter two). * **[`error-handling`](./error-handling)** - Exhibits the use of scoped catchers; contains commented out lines that will cause a launch-time error diff --git a/examples/databases/Cargo.toml b/examples/databases/Cargo.toml index 4efaccbb..1556db32 100644 --- a/examples/databases/Cargo.toml +++ b/examples/databases/Cargo.toml @@ -13,7 +13,11 @@ diesel_migrations = "1.3" [dependencies.sqlx] version = "0.5.1" default-features = false -features = ["runtime-tokio-rustls", "sqlite", "macros", "offline", "migrate"] +features = ["macros", "offline", "migrate"] + +[dependencies.rocket_db_pools] +path = "../../contrib/db_pools/lib/" +features = ["sqlx_sqlite"] [dependencies.rocket_sync_db_pools] path = "../../contrib/sync_db_pools/lib/" diff --git a/examples/databases/src/sqlx.rs b/examples/databases/src/sqlx.rs index bdd07833..4f7bb09c 100644 --- a/examples/databases/src/sqlx.rs +++ b/examples/databases/src/sqlx.rs @@ -1,13 +1,18 @@ -use rocket::{Rocket, State, Build, futures}; +use rocket::{Rocket, Build, futures}; use rocket::fairing::{self, AdHoc}; use rocket::response::status::Created; use rocket::serde::{Serialize, Deserialize, json::Json}; +use rocket_db_pools::{sqlx, Database}; + use futures::stream::TryStreamExt; use futures::future::TryFutureExt; -use sqlx::ConnectOptions; -type Db = sqlx::SqlitePool; +#[derive(Database)] +#[database("sqlx")] +struct Db(sqlx::SqlitePool); + +type Connection = rocket_db_pools::Connection; type Result> = std::result::Result; @@ -21,19 +26,19 @@ struct Post { } #[post("/", data = "")] -async fn create(db: &State, post: Json) -> Result>> { +async fn create(mut db: Connection, post: Json) -> Result>> { // There is no support for `RETURNING`. sqlx::query!("INSERT INTO posts (title, text) VALUES (?, ?)", post.title, post.text) - .execute(&**db) + .execute(&mut *db) .await?; Ok(Created::new("/").body(post)) } #[get("/")] -async fn list(db: &State) -> Result>> { +async fn list(mut db: Connection) -> Result>> { let ids = sqlx::query!("SELECT id FROM posts") - .fetch(&**db) + .fetch(&mut *db) .map_ok(|record| record.id) .try_collect::>() .await?; @@ -42,65 +47,47 @@ async fn list(db: &State) -> Result>> { } #[get("/")] -async fn read(db: &State, id: i64) -> Option> { +async fn read(mut db: Connection, id: i64) -> Option> { sqlx::query!("SELECT id, title, text FROM posts WHERE id = ?", id) - .fetch_one(&**db) + .fetch_one(&mut *db) .map_ok(|r| Json(Post { id: Some(r.id), title: r.title, text: r.text })) .await .ok() } #[delete("/")] -async fn delete(db: &State, id: i64) -> Result> { +async fn delete(mut db: Connection, id: i64) -> Result> { let result = sqlx::query!("DELETE FROM posts WHERE id = ?", id) - .execute(&**db) + .execute(&mut *db) .await?; Ok((result.rows_affected() == 1).then(|| ())) } #[delete("/")] -async fn destroy(db: &State) -> Result<()> { - sqlx::query!("DELETE FROM posts").execute(&**db).await?; +async fn destroy(mut db: Connection) -> Result<()> { + sqlx::query!("DELETE FROM posts").execute(&mut *db).await?; Ok(()) } async fn init_db(rocket: Rocket) -> fairing::Result { - use rocket_sync_db_pools::Config; - - let config = match Config::from("sqlx", &rocket) { - Ok(config) => config, - Err(e) => { - error!("Failed to read SQLx config: {}", e); - return Err(rocket); + match rocket.state::() { + Some(db) => { + if let Err(e) = sqlx::migrate!("db/sqlx/migrations").run(db.pool()).await { + error!("Failed to initialize SQLx database: {}", e); + return Err(rocket); + } + Ok(rocket) } - }; - - let mut opts = sqlx::sqlite::SqliteConnectOptions::new() - .filename(&config.url) - .create_if_missing(true); - - opts.disable_statement_logging(); - let db = match Db::connect_with(opts).await { - Ok(db) => db, - Err(e) => { - error!("Failed to connect to SQLx database: {}", e); - return Err(rocket); - } - }; - - if let Err(e) = sqlx::migrate!("db/sqlx/migrations").run(&db).await { - error!("Failed to initialize SQLx database: {}", e); - return Err(rocket); + None => Err(rocket), } - - Ok(rocket.manage(db)) } pub fn stage() -> AdHoc { AdHoc::on_ignite("SQLx Stage", |rocket| async { rocket + .attach(Db::fairing()) .attach(AdHoc::try_on_ignite("SQLx Database", init_db)) .mount("/sqlx", routes![list, create, read, delete, destroy]) }) diff --git a/scripts/test.sh b/scripts/test.sh index 9e3d6d29..8f92fdb0 100755 --- a/scripts/test.sh +++ b/scripts/test.sh @@ -73,6 +73,16 @@ function indir() { } function test_contrib() { + DB_POOLS_FEATURES=( + deadpool-postgres + deadpool-redis + mongodb + mysql_async + sqlx_mysql + sqlx_postgres + sqlx_sqlite + ) + SYNC_DB_POOLS_FEATURES=( diesel_postgres_pool diesel_sqlite_pool @@ -87,6 +97,11 @@ function test_contrib() { handlebars ) + for feature in "${DB_POOLS_FEATURES[@]}"; do + echo ":: Building and testing db_pools [$feature]..." + $CARGO test -p rocket_db_pools --no-default-features --features $feature $@ + done + for feature in "${SYNC_DB_POOLS_FEATURES[@]}"; do echo ":: Building and testing sync_db_pools [$feature]..." $CARGO test -p rocket_sync_db_pools --no-default-features --features $feature $@