gru/RustBox
1
0 Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity
Files
dc57aa7089d231a844608a72594003151a58d53a
RustBox/src/db/pool.rs
T
Mateusz Gruszczyński dc57aa7089 v0.2.0
2026-10-31 14:36:07 +01:00

2469 lines
87 KiB
Rust
Raw Blame History

use std::fmt;
use std::fs;
use std::future::Future;
use std::ops::Deref;
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex, OnceLock, mpsc};
use rusqlite::OptionalExtension as _;
use rusqlite::types::{ToSql, ToSqlOutput, Value};
use sha2::{Digest, Sha256};
use sqlx::pool::PoolConnection;
use sqlx::{MySql, MySqlPool, PgPool, Postgres, Row as SqlxRow};
use tokio::runtime::{Builder, Handle};
use tokio::sync::Mutex as AsyncMutex;
use crate::config::AppConfig;
pub const BASE_SCHEMA_VERSION: u32 = 46;
const SQLITE_SCHEMA: &str = include_str!("../../migrations/sqlite/schema.sql");
const MYSQL_SCHEMA: &str = include_str!("../../migrations/mysql/schema.sql");
const PGSQL_SCHEMA: &str = include_str!("../../migrations/pgsql/schema.sql");
// The squashed legacy baseline is V046. Runtime migrations start at V047 and
// are stored as real files for every supported SQL engine.
const SQLITE_MIGRATIONS: &[(u32, &str)] = &[(
47,
include_str!("../../migrations/sqlite/V047__blob_keys_and_search_outbox.sql"),
)];
const MYSQL_MIGRATIONS: &[(u32, &str)] = &[(
47,
include_str!("../../migrations/mysql/V047__blob_keys_and_search_outbox.sql"),
)];
const PGSQL_MIGRATIONS: &[(u32, &str)] = &[(
47,
include_str!("../../migrations/pgsql/V047__blob_keys_and_search_outbox.sql"),
)];
const TENANT_TABLES: &[&str] = &[
"calendar_events",
"calendar_settings",
"contact_group_members",
"contact_groups",
"contact_photos",
"contacts",
"display_preferences",
"draft_attachments",
"drafts",
"folder_preferences",
"folders",
"identities",
"known_addresses",
"mail_preferences",
"meeting_templates",
"message_tags",
"messages",
"notification_preferences",
"pgp_keys",
"pgp_preferences",
"response_templates",
"saved_searches",
"search_outbox",
"special_folders",
"tags",
"trusted_senders",
"user_meta",
];
// Parent tables precede children so legacy SQLite data can be copied with foreign keys enabled.
const IMPORT_TABLES: &[&str] = &[
"calendar_events",
"calendar_settings",
"contact_groups",
"contacts",
"contact_group_members",
"contact_photos",
"display_preferences",
"drafts",
"draft_attachments",
"folders",
"folder_preferences",
"identities",
"known_addresses",
"mail_preferences",
"meeting_templates",
"messages",
"tags",
"message_tags",
"notification_preferences",
"pgp_keys",
"pgp_preferences",
"response_templates",
"saved_searches",
"search_outbox",
"special_folders",
"trusted_senders",
"user_meta",
];
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum DatabaseBackend {
Sqlite,
Mysql,
Pgsql,
}
impl DatabaseBackend {
pub fn parse(value: &str) -> Result<Self, String> {
match value.trim().to_ascii_lowercase().as_str() {
"sqlite" => Ok(Self::Sqlite),
"mysql" => Ok(Self::Mysql),
"pgsql" => Ok(Self::Pgsql),
other => Err(format!("unsupported database backend: {other}")),
}
}
}
#[derive(Clone)]
enum ConfiguredDatabase {
Sqlite { path: PathBuf },
Mysql { pool: MySqlPool, runtime: Handle },
Pgsql { pool: PgPool, runtime: Handle },
}
static DATABASE: OnceLock<ConfiguredDatabase> = OnceLock::new();
static DB_RUNTIME: OnceLock<Handle> = OnceLock::new();
fn database_runtime() -> Result<Handle, String> {
if let Some(handle) = DB_RUNTIME.get() {
return Ok(handle.clone());
}
let (tx, rx) = mpsc::sync_channel(1);
std::thread::Builder::new()
.name("rustbox-database-runtime".to_string())
.spawn(move || {
let runtime = Builder::new_multi_thread()
.worker_threads(2)
.enable_all()
.build()
.expect("failed to build RustBox database runtime");
let handle = runtime.handle().clone();
tx.send(handle).expect("database runtime receiver dropped");
runtime.block_on(std::future::pending::<()>());
})
.map_err(|e| format!("failed to start database runtime: {e}"))?;
let handle = rx
.recv()
.map_err(|e| format!("failed to receive database runtime handle: {e}"))?;
let _ = DB_RUNTIME.set(handle.clone());
Ok(handle)
}
fn runtime_wait<T, F>(runtime: &Handle, future: F) -> Result<T, String>
where
T: Send + 'static,
F: Future<Output = Result<T, String>> + Send + 'static,
{
let (tx, rx) = mpsc::sync_channel(1);
runtime.spawn(async move {
let _ = tx.send(future.await);
});
rx.recv()
.map_err(|e| format!("database runtime stopped unexpectedly: {e}"))?
}
fn schema_section(schema: &str, begin: &str, end: &str) -> Result<String, String> {
let start = schema
.find(begin)
.ok_or_else(|| format!("schema marker missing: {begin}"))?
+ begin.len();
let tail = &schema[start..];
let finish = tail
.find(end)
.ok_or_else(|| format!("schema marker missing: {end}"))?;
Ok(tail[..finish].trim().to_string())
}
fn global_schema(schema: &str) -> Result<String, String> {
schema_section(schema, "-- RUSTBOX:GLOBAL-BEGIN", "-- RUSTBOX:GLOBAL-END")
}
pub(crate) fn global_schema_for_backend(backend: DatabaseBackend) -> Result<String, String> {
match backend {
DatabaseBackend::Sqlite => global_schema(SQLITE_SCHEMA),
DatabaseBackend::Mysql => global_schema(MYSQL_SCHEMA),
DatabaseBackend::Pgsql => global_schema(PGSQL_SCHEMA),
}
}
fn tenant_schema(schema: &str, prefix: &str) -> Result<String, String> {
Ok(
schema_section(schema, "-- RUSTBOX:TENANT-BEGIN", "-- RUSTBOX:TENANT-END")?
.replace("{{p}}", prefix),
)
}
fn sqlite_database_path(config: &AppConfig) -> Result<PathBuf, String> {
if let Some(raw) = config
.database_url
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
{
let without_scheme = raw
.strip_prefix("sqlite://")
.or_else(|| raw.strip_prefix("sqlite:"))
.unwrap_or(raw);
let path = without_scheme.split('?').next().unwrap_or_default().trim();
if path.is_empty() || path == ":memory:" {
return Err("SQLite DATABASE_URL must point to a persistent file".to_string());
}
return Ok(PathBuf::from(path));
}
Ok(PathBuf::from(&config.data_dir).join("rustbox.sqlite"))
}
pub async fn configure(config: &AppConfig) -> Result<(), String> {
if DATABASE.get().is_some() {
return Ok(());
}
let backend = DatabaseBackend::parse(&config.database_backend)?;
let configured = match backend {
DatabaseBackend::Sqlite => {
let path = sqlite_database_path(config)?;
if let Some(parent) = path.parent().filter(|p| !p.as_os_str().is_empty()) {
fs::create_dir_all(parent)
.map_err(|e| format!("failed to create SQLite database directory: {e}"))?;
}
let conn = rusqlite::Connection::open(&path)
.map_err(|e| format!("failed to open SQLite database: {e}"))?;
conn.busy_timeout(std::time::Duration::from_secs(5))
.map_err(|e| format!("failed to configure SQLite busy timeout: {e}"))?;
conn.execute_batch("PRAGMA journal_mode=WAL; PRAGMA foreign_keys=ON;")
.map_err(|e| format!("failed to configure SQLite database: {e}"))?;
conn.execute_batch(&global_schema(SQLITE_SCHEMA)?)
.map_err(|e| format!("SQLite global schema initialization failed: {e}"))?;
ConfiguredDatabase::Sqlite { path }
}
DatabaseBackend::Mysql => {
let url = config
.database_url
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.ok_or("DATABASE_URL required for mysql")?
.to_string();
let runtime = database_runtime()?;
let pool = runtime_wait(&runtime, async move {
MySqlPool::connect(&url)
.await
.map_err(|e| format!("failed to connect MySQL database: {e}"))
})?;
let schema = global_schema(MYSQL_SCHEMA)?;
let init_pool = pool.clone();
runtime_wait(&runtime, async move {
sqlx::raw_sql(&schema)
.execute(&init_pool)
.await
.map_err(|e| format!("MySQL global schema initialization failed: {e}"))?;
Ok(())
})?;
ConfiguredDatabase::Mysql { pool, runtime }
}
DatabaseBackend::Pgsql => {
let url = config
.database_url
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.ok_or("DATABASE_URL required for pgsql")?
.to_string();
let runtime = database_runtime()?;
let pool = runtime_wait(&runtime, async move {
PgPool::connect(&url)
.await
.map_err(|e| format!("failed to connect PostgreSQL database: {e}"))
})?;
let schema = global_schema(PGSQL_SCHEMA)?;
let init_pool = pool.clone();
runtime_wait(&runtime, async move {
sqlx::raw_sql(&schema)
.execute(&init_pool)
.await
.map_err(|e| format!("PostgreSQL global schema initialization failed: {e}"))?;
Ok(())
})?;
ConfiguredDatabase::Pgsql { pool, runtime }
}
};
DATABASE
.set(configured)
.map_err(|_| "database already configured".to_string())
}
pub fn configured_backend() -> DatabaseBackend {
match DATABASE.get() {
Some(ConfiguredDatabase::Sqlite { .. }) | None => DatabaseBackend::Sqlite,
Some(ConfiguredDatabase::Mysql { .. }) => DatabaseBackend::Mysql,
Some(ConfiguredDatabase::Pgsql { .. }) => DatabaseBackend::Pgsql,
}
}
pub fn list_account_user_hashes() -> Result<Vec<String>, String> {
match DATABASE.get() {
Some(ConfiguredDatabase::Sqlite { path }) => {
let conn = rusqlite::Connection::open(path)
.map_err(|e| format!("failed to open SQLite account registry: {e}"))?;
let mut stmt = conn
.prepare("SELECT user_hash FROM rustbox_accounts ORDER BY user_hash")
.map_err(|e| format!("failed to prepare account registry query: {e}"))?;
let rows = stmt
.query_map([], |row| row.get::<_, String>(0))
.map_err(|e| format!("failed to query account registry: {e}"))?;
rows.map(|row| row.map_err(|e| format!("failed to read account registry row: {e}")))
.collect()
}
Some(ConfiguredDatabase::Mysql { pool, runtime }) => {
let pool = pool.clone();
runtime_wait(runtime, async move {
sqlx::query_scalar::<_, String>(
"SELECT user_hash FROM rustbox_accounts ORDER BY user_hash",
)
.fetch_all(&pool)
.await
.map_err(|e| format!("failed to query MySQL account registry: {e}"))
})
}
Some(ConfiguredDatabase::Pgsql { pool, runtime }) => {
let pool = pool.clone();
runtime_wait(runtime, async move {
sqlx::query_scalar::<_, String>(
"SELECT user_hash FROM rustbox_accounts ORDER BY user_hash",
)
.fetch_all(&pool)
.await
.map_err(|e| format!("failed to query PostgreSQL account registry: {e}"))
})
}
None => Err("database is not configured".to_string()),
}
}
#[macro_export]
macro_rules! dbq {
($module:ident::$name:ident) => {{
match $crate::db::pool::configured_backend() {
$crate::db::pool::DatabaseBackend::Sqlite => {
$crate::queries::sqlite::queries::$module::$name
}
$crate::db::pool::DatabaseBackend::Mysql => {
$crate::queries::mysql::queries::$module::$name
}
$crate::db::pool::DatabaseBackend::Pgsql => {
$crate::queries::pgsql::queries::$module::$name
}
}
}};
}
#[derive(Debug, Clone)]
pub enum DbError {
QueryReturnedNoRows,
Other(String),
}
impl fmt::Display for DbError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::QueryReturnedNoRows => f.write_str("query returned no rows"),
Self::Other(message) => f.write_str(message),
}
}
}
impl std::error::Error for DbError {}
impl From<rusqlite::Error> for DbError {
fn from(error: rusqlite::Error) -> Self {
match error {
rusqlite::Error::QueryReturnedNoRows => Self::QueryReturnedNoRows,
other => Self::Other(other.to_string()),
}
}
}
pub trait OptionalExtension<T> {
fn optional(self) -> Result<Option<T>, DbError>;
}
impl<T> OptionalExtension<T> for Result<T, DbError> {
fn optional(self) -> Result<Option<T>, DbError> {
match self {
Ok(value) => Ok(Some(value)),
Err(DbError::QueryReturnedNoRows) => Ok(None),
Err(error) => Err(error),
}
}
}
#[derive(Debug, Clone, PartialEq)]
pub enum DbValue {
Text(Option<String>),
I64(Option<i64>),
F64(Option<f64>),
Bytes(Option<Vec<u8>>),
}
impl ToSql for DbValue {
fn to_sql(&self) -> rusqlite::Result<ToSqlOutput<'_>> {
let value = match self {
Self::Text(Some(value)) => Value::Text(value.clone()),
Self::Text(None) => Value::Null,
Self::I64(Some(value)) => Value::Integer(*value),
Self::I64(None) => Value::Null,
Self::F64(Some(value)) => Value::Real(*value),
Self::F64(None) => Value::Null,
Self::Bytes(Some(value)) => Value::Blob(value.clone()),
Self::Bytes(None) => Value::Null,
};
Ok(ToSqlOutput::Owned(value))
}
}
pub trait ToDbValue {
fn to_db_value(&self) -> DbValue;
fn null_db_value() -> DbValue
where
Self: Sized,
{
DbValue::Text(None)
}
}
// Kept explicit rather than clever: the macro above intentionally stays simple for rustfmt/parser compatibility.
impl ToDbValue for i64 {
fn to_db_value(&self) -> DbValue {
DbValue::I64(Some(*self))
}
fn null_db_value() -> DbValue {
DbValue::I64(None)
}
}
impl ToDbValue for i32 {
fn to_db_value(&self) -> DbValue {
DbValue::I64(Some(i64::from(*self)))
}
fn null_db_value() -> DbValue {
DbValue::I64(None)
}
}
impl ToDbValue for i16 {
fn to_db_value(&self) -> DbValue {
DbValue::I64(Some(i64::from(*self)))
}
fn null_db_value() -> DbValue {
DbValue::I64(None)
}
}
impl ToDbValue for u64 {
fn to_db_value(&self) -> DbValue {
DbValue::I64(Some((*self).min(i64::MAX as u64) as i64))
}
fn null_db_value() -> DbValue {
DbValue::I64(None)
}
}
impl ToDbValue for u32 {
fn to_db_value(&self) -> DbValue {
DbValue::I64(Some(i64::from(*self)))
}
fn null_db_value() -> DbValue {
DbValue::I64(None)
}
}
impl ToDbValue for usize {
fn to_db_value(&self) -> DbValue {
DbValue::I64(Some((*self).min(i64::MAX as usize) as i64))
}
fn null_db_value() -> DbValue {
DbValue::I64(None)
}
}
impl ToDbValue for bool {
fn to_db_value(&self) -> DbValue {
DbValue::I64(Some(if *self { 1 } else { 0 }))
}
fn null_db_value() -> DbValue {
DbValue::I64(None)
}
}
impl ToDbValue for f64 {
fn to_db_value(&self) -> DbValue {
DbValue::F64(Some(*self))
}
fn null_db_value() -> DbValue {
DbValue::F64(None)
}
}
impl ToDbValue for String {
fn to_db_value(&self) -> DbValue {
DbValue::Text(Some(self.clone()))
}
fn null_db_value() -> DbValue {
DbValue::Text(None)
}
}
impl ToDbValue for str {
fn to_db_value(&self) -> DbValue {
DbValue::Text(Some(self.to_string()))
}
}
impl ToDbValue for Vec<u8> {
fn to_db_value(&self) -> DbValue {
DbValue::Bytes(Some(self.clone()))
}
fn null_db_value() -> DbValue {
DbValue::Bytes(None)
}
}
impl ToDbValue for [u8] {
fn to_db_value(&self) -> DbValue {
DbValue::Bytes(Some(self.to_vec()))
}
}
impl<T: ToDbValue> ToDbValue for Option<T> {
fn to_db_value(&self) -> DbValue {
self.as_ref()
.map(ToDbValue::to_db_value)
.unwrap_or_else(T::null_db_value)
}
fn null_db_value() -> DbValue {
T::null_db_value()
}
}
impl<T: ToDbValue + ?Sized> ToDbValue for &T {
fn to_db_value(&self) -> DbValue {
(*self).to_db_value()
}
fn null_db_value() -> DbValue
where
Self: Sized,
{
DbValue::Text(None)
}
}
#[macro_export]
macro_rules! db_params {
() => { Vec::<$crate::db::pool::DbValue>::new() };
($($value:expr),+ $(,)?) => {{
vec![$($crate::db::pool::ToDbValue::to_db_value(&$value)),+]
}};
}
pub use crate::db_params as params;
pub trait IntoDbParams {
fn into_db_params(self) -> Vec<DbValue>;
}
impl IntoDbParams for Vec<DbValue> {
fn into_db_params(self) -> Vec<DbValue> {
self
}
}
impl IntoDbParams for &[DbValue] {
fn into_db_params(self) -> Vec<DbValue> {
self.to_vec()
}
}
impl IntoDbParams for [(); 0] {
fn into_db_params(self) -> Vec<DbValue> {
Vec::new()
}
}
// Do not use a const-generic array implementation here. For an empty array (`[]`)
// Rust cannot infer the element type and sees both `[(); 0]` and `[T; 0]` as valid
// `IntoDbParams` candidates. Implement only non-empty array lengths explicitly so
// legacy/rusqlite-style `[]` call sites remain unambiguous.
macro_rules! impl_into_db_params_array {
($($n:expr),+ $(,)?) => {
$(
impl<T: ToDbValue> IntoDbParams for [T; $n] {
fn into_db_params(self) -> Vec<DbValue> {
self.iter().map(ToDbValue::to_db_value).collect()
}
}
)+
};
}
impl_into_db_params_array!(
1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26,
27, 28, 29, 30, 31, 32,
);
#[derive(Clone)]
enum ConnectionInner {
Sqlite(Arc<Mutex<rusqlite::Connection>>),
MysqlPool {
pool: MySqlPool,
runtime: Handle,
},
PgsqlPool {
pool: PgPool,
runtime: Handle,
},
MysqlTx {
conn: Arc<AsyncMutex<PoolConnection<MySql>>>,
runtime: Handle,
},
PgsqlTx {
conn: Arc<AsyncMutex<PoolConnection<Postgres>>>,
runtime: Handle,
},
}
#[derive(Clone)]
pub struct Connection {
inner: ConnectionInner,
prefix: String,
}
impl Connection {
pub fn backend(&self) -> DatabaseBackend {
match &self.inner {
ConnectionInner::Sqlite(_) => DatabaseBackend::Sqlite,
ConnectionInner::MysqlPool { .. } | ConnectionInner::MysqlTx { .. } => {
DatabaseBackend::Mysql
}
ConnectionInner::PgsqlPool { .. } | ConnectionInner::PgsqlTx { .. } => {
DatabaseBackend::Pgsql
}
}
}
fn scoped_sql(&self, sql: &str) -> String {
prefix_tenant_tables(sql, &self.prefix)
}
fn adapted_sql_params<P: IntoDbParams>(&self, sql: &str, params: P) -> (String, Vec<DbValue>) {
let sql = self.scoped_sql(sql);
adapt_sql_params(self.backend(), sql, params.into_db_params())
}
pub fn execute<P: IntoDbParams>(&self, sql: &str, params: P) -> Result<usize, DbError> {
let (sql, params) = self.adapted_sql_params(sql, params);
match &self.inner {
ConnectionInner::Sqlite(conn) => {
let conn = conn
.lock()
.map_err(|_| DbError::Other("SQLite mutex poisoned".into()))?;
conn.execute(&sql, rusqlite::params_from_iter(params.iter()))
.map_err(Into::into)
}
ConnectionInner::MysqlPool { pool, runtime } => {
let pool = pool.clone();
runtime_wait(runtime, async move {
mysql_execute_pool(pool, sql, params).await
})
.map_err(DbError::Other)
}
ConnectionInner::PgsqlPool { pool, runtime } => {
let pool = pool.clone();
runtime_wait(runtime, async move {
pgsql_execute_pool(pool, sql, params).await
})
.map_err(DbError::Other)
}
ConnectionInner::MysqlTx { conn, runtime } => {
let conn = Arc::clone(conn);
runtime_wait(
runtime,
async move { mysql_execute_tx(conn, sql, params).await },
)
.map_err(DbError::Other)
}
ConnectionInner::PgsqlTx { conn, runtime } => {
let conn = Arc::clone(conn);
runtime_wait(
runtime,
async move { pgsql_execute_tx(conn, sql, params).await },
)
.map_err(DbError::Other)
}
}
}
pub fn insert_id<P: IntoDbParams>(&self, sql: &str, params: P) -> Result<i64, DbError> {
let (sql, params) = self.adapted_sql_params(sql, params);
match &self.inner {
ConnectionInner::Sqlite(conn) => {
let conn = conn
.lock()
.map_err(|_| DbError::Other("SQLite mutex poisoned".into()))?;
conn.execute(&sql, rusqlite::params_from_iter(params.iter()))
.map_err(DbError::from)?;
Ok(conn.last_insert_rowid())
}
ConnectionInner::MysqlPool { pool, runtime } => {
let pool = pool.clone();
runtime_wait(runtime, async move {
mysql_insert_id_pool(pool, sql, params).await
})
.map_err(DbError::Other)
}
ConnectionInner::MysqlTx { conn, runtime } => {
let conn = Arc::clone(conn);
runtime_wait(runtime, async move {
mysql_insert_id_tx(conn, sql, params).await
})
.map_err(DbError::Other)
}
ConnectionInner::PgsqlPool { pool, runtime } => {
let pool = pool.clone();
runtime_wait(runtime, async move {
pgsql_insert_id_pool(pool, sql, params).await
})
.map_err(DbError::Other)
}
ConnectionInner::PgsqlTx { conn, runtime } => {
let conn = Arc::clone(conn);
runtime_wait(runtime, async move {
pgsql_insert_id_tx(conn, sql, params).await
})
.map_err(DbError::Other)
}
}
}
pub fn execute_batch(&self, sql: &str) -> Result<(), DbError> {
let sql = self.scoped_sql(sql);
match &self.inner {
ConnectionInner::Sqlite(conn) => {
let conn = conn
.lock()
.map_err(|_| DbError::Other("SQLite mutex poisoned".into()))?;
conn.execute_batch(&sql).map_err(Into::into)
}
ConnectionInner::MysqlPool { pool, runtime } => {
let pool = pool.clone();
runtime_wait(runtime, async move {
sqlx::raw_sql(&sql)
.execute(&pool)
.await
.map_err(|e| e.to_string())?;
Ok(())
})
.map_err(DbError::Other)
}
ConnectionInner::PgsqlPool { pool, runtime } => {
let pool = pool.clone();
runtime_wait(runtime, async move {
sqlx::raw_sql(&sql)
.execute(&pool)
.await
.map_err(|e| e.to_string())?;
Ok(())
})
.map_err(DbError::Other)
}
ConnectionInner::MysqlTx { conn, runtime } => {
let conn = Arc::clone(conn);
runtime_wait(runtime, async move {
mysql_execute_tx(conn, sql, Vec::new()).await.map(|_| ())
})
.map_err(DbError::Other)
}
ConnectionInner::PgsqlTx { conn, runtime } => {
let conn = Arc::clone(conn);
runtime_wait(runtime, async move {
pgsql_execute_tx(conn, sql, Vec::new()).await.map(|_| ())
})
.map_err(DbError::Other)
}
}
}
pub fn query_row<T, P, F>(&self, sql: &str, params: P, map: F) -> Result<T, DbError>
where
P: IntoDbParams,
F: FnOnce(&Row<'_>) -> Result<T, DbError>,
{
let (sql, params) = self.adapted_sql_params(sql, params);
match &self.inner {
ConnectionInner::Sqlite(conn) => {
let conn = conn
.lock()
.map_err(|_| DbError::Other("SQLite mutex poisoned".into()))?;
let mut statement = conn.prepare(&sql).map_err(DbError::from)?;
let mut rows = statement
.query(rusqlite::params_from_iter(params.iter()))
.map_err(DbError::from)?;
let row = rows
.next()
.map_err(DbError::from)?
.ok_or(DbError::QueryReturnedNoRows)?;
map(&Row::Sqlite(row))
}
_ => {
let mut rows = self.fetch_rows_owned(sql, params)?;
let row = rows.drain(..).next().ok_or(DbError::QueryReturnedNoRows)?;
map(&row.as_row())
}
}
}
pub fn prepare(&self, sql: &str) -> Result<Statement, DbError> {
Ok(Statement {
connection: self.clone(),
sql: self.scoped_sql(sql),
})
}
pub fn transaction(&mut self) -> Result<Transaction, DbError> {
self.begin_transaction()
}
pub fn unchecked_transaction(&self) -> Result<Transaction, DbError> {
self.begin_transaction()
}
fn begin_transaction(&self) -> Result<Transaction, DbError> {
let tx_connection = match &self.inner {
ConnectionInner::Sqlite(conn) => {
{
let conn = conn
.lock()
.map_err(|_| DbError::Other("SQLite mutex poisoned".into()))?;
conn.execute_batch("BEGIN IMMEDIATE")
.map_err(DbError::from)?;
}
Connection {
inner: ConnectionInner::Sqlite(Arc::clone(conn)),
prefix: self.prefix.clone(),
}
}
ConnectionInner::MysqlPool { pool, runtime } => {
let pool = pool.clone();
let conn = runtime_wait(runtime, async move {
let mut conn = pool.acquire().await.map_err(|e| e.to_string())?;
sqlx::query("START TRANSACTION")
.execute(&mut *conn)
.await
.map_err(|e| e.to_string())?;
Ok(conn)
})
.map_err(DbError::Other)?;
Connection {
inner: ConnectionInner::MysqlTx {
conn: Arc::new(AsyncMutex::new(conn)),
runtime: runtime.clone(),
},
prefix: self.prefix.clone(),
}
}
ConnectionInner::PgsqlPool { pool, runtime } => {
let pool = pool.clone();
let conn = runtime_wait(runtime, async move {
let mut conn = pool.acquire().await.map_err(|e| e.to_string())?;
sqlx::query("BEGIN")
.execute(&mut *conn)
.await
.map_err(|e| e.to_string())?;
Ok(conn)
})
.map_err(DbError::Other)?;
Connection {
inner: ConnectionInner::PgsqlTx {
conn: Arc::new(AsyncMutex::new(conn)),
runtime: runtime.clone(),
},
prefix: self.prefix.clone(),
}
}
ConnectionInner::MysqlTx { .. } | ConnectionInner::PgsqlTx { .. } => {
return Err(DbError::Other(
"nested database transactions are not supported".into(),
));
}
};
Ok(Transaction {
connection: tx_connection,
finished: false,
})
}
fn fetch_rows_owned(
&self,
sql: String,
params: Vec<DbValue>,
) -> Result<Vec<OwnedRow>, DbError> {
match &self.inner {
ConnectionInner::Sqlite(_) => {
Err(DbError::Other("internal SQLite fetch path mismatch".into()))
}
ConnectionInner::MysqlPool { pool, runtime } => {
let pool = pool.clone();
runtime_wait(
runtime,
async move { mysql_fetch_pool(pool, sql, params).await },
)
.map(|rows| rows.into_iter().map(OwnedRow::Mysql).collect())
.map_err(DbError::Other)
}
ConnectionInner::PgsqlPool { pool, runtime } => {
let pool = pool.clone();
runtime_wait(
runtime,
async move { pgsql_fetch_pool(pool, sql, params).await },
)
.map(|rows| rows.into_iter().map(OwnedRow::Pgsql).collect())
.map_err(DbError::Other)
}
ConnectionInner::MysqlTx { conn, runtime } => {
let conn = Arc::clone(conn);
runtime_wait(
runtime,
async move { mysql_fetch_tx(conn, sql, params).await },
)
.map(|rows| rows.into_iter().map(OwnedRow::Mysql).collect())
.map_err(DbError::Other)
}
ConnectionInner::PgsqlTx { conn, runtime } => {
let conn = Arc::clone(conn);
runtime_wait(
runtime,
async move { pgsql_fetch_tx(conn, sql, params).await },
)
.map(|rows| rows.into_iter().map(OwnedRow::Pgsql).collect())
.map_err(DbError::Other)
}
}
}
}
pub struct Transaction {
connection: Connection,
finished: bool,
}
impl Transaction {
pub fn commit(mut self) -> Result<(), DbError> {
self.connection.execute_batch("COMMIT")?;
self.finished = true;
Ok(())
}
}
impl Deref for Transaction {
type Target = Connection;
fn deref(&self) -> &Self::Target {
&self.connection
}
}
impl Drop for Transaction {
fn drop(&mut self) {
if !self.finished {
let _ = self.connection.execute_batch("ROLLBACK");
}
}
}
pub struct Statement {
connection: Connection,
sql: String,
}
impl Statement {
pub fn query_map<T, P, F>(&mut self, params: P, mut map: F) -> Result<MappedRows<T>, DbError>
where
P: IntoDbParams,
F: FnMut(&Row<'_>) -> Result<T, DbError>,
{
let (sql, params) = adapt_sql_params(
self.connection.backend(),
self.sql.clone(),
params.into_db_params(),
);
let mut output = Vec::new();
match &self.connection.inner {
ConnectionInner::Sqlite(conn) => {
let conn = conn
.lock()
.map_err(|_| DbError::Other("SQLite mutex poisoned".into()))?;
let mut statement = conn.prepare(&sql).map_err(DbError::from)?;
let mut rows = statement
.query(rusqlite::params_from_iter(params.iter()))
.map_err(DbError::from)?;
while let Some(row) = rows.next().map_err(DbError::from)? {
output.push(map(&Row::Sqlite(row)));
}
}
_ => {
for row in self.connection.fetch_rows_owned(sql.clone(), params)? {
output.push(map(&row.as_row()));
}
}
}
Ok(MappedRows {
inner: output.into_iter(),
})
}
}
pub struct MappedRows<T> {
inner: std::vec::IntoIter<Result<T, DbError>>,
}
impl<T> Iterator for MappedRows<T> {
type Item = Result<T, DbError>;
fn next(&mut self) -> Option<Self::Item> {
self.inner.next()
}
}
enum OwnedRow {
Mysql(sqlx::mysql::MySqlRow),
Pgsql(sqlx::postgres::PgRow),
}
impl OwnedRow {
fn as_row(&self) -> Row<'_> {
match self {
Self::Mysql(row) => Row::Mysql(row),
Self::Pgsql(row) => Row::Pgsql(row),
}
}
}
pub enum Row<'a> {
Sqlite(&'a rusqlite::Row<'a>),
Mysql(&'a sqlx::mysql::MySqlRow),
Pgsql(&'a sqlx::postgres::PgRow),
}
pub enum DbColumn<'a> {
Index(usize),
Name(&'a str),
}
pub trait IntoDbColumn<'a> {
fn into_db_column(self) -> DbColumn<'a>;
}
impl<'a> IntoDbColumn<'a> for usize {
fn into_db_column(self) -> DbColumn<'a> {
DbColumn::Index(self)
}
}
impl<'a> IntoDbColumn<'a> for &'a str {
fn into_db_column(self) -> DbColumn<'a> {
DbColumn::Name(self)
}
}
pub trait DbDecode: Sized {
fn decode(row: &Row<'_>, column: DbColumn<'_>) -> Result<Self, DbError>;
}
impl<'a> Row<'a> {
pub fn get<I, T>(&self, index: I) -> Result<T, DbError>
where
I: IntoDbColumn<'a>,
T: DbDecode,
{
T::decode(self, index.into_db_column())
}
}
macro_rules! decode_direct {
($ty:ty) => {
impl DbDecode for $ty {
fn decode(row: &Row<'_>, column: DbColumn<'_>) -> Result<Self, DbError> {
match row {
Row::Sqlite(row) => match column {
DbColumn::Index(i) => row.get::<_, $ty>(i).map_err(DbError::from),
DbColumn::Name(n) => row.get::<_, $ty>(n).map_err(DbError::from),
},
Row::Mysql(row) => match column {
DbColumn::Index(i) => row
.try_get::<$ty, _>(i)
.map_err(|e| DbError::Other(e.to_string())),
DbColumn::Name(n) => row
.try_get::<$ty, _>(n)
.map_err(|e| DbError::Other(e.to_string())),
},
Row::Pgsql(row) => match column {
DbColumn::Index(i) => row
.try_get::<$ty, _>(i)
.map_err(|e| DbError::Other(e.to_string())),
DbColumn::Name(n) => row
.try_get::<$ty, _>(n)
.map_err(|e| DbError::Other(e.to_string())),
},
}
}
}
};
}
decode_direct!(String);
decode_direct!(Option<String>);
decode_direct!(Vec<u8>);
decode_direct!(Option<Vec<u8>>);
decode_direct!(i64);
decode_direct!(Option<i64>);
decode_direct!(f64);
decode_direct!(Option<f64>);
fn mysql_bool_by_index(row: &sqlx::mysql::MySqlRow, index: usize) -> Result<bool, DbError> {
row.try_get::<i64, _>(index)
.map(|value| value != 0)
.or_else(|_| row.try_get::<i32, _>(index).map(|value| value != 0))
.or_else(|_| row.try_get::<i16, _>(index).map(|value| value != 0))
.or_else(|_| row.try_get::<i8, _>(index).map(|value| value != 0))
.map_err(|e| DbError::Other(e.to_string()))
}
fn mysql_bool_by_name(row: &sqlx::mysql::MySqlRow, name: &str) -> Result<bool, DbError> {
row.try_get::<i64, _>(name)
.map(|value| value != 0)
.or_else(|_| row.try_get::<i32, _>(name).map(|value| value != 0))
.or_else(|_| row.try_get::<i16, _>(name).map(|value| value != 0))
.or_else(|_| row.try_get::<i8, _>(name).map(|value| value != 0))
.map_err(|e| DbError::Other(e.to_string()))
}
fn mysql_optional_bool_by_index(
row: &sqlx::mysql::MySqlRow,
index: usize,
) -> Result<Option<bool>, DbError> {
row.try_get::<Option<i64>, _>(index)
.map(|value| value.map(|n| n != 0))
.or_else(|_| {
row.try_get::<Option<i32>, _>(index)
.map(|value| value.map(|n| n != 0))
})
.or_else(|_| {
row.try_get::<Option<i16>, _>(index)
.map(|value| value.map(|n| n != 0))
})
.or_else(|_| {
row.try_get::<Option<i8>, _>(index)
.map(|value| value.map(|n| n != 0))
})
.map_err(|e| DbError::Other(e.to_string()))
}
fn mysql_optional_bool_by_name(
row: &sqlx::mysql::MySqlRow,
name: &str,
) -> Result<Option<bool>, DbError> {
row.try_get::<Option<i64>, _>(name)
.map(|value| value.map(|n| n != 0))
.or_else(|_| {
row.try_get::<Option<i32>, _>(name)
.map(|value| value.map(|n| n != 0))
})
.or_else(|_| {
row.try_get::<Option<i16>, _>(name)
.map(|value| value.map(|n| n != 0))
})
.or_else(|_| {
row.try_get::<Option<i8>, _>(name)
.map(|value| value.map(|n| n != 0))
})
.map_err(|e| DbError::Other(e.to_string()))
}
impl DbDecode for bool {
fn decode(row: &Row<'_>, column: DbColumn<'_>) -> Result<Self, DbError> {
match row {
Row::Sqlite(row) => match column {
DbColumn::Index(i) => row.get::<_, bool>(i).map_err(DbError::from),
DbColumn::Name(n) => row.get::<_, bool>(n).map_err(DbError::from),
},
Row::Mysql(row) => match column {
DbColumn::Index(i) => mysql_bool_by_index(row, i),
DbColumn::Name(n) => mysql_bool_by_name(row, n),
},
Row::Pgsql(row) => match column {
DbColumn::Index(i) => row
.try_get::<bool, _>(i)
.or_else(|_| row.try_get::<i64, _>(i).map(|value| value != 0))
.map_err(|e| DbError::Other(e.to_string())),
DbColumn::Name(n) => row
.try_get::<bool, _>(n)
.or_else(|_| row.try_get::<i64, _>(n).map(|value| value != 0))
.map_err(|e| DbError::Other(e.to_string())),
},
}
}
}
impl DbDecode for Option<bool> {
fn decode(row: &Row<'_>, column: DbColumn<'_>) -> Result<Self, DbError> {
match row {
Row::Sqlite(row) => match column {
DbColumn::Index(i) => row.get::<_, Option<bool>>(i).map_err(DbError::from),
DbColumn::Name(n) => row.get::<_, Option<bool>>(n).map_err(DbError::from),
},
Row::Mysql(row) => match column {
DbColumn::Index(i) => mysql_optional_bool_by_index(row, i),
DbColumn::Name(n) => mysql_optional_bool_by_name(row, n),
},
Row::Pgsql(row) => match column {
DbColumn::Index(i) => row
.try_get::<Option<bool>, _>(i)
.or_else(|_| {
row.try_get::<Option<i64>, _>(i)
.map(|value| value.map(|n| n != 0))
})
.map_err(|e| DbError::Other(e.to_string())),
DbColumn::Name(n) => row
.try_get::<Option<bool>, _>(n)
.or_else(|_| {
row.try_get::<Option<i64>, _>(n)
.map(|value| value.map(|n| n != 0))
})
.map_err(|e| DbError::Other(e.to_string())),
},
}
}
}
macro_rules! decode_via_i64 {
($ty:ty) => {
impl DbDecode for $ty {
fn decode(row: &Row<'_>, column: DbColumn<'_>) -> Result<Self, DbError> {
let value = i64::decode(row, column)?;
<$ty>::try_from(value).map_err(|_| {
DbError::Other(format!(
"database integer {value} does not fit {}",
stringify!($ty)
))
})
}
}
impl DbDecode for Option<$ty> {
fn decode(row: &Row<'_>, column: DbColumn<'_>) -> Result<Self, DbError> {
let value = Option::<i64>::decode(row, column)?;
value
.map(|v| {
<$ty>::try_from(v).map_err(|_| {
DbError::Other(format!(
"database integer {v} does not fit {}",
stringify!($ty)
))
})
})
.transpose()
}
}
};
}
decode_via_i64!(i32);
decode_via_i64!(u32);
decode_via_i64!(u64);
decode_via_i64!(usize);
fn bind_mysql<'q>(
mut query: sqlx::query::Query<'q, MySql, sqlx::mysql::MySqlArguments>,
params: Vec<DbValue>,
) -> sqlx::query::Query<'q, MySql, sqlx::mysql::MySqlArguments> {
for value in params {
query = match value {
DbValue::Text(value) => query.bind(value),
DbValue::I64(value) => query.bind(value),
DbValue::F64(value) => query.bind(value),
DbValue::Bytes(value) => query.bind(value),
};
}
query
}
fn bind_pgsql<'q>(
mut query: sqlx::query::Query<'q, Postgres, sqlx::postgres::PgArguments>,
params: Vec<DbValue>,
) -> sqlx::query::Query<'q, Postgres, sqlx::postgres::PgArguments> {
for value in params {
query = match value {
DbValue::Text(value) => query.bind(value),
DbValue::I64(value) => query.bind(value),
DbValue::F64(value) => query.bind(value),
DbValue::Bytes(value) => query.bind(value),
};
}
query
}
async fn mysql_insert_id_pool(
pool: MySqlPool,
sql: String,
params: Vec<DbValue>,
) -> Result<i64, String> {
let result = bind_mysql(sqlx::query(&sql), params)
.execute(&pool)
.await
.map_err(|e| e.to_string())?;
i64::try_from(result.last_insert_id()).map_err(|_| "MySQL insert id exceeds i64".to_string())
}
async fn mysql_insert_id_tx(
conn: Arc<AsyncMutex<PoolConnection<MySql>>>,
sql: String,
params: Vec<DbValue>,
) -> Result<i64, String> {
let mut conn = conn.lock().await;
let result = bind_mysql(sqlx::query(&sql), params)
.execute(&mut **conn)
.await
.map_err(|e| e.to_string())?;
i64::try_from(result.last_insert_id()).map_err(|_| "MySQL insert id exceeds i64".to_string())
}
fn with_returning_id(sql: String) -> String {
let trimmed = sql.trim_end().trim_end_matches(';').trim_end();
format!("{trimmed} RETURNING id")
}
async fn pgsql_insert_id_pool(
pool: PgPool,
sql: String,
params: Vec<DbValue>,
) -> Result<i64, String> {
let sql = with_returning_id(sql);
let row = bind_pgsql(sqlx::query(&sql), params)
.fetch_one(&pool)
.await
.map_err(|e| e.to_string())?;
row.try_get::<i64, _>(0).map_err(|e| e.to_string())
}
async fn pgsql_insert_id_tx(
conn: Arc<AsyncMutex<PoolConnection<Postgres>>>,
sql: String,
params: Vec<DbValue>,
) -> Result<i64, String> {
let sql = with_returning_id(sql);
let mut conn = conn.lock().await;
let row = bind_pgsql(sqlx::query(&sql), params)
.fetch_one(&mut **conn)
.await
.map_err(|e| e.to_string())?;
row.try_get::<i64, _>(0).map_err(|e| e.to_string())
}
async fn mysql_execute_pool(
pool: MySqlPool,
sql: String,
params: Vec<DbValue>,
) -> Result<usize, String> {
let result = bind_mysql(sqlx::query(&sql), params)
.execute(&pool)
.await
.map_err(|e| e.to_string())?;
Ok(result.rows_affected() as usize)
}
async fn pgsql_execute_pool(
pool: PgPool,
sql: String,
params: Vec<DbValue>,
) -> Result<usize, String> {
let result = bind_pgsql(sqlx::query(&sql), params)
.execute(&pool)
.await
.map_err(|e| e.to_string())?;
Ok(result.rows_affected() as usize)
}
async fn mysql_execute_tx(
conn: Arc<AsyncMutex<PoolConnection<MySql>>>,
sql: String,
params: Vec<DbValue>,
) -> Result<usize, String> {
let mut conn = conn.lock().await;
let result = bind_mysql(sqlx::query(&sql), params)
.execute(&mut **conn)
.await
.map_err(|e| e.to_string())?;
Ok(result.rows_affected() as usize)
}
async fn pgsql_execute_tx(
conn: Arc<AsyncMutex<PoolConnection<Postgres>>>,
sql: String,
params: Vec<DbValue>,
) -> Result<usize, String> {
let mut conn = conn.lock().await;
let result = bind_pgsql(sqlx::query(&sql), params)
.execute(&mut **conn)
.await
.map_err(|e| e.to_string())?;
Ok(result.rows_affected() as usize)
}
async fn mysql_fetch_pool(
pool: MySqlPool,
sql: String,
params: Vec<DbValue>,
) -> Result<Vec<sqlx::mysql::MySqlRow>, String> {
bind_mysql(sqlx::query(&sql), params)
.fetch_all(&pool)
.await
.map_err(|e| e.to_string())
}
async fn pgsql_fetch_pool(
pool: PgPool,
sql: String,
params: Vec<DbValue>,
) -> Result<Vec<sqlx::postgres::PgRow>, String> {
bind_pgsql(sqlx::query(&sql), params)
.fetch_all(&pool)
.await
.map_err(|e| e.to_string())
}
async fn mysql_fetch_tx(
conn: Arc<AsyncMutex<PoolConnection<MySql>>>,
sql: String,
params: Vec<DbValue>,
) -> Result<Vec<sqlx::mysql::MySqlRow>, String> {
let mut conn = conn.lock().await;
bind_mysql(sqlx::query(&sql), params)
.fetch_all(&mut **conn)
.await
.map_err(|e| e.to_string())
}
async fn pgsql_fetch_tx(
conn: Arc<AsyncMutex<PoolConnection<Postgres>>>,
sql: String,
params: Vec<DbValue>,
) -> Result<Vec<sqlx::postgres::PgRow>, String> {
let mut conn = conn.lock().await;
bind_pgsql(sqlx::query(&sql), params)
.fetch_all(&mut **conn)
.await
.map_err(|e| e.to_string())
}
fn adapt_sql_params(
backend: DatabaseBackend,
mut sql: String,
params: Vec<DbValue>,
) -> (String, Vec<DbValue>) {
if backend == DatabaseBackend::Sqlite {
return (sql, params);
}
// Dynamic query builders in the application use SQLite-style ?N placeholders.
// Normalize them here so the same higher-level logic can target all engines.
let chars: Vec<char> = sql.chars().collect();
let mut output = String::with_capacity(sql.len() + 16);
let mut reordered = Vec::new();
let mut sequential = 0usize;
let mut i = 0usize;
let mut quote: Option<char> = None;
while i < chars.len() {
let ch = chars[i];
if let Some(q) = quote {
output.push(ch);
if ch == q {
if i + 1 < chars.len() && chars[i + 1] == q {
output.push(chars[i + 1]);
i += 2;
continue;
}
quote = None;
}
i += 1;
continue;
}
if matches!(ch, '\'' | '"' | '`') {
quote = Some(ch);
output.push(ch);
i += 1;
continue;
}
if ch == '?' {
let mut j = i + 1;
while j < chars.len() && chars[j].is_ascii_digit() {
j += 1;
}
let index = if j > i + 1 {
chars[i + 1..j]
.iter()
.collect::<String>()
.parse::<usize>()
.ok()
.and_then(|value| value.checked_sub(1))
} else {
let value = sequential;
sequential += 1;
Some(value)
};
if backend == DatabaseBackend::Mysql {
output.push('?');
if let Some(index) = index.and_then(|index| params.get(index).map(|_| index)) {
reordered.push(params[index].clone());
}
} else {
let number = index.map(|index| index + 1).unwrap_or(1);
output.push('$');
output.push_str(&number.to_string());
}
i = j;
continue;
}
output.push(ch);
i += 1;
}
sql = output;
if backend == DatabaseBackend::Mysql && !reordered.is_empty() {
// If numbered placeholders were expanded, bind in occurrence order.
// Plain MySQL catalog queries already use '?' and are handled the same way.
return (adapt_dialect_sql(backend, sql), reordered);
}
(adapt_dialect_sql(backend, sql), params)
}
fn adapt_dialect_sql(backend: DatabaseBackend, sql: String) -> String {
match backend {
DatabaseBackend::Sqlite => sql,
DatabaseBackend::Mysql => sql
.replace("datetime('now')", "CURRENT_TIMESTAMP")
.replace("datetime(\"now\")", "CURRENT_TIMESTAMP"),
DatabaseBackend::Pgsql => sql
.replace("datetime('now')", "CURRENT_TIMESTAMP::text")
.replace("datetime(\"now\")", "CURRENT_TIMESTAMP::text")
.replace("instr(", "strpos("),
}
}
fn prefix_tenant_tables(sql: &str, prefix: &str) -> String {
let mut output = String::with_capacity(sql.len() + 32);
let chars: Vec<char> = sql.chars().collect();
let mut i = 0usize;
let mut quote: Option<char> = None;
let mut expect_table = false;
while i < chars.len() {
let ch = chars[i];
if let Some(q) = quote {
output.push(ch);
if ch == q {
if i + 1 < chars.len() && chars[i + 1] == q {
output.push(chars[i + 1]);
i += 2;
continue;
}
quote = None;
}
i += 1;
continue;
}
if matches!(ch, '\'' | '"' | '`') {
quote = Some(ch);
output.push(ch);
i += 1;
continue;
}
if ch.is_ascii_alphabetic() || ch == '_' {
let start = i;
i += 1;
while i < chars.len() && (chars[i].is_ascii_alphanumeric() || chars[i] == '_') {
i += 1;
}
let token: String = chars[start..i].iter().collect();
let upper = token.to_ascii_uppercase();
let mut next = i;
while next < chars.len() && chars[next].is_ascii_whitespace() {
next += 1;
}
let is_qualifier = next < chars.len() && chars[next] == '.';
let is_tenant_table = TENANT_TABLES.contains(&token.as_str());
if is_tenant_table && (expect_table || is_qualifier) {
output.push_str(prefix);
}
output.push_str(&token);
if expect_table {
expect_table = false;
}
if matches!(
upper.as_str(),
"FROM" | "JOIN" | "UPDATE" | "INTO" | "TABLE" | "REFERENCES"
) {
expect_table = true;
}
continue;
}
output.push(ch);
i += 1;
}
output
}
#[derive(Clone)]
struct LegacyTableData {
name: String,
columns: Vec<String>,
rows: Vec<Vec<DbValue>>,
}
fn null_value_for_declared_type(declared: &str) -> DbValue {
let upper = declared.to_ascii_uppercase();
if upper.contains("INT") {
DbValue::I64(None)
} else if upper.contains("REAL") || upper.contains("FLOA") || upper.contains("DOUB") {
DbValue::F64(None)
} else if upper.contains("BLOB") || upper.contains("BINARY") {
DbValue::Bytes(None)
} else {
DbValue::Text(None)
}
}
fn read_legacy_sqlite(path: &Path) -> Result<Vec<LegacyTableData>, String> {
use rusqlite::types::ValueRef;
let conn = rusqlite::Connection::open(path).map_err(|e| {
format!(
"failed to open legacy SQLite database {}: {e}",
path.display()
)
})?;
let version: u32 = conn
.query_row("PRAGMA user_version", [], |row| row.get(0))
.map_err(|e| format!("failed to read legacy SQLite schema version: {e}"))?;
if version < BASE_SCHEMA_VERSION {
return Err(format!(
"legacy SQLite schema version {version} is older than the legacy RustBox baseline {BASE_SCHEMA_VERSION}; upgrade with RustBox 0.1.0 first"
));
}
let latest_supported_sqlite = SQLITE_MIGRATIONS
.last()
.map(|(migration_version, _)| *migration_version)
.unwrap_or(BASE_SCHEMA_VERSION);
if version > latest_supported_sqlite {
return Err(format!(
"legacy SQLite schema version {version} is newer than the latest SQLite schema supported by this build ({latest_supported_sqlite})"
));
}
let mut output = Vec::new();
for &table in IMPORT_TABLES {
let exists: i64 = conn
.query_row(
"SELECT COUNT(*) FROM sqlite_master WHERE type='table' AND name=?1",
[table],
|row| row.get(0),
)
.map_err(|e| format!("failed to inspect legacy table {table}: {e}"))?;
if exists == 0 {
continue;
}
let escaped = table.replace('"', "\"\"");
let mut info = conn
.prepare(&format!("PRAGMA table_info(\"{escaped}\")"))
.map_err(|e| format!("failed to inspect legacy columns for {table}: {e}"))?;
let columns_with_types: Vec<(String, String)> = info
.query_map([], |row| Ok((row.get(1)?, row.get(2)?)))
.map_err(|e| format!("failed to query legacy columns for {table}: {e}"))?
.collect::<rusqlite::Result<Vec<_>>>()
.map_err(|e| format!("failed to read legacy columns for {table}: {e}"))?;
drop(info);
let columns: Vec<String> = columns_with_types
.iter()
.map(|(name, _)| name.clone())
.collect();
let quoted_columns = columns
.iter()
.map(|name| format!("\"{}\"", name.replace('"', "\"\"")))
.collect::<Vec<_>>()
.join(", ");
let mut statement = conn
.prepare(&format!("SELECT {quoted_columns} FROM \"{escaped}\""))
.map_err(|e| format!("failed to prepare legacy read for {table}: {e}"))?;
let mut rows = statement
.query([])
.map_err(|e| format!("failed to query legacy table {table}: {e}"))?;
let mut table_rows = Vec::new();
while let Some(row) = rows
.next()
.map_err(|e| format!("failed to read legacy row from {table}: {e}"))?
{
let mut values = Vec::with_capacity(columns.len());
for (index, (_, declared_type)) in columns_with_types.iter().enumerate() {
let value = match row
.get_ref(index)
.map_err(|e| format!("failed to read legacy value {table}[{index}]: {e}"))?
{
ValueRef::Null => null_value_for_declared_type(declared_type),
ValueRef::Integer(value) => DbValue::I64(Some(value)),
ValueRef::Real(value) => DbValue::F64(Some(value)),
ValueRef::Text(value) => {
DbValue::Text(Some(String::from_utf8_lossy(value).into_owned()))
}
ValueRef::Blob(value) => DbValue::Bytes(Some(value.to_vec())),
};
values.push(value);
}
table_rows.push(values);
}
output.push(LegacyTableData {
name: table.to_string(),
columns,
rows: table_rows,
});
}
Ok(output)
}
fn tenant_seed_sql(backend: DatabaseBackend, prefix: &str) -> String {
match backend {
DatabaseBackend::Sqlite => format!(
"INSERT OR IGNORE INTO {prefix}user_meta (key, value) VALUES ('schema_version', '1');\n\
INSERT OR IGNORE INTO {prefix}notification_preferences (id) VALUES (1);\n\
INSERT OR IGNORE INTO {prefix}mail_preferences (id) VALUES (1);\n\
INSERT OR IGNORE INTO {prefix}pgp_preferences (id) VALUES (1);"
),
DatabaseBackend::Mysql => format!(
"INSERT IGNORE INTO {prefix}user_meta (`key`, value) VALUES ('schema_version', '1');\n\
INSERT IGNORE INTO {prefix}notification_preferences (id) VALUES (1);\n\
INSERT IGNORE INTO {prefix}mail_preferences (id) VALUES (1);\n\
INSERT IGNORE INTO {prefix}pgp_preferences (id) VALUES (1);"
),
DatabaseBackend::Pgsql => format!(
"INSERT INTO {prefix}user_meta (key, value) VALUES ('schema_version', '1') ON CONFLICT (key) DO NOTHING;\n\
INSERT INTO {prefix}notification_preferences (id) VALUES (1) ON CONFLICT (id) DO NOTHING;\n\
INSERT INTO {prefix}mail_preferences (id) VALUES (1) ON CONFLICT (id) DO NOTHING;\n\
INSERT INTO {prefix}pgp_preferences (id) VALUES (1) ON CONFLICT (id) DO NOTHING;"
),
}
}
fn legacy_backup_path(path: &Path) -> PathBuf {
path.with_extension("sqlite.pre-0.1.1.bak")
}
fn backup_legacy_sqlite(path: &Path) {
if path.is_file() {
let backup = legacy_backup_path(path);
if !backup.exists() {
let _ = fs::rename(path, backup);
}
}
}
fn mysql_legacy_insert_sql(prefix: &str, table: &LegacyTableData) -> String {
let columns = table
.columns
.iter()
.map(|name| format!("`{}`", name.replace('`', "``")))
.collect::<Vec<_>>()
.join(", ");
let placeholders = (0..table.columns.len())
.map(|_| "?")
.collect::<Vec<_>>()
.join(", ");
format!(
"INSERT INTO {prefix}{} ({columns}) VALUES ({placeholders})",
table.name
)
}
fn pgsql_legacy_insert_sql(prefix: &str, table: &LegacyTableData) -> String {
let columns = table
.columns
.iter()
.map(|name| format!("\"{}\"", name.replace('"', "\"\"")))
.collect::<Vec<_>>()
.join(", ");
let placeholders = (1..=table.columns.len())
.map(|index| format!("${index}"))
.collect::<Vec<_>>()
.join(", ");
format!(
"INSERT INTO {prefix}{} ({columns}) VALUES ({placeholders})",
table.name
)
}
async fn import_legacy_mysql(
pool: &MySqlPool,
prefix: &str,
tables: Vec<LegacyTableData>,
) -> Result<(), String> {
for table in tables {
let sql = mysql_legacy_insert_sql(prefix, &table);
for row in table.rows {
mysql_execute_pool(pool.clone(), sql.clone(), row)
.await
.map_err(|e| {
format!(
"failed to import legacy table {} into MySQL: {e}",
table.name
)
})?;
}
}
Ok(())
}
async fn import_legacy_pgsql(
pool: &PgPool,
prefix: &str,
tables: Vec<LegacyTableData>,
) -> Result<(), String> {
for table in tables {
let sql = pgsql_legacy_insert_sql(prefix, &table);
for row in table.rows {
pgsql_execute_pool(pool.clone(), sql.clone(), row)
.await
.map_err(|e| {
format!(
"failed to import legacy table {} into PostgreSQL: {e}",
table.name
)
})?;
}
}
for table in [
"identities",
"meeting_templates",
"response_templates",
"saved_searches",
] {
let qualified = format!("{prefix}{table}");
let sql = format!(
"SELECT setval(pg_get_serial_sequence('{qualified}', 'id'), COALESCE(MAX(id), 1), MAX(id) IS NOT NULL) FROM {qualified}"
);
sqlx::query(&sql)
.execute(pool)
.await
.map_err(|e| format!("failed to advance PostgreSQL sequence for {qualified}: {e}"))?;
}
Ok(())
}
fn tenant_prefix(user_hash: &str) -> Result<String, String> {
if user_hash.is_empty() {
return Err("empty user hash for database isolation".to_string());
}
// Production user hashes are SHA-256 hex strings. Tests and legacy callers may
// use symbolic identifiers; hash those first so table prefixes remain safe,
// deterministic and restricted to ASCII hex without weakening isolation.
let normalized = if user_hash.len() >= 24 && user_hash.chars().all(|c| c.is_ascii_hexdigit()) {
user_hash.to_ascii_lowercase()
} else {
format!("{:x}", Sha256::digest(user_hash.as_bytes()))
};
Ok(format!("u_{}__", &normalized[..24]))
}
fn latest_schema_version() -> u32 {
let migrations = match configured_backend() {
DatabaseBackend::Sqlite => SQLITE_MIGRATIONS,
DatabaseBackend::Mysql => MYSQL_MIGRATIONS,
DatabaseBackend::Pgsql => PGSQL_MIGRATIONS,
};
migrations
.last()
.map(|(version, _)| *version)
.unwrap_or(BASE_SCHEMA_VERSION)
}
fn configured_or_default_sqlite(data_dir: &str) -> Result<ConfiguredDatabase, String> {
if let Some(configured) = DATABASE.get() {
return Ok(configured.clone());
}
let path = Path::new(data_dir).join("rustbox.sqlite");
if let Some(parent) = path.parent() {
fs::create_dir_all(parent)
.map_err(|e| format!("failed to create database directory: {e}"))?;
}
let conn = rusqlite::Connection::open(&path)
.map_err(|e| format!("failed to open SQLite database: {e}"))?;
conn.execute_batch("PRAGMA journal_mode=WAL; PRAGMA foreign_keys=ON;")
.map_err(|e| format!("failed to configure SQLite database: {e}"))?;
conn.execute_batch(&global_schema(SQLITE_SCHEMA)?)
.map_err(|e| format!("SQLite global schema initialization failed: {e}"))?;
Ok(ConfiguredDatabase::Sqlite { path })
}
fn sqlite_legacy_insert_sql(prefix: &str, table: &LegacyTableData) -> String {
let columns = table
.columns
.iter()
.map(|name| format!("\"{}\"", name.replace('"', "\"\"")))
.collect::<Vec<_>>()
.join(", ");
let placeholders = (1..=table.columns.len())
.map(|index| format!("?{index}"))
.collect::<Vec<_>>()
.join(", ");
format!(
"INSERT INTO {prefix}{} ({columns}) VALUES ({placeholders})",
table.name
)
}
fn import_legacy_sqlite_data(
tx: &rusqlite::Transaction<'_>,
prefix: &str,
tables: &[LegacyTableData],
) -> Result<(), String> {
for table in tables {
let sql = sqlite_legacy_insert_sql(prefix, table);
for row in &table.rows {
tx.execute(&sql, rusqlite::params_from_iter(row.iter()))
.map_err(|e| format!("failed to import legacy table {}: {e}", table.name))?;
}
}
Ok(())
}
fn ensure_sqlite_tenant(
path: &Path,
data_dir: &str,
user_hash: &str,
prefix: &str,
) -> Result<(), String> {
let mut conn = rusqlite::Connection::open(path)
.map_err(|e| format!("failed to open SQLite database: {e}"))?;
conn.busy_timeout(std::time::Duration::from_secs(5))
.map_err(|e| e.to_string())?;
conn.execute_batch("PRAGMA foreign_keys=ON;")
.map_err(|e| e.to_string())?;
let existing: Option<(String, u32)> = conn
.query_row(
"SELECT table_prefix, schema_version FROM rustbox_accounts WHERE user_hash=?1",
[user_hash],
|row| Ok((row.get(0)?, row.get(1)?)),
)
.optional()
.map_err(|e| e.to_string())?;
if let Some((stored_prefix, version)) = existing {
if stored_prefix != prefix {
return Err("database account-prefix collision detected".to_string());
}
return apply_sqlite_migrations(&conn, prefix, version);
}
let collision: Option<String> = conn
.query_row(
"SELECT user_hash FROM rustbox_accounts WHERE table_prefix=?1",
[prefix],
|row| row.get(0),
)
.optional()
.map_err(|e| e.to_string())?;
if collision.is_some() {
return Err("database account-prefix collision detected".to_string());
}
let legacy_path = Path::new(data_dir).join(user_hash).join("db.sqlite");
let legacy_data = if legacy_path.is_file() {
read_legacy_sqlite(&legacy_path)?
} else {
Vec::new()
};
let tx = conn.transaction().map_err(|e| e.to_string())?;
tx.execute_batch(&tenant_schema(SQLITE_SCHEMA, prefix)?)
.map_err(|e| format!("SQLite tenant schema initialization failed: {e}"))?;
if !legacy_data.is_empty() {
import_legacy_sqlite_data(&tx, prefix, &legacy_data)?;
}
tx.execute_batch(&tenant_seed_sql(DatabaseBackend::Sqlite, prefix))
.map_err(|e| format!("failed to initialize SQLite tenant defaults: {e}"))?;
tx.execute(
"INSERT INTO rustbox_accounts(user_hash, table_prefix, schema_version) VALUES(?1, ?2, ?3)",
rusqlite::params![user_hash, prefix, latest_schema_version()],
)
.map_err(|e| format!("failed to register SQLite account schema: {e}"))?;
tx.commit().map_err(|e| e.to_string())?;
if !legacy_data.is_empty() {
backup_legacy_sqlite(&legacy_path);
}
Ok(())
}
fn apply_sqlite_migrations(
conn: &rusqlite::Connection,
prefix: &str,
current: u32,
) -> Result<(), String> {
if current < BASE_SCHEMA_VERSION {
return Err(format!(
"database schema version {current} is older than the legacy RustBox baseline {BASE_SCHEMA_VERSION}; upgrade with RustBox 0.1.0 first"
));
}
let latest = latest_schema_version();
if current > latest {
return Err(format!(
"database schema version {current} is newer than this RustBox build supports ({latest})"
));
}
for &(version, sql) in SQLITE_MIGRATIONS {
if version <= current {
continue;
}
let tx = conn.unchecked_transaction().map_err(|e| e.to_string())?;
tx.execute_batch(&sql.replace("{{p}}", prefix))
.map_err(|e| format!("SQLite migration V{version:03} failed: {e}"))?;
tx.execute("UPDATE rustbox_accounts SET schema_version=?1, updated_at=datetime('now') WHERE table_prefix=?2", rusqlite::params![version, prefix])
.map_err(|e| e.to_string())?;
tx.commit().map_err(|e| e.to_string())?;
}
Ok(())
}
async fn drop_mysql_tenant(pool: &MySqlPool, prefix: &str) {
let _ = sqlx::query("SET FOREIGN_KEY_CHECKS=0").execute(pool).await;
for table in IMPORT_TABLES.iter().rev() {
let sql = format!("DROP TABLE IF EXISTS {prefix}{table}");
let _ = sqlx::query(&sql).execute(pool).await;
}
let _ = sqlx::query("SET FOREIGN_KEY_CHECKS=1").execute(pool).await;
}
async fn drop_pgsql_tenant(pool: &PgPool, prefix: &str) {
for table in IMPORT_TABLES.iter().rev() {
let sql = format!("DROP TABLE IF EXISTS {prefix}{table} CASCADE");
let _ = sqlx::query(&sql).execute(pool).await;
}
}
fn ensure_server_tenant(
configured: &ConfiguredDatabase,
data_dir: &str,
user_hash: &str,
prefix: &str,
) -> Result<(), String> {
let legacy_path = (!data_dir.trim().is_empty())
.then(|| Path::new(data_dir).join(user_hash).join("db.sqlite"));
let imported_legacy = match configured {
ConfiguredDatabase::Mysql { pool, runtime } => {
let pool = pool.clone();
let user_hash = user_hash.to_string();
let prefix = prefix.to_string();
let tenant = tenant_schema(MYSQL_SCHEMA, &prefix)?;
let legacy_path = legacy_path.clone();
runtime_wait(runtime, async move {
// Serialize first-time provisioning per account across cluster nodes.
// MySQL named locks are session-scoped and are released automatically if the connection dies.
let lock_name = format!("rustbox:{prefix}");
let mut lock_conn = pool
.acquire()
.await
.map_err(|e| format!("failed to acquire MySQL provisioning connection: {e}"))?;
let acquired: Option<i64> = sqlx::query_scalar("SELECT GET_LOCK(?, 10)")
.bind(&lock_name)
.fetch_one(&mut *lock_conn)
.await
.map_err(|e| {
format!("failed to acquire MySQL account provisioning lock: {e}")
})?;
if acquired != Some(1) {
return Err("timed out waiting for MySQL account provisioning lock".to_string());
}
let result: Result<bool, String> = async {
let existing: Option<(String, i64)> = sqlx::query_as("SELECT table_prefix, schema_version FROM rustbox_accounts WHERE user_hash=?")
.bind(&user_hash).fetch_optional(&pool).await.map_err(|e| e.to_string())?;
if let Some((stored, version)) = existing {
if stored != prefix { return Err("database account-prefix collision detected".into()); }
if version < i64::from(BASE_SCHEMA_VERSION) { return Err(format!("database schema version {version} is older than baseline {BASE_SCHEMA_VERSION}")); }
if version > i64::from(latest_schema_version()) { return Err(format!("database schema version {version} is newer than this build")); }
for &(migration_version, migration) in MYSQL_MIGRATIONS {
if i64::from(migration_version) <= version { continue; }
let rendered = migration.replace("{{p}}", &prefix);
sqlx::raw_sql(&rendered).execute(&pool).await.map_err(|e| format!("MySQL migration V{migration_version:03} failed: {e}"))?;
sqlx::query("UPDATE rustbox_accounts SET schema_version=?, updated_at=CURRENT_TIMESTAMP WHERE user_hash=?")
.bind(i64::from(migration_version)).bind(&user_hash).execute(&pool).await.map_err(|e| e.to_string())?;
}
return Ok(false);
}
let collision: Option<String> = sqlx::query_scalar("SELECT user_hash FROM rustbox_accounts WHERE table_prefix=?")
.bind(&prefix).fetch_optional(&pool).await.map_err(|e| e.to_string())?;
if collision.is_some() { return Err("database account-prefix collision detected".into()); }
let legacy_data = if let Some(path) = legacy_path.as_ref().filter(|path| path.is_file()) {
read_legacy_sqlite(path)?
} else {
Vec::new()
};
let imported_legacy = !legacy_data.is_empty();
if let Err(error) = sqlx::raw_sql(&tenant).execute(&pool).await {
drop_mysql_tenant(&pool, &prefix).await;
return Err(format!("MySQL tenant schema initialization failed: {error}"));
}
if let Err(error) = import_legacy_mysql(&pool, &prefix, legacy_data).await {
drop_mysql_tenant(&pool, &prefix).await;
return Err(error);
}
if let Err(error) = sqlx::raw_sql(&tenant_seed_sql(DatabaseBackend::Mysql, &prefix)).execute(&pool).await {
drop_mysql_tenant(&pool, &prefix).await;
return Err(format!("failed to initialize MySQL tenant defaults: {error}"));
}
if let Err(error) = sqlx::query("INSERT INTO rustbox_accounts(user_hash, table_prefix, schema_version) VALUES(?, ?, ?)")
.bind(&user_hash).bind(&prefix).bind(i64::from(latest_schema_version())).execute(&pool).await
{
drop_mysql_tenant(&pool, &prefix).await;
return Err(format!("failed to register MySQL account schema: {error}"));
}
Ok(imported_legacy)
}.await;
let _ = sqlx::query("SELECT RELEASE_LOCK(?)")
.bind(&lock_name)
.execute(&mut *lock_conn)
.await;
result
})?
}
ConfiguredDatabase::Pgsql { pool, runtime } => {
let pool = pool.clone();
let user_hash = user_hash.to_string();
let prefix = prefix.to_string();
let tenant = tenant_schema(PGSQL_SCHEMA, &prefix)?;
let legacy_path = legacy_path.clone();
runtime_wait(runtime, async move {
// PostgreSQL advisory locks are session-scoped, so they work across all RustBox nodes
// that share this database and disappear automatically with the connection.
let lock_name = format!("rustbox:{prefix}");
let mut lock_conn = pool.acquire().await.map_err(|e| {
format!("failed to acquire PostgreSQL provisioning connection: {e}")
})?;
sqlx::query("SELECT pg_advisory_lock(hashtextextended($1, 0))")
.bind(&lock_name)
.execute(&mut *lock_conn)
.await
.map_err(|e| {
format!("failed to acquire PostgreSQL account provisioning lock: {e}")
})?;
let result: Result<bool, String> = async {
let existing: Option<(String, i64)> = sqlx::query_as("SELECT table_prefix, schema_version FROM rustbox_accounts WHERE user_hash=$1")
.bind(&user_hash).fetch_optional(&pool).await.map_err(|e| e.to_string())?;
if let Some((stored, version)) = existing {
if stored != prefix { return Err("database account-prefix collision detected".into()); }
if version < i64::from(BASE_SCHEMA_VERSION) { return Err(format!("database schema version {version} is older than baseline {BASE_SCHEMA_VERSION}")); }
if version > i64::from(latest_schema_version()) { return Err(format!("database schema version {version} is newer than this build")); }
for &(migration_version, migration) in PGSQL_MIGRATIONS {
if i64::from(migration_version) <= version { continue; }
let rendered = migration.replace("{{p}}", &prefix);
sqlx::raw_sql(&rendered).execute(&pool).await.map_err(|e| format!("PostgreSQL migration V{migration_version:03} failed: {e}"))?;
sqlx::query("UPDATE rustbox_accounts SET schema_version=$1, updated_at=CURRENT_TIMESTAMP::text WHERE user_hash=$2")
.bind(i64::from(migration_version)).bind(&user_hash).execute(&pool).await.map_err(|e| e.to_string())?;
}
return Ok(false);
}
let collision: Option<String> = sqlx::query_scalar("SELECT user_hash FROM rustbox_accounts WHERE table_prefix=$1")
.bind(&prefix).fetch_optional(&pool).await.map_err(|e| e.to_string())?;
if collision.is_some() { return Err("database account-prefix collision detected".into()); }
let legacy_data = if let Some(path) = legacy_path.as_ref().filter(|path| path.is_file()) {
read_legacy_sqlite(path)?
} else {
Vec::new()
};
let imported_legacy = !legacy_data.is_empty();
if let Err(error) = sqlx::raw_sql(&tenant).execute(&pool).await {
drop_pgsql_tenant(&pool, &prefix).await;
return Err(format!("PostgreSQL tenant schema initialization failed: {error}"));
}
if let Err(error) = import_legacy_pgsql(&pool, &prefix, legacy_data).await {
drop_pgsql_tenant(&pool, &prefix).await;
return Err(error);
}
if let Err(error) = sqlx::raw_sql(&tenant_seed_sql(DatabaseBackend::Pgsql, &prefix)).execute(&pool).await {
drop_pgsql_tenant(&pool, &prefix).await;
return Err(format!("failed to initialize PostgreSQL tenant defaults: {error}"));
}
if let Err(error) = sqlx::query("INSERT INTO rustbox_accounts(user_hash, table_prefix, schema_version) VALUES($1, $2, $3)")
.bind(&user_hash).bind(&prefix).bind(i64::from(latest_schema_version())).execute(&pool).await
{
drop_pgsql_tenant(&pool, &prefix).await;
return Err(format!("failed to register PostgreSQL account schema: {error}"));
}
Ok(imported_legacy)
}.await;
let _ = sqlx::query("SELECT pg_advisory_unlock(hashtextextended($1, 0))")
.bind(&lock_name)
.execute(&mut *lock_conn)
.await;
result
})?
}
ConfiguredDatabase::Sqlite { .. } => false,
};
if imported_legacy {
if let Some(path) = legacy_path.as_ref() {
backup_legacy_sqlite(path);
}
}
Ok(())
}
pub fn open_user_db(data_dir: &str, user_hash: &str) -> Result<Connection, String> {
let prefix = tenant_prefix(user_hash)?;
let configured = configured_or_default_sqlite(data_dir)?;
match configured {
ConfiguredDatabase::Sqlite { path } => {
ensure_sqlite_tenant(&path, data_dir, user_hash, &prefix)?;
let conn = rusqlite::Connection::open(path)
.map_err(|e| format!("failed to open SQLite database: {e}"))?;
conn.busy_timeout(std::time::Duration::from_secs(5))
.map_err(|e| e.to_string())?;
conn.execute_batch("PRAGMA journal_mode=WAL; PRAGMA foreign_keys=ON;")
.map_err(|e| e.to_string())?;
Ok(Connection {
inner: ConnectionInner::Sqlite(Arc::new(Mutex::new(conn))),
prefix,
})
}
ConfiguredDatabase::Mysql { pool, runtime } => {
let configured = ConfiguredDatabase::Mysql {
pool: pool.clone(),
runtime: runtime.clone(),
};
ensure_server_tenant(&configured, data_dir, user_hash, &prefix)?;
Ok(Connection {
inner: ConnectionInner::MysqlPool { pool, runtime },
prefix,
})
}
ConfiguredDatabase::Pgsql { pool, runtime } => {
let configured = ConfiguredDatabase::Pgsql {
pool: pool.clone(),
runtime: runtime.clone(),
};
ensure_server_tenant(&configured, data_dir, user_hash, &prefix)?;
Ok(Connection {
inner: ConnectionInner::PgsqlPool { pool, runtime },
prefix,
})
}
}
}
#[derive(Clone)]
pub(crate) struct UserDb {
data_dir: String,
user_hash: String,
}
impl UserDb {
pub(crate) fn new(data_dir: &str, user_hash: &str) -> Self {
Self {
data_dir: data_dir.into(),
user_hash: user_hash.into(),
}
}
pub(crate) async fn run<T: Send + 'static>(
&self,
work: impl FnOnce(&Connection) -> Result<T, String> + Send + 'static,
) -> Result<T, String> {
let context = self.clone();
tokio::task::spawn_blocking(move || {
let conn = open_user_db(&context.data_dir, &context.user_hash)?;
work(&conn)
})
.await
.map_err(|e| format!("Database task failed: {e}"))?
}
}
#[cfg(test)]
pub fn open_test_db() -> Connection {
let user_hash = format!("{:x}", Sha256::digest(b"rustbox-test"));
let raw = rusqlite::Connection::open_in_memory().expect("Failed to open in-memory SQLite");
raw.execute_batch("PRAGMA foreign_keys=ON;").unwrap();
raw.execute_batch(&global_schema(SQLITE_SCHEMA).unwrap())
.unwrap();
let prefix = tenant_prefix(&user_hash).unwrap();
raw.execute_batch(&tenant_schema(SQLITE_SCHEMA, &prefix).unwrap())
.unwrap();
raw.execute_batch(&tenant_seed_sql(DatabaseBackend::Sqlite, &prefix))
.unwrap();
raw.execute(
"INSERT INTO rustbox_accounts(user_hash, table_prefix, schema_version) VALUES(?1, ?2, ?3)",
rusqlite::params![user_hash, prefix, latest_schema_version()],
)
.unwrap();
Connection {
inner: ConnectionInner::Sqlite(Arc::new(Mutex::new(raw))),
prefix,
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn table_prefix_is_short_and_deterministic() {
let hash = "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef";
assert_eq!(tenant_prefix(hash).unwrap(), "u_0123456789abcdef01234567__");
}
#[test]
fn symbolic_user_hash_gets_safe_deterministic_prefix() {
let first = tenant_prefix("header-race").unwrap();
let second = tenant_prefix("header-race").unwrap();
assert_eq!(first, second);
assert!(first.starts_with("u_"));
assert!(first.ends_with("__"));
assert_eq!(first.len(), 28);
assert!(first[2..26].chars().all(|c| c.is_ascii_hexdigit()));
}
#[test]
fn tenant_rewriter_skips_string_literals() {
let sql = prefix_tenant_tables("SELECT * FROM contacts WHERE name='contacts'", "u_test__");
assert_eq!(sql, "SELECT * FROM u_test__contacts WHERE name='contacts'");
}
#[test]
fn tenant_rewriter_does_not_prefix_same_named_columns() {
let sql = prefix_tenant_tables(
"SELECT folders, contacts.email FROM notification_preferences JOIN contacts ON 1=1",
"u_test__",
);
assert_eq!(
sql,
"SELECT folders, u_test__contacts.email FROM u_test__notification_preferences JOIN u_test__contacts ON 1=1"
);
}
#[test]
fn test_schema_contains_sessions_and_user_tables() {
let conn = open_test_db();
let tables: Vec<String> = conn
.prepare("SELECT name FROM sqlite_master WHERE type='table' ORDER BY name")
.unwrap()
.query_map([], |row| row.get(0))
.unwrap()
.filter_map(Result::ok)
.collect();
assert!(tables.contains(&"sessions".to_string()));
assert!(tables.iter().any(|name| name.ends_with("__contacts")));
assert!(tables.iter().any(|name| name.ends_with("__messages")));
}
#[test]
fn test_connection_scopes_queries() {
let conn = open_test_db();
conn.execute(
crate::dbq!(db_trusted_senders::Q_002),
params!["alice@example.com"],
)
.unwrap();
let value: String = conn
.query_row("SELECT email FROM trusted_senders", [], |row| row.get(0))
.unwrap();
assert_eq!(value, "alice@example.com");
}
#[test]
fn mysql_numbered_placeholders_are_reordered_and_duplicated() {
let params = vec![DbValue::Text(Some("needle".into())), DbValue::I64(Some(25))];
let (sql, reordered) = adapt_sql_params(
DatabaseBackend::Mysql,
"SELECT ?1, ?2, ?1".to_string(),
params,
);
assert_eq!(sql, "SELECT ?, ?, ?");
assert_eq!(
reordered,
vec![
DbValue::Text(Some("needle".into())),
DbValue::I64(Some(25)),
DbValue::Text(Some("needle".into())),
]
);
}
#[test]
fn pgsql_numbered_placeholders_preserve_parameter_numbers() {
let params = vec![DbValue::Text(Some("needle".into())), DbValue::I64(Some(25))];
let (sql, unchanged) = adapt_sql_params(
DatabaseBackend::Pgsql,
"SELECT ?2, ?1, ?2".to_string(),
params.clone(),
);
assert_eq!(sql, "SELECT $2, $1, $2");
assert_eq!(unchanged, params);
}
#[test]
fn tenant_and_import_table_sets_match() {
let mut tenant = TENANT_TABLES.to_vec();
let mut imported = IMPORT_TABLES.to_vec();
tenant.sort_unstable();
imported.sort_unstable();
assert_eq!(tenant, imported);
}
#[test]
fn legacy_sqlite_v046_is_imported_and_backed_up() {
let temp = tempfile::tempdir().unwrap();
let data_dir = temp.path();
let user_hash = format!("{:x}", Sha256::digest(b"legacy-user"));
let legacy_dir = data_dir.join(&user_hash);
fs::create_dir_all(&legacy_dir).unwrap();
let legacy_path = legacy_dir.join("db.sqlite");
{
let legacy = rusqlite::Connection::open(&legacy_path).unwrap();
legacy
.execute_batch(&tenant_schema(SQLITE_SCHEMA, "").unwrap())
.unwrap();
legacy
.execute_batch(&tenant_seed_sql(DatabaseBackend::Sqlite, ""))
.unwrap();
legacy.execute_batch("PRAGMA user_version=46;").unwrap();
legacy.execute(
"INSERT INTO display_preferences(id, theme, accent_color) VALUES(1, 'dark', '#123456')",
[],
).unwrap();
legacy
.execute(
"INSERT INTO trusted_senders(email) VALUES(?1)",
["legacy@example.com"],
)
.unwrap();
}
let central_path = data_dir.join("rustbox.sqlite");
{
let central = rusqlite::Connection::open(&central_path).unwrap();
central
.execute_batch(&global_schema(SQLITE_SCHEMA).unwrap())
.unwrap();
}
let prefix = tenant_prefix(&user_hash).unwrap();
ensure_sqlite_tenant(
&central_path,
data_dir.to_str().unwrap(),
&user_hash,
&prefix,
)
.unwrap();
let central = rusqlite::Connection::open(&central_path).unwrap();
let theme: String = central
.query_row(
&format!("SELECT theme FROM {prefix}display_preferences WHERE id=1"),
[],
|row| row.get(0),
)
.unwrap();
let accent: String = central
.query_row(
&format!("SELECT accent_color FROM {prefix}display_preferences WHERE id=1"),
[],
|row| row.get(0),
)
.unwrap();
let trusted: String = central
.query_row(
&format!("SELECT email FROM {prefix}trusted_senders"),
[],
|row| row.get(0),
)
.unwrap();
assert_eq!(theme, "dark");
assert_eq!(accent, "#123456");
assert_eq!(trusted, "legacy@example.com");
assert!(!legacy_path.exists());
assert!(legacy_backup_path(&legacy_path).exists());
}
}
Reference in new issue View Git Blame Copy Permalink