vw_small

Hardened fork of Vaultwarden (https://github.com/dani-garcia/vaultwarden) with fewer features.
git clone https://git.philomathiclife.com/repos/vw_small
Log | Files | Refs | README

commit b617ffd2af71d81c24955c00c1f378325613dec0
parent 3abf173d8954c982383285b29b39d15043e974cb
Author: Jeremy Lin <jeremy.lin@gmail.com>
Date:   Tue, 26 Apr 2022 17:50:20 -0700

Add support for database connection init statements

This is probably mainly useful for running connection-scoped pragma statements.

Diffstat:
M.env.template | 11+++++++++++
Msrc/config.rs | 11++++++++++-
Msrc/db/mod.rs | 29+++++++++++++++++++++++++++--
3 files changed, 48 insertions(+), 3 deletions(-)

diff --git a/.env.template b/.env.template @@ -29,6 +29,17 @@ ## Define the size of the connection pool used for connecting to the database. # DATABASE_MAX_CONNS=10 +## Database connection initialization +## Allows SQL statements to be run whenever a new database connection is created. +## For example, this can be used to run connection-scoped pragma statements. +## +## Statements to run when creating a new SQLite connection +# SQLITE_CONN_INIT="" +## Statements to run when creating a new MySQL connection +# MYSQL_CONN_INIT="" +## Statements to run when creating a new PostgreSQL connection +# POSTGRESQL_CONN_INIT="" + ## Individual folders, these override %DATA_FOLDER% # RSA_KEY_FILENAME=data/rsa_key # ICON_CACHE_FOLDER=data/icon_cache diff --git a/src/config.rs b/src/config.rs @@ -515,11 +515,20 @@ make_config! { db_connection_retries: u32, false, def, 15; /// Timeout when aquiring database connection - database_timeout: u64, false, def, 30; + database_timeout: u64, false, def, 30; /// Database connection pool size database_max_conns: u32, false, def, 10; + /// SQLite connection init |> Statements to run when creating a new SQLite connection + sqlite_conn_init: String, false, def, "".to_string(); + + /// MySQL connection init |> Statements to run when creating a new MySQL connection + mysql_conn_init: String, false, def, "".to_string(); + + /// PostgreSQL connection init |> Statements to run when creating a new PostgreSQL connection + postgresql_conn_init: String, false, def, "".to_string(); + /// Bypass admin page security (Know the risks!) |> Disables the Admin Token for the admin page so you may use your own auth in-front disable_admin_token: bool, true, def, false; diff --git a/src/db/mod.rs b/src/db/mod.rs @@ -1,6 +1,10 @@ use std::{sync::Arc, time::Duration}; -use diesel::r2d2::{ConnectionManager, Pool, PooledConnection}; +use diesel::{ + connection::SimpleConnection, + r2d2::{ConnectionManager, CustomizeConnection, Pool, PooledConnection}, +}; + use rocket::{ http::Status, outcome::IntoOutcome, @@ -62,6 +66,23 @@ macro_rules! generate_connections { #[allow(non_camel_case_types)] pub enum DbConnInner { $( #[cfg($name)] $name(PooledConnection<ConnectionManager< $ty >>), )+ } + #[derive(Debug)] + pub struct DbConnOptions { + pub init_stmts: String, + } + + $( // Based on <https://stackoverflow.com/a/57717533>. + #[cfg($name)] + impl CustomizeConnection<$ty, diesel::r2d2::Error> for DbConnOptions { + fn on_acquire(&self, conn: &mut $ty) -> Result<(), diesel::r2d2::Error> { + (|| { + if !self.init_stmts.is_empty() { + conn.batch_execute(&self.init_stmts)?; + } + Ok(()) + })().map_err(diesel::r2d2::Error::QueryError) + } + })+ #[derive(Clone)] pub struct DbPool { @@ -103,7 +124,8 @@ macro_rules! generate_connections { } impl DbPool { - // For the given database URL, guess it's type, run migrations create pool and return it + // For the given database URL, guess its type, run migrations, create pool, and return it + #[allow(clippy::diverging_sub_expression)] pub fn from_config() -> Result<Self, Error> { let url = CONFIG.database_url(); let conn_type = DbConnType::from_url(&url)?; @@ -117,6 +139,9 @@ macro_rules! generate_connections { let pool = Pool::builder() .max_size(CONFIG.database_max_conns()) .connection_timeout(Duration::from_secs(CONFIG.database_timeout())) + .connection_customizer(Box::new(DbConnOptions{ + init_stmts: paste::paste!{ CONFIG. [< $name _conn_init >] () } + })) .build(manager) .map_res("Failed to create pool")?; return Ok(DbPool {