From b617ffd2af71d81c24955c00c1f378325613dec0 Mon Sep 17 00:00:00 2001
From: Jeremy Lin <jeremy.lin@gmail.com>
Date: Tue, 26 Apr 2022 17:50:20 -0700
Subject: [PATCH] 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<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 {