[Mig] storage to SQLite pool
[Clean] split message handler dispatch
This commit is contained in:
Alex Emmet
parent
9b20d159b5
commit
3be1d9f308
18 changed files with 1949 additions and 1700 deletions
|
|
@ -1,88 +1,207 @@
|
|||
//! Database helper utilities.
|
||||
//!
|
||||
//! This module provides small helpers to open/init sqlite databases and to
|
||||
//! create a shared (Arc<Mutex<Connection>>) connection wrapper callers can
|
||||
//! reuse. The goal is to centralize the "open and initialize" logic and
|
||||
//! provide small convenience helpers used by other util modules.
|
||||
|
||||
use iota_util::file_util::get_directory;
|
||||
use rusqlite::{Connection, Error as RusqliteError};
|
||||
use once_cell::sync::Lazy;
|
||||
use r2d2::ManageConnection;
|
||||
use rusqlite::Connection;
|
||||
use std::path::PathBuf;
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
/// Returns the file path for a named DB inside the application's data directory.
|
||||
///
|
||||
/// Arguments:
|
||||
/// - `db_name` : name of the DB (without extension). Example: `"messages"`.
|
||||
pub fn db_file_path(db_name: &str) -> String {
|
||||
use crate::storage_error::StorageError;
|
||||
|
||||
const DB_NAME: &str = "messages";
|
||||
|
||||
/// A simple r2d2 manager for rusqlite connections.
|
||||
pub struct SqliteManager;
|
||||
|
||||
impl ManageConnection for SqliteManager {
|
||||
type Connection = Connection;
|
||||
type Error = rusqlite::Error;
|
||||
|
||||
fn connect(&self) -> Result<Connection, rusqlite::Error> {
|
||||
let path = db_file_path(DB_NAME);
|
||||
let conn = Connection::open(path)?;
|
||||
conn.execute_batch("PRAGMA journal_mode = WAL; PRAGMA synchronous = NORMAL;")?;
|
||||
conn.busy_timeout(Duration::from_millis(250))?;
|
||||
Ok(conn)
|
||||
}
|
||||
|
||||
fn is_valid(&self, conn: &mut Connection) -> Result<(), rusqlite::Error> {
|
||||
conn.execute_batch("SELECT 1")
|
||||
}
|
||||
|
||||
fn has_broken(&self, _conn: &mut Connection) -> bool {
|
||||
false
|
||||
}
|
||||
}
|
||||
|
||||
static POOL: Lazy<Arc<r2d2::Pool<SqliteManager>>> = Lazy::new(|| {
|
||||
let manager = SqliteManager;
|
||||
let pool = r2d2::Pool::builder()
|
||||
.max_size(8)
|
||||
.build(manager)
|
||||
.expect("Failed to create database connection pool");
|
||||
run_migrations(&pool).expect("Failed to run database migrations");
|
||||
Arc::new(pool)
|
||||
});
|
||||
|
||||
pub fn pool() -> Arc<r2d2::Pool<SqliteManager>> {
|
||||
POOL.clone()
|
||||
}
|
||||
|
||||
pub fn with_db<T, F>(f: F) -> Result<T, StorageError>
|
||||
where
|
||||
F: FnOnce(&Connection) -> Result<T, StorageError>,
|
||||
{
|
||||
let conn = POOL.get().map_err(|e| StorageError::Pool(e.to_string()))?;
|
||||
f(&conn)
|
||||
}
|
||||
|
||||
fn db_file_path(db_name: &str) -> String {
|
||||
let mut p = PathBuf::from(get_directory());
|
||||
p.push(format!("{db_name}.sqlite3"));
|
||||
p.to_string_lossy().to_string()
|
||||
}
|
||||
|
||||
/// Open a sqlite connection to the named DB file (no initialization).
|
||||
///
|
||||
/// Arguments:
|
||||
/// - `db_name`: name of the DB (without extension).
|
||||
pub fn open_connection(db_name: &str) -> Result<Connection, RusqliteError> {
|
||||
fn run_migrations(pool: &r2d2::Pool<SqliteManager>) -> Result<(), StorageError> {
|
||||
let conn = pool.get().map_err(|e| StorageError::Pool(e.to_string()))?;
|
||||
let current_version: i64 = conn
|
||||
.pragma_query_value(None, "user_version", |r| r.get(0))
|
||||
.unwrap_or(0);
|
||||
|
||||
if current_version < 1 {
|
||||
conn.execute_batch(
|
||||
r#"
|
||||
CREATE TABLE IF NOT EXISTS messages (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
storage_owner INTEGER NOT NULL,
|
||||
external_user INTEGER NOT NULL,
|
||||
message_time INTEGER NOT NULL,
|
||||
content TEXT NOT NULL,
|
||||
sent_by_self INTEGER NOT NULL,
|
||||
message_state TEXT NOT NULL,
|
||||
height INTEGER NOT NULL DEFAULT 0
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS idx_messages_lookup
|
||||
ON messages (storage_owner, external_user, message_time DESC);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS contacts (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
storage_owner INTEGER NOT NULL,
|
||||
user_id INTEGER NOT NULL,
|
||||
user_name TEXT,
|
||||
last_message_at INTEGER,
|
||||
UNIQUE(storage_owner, user_id)
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS idx_contacts_owner
|
||||
ON contacts (storage_owner, last_message_at DESC, user_id ASC);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS communities (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
storage_owner INTEGER NOT NULL,
|
||||
address TEXT NOT NULL,
|
||||
title TEXT NOT NULL,
|
||||
position TEXT NOT NULL,
|
||||
UNIQUE(storage_owner, address)
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS idx_communities_owner
|
||||
ON communities (storage_owner);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS users (
|
||||
user_id INTEGER PRIMARY KEY,
|
||||
username TEXT NOT NULL UNIQUE,
|
||||
public_key TEXT NOT NULL,
|
||||
private_key_hash TEXT NOT NULL,
|
||||
reset_token TEXT NOT NULL,
|
||||
created_at INTEGER NOT NULL,
|
||||
display_name TEXT
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS trusted_apps (
|
||||
user_id INTEGER NOT NULL,
|
||||
app_id TEXT NOT NULL,
|
||||
app_secret TEXT NOT NULL,
|
||||
PRIMARY KEY (user_id, app_id)
|
||||
);
|
||||
|
||||
PRAGMA user_version = 1;
|
||||
"#,
|
||||
)?;
|
||||
}
|
||||
|
||||
if current_version < 2 {
|
||||
conn.execute_batch(
|
||||
r#"
|
||||
ALTER TABLE messages ADD COLUMN height INTEGER NOT NULL DEFAULT 0;
|
||||
PRAGMA user_version = 2;
|
||||
"#,
|
||||
)?;
|
||||
}
|
||||
|
||||
if current_version < 3 {
|
||||
conn.execute_batch(
|
||||
r#"
|
||||
ALTER TABLE messages ADD COLUMN reply_to INTEGER;
|
||||
PRAGMA user_version = 3;
|
||||
"#,
|
||||
)?;
|
||||
}
|
||||
|
||||
if current_version < 4 {
|
||||
conn.execute_batch(
|
||||
r#"
|
||||
ALTER TABLE messages ADD COLUMN edited_count INTEGER NOT NULL DEFAULT 0;
|
||||
ALTER TABLE messages ADD COLUMN deleted_by_external INTEGER NOT NULL DEFAULT 0;
|
||||
|
||||
CREATE TABLE IF NOT EXISTS message_edits (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
message_id INTEGER NOT NULL REFERENCES messages(id),
|
||||
content_before TEXT NOT NULL,
|
||||
content_after TEXT NOT NULL,
|
||||
edited_at INTEGER NOT NULL,
|
||||
edited_by INTEGER NOT NULL
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS idx_message_edits_msg
|
||||
ON message_edits (message_id, edited_at DESC);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS reactions (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
message_id INTEGER NOT NULL REFERENCES messages(id),
|
||||
user_id INTEGER NOT NULL,
|
||||
reaction TEXT NOT NULL,
|
||||
created_at INTEGER NOT NULL,
|
||||
UNIQUE(message_id, user_id, reaction)
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS idx_reactions_msg
|
||||
ON reactions (message_id, reaction);
|
||||
|
||||
PRAGMA user_version = 4;
|
||||
"#,
|
||||
)?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn open_connection(db_name: &str) -> Result<Connection, rusqlite::Error> {
|
||||
let path = db_file_path(db_name);
|
||||
Connection::open(path)
|
||||
}
|
||||
|
||||
/// Open a connection and immediately run `init_sql` via `execute_batch`.
|
||||
///
|
||||
/// Arguments:
|
||||
/// - `db_name`: name of the DB (without extension).
|
||||
/// - `init_sql`: SQL statements to initialize schema & PRAGMAs (can be multiple).
|
||||
pub fn open_and_init(db_name: &str, init_sql: &str) -> Result<Connection, RusqliteError> {
|
||||
let conn = open_connection(db_name)?;
|
||||
conn.execute_batch(init_sql)?;
|
||||
Ok(conn)
|
||||
}
|
||||
|
||||
/// Create a shared, Arc<Mutex<Connection>> initialized with the given SQL.
|
||||
///
|
||||
/// This is a convenience wrapper that returns an owned Arc<Mutex<Connection>>
|
||||
/// so caller modules can store it in a `static` or pass it around.
|
||||
///
|
||||
/// Arguments:
|
||||
/// - `db_name`: DB name (without extension).
|
||||
/// - `init_sql`: init SQL (eg PRAGMA + CREATE TABLE statements).
|
||||
pub fn create_shared_connection(
|
||||
db_name: &str,
|
||||
init_sql: &str,
|
||||
) -> Result<Arc<Mutex<Connection>>, String> {
|
||||
match open_and_init(db_name, init_sql) {
|
||||
Ok(conn) => {
|
||||
// Configure some sensible defaults for concurrency
|
||||
// Attempt to set a busy timeout to reduce SQLITE_BUSY failures.
|
||||
let _ = conn.busy_timeout(Duration::from_millis(250));
|
||||
Ok(Arc::new(Mutex::new(conn)))
|
||||
}
|
||||
Err(e) => Err(format!("Failed to open/init DB '{}': {}", db_name, e)),
|
||||
}
|
||||
) -> Result<Arc<std::sync::Mutex<Connection>>, String> {
|
||||
let path = db_file_path(db_name);
|
||||
let conn = Connection::open(path).map_err(|e| e.to_string())?;
|
||||
conn.execute_batch(init_sql).map_err(|e| e.to_string())?;
|
||||
let _ = conn.busy_timeout(Duration::from_millis(250));
|
||||
Ok(Arc::new(std::sync::Mutex::new(conn)))
|
||||
}
|
||||
|
||||
/// Acquire the Connection from an Arc<Mutex<Connection>> and run the provided
|
||||
/// closure. Converts rusqlite::Error into a String on error.
|
||||
///
|
||||
/// Arguments:
|
||||
/// - `shared`: Arc<Mutex<Connection>>
|
||||
/// - `f`: closure that receives &Connection and returns Result<T, RusqliteError>
|
||||
///
|
||||
/// Returns Ok(T) or Err(String).
|
||||
pub fn with_conn<T, F>(shared: &Arc<Mutex<Connection>>, f: F) -> Result<T, String>
|
||||
pub fn with_conn<T, F>(shared: &Arc<std::sync::Mutex<Connection>>, f: F) -> Result<T, String>
|
||||
where
|
||||
F: FnOnce(&Connection) -> Result<T, RusqliteError>,
|
||||
F: FnOnce(&Connection) -> Result<T, rusqlite::Error>,
|
||||
{
|
||||
// When invoked from within an async runtime (such as Tokio), taking a blocking
|
||||
// std::sync::Mutex lock on the runtime thread can cause deadlocks or permanent
|
||||
// awaits. Detect whether we're running inside a Tokio runtime and, if so,
|
||||
// execute the blocking lock + database closure using Tokio's blocking helper.
|
||||
//
|
||||
// The blocking section returns Result<T, String> so we can propagate errors
|
||||
// in the same form as before.
|
||||
if tokio::runtime::Handle::try_current().is_ok() {
|
||||
tokio::task::block_in_place(|| {
|
||||
let guard = shared
|
||||
|
|
@ -98,97 +217,7 @@ where
|
|||
}
|
||||
}
|
||||
|
||||
/// Initialize a general-purpose messages+contacts DB and return a shared
|
||||
/// connection. This helper creates a single DB file that can contain multiple
|
||||
/// tables (messages, contacts, ...). The SQL here is conservative and intended
|
||||
/// to be safe if called multiple times.
|
||||
///
|
||||
/// Callers may prefer to call `create_shared_connection("messages", INIT_SQL)`
|
||||
/// directly, but this convenience is useful for code that expects both tables.
|
||||
pub fn create_general_messages_db() -> Result<Arc<Mutex<Connection>>, String> {
|
||||
// Keep PRAGMA and schema in one multi-statement string so callers only
|
||||
// need to call a single execute_batch.
|
||||
const INIT_SQL: &str = r#"
|
||||
PRAGMA journal_mode = WAL;
|
||||
PRAGMA synchronous = NORMAL;
|
||||
|
||||
CREATE TABLE IF NOT EXISTS messages (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
storage_owner INTEGER NOT NULL,
|
||||
external_user INTEGER NOT NULL,
|
||||
message_time INTEGER NOT NULL,
|
||||
content TEXT NOT NULL,
|
||||
sent_by_self INTEGER NOT NULL,
|
||||
message_state TEXT NOT NULL,
|
||||
height INTEGER NOT NULL DEFAULT 0
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_messages_lookup
|
||||
ON messages (storage_owner, external_user, message_time DESC);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS contacts (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
storage_owner INTEGER NOT NULL,
|
||||
user_id INTEGER NOT NULL,
|
||||
user_name TEXT,
|
||||
last_message_at INTEGER,
|
||||
UNIQUE(storage_owner, user_id)
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_contacts_owner
|
||||
ON contacts (storage_owner, last_message_at DESC, user_id ASC);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS communities (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
storage_owner INTEGER NOT NULL,
|
||||
address TEXT NOT NULL,
|
||||
title TEXT NOT NULL,
|
||||
position TEXT NOT NULL,
|
||||
UNIQUE(storage_owner, address)
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_communities_owner
|
||||
ON communities (storage_owner);
|
||||
"#;
|
||||
|
||||
match create_shared_connection("messages", INIT_SQL) {
|
||||
Ok(shared_conn) => {
|
||||
// Attempt to add the height column for backwards compatibility.
|
||||
// This will fail if the column already exists, which is expected.
|
||||
let _ = with_conn(&shared_conn, |conn| {
|
||||
let _ = conn.execute(
|
||||
"ALTER TABLE messages ADD COLUMN height INTEGER NOT NULL DEFAULT 0",
|
||||
[],
|
||||
);
|
||||
Ok(())
|
||||
});
|
||||
// Attempt to add the reply_to column for backwards compatibility.
|
||||
// This will fail if the column already exists, which is expected.
|
||||
let _ = with_conn(&shared_conn, |conn| {
|
||||
let _ = conn.execute(
|
||||
"ALTER TABLE messages ADD COLUMN reply_to INTEGER",
|
||||
[],
|
||||
);
|
||||
Ok(())
|
||||
});
|
||||
Ok(shared_conn)
|
||||
}
|
||||
Err(e) => Err(e),
|
||||
}
|
||||
/// Legacy - kept for e2ee_storage which uses its own DB.
|
||||
pub fn create_general_messages_db() -> Result<Arc<std::sync::Mutex<Connection>>, String> {
|
||||
create_shared_connection(DB_NAME, "")
|
||||
}
|
||||
|
||||
/*
|
||||
Example usage:
|
||||
|
||||
// In some util module (at init time, e.g. lazy_static or LazyLock)
|
||||
static MESSAGES_DB: LazyLock<Arc<Mutex<Connection>>> = LazyLock::new(|| {
|
||||
create_general_messages_db().expect("failed to create messages DB")
|
||||
});
|
||||
|
||||
// Later, to run a query:
|
||||
let res: Result<Vec<MyRow>, String> = with_conn(&MESSAGES_DB, |conn| {
|
||||
let mut stmt = conn.prepare("SELECT ...")?;
|
||||
let rows = stmt.query_map(...)?;
|
||||
// collect and return Ok(...)
|
||||
});
|
||||
*/
|
||||
|
|
|
|||
Loading…
Reference in a new issue