mirror of
https://github.com/dani-garcia/vaultwarden.wiki.git
synced 2026-07-29 21:45:06 +03:00
Use Diesels MultiConnections Derive (#6279)
* Use Diesels MultiConnections Derive With this PR we remove almost all custom macro's to create the multiple database type code. This is now handled by Diesel it self. This removed the need of the following functions/macro's: - `db_object!` - `::to_db` - `.from_db()` It is also possible to just use one schema instead of multiple per type. Also done: - Refactored the SQLite backup function - Some formatting of queries so every call is one a separate line, this looks a bit better - Declare `conn` as mut inside each `db_run!` instead of having to declare it as `mut` in functions or calls - Added an `ACTIVE_DB_TYPE` static which holds the currently active database type - Removed `diesel_logger` crate and use Diesel's `set_default_instrumentation()` If you want debug queries you can now simply change the log level of `vaultwarden::db::query_logger` - Use PostgreSQL v17 in the Alpine images to match the Debian Trixie version - Optimized the Workflows since `diesel_logger` isn't needed anymore And on the extra plus-side, this lowers the compile-time and binary size too. Signed-off-by: BlackDex <black.dex@gmail.com> * Adjust query_logger and some other small items Signed-off-by: BlackDex <black.dex@gmail.com> * Remove macro, replaced with an function Signed-off-by: BlackDex <black.dex@gmail.com> * Implement custom connection manager Signed-off-by: BlackDex <black.dex@gmail.com> * Updated some crates to keep up2date Signed-off-by: BlackDex <black.dex@gmail.com> * Small adjustment Signed-off-by: BlackDex <black.dex@gmail.com> * crate updates Signed-off-by: BlackDex <black.dex@gmail.com> * Update crates Signed-off-by: BlackDex <black.dex@gmail.com> --------- Signed-off-by: BlackDex <black.dex@gmail.com>
This commit is contained in:
Mathijs van Veluw
+303
-312
@@ -1,8 +1,14 @@
|
||||
use std::{sync::Arc, time::Duration};
|
||||
mod query_logger;
|
||||
|
||||
use std::{
|
||||
sync::{Arc, OnceLock},
|
||||
time::Duration,
|
||||
};
|
||||
|
||||
use diesel::{
|
||||
connection::SimpleConnection,
|
||||
r2d2::{ConnectionManager, CustomizeConnection, Pool, PooledConnection},
|
||||
r2d2::{CustomizeConnection, Pool, PooledConnection},
|
||||
Connection, RunQueryDsl,
|
||||
};
|
||||
|
||||
use rocket::{
|
||||
@@ -21,20 +27,7 @@ use crate::{
|
||||
CONFIG,
|
||||
};
|
||||
|
||||
#[cfg(sqlite)]
|
||||
#[path = "schemas/sqlite/schema.rs"]
|
||||
pub mod __sqlite_schema;
|
||||
|
||||
#[cfg(mysql)]
|
||||
#[path = "schemas/mysql/schema.rs"]
|
||||
pub mod __mysql_schema;
|
||||
|
||||
#[cfg(postgresql)]
|
||||
#[path = "schemas/postgresql/schema.rs"]
|
||||
pub mod __postgresql_schema;
|
||||
|
||||
// These changes are based on Rocket 0.5-rc wrapper of Diesel: https://github.com/SergioBenitez/Rocket/blob/v0.5-rc/contrib/sync_db_pools
|
||||
|
||||
// A wrapper around spawn_blocking that propagates panics to the calling code.
|
||||
pub async fn run_blocking<F, R>(job: F) -> R
|
||||
where
|
||||
@@ -51,162 +44,230 @@ where
|
||||
}
|
||||
|
||||
// This is used to generate the main DbConn and DbPool enums, which contain one variant for each database supported
|
||||
macro_rules! generate_connections {
|
||||
( $( $name:ident: $ty:ty ),+ ) => {
|
||||
#[allow(non_camel_case_types, dead_code)]
|
||||
#[derive(Eq, PartialEq)]
|
||||
pub enum DbConnType { $( $name, )+ }
|
||||
|
||||
pub struct DbConn {
|
||||
conn: Arc<Mutex<Option<DbConnInner>>>,
|
||||
permit: Option<OwnedSemaphorePermit>,
|
||||
}
|
||||
|
||||
#[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).map_err(diesel::r2d2::Error::QueryError)?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
})+
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct DbPool {
|
||||
// This is an 'Option' so that we can drop the pool in a 'spawn_blocking'.
|
||||
pool: Option<DbPoolInner>,
|
||||
semaphore: Arc<Semaphore>
|
||||
}
|
||||
|
||||
#[allow(non_camel_case_types)]
|
||||
#[derive(Clone)]
|
||||
pub enum DbPoolInner { $( #[cfg($name)] $name(Pool<ConnectionManager< $ty >>), )+ }
|
||||
|
||||
impl Drop for DbConn {
|
||||
fn drop(&mut self) {
|
||||
let conn = Arc::clone(&self.conn);
|
||||
let permit = self.permit.take();
|
||||
|
||||
// Since connection can't be on the stack in an async fn during an
|
||||
// await, we have to spawn a new blocking-safe thread...
|
||||
tokio::task::spawn_blocking(move || {
|
||||
// And then re-enter the runtime to wait on the async mutex, but in a blocking fashion.
|
||||
let mut conn = tokio::runtime::Handle::current().block_on(conn.lock_owned());
|
||||
|
||||
if let Some(conn) = conn.take() {
|
||||
drop(conn);
|
||||
}
|
||||
|
||||
// Drop permit after the connection is dropped
|
||||
drop(permit);
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for DbPool {
|
||||
fn drop(&mut self) {
|
||||
let pool = self.pool.take();
|
||||
tokio::task::spawn_blocking(move || drop(pool));
|
||||
}
|
||||
}
|
||||
|
||||
impl DbPool {
|
||||
// For the given database URL, guess its type, run migrations, create pool, and return it
|
||||
pub fn from_config() -> Result<Self, Error> {
|
||||
let url = CONFIG.database_url();
|
||||
let conn_type = DbConnType::from_url(&url)?;
|
||||
|
||||
match conn_type { $(
|
||||
DbConnType::$name => {
|
||||
#[cfg($name)]
|
||||
{
|
||||
pastey::paste!{ [< $name _migrations >]::run_migrations()?; }
|
||||
let manager = ConnectionManager::new(&url);
|
||||
let pool = Pool::builder()
|
||||
.max_size(CONFIG.database_max_conns())
|
||||
.min_idle(Some(CONFIG.database_min_conns()))
|
||||
.idle_timeout(Some(Duration::from_secs(CONFIG.database_idle_timeout())))
|
||||
.connection_timeout(Duration::from_secs(CONFIG.database_timeout()))
|
||||
.connection_customizer(Box::new(DbConnOptions{
|
||||
init_stmts: conn_type.get_init_stmts()
|
||||
}))
|
||||
.build(manager)
|
||||
.map_res("Failed to create pool")?;
|
||||
Ok(DbPool {
|
||||
pool: Some(DbPoolInner::$name(pool)),
|
||||
semaphore: Arc::new(Semaphore::new(CONFIG.database_max_conns() as usize)),
|
||||
})
|
||||
}
|
||||
#[cfg(not($name))]
|
||||
unreachable!("Trying to use a DB backend when it's feature is disabled")
|
||||
},
|
||||
)+ }
|
||||
}
|
||||
// Get a connection from the pool
|
||||
pub async fn get(&self) -> Result<DbConn, Error> {
|
||||
let duration = Duration::from_secs(CONFIG.database_timeout());
|
||||
let permit = match timeout(duration, Arc::clone(&self.semaphore).acquire_owned()).await {
|
||||
Ok(p) => p.expect("Semaphore should be open"),
|
||||
Err(_) => {
|
||||
err!("Timeout waiting for database connection");
|
||||
}
|
||||
};
|
||||
|
||||
match self.pool.as_ref().expect("DbPool.pool should always be Some()") { $(
|
||||
#[cfg($name)]
|
||||
DbPoolInner::$name(p) => {
|
||||
let pool = p.clone();
|
||||
let c = run_blocking(move || pool.get_timeout(duration)).await.map_res("Error retrieving connection from pool")?;
|
||||
|
||||
Ok(DbConn {
|
||||
conn: Arc::new(Mutex::new(Some(DbConnInner::$name(c)))),
|
||||
permit: Some(permit)
|
||||
})
|
||||
},
|
||||
)+ }
|
||||
}
|
||||
}
|
||||
};
|
||||
#[derive(diesel::MultiConnection)]
|
||||
pub enum DbConnInner {
|
||||
#[cfg(mysql)]
|
||||
Mysql(diesel::mysql::MysqlConnection),
|
||||
#[cfg(postgresql)]
|
||||
Postgresql(diesel::pg::PgConnection),
|
||||
#[cfg(sqlite)]
|
||||
Sqlite(diesel::sqlite::SqliteConnection),
|
||||
}
|
||||
|
||||
#[cfg(not(query_logger))]
|
||||
generate_connections! {
|
||||
sqlite: diesel::sqlite::SqliteConnection,
|
||||
mysql: diesel::mysql::MysqlConnection,
|
||||
postgresql: diesel::pg::PgConnection
|
||||
/// Custom connection manager that implements manual connection establishment
|
||||
pub struct DbConnManager {
|
||||
database_url: String,
|
||||
}
|
||||
|
||||
#[cfg(query_logger)]
|
||||
generate_connections! {
|
||||
sqlite: diesel_logger::LoggingConnection<diesel::sqlite::SqliteConnection>,
|
||||
mysql: diesel_logger::LoggingConnection<diesel::mysql::MysqlConnection>,
|
||||
postgresql: diesel_logger::LoggingConnection<diesel::pg::PgConnection>
|
||||
impl DbConnManager {
|
||||
pub fn new(database_url: &str) -> Self {
|
||||
Self {
|
||||
database_url: database_url.to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
fn establish_connection(&self) -> Result<DbConnInner, diesel::r2d2::Error> {
|
||||
match DbConnType::from_url(&self.database_url) {
|
||||
#[cfg(mysql)]
|
||||
Ok(DbConnType::Mysql) => {
|
||||
let conn = diesel::mysql::MysqlConnection::establish(&self.database_url)?;
|
||||
Ok(DbConnInner::Mysql(conn))
|
||||
}
|
||||
#[cfg(postgresql)]
|
||||
Ok(DbConnType::Postgresql) => {
|
||||
let conn = diesel::pg::PgConnection::establish(&self.database_url)?;
|
||||
Ok(DbConnInner::Postgresql(conn))
|
||||
}
|
||||
#[cfg(sqlite)]
|
||||
Ok(DbConnType::Sqlite) => {
|
||||
let conn = diesel::sqlite::SqliteConnection::establish(&self.database_url)?;
|
||||
Ok(DbConnInner::Sqlite(conn))
|
||||
}
|
||||
|
||||
Err(e) => Err(diesel::r2d2::Error::ConnectionError(diesel::ConnectionError::InvalidConnectionUrl(
|
||||
format!("Unable to estabilsh a connection: {e:?}"),
|
||||
))),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl diesel::r2d2::ManageConnection for DbConnManager {
|
||||
type Connection = DbConnInner;
|
||||
type Error = diesel::r2d2::Error;
|
||||
|
||||
fn connect(&self) -> Result<Self::Connection, Self::Error> {
|
||||
self.establish_connection()
|
||||
}
|
||||
|
||||
fn is_valid(&self, conn: &mut Self::Connection) -> Result<(), Self::Error> {
|
||||
use diesel::r2d2::R2D2Connection;
|
||||
conn.ping().map_err(diesel::r2d2::Error::QueryError)
|
||||
}
|
||||
|
||||
fn has_broken(&self, conn: &mut Self::Connection) -> bool {
|
||||
use diesel::r2d2::R2D2Connection;
|
||||
conn.is_broken()
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Eq, PartialEq)]
|
||||
pub enum DbConnType {
|
||||
#[cfg(mysql)]
|
||||
Mysql,
|
||||
#[cfg(postgresql)]
|
||||
Postgresql,
|
||||
#[cfg(sqlite)]
|
||||
Sqlite,
|
||||
}
|
||||
|
||||
pub static ACTIVE_DB_TYPE: OnceLock<DbConnType> = OnceLock::new();
|
||||
|
||||
pub struct DbConn {
|
||||
conn: Arc<Mutex<Option<PooledConnection<DbConnManager>>>>,
|
||||
permit: Option<OwnedSemaphorePermit>,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct DbConnOptions {
|
||||
pub init_stmts: String,
|
||||
}
|
||||
|
||||
impl CustomizeConnection<DbConnInner, diesel::r2d2::Error> for DbConnOptions {
|
||||
fn on_acquire(&self, conn: &mut DbConnInner) -> Result<(), diesel::r2d2::Error> {
|
||||
if !self.init_stmts.is_empty() {
|
||||
conn.batch_execute(&self.init_stmts).map_err(diesel::r2d2::Error::QueryError)?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct DbPool {
|
||||
// This is an 'Option' so that we can drop the pool in a 'spawn_blocking'.
|
||||
pool: Option<Pool<DbConnManager>>,
|
||||
semaphore: Arc<Semaphore>,
|
||||
}
|
||||
|
||||
impl Drop for DbConn {
|
||||
fn drop(&mut self) {
|
||||
let conn = Arc::clone(&self.conn);
|
||||
let permit = self.permit.take();
|
||||
|
||||
// Since connection can't be on the stack in an async fn during an
|
||||
// await, we have to spawn a new blocking-safe thread...
|
||||
tokio::task::spawn_blocking(move || {
|
||||
// And then re-enter the runtime to wait on the async mutex, but in a blocking fashion.
|
||||
let mut conn = tokio::runtime::Handle::current().block_on(conn.lock_owned());
|
||||
|
||||
if let Some(conn) = conn.take() {
|
||||
drop(conn);
|
||||
}
|
||||
|
||||
// Drop permit after the connection is dropped
|
||||
drop(permit);
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for DbPool {
|
||||
fn drop(&mut self) {
|
||||
let pool = self.pool.take();
|
||||
// Only use spawn_blocking if the Tokio runtime is still available
|
||||
// Otherwise the pool will be dropped on the current thread
|
||||
if let Ok(handle) = tokio::runtime::Handle::try_current() {
|
||||
handle.spawn_blocking(move || drop(pool));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl DbPool {
|
||||
// For the given database URL, guess its type, run migrations, create pool, and return it
|
||||
pub fn from_config() -> Result<Self, Error> {
|
||||
let db_url = CONFIG.database_url();
|
||||
let conn_type = DbConnType::from_url(&db_url)?;
|
||||
|
||||
// Only set the default instrumentation if the log level is specifically set to either warn, info or debug
|
||||
if log_enabled!(target: "vaultwarden::db::query_logger", log::Level::Warn)
|
||||
|| log_enabled!(target: "vaultwarden::db::query_logger", log::Level::Info)
|
||||
|| log_enabled!(target: "vaultwarden::db::query_logger", log::Level::Debug)
|
||||
{
|
||||
drop(diesel::connection::set_default_instrumentation(query_logger::simple_logger));
|
||||
}
|
||||
|
||||
match conn_type {
|
||||
#[cfg(mysql)]
|
||||
DbConnType::Mysql => {
|
||||
mysql_migrations::run_migrations(&db_url)?;
|
||||
}
|
||||
#[cfg(postgresql)]
|
||||
DbConnType::Postgresql => {
|
||||
postgresql_migrations::run_migrations(&db_url)?;
|
||||
}
|
||||
#[cfg(sqlite)]
|
||||
DbConnType::Sqlite => {
|
||||
sqlite_migrations::run_migrations(&db_url)?;
|
||||
}
|
||||
}
|
||||
|
||||
let max_conns = CONFIG.database_max_conns();
|
||||
let manager = DbConnManager::new(&db_url);
|
||||
let pool = Pool::builder()
|
||||
.max_size(max_conns)
|
||||
.min_idle(Some(CONFIG.database_min_conns()))
|
||||
.idle_timeout(Some(Duration::from_secs(CONFIG.database_idle_timeout())))
|
||||
.connection_timeout(Duration::from_secs(CONFIG.database_timeout()))
|
||||
.connection_customizer(Box::new(DbConnOptions {
|
||||
init_stmts: conn_type.get_init_stmts(),
|
||||
}))
|
||||
.build(manager)
|
||||
.map_res("Failed to create pool")?;
|
||||
|
||||
// Set a global to determine the database more easily throughout the rest of the code
|
||||
if ACTIVE_DB_TYPE.set(conn_type).is_err() {
|
||||
error!("Tried to set the active database connection type more than once.")
|
||||
}
|
||||
|
||||
Ok(DbPool {
|
||||
pool: Some(pool),
|
||||
semaphore: Arc::new(Semaphore::new(max_conns as usize)),
|
||||
})
|
||||
}
|
||||
|
||||
// Get a connection from the pool
|
||||
pub async fn get(&self) -> Result<DbConn, Error> {
|
||||
let duration = Duration::from_secs(CONFIG.database_timeout());
|
||||
let permit = match timeout(duration, Arc::clone(&self.semaphore).acquire_owned()).await {
|
||||
Ok(p) => p.expect("Semaphore should be open"),
|
||||
Err(_) => {
|
||||
err!("Timeout waiting for database connection");
|
||||
}
|
||||
};
|
||||
|
||||
let p = self.pool.as_ref().expect("DbPool.pool should always be Some()");
|
||||
let pool = p.clone();
|
||||
let c =
|
||||
run_blocking(move || pool.get_timeout(duration)).await.map_res("Error retrieving connection from pool")?;
|
||||
Ok(DbConn {
|
||||
conn: Arc::new(Mutex::new(Some(c))),
|
||||
permit: Some(permit),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl DbConnType {
|
||||
pub fn from_url(url: &str) -> Result<DbConnType, Error> {
|
||||
pub fn from_url(url: &str) -> Result<Self, Error> {
|
||||
// Mysql
|
||||
if url.starts_with("mysql:") {
|
||||
if url.len() > 6 && &url[..6] == "mysql:" {
|
||||
#[cfg(mysql)]
|
||||
return Ok(DbConnType::mysql);
|
||||
return Ok(DbConnType::Mysql);
|
||||
|
||||
#[cfg(not(mysql))]
|
||||
err!("`DATABASE_URL` is a MySQL URL, but the 'mysql' feature is not enabled")
|
||||
|
||||
// Postgres
|
||||
} else if url.starts_with("postgresql:") || url.starts_with("postgres:") {
|
||||
// Postgresql
|
||||
} else if url.len() > 11 && (&url[..11] == "postgresql:" || &url[..9] == "postgres:") {
|
||||
#[cfg(postgresql)]
|
||||
return Ok(DbConnType::postgresql);
|
||||
return Ok(DbConnType::Postgresql);
|
||||
|
||||
#[cfg(not(postgresql))]
|
||||
err!("`DATABASE_URL` is a PostgreSQL URL, but the 'postgresql' feature is not enabled")
|
||||
@@ -214,7 +275,7 @@ impl DbConnType {
|
||||
//Sqlite
|
||||
} else {
|
||||
#[cfg(sqlite)]
|
||||
return Ok(DbConnType::sqlite);
|
||||
return Ok(DbConnType::Sqlite);
|
||||
|
||||
#[cfg(not(sqlite))]
|
||||
err!("`DATABASE_URL` looks like a SQLite URL, but 'sqlite' feature is not enabled")
|
||||
@@ -232,175 +293,106 @@ impl DbConnType {
|
||||
|
||||
pub fn default_init_stmts(&self) -> String {
|
||||
match self {
|
||||
Self::sqlite => "PRAGMA busy_timeout = 5000; PRAGMA synchronous = NORMAL;".to_string(),
|
||||
Self::mysql => String::new(),
|
||||
Self::postgresql => String::new(),
|
||||
#[cfg(mysql)]
|
||||
Self::Mysql => String::new(),
|
||||
#[cfg(postgresql)]
|
||||
Self::Postgresql => String::new(),
|
||||
#[cfg(sqlite)]
|
||||
Self::Sqlite => "PRAGMA busy_timeout = 5000; PRAGMA synchronous = NORMAL;".to_string(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl DbConn {
|
||||
pub async fn run<F, R>(&self, f: F) -> R
|
||||
where
|
||||
F: FnOnce(&mut DbConnInner) -> R + Send,
|
||||
R: Send + 'static,
|
||||
{
|
||||
let conn = Arc::clone(&self.conn);
|
||||
let mut conn = conn.lock_owned().await;
|
||||
let conn = conn.as_mut().expect("Internal invariant broken: self.conn is Some");
|
||||
|
||||
// Run blocking can't be used due to the 'static limitation, use block_in_place instead
|
||||
tokio::task::block_in_place(move || f(conn))
|
||||
}
|
||||
}
|
||||
|
||||
#[macro_export]
|
||||
macro_rules! db_run {
|
||||
// Same for all dbs
|
||||
( $conn:ident: $body:block ) => {
|
||||
db_run! { $conn: sqlite, mysql, postgresql $body }
|
||||
$conn.run(move |$conn| $body).await
|
||||
};
|
||||
|
||||
( @raw $conn:ident: $body:block ) => {
|
||||
db_run! { @raw $conn: sqlite, mysql, postgresql $body }
|
||||
};
|
||||
|
||||
// Different code for each db
|
||||
( $conn:ident: $( $($db:ident),+ $body:block )+ ) => {{
|
||||
#[allow(unused)] use diesel::prelude::*;
|
||||
#[allow(unused)] use $crate::db::FromDb;
|
||||
|
||||
let conn = $conn.conn.clone();
|
||||
let mut conn = conn.lock_owned().await;
|
||||
match conn.as_mut().expect("internal invariant broken: self.connection is Some") {
|
||||
( $conn:ident: $( $($db:ident),+ $body:block )+ ) => {
|
||||
$conn.run(move |$conn| {
|
||||
match $conn {
|
||||
$($(
|
||||
#[cfg($db)]
|
||||
$crate::db::DbConnInner::$db($conn) => {
|
||||
pastey::paste! {
|
||||
#[allow(unused)] use $crate::db::[<__ $db _schema>]::{self as schema, *};
|
||||
#[allow(unused)] use [<__ $db _model>]::*;
|
||||
}
|
||||
|
||||
tokio::task::block_in_place(move || { $body }) // Run blocking can't be used due to the 'static limitation, use block_in_place instead
|
||||
pastey::paste!(&mut $crate::db::DbConnInner::[<$db:camel>](ref mut $conn)) => {
|
||||
$body
|
||||
},
|
||||
)+)+
|
||||
}
|
||||
}};
|
||||
|
||||
( @raw $conn:ident: $( $($db:ident),+ $body:block )+ ) => {{
|
||||
#[allow(unused)] use diesel::prelude::*;
|
||||
#[allow(unused)] use $crate::db::FromDb;
|
||||
|
||||
let conn = $conn.conn.clone();
|
||||
let mut conn = conn.lock_owned().await;
|
||||
match conn.as_mut().expect("internal invariant broken: self.connection is Some") {
|
||||
$($(
|
||||
#[cfg($db)]
|
||||
$crate::db::DbConnInner::$db($conn) => {
|
||||
pastey::paste! {
|
||||
#[allow(unused)] use $crate::db::[<__ $db _schema>]::{self as schema, *};
|
||||
// @ RAW: #[allow(unused)] use [<__ $db _model>]::*;
|
||||
}
|
||||
|
||||
tokio::task::block_in_place(move || { $body }) // Run blocking can't be used due to the 'static limitation, use block_in_place instead
|
||||
},
|
||||
)+)+
|
||||
}
|
||||
}};
|
||||
}
|
||||
|
||||
pub trait FromDb {
|
||||
type Output;
|
||||
#[allow(clippy::wrong_self_convention)]
|
||||
fn from_db(self) -> Self::Output;
|
||||
}
|
||||
|
||||
impl<T: FromDb> FromDb for Vec<T> {
|
||||
type Output = Vec<T::Output>;
|
||||
#[inline(always)]
|
||||
fn from_db(self) -> Self::Output {
|
||||
self.into_iter().map(FromDb::from_db).collect()
|
||||
}
|
||||
}
|
||||
|
||||
impl<T: FromDb> FromDb for Option<T> {
|
||||
type Output = Option<T::Output>;
|
||||
#[inline(always)]
|
||||
fn from_db(self) -> Self::Output {
|
||||
self.map(FromDb::from_db)
|
||||
}
|
||||
}
|
||||
|
||||
// For each struct eg. Cipher, we create a CipherDb inside a module named __$db_model (where $db is sqlite, mysql or postgresql),
|
||||
// to implement the Diesel traits. We also provide methods to convert between them and the basic structs. Later, that module will be auto imported when using db_run!
|
||||
#[macro_export]
|
||||
macro_rules! db_object {
|
||||
( $(
|
||||
$( #[$attr:meta] )*
|
||||
pub struct $name:ident {
|
||||
$( $( #[$field_attr:meta] )* $vis:vis $field:ident : $typ:ty ),+
|
||||
$(,)?
|
||||
}
|
||||
)+ ) => {
|
||||
// Create the normal struct, without attributes
|
||||
$( pub struct $name { $( /*$( #[$field_attr] )**/ $vis $field : $typ, )+ } )+
|
||||
|
||||
#[cfg(sqlite)]
|
||||
pub mod __sqlite_model { $( db_object! { @db sqlite | $( #[$attr] )* | $name | $( $( #[$field_attr] )* $field : $typ ),+ } )+ }
|
||||
#[cfg(mysql)]
|
||||
pub mod __mysql_model { $( db_object! { @db mysql | $( #[$attr] )* | $name | $( $( #[$field_attr] )* $field : $typ ),+ } )+ }
|
||||
#[cfg(postgresql)]
|
||||
pub mod __postgresql_model { $( db_object! { @db postgresql | $( #[$attr] )* | $name | $( $( #[$field_attr] )* $field : $typ ),+ } )+ }
|
||||
};
|
||||
|
||||
( @db $db:ident | $( #[$attr:meta] )* | $name:ident | $( $( #[$field_attr:meta] )* $vis:vis $field:ident : $typ:ty),+) => {
|
||||
pastey::paste! {
|
||||
#[allow(unused)] use super::*;
|
||||
#[allow(unused)] use diesel::prelude::*;
|
||||
#[allow(unused)] use $crate::db::[<__ $db _schema>]::*;
|
||||
|
||||
$( #[$attr] )*
|
||||
pub struct [<$name Db>] { $(
|
||||
$( #[$field_attr] )* $vis $field : $typ,
|
||||
)+ }
|
||||
|
||||
impl [<$name Db>] {
|
||||
#[allow(clippy::wrong_self_convention)]
|
||||
#[inline(always)] pub fn to_db(x: &super::$name) -> Self { Self { $( $field: x.$field.clone(), )+ } }
|
||||
}
|
||||
|
||||
impl $crate::db::FromDb for [<$name Db>] {
|
||||
type Output = super::$name;
|
||||
#[allow(clippy::wrong_self_convention)]
|
||||
#[inline(always)] fn from_db(self) -> Self::Output { super::$name { $( $field: self.$field, )+ } }
|
||||
}
|
||||
}
|
||||
)+)+}
|
||||
}).await
|
||||
};
|
||||
}
|
||||
|
||||
pub mod schema;
|
||||
|
||||
// Reexport the models, needs to be after the macros are defined so it can access them
|
||||
pub mod models;
|
||||
|
||||
/// Creates a back-up of the sqlite database
|
||||
/// MySQL/MariaDB and PostgreSQL are not supported.
|
||||
pub async fn backup_database(conn: &mut DbConn) -> Result<String, Error> {
|
||||
db_run! {@raw conn:
|
||||
postgresql, mysql {
|
||||
let _ = conn;
|
||||
err!("PostgreSQL and MySQL/MariaDB do not support this backup feature");
|
||||
}
|
||||
sqlite {
|
||||
let db_url = CONFIG.database_url();
|
||||
let db_path = std::path::Path::new(&db_url).parent().unwrap();
|
||||
let backup_file = db_path
|
||||
.join(format!("db_{}.sqlite3", chrono::Utc::now().format("%Y%m%d_%H%M%S")))
|
||||
.to_string_lossy()
|
||||
.into_owned();
|
||||
diesel::sql_query(format!("VACUUM INTO '{backup_file}'")).execute(conn)?;
|
||||
Ok(backup_file)
|
||||
#[cfg(sqlite)]
|
||||
pub fn backup_sqlite() -> Result<String, Error> {
|
||||
use diesel::Connection;
|
||||
use std::{fs::File, io::Write};
|
||||
|
||||
let db_url = CONFIG.database_url();
|
||||
if DbConnType::from_url(&CONFIG.database_url()).map(|t| t == DbConnType::Sqlite).unwrap_or(false) {
|
||||
// Since we do not allow any schema for sqlite database_url's like `file:` or `sqlite:` to be set, we can assume here it isn't
|
||||
// This way we can set a readonly flag on the opening mode without issues.
|
||||
let mut conn = diesel::sqlite::SqliteConnection::establish(&format!("sqlite://{db_url}?mode=ro"))?;
|
||||
|
||||
let db_path = std::path::Path::new(&db_url).parent().unwrap();
|
||||
let backup_file = db_path
|
||||
.join(format!("db_{}.sqlite3", chrono::Utc::now().format("%Y%m%d_%H%M%S")))
|
||||
.to_string_lossy()
|
||||
.into_owned();
|
||||
|
||||
match File::create(backup_file.clone()) {
|
||||
Ok(mut f) => {
|
||||
let serialized_db = conn.serialize_database_to_buffer();
|
||||
f.write_all(serialized_db.as_slice()).expect("Error writing SQLite backup");
|
||||
Ok(backup_file)
|
||||
}
|
||||
Err(e) => {
|
||||
err_silent!(format!("Unable to save SQLite backup: {e:?}"))
|
||||
}
|
||||
}
|
||||
} else {
|
||||
err_silent!("The database type is not SQLite. Backups only works for SQLite databases")
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(not(sqlite))]
|
||||
pub fn backup_sqlite() -> Result<String, Error> {
|
||||
err_silent!("The database type is not SQLite. Backups only works for SQLite databases")
|
||||
}
|
||||
|
||||
/// Get the SQL Server version
|
||||
pub async fn get_sql_server_version(conn: &mut DbConn) -> String {
|
||||
db_run! {@raw conn:
|
||||
postgresql, mysql {
|
||||
define_sql_function!{
|
||||
fn version() -> diesel::sql_types::Text;
|
||||
}
|
||||
diesel::select(version()).get_result::<String>(conn).unwrap_or_else(|_| "Unknown".to_string())
|
||||
pub async fn get_sql_server_version(conn: &DbConn) -> String {
|
||||
db_run! { conn:
|
||||
postgresql,mysql {
|
||||
diesel::select(diesel::dsl::sql::<diesel::sql_types::Text>("version();"))
|
||||
.get_result::<String>(conn)
|
||||
.unwrap_or_else(|_| "Unknown".to_string())
|
||||
}
|
||||
sqlite {
|
||||
define_sql_function!{
|
||||
fn sqlite_version() -> diesel::sql_types::Text;
|
||||
}
|
||||
diesel::select(sqlite_version()).get_result::<String>(conn).unwrap_or_else(|_| "Unknown".to_string())
|
||||
diesel::select(diesel::dsl::sql::<diesel::sql_types::Text>("sqlite_version();"))
|
||||
.get_result::<String>(conn)
|
||||
.unwrap_or_else(|_| "Unknown".to_string())
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -428,16 +420,14 @@ impl<'r> FromRequest<'r> for DbConn {
|
||||
// https://docs.rs/diesel_migrations/*/diesel_migrations/macro.embed_migrations.html
|
||||
#[cfg(sqlite)]
|
||||
mod sqlite_migrations {
|
||||
use diesel::{Connection, RunQueryDsl};
|
||||
use diesel_migrations::{EmbeddedMigrations, MigrationHarness};
|
||||
pub const MIGRATIONS: EmbeddedMigrations = embed_migrations!("migrations/sqlite");
|
||||
|
||||
pub fn run_migrations() -> Result<(), super::Error> {
|
||||
use diesel::{Connection, RunQueryDsl};
|
||||
let url = crate::CONFIG.database_url();
|
||||
|
||||
pub fn run_migrations(db_url: &str) -> Result<(), super::Error> {
|
||||
// Establish a connection to the sqlite database (this will create a new one, if it does
|
||||
// not exist, and exit if there is an error).
|
||||
let mut connection = diesel::sqlite::SqliteConnection::establish(&url)?;
|
||||
let mut connection = diesel::sqlite::SqliteConnection::establish(db_url)?;
|
||||
|
||||
// Run the migrations after successfully establishing a connection
|
||||
// Disable Foreign Key Checks during migration
|
||||
@@ -458,15 +448,15 @@ mod sqlite_migrations {
|
||||
|
||||
#[cfg(mysql)]
|
||||
mod mysql_migrations {
|
||||
use diesel::{Connection, RunQueryDsl};
|
||||
use diesel_migrations::{EmbeddedMigrations, MigrationHarness};
|
||||
pub const MIGRATIONS: EmbeddedMigrations = embed_migrations!("migrations/mysql");
|
||||
|
||||
pub fn run_migrations() -> Result<(), super::Error> {
|
||||
use diesel::{Connection, RunQueryDsl};
|
||||
pub fn run_migrations(db_url: &str) -> Result<(), super::Error> {
|
||||
// Make sure the database is up to date (create if it doesn't exist, or run the migrations)
|
||||
let mut connection = diesel::mysql::MysqlConnection::establish(&crate::CONFIG.database_url())?;
|
||||
// Disable Foreign Key Checks during migration
|
||||
let mut connection = diesel::mysql::MysqlConnection::establish(db_url)?;
|
||||
|
||||
// Disable Foreign Key Checks during migration
|
||||
// Scoped to a connection/session.
|
||||
diesel::sql_query("SET FOREIGN_KEY_CHECKS = 0")
|
||||
.execute(&mut connection)
|
||||
@@ -479,13 +469,14 @@ mod mysql_migrations {
|
||||
|
||||
#[cfg(postgresql)]
|
||||
mod postgresql_migrations {
|
||||
use diesel::Connection;
|
||||
use diesel_migrations::{EmbeddedMigrations, MigrationHarness};
|
||||
pub const MIGRATIONS: EmbeddedMigrations = embed_migrations!("migrations/postgresql");
|
||||
|
||||
pub fn run_migrations() -> Result<(), super::Error> {
|
||||
use diesel::Connection;
|
||||
pub fn run_migrations(db_url: &str) -> Result<(), super::Error> {
|
||||
// Make sure the database is up to date (create if it doesn't exist, or run the migrations)
|
||||
let mut connection = diesel::pg::PgConnection::establish(&crate::CONFIG.database_url())?;
|
||||
let mut connection = diesel::pg::PgConnection::establish(db_url)?;
|
||||
|
||||
connection.run_pending_migrations(MIGRATIONS).expect("Error running migrations");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user