From b617ffd2af71d81c24955c00c1f378325613dec0 Mon Sep 17 00:00:00 2001 From: Jeremy Lin Date: Tue, 26 Apr 2022 17:50:20 -0700 Subject: [PATCH 1/3] Add support for database connection init statements This is probably mainly useful for running connection-scoped pragma statements. --- .env.template | 11 +++++++++++ src/config.rs | 11 ++++++++++- src/db/mod.rs | 29 +++++++++++++++++++++++++++-- 3 files changed, 48 insertions(+), 3 deletions(-) diff --git a/.env.template b/.env.template index 3c8a5ebb..68c487f4 100644 --- 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 index 2cef76e2..cd6dc138 100644 --- 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 index 6fcb63e5..b382260f 100644 --- 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>), )+ } + #[derive(Debug)] + pub struct DbConnOptions { + pub init_stmts: String, + } + + $( // Based on . + #[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 { 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 { From 78d07e2fda8816fb41b91de2e84f8f4b3df8180e Mon Sep 17 00:00:00 2001 From: Jeremy Lin Date: Tue, 26 Apr 2022 17:55:19 -0700 Subject: [PATCH 2/3] Add default connection-scoped pragmas for SQLite `PRAGMA busy_timeout = 5000` tells SQLite to keep trying for up to 5000 ms when there is lock contention, rather than aborting immediately. This should hopefully prevent the vast majority of "database is locked" panics observed since the async transition. `PRAGMA synchronous = NORMAL` trades better performance for a small potential loss in durability (the default is `FULL`). The SQLite docs recommend `NORMAL` as "a good choice for most applications running in WAL mode". --- .env.template | 2 +- src/config.rs | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.env.template b/.env.template index 68c487f4..5ff34db7 100644 --- a/.env.template +++ b/.env.template @@ -34,7 +34,7 @@ ## For example, this can be used to run connection-scoped pragma statements. ## ## Statements to run when creating a new SQLite connection -# SQLITE_CONN_INIT="" +# SQLITE_CONN_INIT="PRAGMA busy_timeout = 5000; PRAGMA synchronous = NORMAL;" ## Statements to run when creating a new MySQL connection # MYSQL_CONN_INIT="" ## Statements to run when creating a new PostgreSQL connection diff --git a/src/config.rs b/src/config.rs index cd6dc138..eaabdf6b 100644 --- a/src/config.rs +++ b/src/config.rs @@ -521,7 +521,7 @@ make_config! { 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(); + sqlite_conn_init: String, false, def, "PRAGMA busy_timeout = 5000; PRAGMA synchronous = NORMAL;".to_string(); /// MySQL connection init |> Statements to run when creating a new MySQL connection mysql_conn_init: String, false, def, "".to_string(); From 542a73cc6edb4770755663df74b5e052019f425c Mon Sep 17 00:00:00 2001 From: Jeremy Lin Date: Fri, 29 Apr 2022 00:26:49 -0700 Subject: [PATCH 3/3] Switch to a single config option for database connection init The main pro is less config options, while the main con is less clarity in what the defaults are for the various database types. --- .env.template | 14 ++++++-------- src/config.rs | 10 ++-------- src/db/mod.rs | 19 ++++++++++++++++++- 3 files changed, 26 insertions(+), 17 deletions(-) diff --git a/.env.template b/.env.template index 5ff34db7..a835200a 100644 --- a/.env.template +++ b/.env.template @@ -31,14 +31,12 @@ ## 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="PRAGMA busy_timeout = 5000; PRAGMA synchronous = NORMAL;" -## Statements to run when creating a new MySQL connection -# MYSQL_CONN_INIT="" -## Statements to run when creating a new PostgreSQL connection -# POSTGRESQL_CONN_INIT="" +## This is mainly useful for connection-scoped pragmas. +## If empty, a database-specific default is used: +## - SQLite: "PRAGMA busy_timeout = 5000; PRAGMA synchronous = NORMAL;" +## - MySQL: "" +## - PostgreSQL: "" +# DATABASE_CONN_INIT="" ## Individual folders, these override %DATA_FOLDER% # RSA_KEY_FILENAME=data/rsa_key diff --git a/src/config.rs b/src/config.rs index eaabdf6b..cd90caa1 100644 --- a/src/config.rs +++ b/src/config.rs @@ -520,14 +520,8 @@ make_config! { /// 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, "PRAGMA busy_timeout = 5000; PRAGMA synchronous = NORMAL;".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(); + /// Database connection init |> SQL statements to run when creating a new database connection, mainly useful for connection-scoped pragmas. If empty, a database-specific default is used. + database_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 index b382260f..41560827 100644 --- a/src/db/mod.rs +++ b/src/db/mod.rs @@ -140,7 +140,7 @@ macro_rules! generate_connections { .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 >] () } + init_stmts: conn_type.get_init_stmts() })) .build(manager) .map_res("Failed to create pool")?; @@ -215,6 +215,23 @@ impl DbConnType { err!("`DATABASE_URL` looks like a SQLite URL, but 'sqlite' feature is not enabled") } } + + pub fn get_init_stmts(&self) -> String { + let init_stmts = CONFIG.database_conn_init(); + if !init_stmts.is_empty() { + init_stmts + } else { + self.default_init_stmts() + } + } + + pub fn default_init_stmts(&self) -> String { + match self { + Self::sqlite => "PRAGMA busy_timeout = 5000; PRAGMA synchronous = NORMAL;".to_string(), + Self::mysql => "".to_string(), + Self::postgresql => "".to_string(), + } + } } #[macro_export]