MirrorHub/bitwarden_rs
0
0
Fork 0
mirror of https://github.com/dani-garcia/vaultwarden synced 2024-11-05 07:08:59 +01:00
Code Issues Releases Wiki Activity

Add support for database connection init statements

Browse source 
This is probably mainly useful for running connection-scoped pragma statements.
This commit is contained in:
Jeremy Lin 2022-04-26 17:50:20 -07:00
parent 3abf173d89
commit b617ffd2af
3 changed files with 48 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
.env.template
View file

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

11
src/config.rs
View file

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

29
src/db/mod.rs
View file

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