/* * Copyright (c) 2026, Tobias Müller git@tsmr.eu * */ use crate::error::{Result, TwonlyError}; use sqlx::sqlite::{SqliteConnectOptions, SqlitePoolOptions, UpdateHookResult}; use sqlx::{AssertSqlSafe, Column, ConnectOptions, Row, SqlitePool, TypeInfo, ValueRef}; use std::collections::BTreeSet; use std::str::FromStr; use std::sync::{Arc, Mutex}; use std::time::Duration; use tokio::sync::broadcast; mod legacy_import; pub mod tables; pub const APP_DATABASE_FILE: &str = "app_db.sqlite"; pub const APP_SCHEMA_VERSION: i64 = 6; /// Tables imported from the legacy Drift database. Every entry must exist in /// Drift schema 25, because a missing table aborts the whole import. Rust-only /// tables such as `notification_outbox` are deliberately absent: they have no /// legacy counterpart, and importing stale rows would replay old notifications. pub const APPLICATION_TABLES: &[&str] = &[ "contacts", "groups", "media_files", "messages", "message_histories", "reactions", "group_members", "receipts", "received_receipts", "message_actions", "group_histories", "key_verifications", "verification_tokens", "user_discovery_announced_users", "user_discovery_user_relations", "user_discovery_other_promotions", "user_discovery_own_promotions", "user_discovery_shares", "shortcuts", "shortcut_members", "labels", "contact_labels", ]; #[derive(Clone, Debug, PartialEq, Eq)] pub struct DatabaseChange { pub tables: BTreeSet, } pub struct AppDatabase { pub pool: SqlitePool, changes: broadcast::Sender, } impl AppDatabase { pub async fn new(db_path: &str, encryption_key: Option<&str>, read_only: bool) -> Result { let db_url = format!("sqlite://{db_path}"); let mut options = SqliteConnectOptions::from_str(&format!("{db_url}?mode=rwc"))? .create_if_missing(!read_only) .journal_mode(sqlx::sqlite::SqliteJournalMode::Delete) .foreign_keys(true) .read_only(read_only) .busy_timeout(Duration::from_secs(30)) .pragma("synchronous", "FULL") .log_statements(tracing::log::LevelFilter::Off) .log_slow_statements(tracing::log::LevelFilter::Warn, Duration::from_millis(500)); if let Some(key) = encryption_key { // Migrates a database still encrypted with the old passphrase-derived // key before the pool opens it. See `database::cipher`. let key_pragma = crate::database::cipher::key_pragma(db_path, key, read_only).await?; options = options.pragma("key", key_pragma); } let (changes, _) = broadcast::channel(256); // SQLite itself reports which tables a statement touched, so Drift's // query streams stay correct without any write site having to announce // what it changed. Rows are collected as they are written and only // published once the transaction commits, so a rollback never reaches // the UI. // // The one write this cannot see is `DELETE FROM ` with no // `WHERE`: SQLite's truncate optimization drops the rows without // invoking the update hook. Give such a delete a `WHERE 1` so it takes // the ordinary path. let pending: Arc>> = Arc::default(); let hook_pending = pending.clone(); let hook_changes = changes.clone(); let pool = SqlitePoolOptions::new() // The compatibility executor uses statement-based transactions. // Keeping one connection guarantees BEGIN, all statements, and // COMMIT are executed on that same native connection. .max_connections(1) .acquire_timeout(Duration::from_secs(30)) .after_connect(move |connection, _meta| { let pending = hook_pending.clone(); let changes = hook_changes.clone(); Box::pin(async move { let mut handle = connection.lock_handle().await?; let updated = pending.clone(); handle.set_update_hook(move |result: UpdateHookResult| { if let Ok(mut tables) = updated.lock() { tables.insert(result.table.to_owned()); } }); let committed = pending.clone(); handle.set_commit_hook(move || { let tables = committed .lock() .map(|mut tables| std::mem::take(&mut *tables)) .unwrap_or_default(); if !tables.is_empty() { let _ = changes.send(DatabaseChange { tables }); } // Never veto the commit. true }); handle.set_rollback_hook(move || { if let Ok(mut tables) = pending.lock() { tables.clear(); } }); Ok(()) }) }) .connect_with(options) .await?; Ok(Self { pool, changes }) } pub async fn run_migrations(&self) -> Result<()> { sqlx::migrate!("./src/database/app/migrations") .run(&self.pool) .await .map_err(|error| { TwonlyError::Generic(format!("App database migration failed: {error}")) })?; sqlx::query!( r#" INSERT INTO app_metadata(key, value) VALUES('schema_version', ?) ON CONFLICT(key) DO UPDATE SET value = excluded.value "#, APP_SCHEMA_VERSION.to_string() ) .execute(&self.pool) .await?; Ok(()) } pub fn subscribe(&self) -> broadcast::Receiver { self.changes.subscribe() } pub async fn raw_select(&self, statement: String, arguments: Vec) -> Result { let rows = bind_arguments(sqlx::query(AssertSqlSafe(statement)), arguments)? .fetch_all(&self.pool) .await?; let columns = rows .first() .map(|row| { row.columns() .iter() .map(|column| column.name().to_owned()) .collect() }) .unwrap_or_default(); let mut output_rows = Vec::with_capacity(rows.len()); for row in rows { let mut values = Vec::with_capacity(row.len()); for index in 0..row.len() { let raw = row.try_get_raw(index)?; let value = if raw.is_null() { SqlValue::null() } else { match raw.type_info().name() { "INTEGER" | "INT" | "BOOLEAN" | "DATETIME" => { SqlValue::integer(row.try_get(index)?) } "REAL" | "FLOAT" | "DOUBLE" => SqlValue::real(row.try_get(index)?), "BLOB" => SqlValue::blob(row.try_get(index)?), _ => SqlValue::text(row.try_get(index)?), } }; values.push(value); } output_rows.push(SqlRow { values }); } Ok(SqlRows { columns, rows: output_rows, }) } pub async fn raw_execute( &self, statement: String, arguments: Vec, ) -> Result { let result = bind_arguments(sqlx::query(AssertSqlSafe(statement)), arguments)? .execute(&self.pool) .await?; Ok(SqlExecutionResult { affected_rows: result.rows_affected() as i64, last_insert_row_id: result.last_insert_rowid(), }) } pub async fn create_backup(&self, output_path: &str, encryption_key: &str) -> Result<()> { let mut connection = self.pool.acquire().await?; // SQLCipher-only syntax cannot be described by the SQLite database // used by SQLx during compile-time query preparation. sqlx::query(r#"ATTACH DATABASE ? AS backup KEY ?"#) .bind(output_path) .bind(encryption_key) .execute(&mut *connection) .await .map_err(|error| TwonlyError::Generic(format!("Attach app backup failed: {error}")))?; let export = sqlx::query(r#"SELECT sqlcipher_export('backup')"#) .execute(&mut *connection) .await; let detach = sqlx::query(r#"DETACH DATABASE backup"#) .execute(&mut *connection) .await; export.map_err(|error| TwonlyError::Generic(format!("App export failed: {error}")))?; detach .map_err(|error| TwonlyError::Generic(format!("Detach app backup failed: {error}")))?; Ok(()) } } #[derive(Clone, Debug, PartialEq)] pub struct SqlValue { /// 0 = null, 1 = integer, 2 = real, 3 = text, 4 = blob. pub kind: u8, pub integer_value: Option, pub real_value: Option, pub text_value: Option, pub blob_value: Option>, } impl SqlValue { fn null() -> Self { Self { kind: 0, integer_value: None, real_value: None, text_value: None, blob_value: None, } } fn integer(value: i64) -> Self { Self { kind: 1, integer_value: Some(value), real_value: None, text_value: None, blob_value: None, } } fn real(value: f64) -> Self { Self { kind: 2, integer_value: None, real_value: Some(value), text_value: None, blob_value: None, } } fn text(value: String) -> Self { Self { kind: 3, integer_value: None, real_value: None, text_value: Some(value), blob_value: None, } } fn blob(value: Vec) -> Self { Self { kind: 4, integer_value: None, real_value: None, text_value: None, blob_value: Some(value), } } } #[derive(Clone, Debug, PartialEq)] pub struct SqlRows { pub columns: Vec, pub rows: Vec, } #[derive(Clone, Debug, PartialEq)] pub struct SqlRow { pub values: Vec, } #[derive(Clone, Debug, PartialEq, Eq)] pub struct SqlExecutionResult { pub affected_rows: i64, pub last_insert_row_id: i64, } fn bind_arguments<'q>( mut query: sqlx::query::Query<'q, sqlx::Sqlite, sqlx::sqlite::SqliteArguments>, arguments: Vec, ) -> Result> { for argument in arguments { query = match argument { SqlValue { kind: 0, .. } => query.bind(Option::::None), SqlValue { kind: 1, integer_value: Some(value), .. } => query.bind(value), SqlValue { kind: 2, real_value: Some(value), .. } => query.bind(value), SqlValue { kind: 3, text_value: Some(value), .. } => query.bind(value), SqlValue { kind: 4, blob_value: Some(value), .. } => query.bind(value), invalid => { return Err(TwonlyError::Generic(format!( "Invalid SQL value received from bridge: {invalid:?}" ))) } }; } Ok(query) } #[derive(Clone, Debug, PartialEq, Eq)] pub struct MigrationReport { pub legacy_version: i64, pub tables: Vec, } #[derive(Clone, Debug, PartialEq, Eq)] pub struct TableMigrationCount { pub table: String, pub rows: i64, } #[cfg(test)] mod change_notification_tests { use super::*; use tempfile::tempdir; use tokio::sync::broadcast::error::TryRecvError; async fn open() -> (tempfile::TempDir, AppDatabase) { let directory = tempdir().unwrap(); let path = directory.path().join("app_db.sqlite"); let database = AppDatabase::new(path.to_str().unwrap(), None, false) .await .unwrap(); database.run_migrations().await.unwrap(); (directory, database) } fn insert(key: &str) -> String { format!("INSERT INTO app_metadata(key, value) VALUES('{key}', '1')") } #[tokio::test] async fn committed_writes_report_their_tables() { let (_dir, database) = open().await; let mut changes = database.subscribe(); database .raw_execute(insert("hook"), Vec::new()) .await .unwrap(); let change = changes.try_recv().unwrap(); assert!(change.tables.contains("app_metadata")); } #[tokio::test] async fn rolled_back_writes_report_nothing() { let (_dir, database) = open().await; let mut changes = database.subscribe(); database .raw_execute("BEGIN".to_owned(), Vec::new()) .await .unwrap(); database .raw_execute(insert("discarded"), Vec::new()) .await .unwrap(); database .raw_execute("ROLLBACK".to_owned(), Vec::new()) .await .unwrap(); assert_eq!(changes.try_recv().unwrap_err(), TryRecvError::Empty); } #[tokio::test] async fn a_transaction_reports_every_table_once_on_commit() { let (_dir, database) = open().await; let mut changes = database.subscribe(); database .raw_execute("BEGIN".to_owned(), Vec::new()) .await .unwrap(); database .raw_execute(insert("first"), Vec::new()) .await .unwrap(); database .raw_execute(insert("second"), Vec::new()) .await .unwrap(); // Nothing may reach the UI before the transaction commits. assert_eq!(changes.try_recv().unwrap_err(), TryRecvError::Empty); database .raw_execute("COMMIT".to_owned(), Vec::new()) .await .unwrap(); let change = changes.try_recv().unwrap(); assert_eq!(change.tables, BTreeSet::from(["app_metadata".to_owned()])); assert_eq!(changes.try_recv().unwrap_err(), TryRecvError::Empty); } }