diff --git a/Cargo.lock b/Cargo.lock index afea57811..b48371bc9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1855,6 +1855,7 @@ dependencies = [ "redis", "regex", "reqwest 0.12.28", + "rusqlite", "russh", "rust_decimal", "rustls 0.23.40", @@ -2183,7 +2184,7 @@ dependencies = [ "comfy-table", "fallible-iterator", "fallible-streaming-iterator", - "hashlink", + "hashlink 0.10.0", "libduckdb-sys", "num-integer", "rust_decimal", @@ -3147,6 +3148,9 @@ name = "hashbrown" version = "0.14.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e5274423e17b7c9fc20b6e7e208532f9b19825d82dfd615708b70edd83df41f1" +dependencies = [ + "ahash 0.8.12", +] [[package]] name = "hashbrown" @@ -3176,6 +3180,15 @@ version = "0.17.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4f467dd6dccf739c208452f8014c75c18bb8301b050ad1cfb27153803edb0f51" +[[package]] +name = "hashlink" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ba4ff7128dee98c7dc9794b6a411377e1404dba1c97deb8d1a55297bd25d8af" +dependencies = [ + "hashbrown 0.14.5", +] + [[package]] name = "hashlink" version = "0.10.0" @@ -6228,6 +6241,20 @@ dependencies = [ "zeroize", ] +[[package]] +name = "rusqlite" +version = "0.32.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7753b721174eb8ff87a9a0e799e2d7bc3749323e773db92e0984debb00019d6e" +dependencies = [ + "bitflags 2.11.1", + "fallible-iterator", + "fallible-streaming-iterator", + "hashlink 0.9.1", + "libsqlite3-sys", + "smallvec", +] + [[package]] name = "russh" version = "0.60.2" @@ -7291,7 +7318,7 @@ dependencies = [ "futures-io", "futures-util", "hashbrown 0.15.5", - "hashlink", + "hashlink 0.10.0", "indexmap 2.14.0", "log", "memchr", diff --git a/crates/dbx-core/Cargo.toml b/crates/dbx-core/Cargo.toml index 43e1373e3..db57582d3 100644 --- a/crates/dbx-core/Cargo.toml +++ b/crates/dbx-core/Cargo.toml @@ -19,7 +19,8 @@ chrono = { version = "0.4", features = ["serde"] } rust_decimal = { version = "1", features = ["serde"] } anyhow = "1" uuid = { version = "1", features = ["v4", "serde"] } -sqlx = { version = "0.8", features = ["runtime-tokio", "tls-rustls", "postgres", "sqlite", "json", "chrono", "uuid", "rust_decimal"] } +sqlx = { version = "0.8", features = ["runtime-tokio", "tls-rustls", "postgres", "json", "chrono", "uuid", "rust_decimal"] } +rusqlite = { version = "0.32", features = ["bundled"] } mysql_async = { version = "0.36", default-features = false, features = ["default-rustls", "client_ed25519", "chrono", "rust_decimal"] } sqlparser = "0.62.0" redis = { version = "0.32.2", features = ["tokio-comp", "tls-rustls", "tokio-rustls-comp"] } diff --git a/crates/dbx-core/src/connection.rs b/crates/dbx-core/src/connection.rs index 0b61aa0b5..64e74128f 100644 --- a/crates/dbx-core/src/connection.rs +++ b/crates/dbx-core/src/connection.rs @@ -42,7 +42,7 @@ pub enum MysqlMode { pub enum PoolKind { Mysql(db::mysql::MySqlPool, MysqlMode), Postgres(sqlx::postgres::PgPool), - Sqlite(sqlx::sqlite::SqlitePool), + Sqlite(db::sqlite::SqliteHandle), Redis(tokio::sync::Mutex), DuckDb(Arc>), MongoDb(mongodb::Client), @@ -1098,7 +1098,7 @@ mod tests { #[tokio::test] async fn remove_connection_pools_clears_base_and_database_scoped_pools() { let (state, dir) = test_app_state().await; - let pool = sqlx::sqlite::SqlitePoolOptions::new().max_connections(1).connect(":memory:").await.unwrap(); + let pool = crate::db::sqlite::connect_path(":memory:").await.unwrap(); { let mut conns = state.connections.write().await; diff --git a/crates/dbx-core/src/db/sqlite.rs b/crates/dbx-core/src/db/sqlite.rs index 5e9de23da..f9ec70554 100644 --- a/crates/dbx-core/src/db/sqlite.rs +++ b/crates/dbx-core/src/db/sqlite.rs @@ -1,36 +1,82 @@ -use futures::StreamExt; -use sqlx::sqlite::{SqliteConnectOptions, SqlitePool, SqlitePoolOptions, SqliteRow}; -use sqlx::{Column, Executor, Row}; -use std::str::FromStr; -use std::time::{Duration, Instant}; +use base64::prelude::{Engine as _, BASE64_STANDARD}; +use rusqlite::types::ValueRef; +use rusqlite::{Connection, OpenFlags}; +use std::path::Path; +use std::sync::{Arc, Mutex}; +use std::time::Instant; use super::file_validator::validate_file_path; use crate::sql::starts_with_executable_sql_keyword; use crate::types::{ColumnInfo, DatabaseInfo, ForeignKeyInfo, IndexInfo, QueryResult, TableInfo, TriggerInfo}; -pub async fn connect_path(path: &str) -> Result { +#[derive(Clone)] +pub struct SqliteHandle { + conn: Arc>, +} + +impl SqliteHandle { + pub fn with_connection(&self, f: F) -> Result + where + F: FnOnce(&mut Connection) -> Result, + { + let mut conn = self.conn.lock().map_err(|e| e.to_string())?; + f(&mut conn) + } +} + +pub async fn connect_path(path: &str) -> Result { + connect_path_with_options(path, false).await +} + +pub async fn connect_path_create_if_missing(path: &str) -> Result { + connect_path_with_options(path, true).await +} + +async fn connect_path_with_options(path: &str, create_if_missing: bool) -> Result { + let path = path.to_string(); + tokio::task::spawn_blocking(move || open_sqlite_handle(&path, create_if_missing)) + .await + .map_err(|e| e.to_string())? +} + +fn open_sqlite_handle(path: &str, create_if_missing: bool) -> Result { let is_memory = is_memory_database_path(path); - if !is_memory { + if !is_memory && !create_if_missing { validate_file_path(path, is_network_path)?; } - let mut options = if is_memory { - SqliteConnectOptions::from_str("sqlite::memory:").map_err(|e| format!("SQLite connection failed: {e}"))? - } else { - SqliteConnectOptions::new().filename(path).create_if_missing(false) - }; - - if is_network_path(path) { - options = options.vfs("unix-nolock"); + if !is_memory && create_if_missing { + ensure_parent_dir(path)?; } - SqlitePoolOptions::new() - .max_connections(if is_memory { 1 } else { 5 }) - .acquire_timeout(Duration::from_secs(10)) - .idle_timeout(Duration::from_secs(300)) - .connect_with(options) - .await - .map_err(|e| format!("SQLite connection failed: {e}")) + let conn = if is_memory { + Connection::open_in_memory().map_err(|e| format!("SQLite connection failed: {e}"))? + } else { + let mut flags = OpenFlags::SQLITE_OPEN_READ_WRITE; + if create_if_missing { + flags |= OpenFlags::SQLITE_OPEN_CREATE; + } + if is_network_path(path) { + flags |= OpenFlags::SQLITE_OPEN_URI; + Connection::open_with_flags(format!("file:{}?vfs=unix-nolock", path), flags) + .map_err(|e| format!("SQLite connection failed: {e}"))? + } else { + Connection::open_with_flags(path, flags).map_err(|e| format!("SQLite connection failed: {e}"))? + } + }; + + conn.busy_timeout(std::time::Duration::from_secs(10)).map_err(|e| e.to_string())?; + + Ok(SqliteHandle { conn: Arc::new(Mutex::new(conn)) }) +} + +fn ensure_parent_dir(path: &str) -> Result<(), String> { + if let Some(parent) = Path::new(path).parent() { + if !parent.as_os_str().is_empty() { + std::fs::create_dir_all(parent).map_err(|e| e.to_string())?; + } + } + Ok(()) } fn is_network_path(path: &str) -> bool { @@ -107,8 +153,6 @@ mod tests { async fn view_with_if_function_works_after_normalization() { let pool = connect_path(":memory:").await.expect("connect in-memory SQLite"); - // Create a view that uses if() — this succeeds because CREATE VIEW - // just stores the SQL text without evaluating it execute_query(&pool, "CREATE TABLE t (x INTEGER); INSERT INTO t VALUES (1), (2), (3);") .await .expect("create and populate table"); @@ -117,7 +161,6 @@ mod tests { .await .expect("create view"); - // Query the view — this must go through normalize_sqlite_sql let result = execute_query(&pool, "SELECT * FROM v ORDER BY x").await.expect("query view"); assert_eq!(result.rows.len(), 3); @@ -129,7 +172,6 @@ mod tests { async fn if_rewrite_works_in_direct_query() { let pool = connect_path(":memory:").await.expect("connect in-memory SQLite"); - // if() is not a built-in SQLite function — the normalizer must rewrite it to IIF() let result = execute_query(&pool, "SELECT if(1 = 1, 'yes', 'no') AS answer") .await .expect("if() should be rewritten to IIF()"); @@ -165,138 +207,176 @@ mod tests { } } -pub async fn list_databases(_pool: &SqlitePool) -> Result, String> { +pub async fn list_databases(_pool: &SqliteHandle) -> Result, String> { Ok(vec![DatabaseInfo { name: "main".to_string() }]) } -pub async fn list_tables(pool: &SqlitePool, _schema: &str) -> Result, String> { - let rows: Vec = sqlx::query( - "SELECT name, type FROM sqlite_master WHERE type IN ('table', 'view') AND name NOT LIKE 'sqlite_%' ORDER BY name", - ) - .fetch_all(pool) +pub async fn list_tables(pool: &SqliteHandle, _schema: &str) -> Result, String> { + let pool = pool.clone(); + tokio::task::spawn_blocking(move || { + pool.with_connection(|conn| { + let mut stmt = conn + .prepare( + "SELECT name, type FROM sqlite_master \ + WHERE type IN ('table', 'view') AND name NOT LIKE 'sqlite_%' ORDER BY name", + ) + .map_err(|e| e.to_string())?; + let rows = stmt + .query_map([], |row| { + let table_type: String = row.get(1)?; + Ok(TableInfo { + name: row.get(0)?, + table_type: if table_type == "view" { "VIEW".to_string() } else { "BASE TABLE".to_string() }, + comment: None, + }) + }) + .map_err(|e| e.to_string())?; + rows.collect::, _>>().map_err(|e| e.to_string()) + }) + }) .await - .map_err(|e| e.to_string())?; + .map_err(|e| e.to_string())? +} - Ok(rows - .iter() - .map(|row| { - let t: String = row.get("type"); - TableInfo { - name: row.get::("name"), - table_type: if t == "view" { "VIEW".to_string() } else { "BASE TABLE".to_string() }, - comment: None, +pub async fn get_columns(pool: &SqliteHandle, _schema: &str, table: &str) -> Result, String> { + let pool = pool.clone(); + let table = table.to_string(); + tokio::task::spawn_blocking(move || { + let sql = format!("PRAGMA table_info(\"{}\")", table.replace('"', "\"\"")); + pool.with_connection(|conn| { + let mut stmt = conn.prepare(&sql).map_err(|e| e.to_string())?; + let rows = stmt + .query_map([], |row| { + Ok(ColumnInfo { + name: row.get("name")?, + data_type: row.get("type")?, + is_nullable: row.get::<_, i32>("notnull")? == 0, + column_default: row.get("dflt_value")?, + is_primary_key: row.get::<_, i32>("pk")? > 0, + extra: None, + comment: None, + numeric_precision: None, + numeric_scale: None, + character_maximum_length: None, + }) + }) + .map_err(|e| e.to_string())?; + rows.collect::, _>>().map_err(|e| e.to_string()) + }) + }) + .await + .map_err(|e| e.to_string())? +} + +pub async fn list_indexes(pool: &SqliteHandle, _schema: &str, table: &str) -> Result, String> { + let pool = pool.clone(); + let table = table.to_string(); + tokio::task::spawn_blocking(move || { + let safe_table = table.replace('"', "\"\""); + pool.with_connection(|conn| { + let mut stmt = conn.prepare(&format!("PRAGMA index_list(\"{safe_table}\")")).map_err(|e| e.to_string())?; + let idx_rows = stmt + .query_map([], |row| { + Ok(( + row.get::<_, String>("name")?, + row.get::<_, i32>("unique")? != 0, + row.get::<_, String>("origin")?, + )) + }) + .map_err(|e| e.to_string())? + .collect::, _>>() + .map_err(|e| e.to_string())?; + + let mut indexes = Vec::new(); + for (name, is_unique, origin) in idx_rows { + let safe_name = name.replace('"', "\"\""); + let mut col_stmt = + conn.prepare(&format!("PRAGMA index_info(\"{safe_name}\")")).map_err(|e| e.to_string())?; + let columns = col_stmt + .query_map([], |row| row.get::<_, String>("name")) + .map_err(|e| e.to_string())? + .collect::, _>>() + .map_err(|e| e.to_string())?; + + indexes.push(IndexInfo { + name, + columns, + is_unique, + is_primary: origin == "pk", + filter: None, + index_type: None, + included_columns: None, + comment: None, + }); } + Ok(indexes) }) - .collect()) + }) + .await + .map_err(|e| e.to_string())? } -pub async fn get_columns(pool: &SqlitePool, _schema: &str, table: &str) -> Result, String> { - let rows: Vec = - sqlx::query(&format!("PRAGMA table_info(\"{}\")", table)).fetch_all(pool).await.map_err(|e| e.to_string())?; - - Ok(rows - .iter() - .map(|row| ColumnInfo { - name: row.get::("name"), - data_type: row.get::("type"), - is_nullable: row.get::("notnull") == 0, - column_default: row.get::, _>("dflt_value"), - is_primary_key: row.get::("pk") > 0, - extra: None, - comment: None, - numeric_precision: None, - numeric_scale: None, - character_maximum_length: None, +pub async fn list_foreign_keys(pool: &SqliteHandle, _schema: &str, table: &str) -> Result, String> { + let pool = pool.clone(); + let table = table.to_string(); + tokio::task::spawn_blocking(move || { + let sql = format!("PRAGMA foreign_key_list(\"{}\")", table.replace('"', "\"\"")); + pool.with_connection(|conn| { + let mut stmt = conn.prepare(&sql).map_err(|e| e.to_string())?; + let rows = stmt + .query_map([], |row| { + Ok(ForeignKeyInfo { + name: format!("fk_{}", row.get::<_, i32>("id")?), + column: row.get("from")?, + ref_table: row.get("table")?, + ref_column: row.get("to")?, + }) + }) + .map_err(|e| e.to_string())?; + rows.collect::, _>>().map_err(|e| e.to_string()) }) - .collect()) + }) + .await + .map_err(|e| e.to_string())? } -pub async fn list_indexes(pool: &SqlitePool, _schema: &str, table: &str) -> Result, String> { - let safe_table = table.replace('"', "\"\""); - let idx_rows: Vec = sqlx::query(&format!("PRAGMA index_list(\"{safe_table}\")")) - .fetch_all(pool) - .await - .map_err(|e| e.to_string())?; - - let mut indexes = Vec::new(); - for idx_row in &idx_rows { - let name: String = idx_row.get("name"); - let is_unique: bool = idx_row.get::("unique") != 0; - let origin: String = idx_row.get::("origin"); - let is_primary = origin == "pk"; - - let safe_name = name.replace('"', "\"\""); - let col_rows: Vec = sqlx::query(&format!("PRAGMA index_info(\"{safe_name}\")")) - .fetch_all(pool) - .await - .map_err(|e| e.to_string())?; - - let columns: Vec = col_rows.iter().map(|r| r.get::("name")).collect(); - - indexes.push(IndexInfo { - name, - columns, - is_unique, - is_primary, - filter: None, - index_type: None, - included_columns: None, - comment: None, - }); - } - Ok(indexes) -} - -pub async fn list_foreign_keys(pool: &SqlitePool, _schema: &str, table: &str) -> Result, String> { - let rows: Vec = sqlx::query(&format!("PRAGMA foreign_key_list(\"{}\")", table)) - .fetch_all(pool) - .await - .map_err(|e| e.to_string())?; - - Ok(rows - .iter() - .map(|row| ForeignKeyInfo { - name: format!("fk_{}", row.get::("id")), - column: row.get::("from"), - ref_table: row.get::("table"), - ref_column: row.get::("to"), +pub async fn list_triggers(pool: &SqliteHandle, _schema: &str, table: &str) -> Result, String> { + let pool = pool.clone(); + let table = table.to_string(); + tokio::task::spawn_blocking(move || { + pool.with_connection(|conn| { + let mut stmt = conn + .prepare("SELECT name, sql FROM sqlite_master WHERE type = 'trigger' AND tbl_name = ? ORDER BY name") + .map_err(|e| e.to_string())?; + let rows = stmt + .query_map([table], |row| { + let sql_text: Option = row.get("sql")?; + let upper = sql_text.unwrap_or_default().to_uppercase(); + let timing = if upper.contains("BEFORE") { + "BEFORE" + } else if upper.contains("AFTER") { + "AFTER" + } else { + "INSTEAD OF" + }; + let event = if upper.contains("INSERT") { + "INSERT" + } else if upper.contains("UPDATE") { + "UPDATE" + } else { + "DELETE" + }; + Ok(TriggerInfo { name: row.get("name")?, event: event.to_string(), timing: timing.to_string() }) + }) + .map_err(|e| e.to_string())?; + rows.collect::, _>>().map_err(|e| e.to_string()) }) - .collect()) + }) + .await + .map_err(|e| e.to_string())? } -pub async fn list_triggers(pool: &SqlitePool, _schema: &str, table: &str) -> Result, String> { - let rows: Vec = - sqlx::query("SELECT name, sql FROM sqlite_master WHERE type = 'trigger' AND tbl_name = ? ORDER BY name") - .bind(table) - .fetch_all(pool) - .await - .map_err(|e| e.to_string())?; - - Ok(rows - .iter() - .map(|row| { - let sql_text: String = row.get::, _>("sql").unwrap_or_default(); - let upper = sql_text.to_uppercase(); - let timing = if upper.contains("BEFORE") { - "BEFORE" - } else if upper.contains("AFTER") { - "AFTER" - } else { - "INSTEAD OF" - }; - let event = if upper.contains("INSERT") { - "INSERT" - } else if upper.contains("UPDATE") { - "UPDATE" - } else { - "DELETE" - }; - TriggerInfo { name: row.get::("name"), event: event.to_string(), timing: timing.to_string() } - }) - .collect()) -} - -pub async fn execute_query(pool: &SqlitePool, sql: &str) -> Result { +pub async fn execute_query(pool: &SqliteHandle, sql: &str) -> Result { execute_query_with_max_rows(pool, sql, None).await } @@ -304,13 +384,8 @@ fn query_result_row_limit(max_rows: Option) -> usize { max_rows.unwrap_or(crate::query::MAX_ROWS).max(1) } -/// Function-name rewrites for SQLite compatibility. -/// Keys are lowercase source names; values are replacement names. -/// Applied only at word boundaries followed by optional whitespace and `(`. const SQLITE_FUNCTION_ALIASES: &[(&str, &str)] = &[("if", "IIF"), ("substring", "substr")]; -/// Rewrites known non-SQLite function names (e.g. `if()` → `IIF()`, `substring()` → `substr()`) -/// before sending SQL to SQLite. Avoids modifying string literals, comments, and identifiers. fn normalize_sqlite_sql(sql: &str) -> String { let mut result = String::with_capacity(sql.len()); let chars: Vec = sql.chars().collect(); @@ -397,71 +472,76 @@ fn normalize_sqlite_sql(sql: &str) -> String { } pub async fn execute_query_with_max_rows( - pool: &SqlitePool, + pool: &SqliteHandle, sql: &str, max_rows: Option, ) -> Result { + let pool = pool.clone(); + let sql = normalize_sqlite_sql(sql); + tokio::task::spawn_blocking(move || execute_query_blocking(&pool, &sql, max_rows)) + .await + .map_err(|e| e.to_string())? +} + +fn execute_query_blocking(pool: &SqliteHandle, sql: &str, max_rows: Option) -> Result { let start = Instant::now(); let row_limit = query_result_row_limit(max_rows); - let sql = normalize_sqlite_sql(sql); - if starts_with_executable_sql_keyword(&sql, &["SELECT", "PRAGMA", "EXPLAIN", "WITH"]) { - let desc = pool.describe(&sql).await.map_err(|e| e.to_string())?; - let columns: Vec = desc.columns().iter().map(|c| c.name().to_string()).collect(); + pool.with_connection(|conn| { + if starts_with_executable_sql_keyword(sql, &["SELECT", "PRAGMA", "EXPLAIN", "WITH"]) { + let mut stmt = conn.prepare(sql).map_err(|e| e.to_string())?; + let columns = stmt.column_names().iter().map(|name| name.to_string()).collect::>(); + let mut rows = stmt.query([]).map_err(|e| e.to_string())?; + let mut result_rows = Vec::new(); - let mut stream = sqlx::query(&sql).fetch(pool); - let mut result_rows: Vec> = Vec::new(); - - while let Some(row) = stream.next().await { - let row = row.map_err(|e| e.to_string())?; - result_rows.push( - (0..row.len()) - .map(|i| { - row.try_get::(i) - .map(serde_json::Value::String) - .or_else(|_| row.try_get::(i).map(super::safe_i64_to_json)) - .or_else(|_| { - row.try_get::(i).map(|v| { - serde_json::Number::from_f64(v) - .map(serde_json::Value::Number) - .unwrap_or(serde_json::Value::Null) - }) - }) - .or_else(|_| row.try_get::(i).map(serde_json::Value::Bool)) - .unwrap_or(serde_json::Value::Null) - }) - .collect(), - ); - if result_rows.len() > row_limit { - break; + while let Some(row) = rows.next().map_err(|e| e.to_string())? { + let mut values = Vec::with_capacity(columns.len()); + for i in 0..columns.len() { + values.push(value_ref_to_json(row.get_ref(i).map_err(|e| e.to_string())?)); + } + result_rows.push(values); + if result_rows.len() > row_limit { + break; + } } + + let truncated = result_rows.len() > row_limit; + if truncated { + result_rows.truncate(row_limit); + } + + Ok(QueryResult { + columns, + rows: result_rows, + affected_rows: 0, + execution_time_ms: start.elapsed().as_millis(), + truncated, + session_id: None, + has_more: false, + }) + } else { + conn.execute_batch(sql).map_err(|e| e.to_string())?; + Ok(QueryResult { + columns: vec![], + rows: vec![], + affected_rows: conn.changes(), + execution_time_ms: start.elapsed().as_millis(), + truncated: false, + session_id: None, + has_more: false, + }) } + }) +} - let truncated = result_rows.len() > row_limit; - if truncated { - result_rows.truncate(row_limit); +fn value_ref_to_json(value: ValueRef<'_>) -> serde_json::Value { + match value { + ValueRef::Null => serde_json::Value::Null, + ValueRef::Integer(v) => super::safe_i64_to_json(v), + ValueRef::Real(v) => { + serde_json::Number::from_f64(v).map(serde_json::Value::Number).unwrap_or(serde_json::Value::Null) } - - Ok(QueryResult { - columns, - rows: result_rows, - affected_rows: 0, - execution_time_ms: start.elapsed().as_millis(), - truncated, - session_id: None, - has_more: false, - }) - } else { - let result = sqlx::query(&sql).execute(pool).await.map_err(|e| e.to_string())?; - - Ok(QueryResult { - columns: vec![], - rows: vec![], - affected_rows: result.rows_affected(), - execution_time_ms: start.elapsed().as_millis(), - truncated: false, - session_id: None, - has_more: false, - }) + ValueRef::Text(v) => serde_json::Value::String(String::from_utf8_lossy(v).to_string()), + ValueRef::Blob(v) => serde_json::Value::String(BASE64_STANDARD.encode(v)), } } diff --git a/crates/dbx-core/src/query.rs b/crates/dbx-core/src/query.rs index 3352d37f3..4ccdd0e05 100644 --- a/crates/dbx-core/src/query.rs +++ b/crates/dbx-core/src/query.rs @@ -866,13 +866,12 @@ pub async fn execute_statements( } /// Execute multiple SQL statements within a single transaction. -/// For sqlx-based pools (Postgres/MySQL/SQLite), uses the Transaction API to -/// guarantee all statements run on the same physical connection. -/// For custom drivers (ClickHouse/SqlServer/Agent), uses explicit -/// BEGIN/COMMIT/ROLLBACK on the already-single-connection client. +/// For pooled drivers (Postgres/MySQL), uses the driver transaction API. +/// For SQLite and already-single-connection drivers (ClickHouse/SqlServer/Agent), +/// uses explicit BEGIN/COMMIT/ROLLBACK on the shared connection. /// For databases that don't support explicit transactions (Redis, MongoDB, Oracle), /// executes statements sequentially without transaction. -/// If BEGIN fails, returns an error — no silent fallback to auto-commit. +/// If BEGIN fails, returns an error instead of silently falling back to auto-commit. pub async fn execute_statements_in_transaction( state: &AppState, connection_id: &str, @@ -919,7 +918,7 @@ pub async fn execute_statements_in_transaction( enum TxPath { Pg(sqlx::postgres::PgPool), Mysql(mysql_async::Pool, bool), - Sqlite(sqlx::sqlite::SqlitePool), + Sqlite(db::sqlite::SqliteHandle), Explicit, None, } @@ -993,32 +992,38 @@ async fn exec_tx_mysql_inner( } async fn exec_tx_sqlite_inner( - pool: sqlx::sqlite::SqlitePool, + pool: db::sqlite::SqliteHandle, statements: &[String], start: std::time::Instant, ) -> Result { - let mut conn = pool.acquire().await.map_err(|e| format!("Failed to acquire connection: {}", e))?; - sqlx::query("BEGIN").execute(&mut *conn).await.map_err(|e| format!("Failed to begin transaction: {}", e))?; - let mut total_affected: u64 = 0; - for (i, sql) in statements.iter().enumerate() { - match sqlx::query(sql).execute(&mut *conn).await { - Ok(r) => total_affected += r.rows_affected(), - Err(e) => { - let _ = sqlx::query("ROLLBACK").execute(&mut *conn).await; - return Err(format!("Statement {} failed: {}", i + 1, e)); + let statements = statements.to_vec(); + tokio::task::spawn_blocking(move || { + pool.with_connection(|conn| { + conn.execute_batch("BEGIN").map_err(|e| format!("Failed to begin transaction: {}", e))?; + let mut total_affected: u64 = 0; + for (i, sql) in statements.iter().enumerate() { + match conn.execute_batch(sql) { + Ok(_) => total_affected += conn.changes(), + Err(e) => { + let _ = conn.execute_batch("ROLLBACK"); + return Err(format!("Statement {} failed: {}", i + 1, e)); + } + } } - } - } - sqlx::query("COMMIT").execute(&mut *conn).await.map_err(|e| format!("COMMIT failed: {}", e))?; - Ok(db::QueryResult { - columns: vec![], - rows: vec![], - affected_rows: total_affected, - execution_time_ms: start.elapsed().as_millis(), - truncated: false, - session_id: None, - has_more: false, + conn.execute_batch("COMMIT").map_err(|e| format!("COMMIT failed: {}", e))?; + Ok(db::QueryResult { + columns: vec![], + rows: vec![], + affected_rows: total_affected, + execution_time_ms: start.elapsed().as_millis(), + truncated: false, + session_id: None, + has_more: false, + }) + }) }) + .await + .map_err(|e| e.to_string())? } async fn exec_tx_explicit_inner( diff --git a/crates/dbx-core/src/schema.rs b/crates/dbx-core/src/schema.rs index 8473437ad..592aed807 100644 --- a/crates/dbx-core/src/schema.rs +++ b/crates/dbx-core/src/schema.rs @@ -1103,14 +1103,19 @@ pub async fn mysql_ddl(pool: &db::mysql::MySqlPool, table: &str) -> Result Result { - use sqlx::Row; - let row: sqlx::sqlite::SqliteRow = sqlx::query("SELECT sql FROM sqlite_master WHERE type='table' AND name=?") - .bind(table) - .fetch_one(pool) - .await - .map_err(|e| e.to_string())?; - row.try_get::(0).map_err(|e| e.to_string()) +pub async fn sqlite_ddl(pool: &db::sqlite::SqliteHandle, table: &str) -> Result { + let pool = pool.clone(); + let table = table.to_string(); + tokio::task::spawn_blocking(move || { + pool.with_connection(|conn| { + conn.query_row("SELECT sql FROM sqlite_master WHERE type='table' AND name=?1", [table], |row| { + row.get::<_, String>(0) + }) + .map_err(|e| e.to_string()) + }) + }) + .await + .map_err(|e| e.to_string())? } pub async fn pg_ddl(pool: &sqlx::postgres::PgPool, schema: &str, table: &str) -> Result { diff --git a/crates/dbx-core/src/sql.rs b/crates/dbx-core/src/sql.rs index df39b40da..6ef303107 100644 --- a/crates/dbx-core/src/sql.rs +++ b/crates/dbx-core/src/sql.rs @@ -220,9 +220,9 @@ impl SqlStatementSplitter { if let Some(new_delim) = parse_delimiter_command(last_line) { self.custom_delimiter = if new_delim == ";" { None } else { Some(new_delim.to_string()) }; if last_line_start > 0 { - let before = &self.buffer[..last_line_start]; - if let Some(statement) = executable_sql_slice(before, self.options) { - statements.push(statement.to_string()); + let before = self.buffer[..last_line_start].trim(); + if has_executable_sql_with_options(before, self.options) { + statements.push(before.to_string()); } } self.buffer.clear(); @@ -252,8 +252,8 @@ impl SqlStatementSplitter { let last_line = trimmed.rsplit('\n').next().unwrap_or(trimmed).trim(); if parse_delimiter_command(last_line).is_some() { let before = trimmed.rsplitn(2, '\n').nth(1).unwrap_or("").trim(); - if let Some(statement) = executable_sql_slice(before, self.options) { - statements.push(statement.to_string()); + if has_executable_sql_with_options(before, self.options) { + statements.push(before.to_string()); } self.buffer.clear(); } else if let Some(ref delim) = self.custom_delimiter { @@ -266,7 +266,8 @@ impl SqlStatementSplitter { } fn push_current_statement(&mut self, statements: &mut Vec) { - if let Some(statement) = executable_sql_slice(&self.buffer, self.options) { + let statement = self.buffer.trim(); + if has_executable_sql_with_options(statement, self.options) { statements.push(statement.to_string()); } self.buffer.clear(); @@ -1414,7 +1415,7 @@ SELECT 2;"; let sql = "SELECT 1; # mysql comment\n\nSELECT 2 # trailing comment"; assert_eq!( split_sql_statements_for_database(sql, DatabaseType::Mysql), - vec!["SELECT 1", "SELECT 2 # trailing comment"] + vec!["SELECT 1", "# mysql comment\n\nSELECT 2 # trailing comment"] ); } diff --git a/crates/dbx-core/src/storage.rs b/crates/dbx-core/src/storage.rs index fd7083eac..2b0d87e08 100644 --- a/crates/dbx-core/src/storage.rs +++ b/crates/dbx-core/src/storage.rs @@ -1,16 +1,17 @@ +use std::collections::{HashMap, HashSet}; use std::path::Path; -use std::str::FromStr; +use rusqlite::{params, params_from_iter, Connection, OptionalExtension, ToSql}; use serde::{Deserialize, Serialize}; -use sqlx::sqlite::{SqliteConnectOptions, SqlitePool, SqlitePoolOptions}; use crate::ai::{AiChatMessage, AiConfig, AiConversation}; -use crate::history::HistoryEntry; +use crate::db::sqlite::{connect_path_create_if_missing, SqliteHandle}; +use crate::history::{HistoryEntry, MAX_HISTORY}; use crate::models::connection::ConnectionConfig; use crate::saved_sql::{SavedSqlFile, SavedSqlFolder, SavedSqlLibrary}; pub struct Storage { - db: SqlitePool, + db: SqliteHandle, } #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] @@ -98,27 +99,36 @@ const SCHEMA_STATEMENTS: &[&str] = &[ )", ]; -// --------------------------------------------------------------------------- -// Construction / schema -// --------------------------------------------------------------------------- - impl Storage { pub async fn open(db_path: &Path) -> Result { - let url = format!("sqlite:{}?mode=rwc", db_path.display()); - let options = SqliteConnectOptions::from_str(&url).map_err(|e| e.to_string())?.create_if_missing(true); - let pool = - SqlitePoolOptions::new().max_connections(5).connect_with(options).await.map_err(|e| e.to_string())?; + let db_path = db_path.to_string_lossy().to_string(); + let db = connect_path_create_if_missing(&db_path).await?; + let storage = Self { db }; + storage.init_schema().await?; + Ok(storage) + } - for statement in SCHEMA_STATEMENTS { - sqlx::query(statement).execute(&pool).await.map_err(|e| e.to_string())?; - } - ensure_history_columns(&pool).await?; + async fn init_schema(&self) -> Result<(), String> { + self.db.with_connection(|conn| { + for statement in SCHEMA_STATEMENTS { + conn.execute(statement, []).map_err(|e| e.to_string())?; + } + ensure_history_columns_sync(conn)?; + Ok(()) + }) + } - Ok(Self { db: pool }) + async fn with_conn(&self, f: F) -> Result + where + T: Send + 'static, + F: FnOnce(&mut Connection) -> Result + Send + 'static, + { + let db = self.db.clone(); + tokio::task::spawn_blocking(move || db.with_connection(f)).await.map_err(|e| e.to_string())? } } -async fn ensure_history_columns(pool: &SqlitePool) -> Result<(), String> { +fn ensure_history_columns_sync(conn: &Connection) -> Result<(), String> { const COLUMNS: &[(&str, &str)] = &[ ("activity_kind", "TEXT NOT NULL DEFAULT 'query'"), ("connection_id", "TEXT NOT NULL DEFAULT ''"), @@ -129,169 +139,156 @@ async fn ensure_history_columns(pool: &SqlitePool) -> Result<(), String> { ("details_json", "TEXT"), ]; - let rows: Vec<(String,)> = sqlx::query_as("SELECT name FROM pragma_table_info('history')") - .fetch_all(pool) - .await + let mut stmt = conn.prepare("SELECT name FROM pragma_table_info('history')").map_err(|e| e.to_string())?; + let existing = stmt + .query_map([], |row| row.get::<_, String>(0)) + .map_err(|e| e.to_string())? + .collect::, _>>() .map_err(|e| e.to_string())?; - let existing: std::collections::HashSet = rows.into_iter().map(|(name,)| name).collect(); + for (name, definition) in COLUMNS { if existing.contains(*name) { continue; } - sqlx::query(&format!("ALTER TABLE history ADD COLUMN {name} {definition}")) - .execute(pool) - .await - .map_err(|e| e.to_string())?; + conn.execute(&format!("ALTER TABLE history ADD COLUMN {name} {definition}"), []).map_err(|e| e.to_string())?; } Ok(()) } -// --------------------------------------------------------------------------- // History -// --------------------------------------------------------------------------- - -#[derive(sqlx::FromRow)] -struct HistoryRow { - id: String, - connection_id: String, - connection_name: String, - database: String, - sql_text: String, - executed_at: String, - execution_time_ms: i64, - success: bool, - error: Option, - activity_kind: String, - operation: String, - target: String, - affected_rows: Option, - rollback_sql: Option, - details_json: Option, -} impl Storage { pub async fn save_history_entry(&self, entry: &HistoryEntry) -> Result<(), String> { - sqlx::query( - "INSERT OR REPLACE INTO history \ - (id, connection_name, database, sql_text, executed_at, execution_time_ms, success, error, \ - activity_kind, connection_id, operation, target, affected_rows, rollback_sql, details_json) \ - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", - ) - .bind(&entry.id) - .bind(&entry.connection_name) - .bind(&entry.database) - .bind(&entry.sql) - .bind(&entry.executed_at) - .bind(entry.execution_time_ms as i64) - .bind(entry.success) - .bind(&entry.error) - .bind(&entry.activity_kind) - .bind(&entry.connection_id) - .bind(&entry.operation) - .bind(&entry.target) - .bind(entry.affected_rows) - .bind(&entry.rollback_sql) - .bind(&entry.details_json) - .execute(&self.db) - .await - .map_err(|e| e.to_string())?; + let entry = entry.clone(); + self.with_conn(move |conn| { + conn.execute( + "INSERT OR REPLACE INTO history \ + (id, connection_name, database, sql_text, executed_at, execution_time_ms, success, error, \ + activity_kind, connection_id, operation, target, affected_rows, rollback_sql, details_json) \ + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + params![ + entry.id, + entry.connection_name, + entry.database, + entry.sql, + entry.executed_at, + entry.execution_time_ms as i64, + entry.success, + entry.error, + entry.activity_kind, + entry.connection_id, + entry.operation, + entry.target, + entry.affected_rows, + entry.rollback_sql, + entry.details_json + ], + ) + .map_err(|e| e.to_string())?; - // Keep at most MAX_HISTORY entries - sqlx::query( - "DELETE FROM history WHERE id NOT IN \ - (SELECT id FROM history ORDER BY executed_at DESC LIMIT 1000)", - ) - .execute(&self.db) + conn.execute( + "DELETE FROM history WHERE id NOT IN \ + (SELECT id FROM history ORDER BY executed_at DESC LIMIT ?1)", + [MAX_HISTORY as i64], + ) + .map_err(|e| e.to_string())?; + Ok(()) + }) .await - .map_err(|e| e.to_string())?; - - Ok(()) } pub async fn load_history_entries(&self, limit: usize, offset: usize) -> Result, String> { - let rows: Vec = sqlx::query_as( - "SELECT id, connection_name, database, sql_text, executed_at, \ - execution_time_ms, success, error, activity_kind, connection_id, operation, target, \ - affected_rows, rollback_sql, details_json \ - FROM history ORDER BY executed_at DESC LIMIT ? OFFSET ?", - ) - .bind(limit as i64) - .bind(offset as i64) - .fetch_all(&self.db) + self.with_conn(move |conn| { + let mut stmt = conn + .prepare( + "SELECT id, connection_name, database, sql_text, executed_at, execution_time_ms, success, \ + error, activity_kind, connection_id, operation, target, affected_rows, rollback_sql, details_json \ + FROM history ORDER BY executed_at DESC LIMIT ?1 OFFSET ?2", + ) + .map_err(|e| e.to_string())?; + let rows = stmt + .query_map(params![limit as i64, offset as i64], |row| { + Ok(HistoryEntry { + id: row.get(0)?, + connection_name: row.get(1)?, + database: row.get(2)?, + sql: row.get(3)?, + executed_at: row.get(4)?, + execution_time_ms: row.get::<_, i64>(5)? as u128, + success: row.get(6)?, + error: row.get(7)?, + activity_kind: { + let value: String = row.get(8)?; + if value.is_empty() { + "query".to_string() + } else { + value + } + }, + connection_id: row.get(9)?, + operation: row.get(10)?, + target: row.get(11)?, + affected_rows: row.get(12)?, + rollback_sql: row.get(13)?, + details_json: row.get(14)?, + }) + }) + .map_err(|e| e.to_string())?; + rows.collect::, _>>().map_err(|e| e.to_string()) + }) .await - .map_err(|e| e.to_string())?; - - Ok(rows - .into_iter() - .map(|r| HistoryEntry { - id: r.id, - connection_id: r.connection_id, - connection_name: r.connection_name, - database: r.database, - sql: r.sql_text, - executed_at: r.executed_at, - execution_time_ms: r.execution_time_ms as u128, - success: r.success, - error: r.error, - activity_kind: if r.activity_kind.is_empty() { "query".to_string() } else { r.activity_kind }, - operation: r.operation, - target: r.target, - affected_rows: r.affected_rows, - rollback_sql: r.rollback_sql, - details_json: r.details_json, - }) - .collect()) } pub async fn clear_history(&self) -> Result<(), String> { - sqlx::query("DELETE FROM history").execute(&self.db).await.map_err(|e| e.to_string())?; - Ok(()) + self.with_conn(|conn| conn.execute("DELETE FROM history", []).map(|_| ()).map_err(|e| e.to_string())).await } pub async fn delete_history_entry(&self, id: &str) -> Result<(), String> { - sqlx::query("DELETE FROM history WHERE id = ?").bind(id).execute(&self.db).await.map_err(|e| e.to_string())?; - Ok(()) + let id = id.to_string(); + self.with_conn(move |conn| { + conn.execute("DELETE FROM history WHERE id = ?1", [id]).map(|_| ()).map_err(|e| e.to_string()) + }) + .await } } -// --------------------------------------------------------------------------- // AI Config -// --------------------------------------------------------------------------- impl Storage { pub async fn save_ai_config(&self, config: &AiConfig) -> Result<(), String> { let json = serde_json::to_string(config).map_err(|e| e.to_string())?; - sqlx::query("INSERT OR REPLACE INTO ai_config (id, config_json) VALUES (1, ?)") - .bind(&json) - .execute(&self.db) - .await - .map_err(|e| e.to_string())?; - Ok(()) + self.with_conn(move |conn| { + conn.execute("INSERT OR REPLACE INTO ai_config (id, config_json) VALUES (1, ?1)", [json]) + .map(|_| ()) + .map_err(|e| e.to_string()) + }) + .await } pub async fn load_ai_config(&self) -> Result, String> { - let row: Option<(String,)> = sqlx::query_as("SELECT config_json FROM ai_config WHERE id = 1") - .fetch_optional(&self.db) - .await - .map_err(|e| e.to_string())?; - match row { - Some((json,)) => serde_json::from_str(&json).map(Some).map_err(|e| e.to_string()), - None => Ok(None), - } + let json: Option = self + .with_conn(|conn| { + conn.query_row("SELECT config_json FROM ai_config WHERE id = 1", [], |row| row.get(0)) + .optional() + .map_err(|e| e.to_string()) + }) + .await?; + json.map(|value| serde_json::from_str(&value).map_err(|e| e.to_string())).transpose() } } -// --------------------------------------------------------------------------- // App Settings -// --------------------------------------------------------------------------- impl Storage { async fn load_app_settings_json(&self) -> Result, String> { - let row: Option<(String,)> = sqlx::query_as("SELECT settings_json FROM app_settings WHERE id = 1") - .fetch_optional(&self.db) - .await - .map_err(|e| e.to_string())?; - let Some((json,)) = row else { + let json: Option = self + .with_conn(|conn| { + conn.query_row("SELECT settings_json FROM app_settings WHERE id = 1", [], |row| row.get(0)) + .optional() + .map_err(|e| e.to_string()) + }) + .await?; + let Some(json) = json else { return Ok(serde_json::Map::new()); }; match serde_json::from_str::(&json).map_err(|e| e.to_string())? { @@ -305,12 +302,12 @@ impl Storage { settings: &serde_json::Map, ) -> Result<(), String> { let json = serde_json::Value::Object(settings.clone()).to_string(); - sqlx::query("INSERT OR REPLACE INTO app_settings (id, settings_json) VALUES (1, ?)") - .bind(&json) - .execute(&self.db) - .await - .map_err(|e| e.to_string())?; - Ok(()) + self.with_conn(move |conn| { + conn.execute("INSERT OR REPLACE INTO app_settings (id, settings_json) VALUES (1, ?1)", [json]) + .map(|_| ()) + .map_err(|e| e.to_string()) + }) + .await } pub async fn save_password_hash(&self, hash: &str) -> Result<(), String> { @@ -343,7 +340,7 @@ impl Storage { pub async fn save_pinned_tree_node_ids(&self, ids: &[String]) -> Result<(), String> { let mut settings = self.load_app_settings_json().await?; - let values = ids.iter().map(|id| serde_json::Value::String(id.clone())).collect::>(); + let values = ids.iter().map(|id| serde_json::Value::String(id.clone())).collect::>(); settings.insert("pinned_tree_node_ids".to_string(), serde_json::Value::Array(values)); self.save_app_settings_json(&settings).await } @@ -360,155 +357,140 @@ impl Storage { } } -// --------------------------------------------------------------------------- // AI Conversations -// --------------------------------------------------------------------------- - -#[derive(sqlx::FromRow)] -struct AiConversationRow { - id: String, - title: String, - connection_name: String, - database: String, - messages_json: String, - created_at: String, - updated_at: String, -} impl Storage { pub async fn save_ai_conversation(&self, conv: &AiConversation) -> Result<(), String> { + let conv = conv.clone(); let messages_json = serde_json::to_string(&conv.messages).map_err(|e| e.to_string())?; - sqlx::query( - "INSERT OR REPLACE INTO ai_conversations \ - (id, title, connection_name, database, messages_json, created_at, updated_at) \ - VALUES (?, ?, ?, ?, ?, ?, ?)", - ) - .bind(&conv.id) - .bind(&conv.title) - .bind(&conv.connection_name) - .bind(&conv.database) - .bind(&messages_json) - .bind(&conv.created_at) - .bind(&conv.updated_at) - .execute(&self.db) - .await - .map_err(|e| e.to_string())?; + self.with_conn(move |conn| { + conn.execute( + "INSERT OR REPLACE INTO ai_conversations \ + (id, title, connection_name, database, messages_json, created_at, updated_at) \ + VALUES (?, ?, ?, ?, ?, ?, ?)", + params![ + conv.id, + conv.title, + conv.connection_name, + conv.database, + messages_json, + conv.created_at, + conv.updated_at + ], + ) + .map_err(|e| e.to_string())?; - // Keep at most 50 conversations - sqlx::query( - "DELETE FROM ai_conversations WHERE id NOT IN \ - (SELECT id FROM ai_conversations ORDER BY updated_at DESC LIMIT 50)", - ) - .execute(&self.db) + conn.execute( + "DELETE FROM ai_conversations WHERE id NOT IN \ + (SELECT id FROM ai_conversations ORDER BY updated_at DESC LIMIT 50)", + [], + ) + .map_err(|e| e.to_string())?; + Ok(()) + }) .await - .map_err(|e| e.to_string())?; - - Ok(()) } pub async fn load_ai_conversations(&self) -> Result, String> { - let rows: Vec = sqlx::query_as( - "SELECT id, title, connection_name, database, messages_json, \ - created_at, updated_at \ - FROM ai_conversations ORDER BY updated_at DESC", - ) - .fetch_all(&self.db) - .await - .map_err(|e| e.to_string())?; - - rows.into_iter() - .map(|r| { - let messages: Vec = serde_json::from_str(&r.messages_json).map_err(|e| e.to_string())?; - Ok(AiConversation { - id: r.id, - title: r.title, - connection_name: r.connection_name, - database: r.database, - messages, - created_at: r.created_at, - updated_at: r.updated_at, + self.with_conn(|conn| { + let mut stmt = conn + .prepare( + "SELECT id, title, connection_name, database, messages_json, created_at, updated_at \ + FROM ai_conversations ORDER BY updated_at DESC", + ) + .map_err(|e| e.to_string())?; + let rows = stmt + .query_map([], |row| { + let messages_json: String = row.get(4)?; + let messages: Vec = + serde_json::from_str(&messages_json).map_err(map_from_sql_err)?; + Ok(AiConversation { + id: row.get(0)?, + title: row.get(1)?, + connection_name: row.get(2)?, + database: row.get(3)?, + messages, + created_at: row.get(5)?, + updated_at: row.get(6)?, + }) }) - }) - .collect() + .map_err(|e| e.to_string())?; + rows.collect::, _>>().map_err(|e| e.to_string()) + }) + .await } pub async fn delete_ai_conversation(&self, id: &str) -> Result<(), String> { - sqlx::query("DELETE FROM ai_conversations WHERE id = ?") - .bind(id) - .execute(&self.db) - .await - .map_err(|e| e.to_string())?; - Ok(()) + let id = id.to_string(); + self.with_conn(move |conn| { + conn.execute("DELETE FROM ai_conversations WHERE id = ?1", [id]).map(|_| ()).map_err(|e| e.to_string()) + }) + .await } } -// --------------------------------------------------------------------------- -// Connections (with inline secrets) -// --------------------------------------------------------------------------- +// Connections impl Storage { pub async fn save_connections(&self, configs: &[ConnectionConfig]) -> Result<(), String> { - let mut tx = self.db.begin().await.map_err(|e| e.to_string())?; + let configs = configs.to_vec(); + self.with_conn(move |conn| { + let tx = conn.transaction().map_err(|e| e.to_string())?; + tx.execute("DELETE FROM connections", []).map_err(|e| e.to_string())?; - sqlx::query("DELETE FROM connections").execute(&mut *tx).await.map_err(|e| e.to_string())?; + for config in &configs { + let config = config.canonicalized(); + let config_id = config.id.clone(); + let mut sanitized = config.clone(); + sanitized.password = String::new(); + sanitized.ssh_password = String::new(); + sanitized.ssh_key_passphrase = String::new(); + sanitized.proxy_password = String::new(); + sanitized.connection_string = None; + let json = serde_json::to_string(&sanitized).map_err(|e| e.to_string())?; - for config in configs { - let config = config.canonicalized(); - // Store config without secrets - let mut sanitized = config.clone(); - sanitized.password = String::new(); - sanitized.ssh_password = String::new(); - sanitized.ssh_key_passphrase = String::new(); - sanitized.proxy_password = String::new(); - sanitized.connection_string = None; - let json = serde_json::to_string(&sanitized).map_err(|e| e.to_string())?; - - sqlx::query("INSERT INTO connections (id, config_json) VALUES (?, ?)") - .bind(&config.id) - .bind(&json) - .execute(&mut *tx) - .await - .map_err(|e| e.to_string())?; - - // Store secrets - persist_secret_in_tx(&mut tx, &config.id, "password", &config.password).await?; - persist_secret_in_tx(&mut tx, &config.id, "ssh_password", &config.ssh_password).await?; - persist_secret_in_tx(&mut tx, &config.id, "ssh_key_passphrase", &config.ssh_key_passphrase).await?; - persist_secret_in_tx(&mut tx, &config.id, "proxy_password", &config.proxy_password).await?; - if let Some(cs) = &config.connection_string { - persist_secret_in_tx(&mut tx, &config.id, "connection_string", cs).await?; - } else { - sqlx::query("DELETE FROM connection_secrets WHERE connection_id = ? AND key = ?") - .bind(&config.id) - .bind("connection_string") - .execute(&mut *tx) - .await + tx.execute("INSERT INTO connections (id, config_json) VALUES (?1, ?2)", params![config_id, json]) .map_err(|e| e.to_string())?; - } - } - // Remove secrets for connections that no longer exist - if configs.is_empty() { - sqlx::query("DELETE FROM connection_secrets").execute(&mut *tx).await.map_err(|e| e.to_string())?; - } else { - let placeholders: Vec<&str> = configs.iter().map(|_| "?").collect(); - let sql = format!("DELETE FROM connection_secrets WHERE connection_id NOT IN ({})", placeholders.join(",")); - let mut query = sqlx::query(&sql); - for config in configs { - query = query.bind(&config.id); + persist_secret_in_tx(&tx, &config.id, "password", &config.password)?; + persist_secret_in_tx(&tx, &config.id, "ssh_password", &config.ssh_password)?; + persist_secret_in_tx(&tx, &config.id, "ssh_key_passphrase", &config.ssh_key_passphrase)?; + persist_secret_in_tx(&tx, &config.id, "proxy_password", &config.proxy_password)?; + if let Some(cs) = &config.connection_string { + persist_secret_in_tx(&tx, &config.id, "connection_string", cs)?; + } else { + tx.execute( + "DELETE FROM connection_secrets WHERE connection_id = ?1 AND key = ?2", + params![config.id, "connection_string"], + ) + .map_err(|e| e.to_string())?; + } } - query.execute(&mut *tx).await.map_err(|e| e.to_string())?; - } - tx.commit().await.map_err(|e| e.to_string())?; - Ok(()) + if configs.is_empty() { + tx.execute("DELETE FROM connection_secrets", []).map_err(|e| e.to_string())?; + } else { + let placeholders = vec!["?"; configs.len()].join(","); + let sql = format!("DELETE FROM connection_secrets WHERE connection_id NOT IN ({placeholders})"); + let ids = configs.iter().map(|config| &config.id as &dyn ToSql); + tx.execute(&sql, params_from_iter(ids)).map_err(|e| e.to_string())?; + } + + tx.commit().map_err(|e| e.to_string()) + }) + .await } pub async fn load_connections(&self) -> Result, String> { - let rows: Vec<(String, String)> = sqlx::query_as("SELECT id, config_json FROM connections") - .fetch_all(&self.db) - .await - .map_err(|e| e.to_string())?; + let rows: Vec<(String, String)> = self + .with_conn(|conn| { + let mut stmt = conn.prepare("SELECT id, config_json FROM connections").map_err(|e| e.to_string())?; + let rows = stmt + .query_map([], |row| Ok((row.get::<_, String>(0)?, row.get::<_, String>(1)?))) + .map_err(|e| e.to_string())?; + rows.collect::, _>>().map_err(|e| e.to_string()) + }) + .await?; let mut configs = Vec::new(); for (id, json) in rows { @@ -524,273 +506,252 @@ impl Storage { } } -// --------------------------------------------------------------------------- -// Saved SQL Library -// --------------------------------------------------------------------------- - -#[derive(sqlx::FromRow)] -struct SavedSqlFolderRow { - id: String, - connection_id: String, - name: String, - created_at: String, - updated_at: String, -} - -#[derive(sqlx::FromRow)] -struct SavedSqlFileRow { - id: String, - connection_id: String, - folder_id: Option, - name: String, - database_name: String, - schema_name: Option, - sql_text: String, - created_at: String, - updated_at: String, -} - -impl From for SavedSqlFolder { - fn from(row: SavedSqlFolderRow) -> Self { - Self { - id: row.id, - connection_id: row.connection_id, - name: row.name, - created_at: row.created_at, - updated_at: row.updated_at, - } - } -} - -impl From for SavedSqlFile { - fn from(row: SavedSqlFileRow) -> Self { - Self { - id: row.id, - connection_id: row.connection_id, - folder_id: row.folder_id, - name: row.name, - database: row.database_name, - schema: row.schema_name, - sql: row.sql_text, - created_at: row.created_at, - updated_at: row.updated_at, - } - } -} +// Saved SQL impl Storage { pub async fn load_saved_sql_library(&self) -> Result { - let folder_rows: Vec = sqlx::query_as( - "SELECT id, connection_id, name, created_at, updated_at \ - FROM saved_sql_folders ORDER BY connection_id, name COLLATE NOCASE", - ) - .fetch_all(&self.db) - .await - .map_err(|e| e.to_string())?; + self.with_conn(|conn| { + let mut folder_stmt = conn + .prepare( + "SELECT id, connection_id, name, created_at, updated_at \ + FROM saved_sql_folders ORDER BY connection_id, name COLLATE NOCASE", + ) + .map_err(|e| e.to_string())?; + let folders = folder_stmt + .query_map([], |row| { + Ok(SavedSqlFolder { + id: row.get(0)?, + connection_id: row.get(1)?, + name: row.get(2)?, + created_at: row.get(3)?, + updated_at: row.get(4)?, + }) + }) + .map_err(|e| e.to_string())? + .collect::, _>>() + .map_err(|e| e.to_string())?; - let file_rows: Vec = sqlx::query_as( - "SELECT id, connection_id, folder_id, name, database_name, schema_name, sql_text, created_at, updated_at \ - FROM saved_sql_files ORDER BY connection_id, folder_id, name COLLATE NOCASE", - ) - .fetch_all(&self.db) - .await - .map_err(|e| e.to_string())?; + let mut file_stmt = conn + .prepare( + "SELECT id, connection_id, folder_id, name, database_name, schema_name, sql_text, created_at, updated_at \ + FROM saved_sql_files ORDER BY connection_id, folder_id, name COLLATE NOCASE", + ) + .map_err(|e| e.to_string())?; + let files = file_stmt + .query_map([], |row| { + Ok(SavedSqlFile { + id: row.get(0)?, + connection_id: row.get(1)?, + folder_id: row.get(2)?, + name: row.get(3)?, + database: row.get(4)?, + schema: row.get(5)?, + sql: row.get(6)?, + created_at: row.get(7)?, + updated_at: row.get(8)?, + }) + }) + .map_err(|e| e.to_string())? + .collect::, _>>() + .map_err(|e| e.to_string())?; - Ok(SavedSqlLibrary { - folders: folder_rows.into_iter().map(Into::into).collect(), - files: file_rows.into_iter().map(Into::into).collect(), + Ok(SavedSqlLibrary { folders, files }) }) + .await } pub async fn save_saved_sql_folder(&self, folder: &SavedSqlFolder) -> Result<(), String> { - sqlx::query( - "INSERT INTO saved_sql_folders (id, connection_id, name, created_at, updated_at) \ - VALUES (?, ?, ?, ?, ?) \ - ON CONFLICT(id) DO UPDATE SET \ - connection_id = excluded.connection_id, \ - name = excluded.name, \ - updated_at = excluded.updated_at", - ) - .bind(&folder.id) - .bind(&folder.connection_id) - .bind(&folder.name) - .bind(&folder.created_at) - .bind(&folder.updated_at) - .execute(&self.db) + let folder = folder.clone(); + self.with_conn(move |conn| { + conn.execute( + "INSERT INTO saved_sql_folders (id, connection_id, name, created_at, updated_at) \ + VALUES (?, ?, ?, ?, ?) \ + ON CONFLICT(id) DO UPDATE SET \ + connection_id = excluded.connection_id, \ + name = excluded.name, \ + updated_at = excluded.updated_at", + params![folder.id, folder.connection_id, folder.name, folder.created_at, folder.updated_at], + ) + .map(|_| ()) + .map_err(|e| e.to_string()) + }) .await - .map_err(|e| e.to_string())?; - Ok(()) } pub async fn delete_saved_sql_folder(&self, id: &str) -> Result<(), String> { - let mut tx = self.db.begin().await.map_err(|e| e.to_string())?; - sqlx::query("DELETE FROM saved_sql_files WHERE folder_id = ?") - .bind(id) - .execute(&mut *tx) - .await - .map_err(|e| e.to_string())?; - sqlx::query("DELETE FROM saved_sql_folders WHERE id = ?") - .bind(id) - .execute(&mut *tx) - .await - .map_err(|e| e.to_string())?; - tx.commit().await.map_err(|e| e.to_string())?; - Ok(()) + let id = id.to_string(); + self.with_conn(move |conn| { + let tx = conn.transaction().map_err(|e| e.to_string())?; + tx.execute("DELETE FROM saved_sql_files WHERE folder_id = ?1", [id.as_str()]).map_err(|e| e.to_string())?; + tx.execute("DELETE FROM saved_sql_folders WHERE id = ?1", [id.as_str()]).map_err(|e| e.to_string())?; + tx.commit().map_err(|e| e.to_string()) + }) + .await } pub async fn save_saved_sql_file(&self, file: &SavedSqlFile) -> Result<(), String> { - sqlx::query( - "INSERT INTO saved_sql_files \ - (id, connection_id, folder_id, name, database_name, schema_name, sql_text, created_at, updated_at) \ - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) \ - ON CONFLICT(id) DO UPDATE SET \ - connection_id = excluded.connection_id, \ - folder_id = excluded.folder_id, \ - name = excluded.name, \ - database_name = excluded.database_name, \ - schema_name = excluded.schema_name, \ - sql_text = excluded.sql_text, \ - updated_at = excluded.updated_at", - ) - .bind(&file.id) - .bind(&file.connection_id) - .bind(&file.folder_id) - .bind(&file.name) - .bind(&file.database) - .bind(&file.schema) - .bind(&file.sql) - .bind(&file.created_at) - .bind(&file.updated_at) - .execute(&self.db) + let file = file.clone(); + self.with_conn(move |conn| { + conn.execute( + "INSERT INTO saved_sql_files \ + (id, connection_id, folder_id, name, database_name, schema_name, sql_text, created_at, updated_at) \ + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) \ + ON CONFLICT(id) DO UPDATE SET \ + connection_id = excluded.connection_id, \ + folder_id = excluded.folder_id, \ + name = excluded.name, \ + database_name = excluded.database_name, \ + schema_name = excluded.schema_name, \ + sql_text = excluded.sql_text, \ + updated_at = excluded.updated_at", + params![ + file.id, + file.connection_id, + file.folder_id, + file.name, + file.database, + file.schema, + file.sql, + file.created_at, + file.updated_at + ], + ) + .map(|_| ()) + .map_err(|e| e.to_string()) + }) .await - .map_err(|e| e.to_string())?; - Ok(()) } pub async fn delete_saved_sql_file(&self, id: &str) -> Result<(), String> { - sqlx::query("DELETE FROM saved_sql_files WHERE id = ?") - .bind(id) - .execute(&self.db) - .await - .map_err(|e| e.to_string())?; - Ok(()) + let id = id.to_string(); + self.with_conn(move |conn| { + conn.execute("DELETE FROM saved_sql_files WHERE id = ?1", [id]).map(|_| ()).map_err(|e| e.to_string()) + }) + .await } } -// --------------------------------------------------------------------------- // Secrets -// --------------------------------------------------------------------------- impl Storage { pub async fn get_secret(&self, connection_id: &str, key: &str) -> Result, String> { - let row: Option<(String,)> = - sqlx::query_as("SELECT secret FROM connection_secrets WHERE connection_id = ? AND key = ?") - .bind(connection_id) - .bind(key) - .fetch_optional(&self.db) - .await - .map_err(|e| e.to_string())?; - Ok(row.map(|(s,)| s)) + let connection_id = connection_id.to_string(); + let key = key.to_string(); + self.with_conn(move |conn| { + conn.query_row( + "SELECT secret FROM connection_secrets WHERE connection_id = ?1 AND key = ?2", + params![connection_id, key], + |row| row.get(0), + ) + .optional() + .map_err(|e| e.to_string()) + }) + .await } pub async fn set_secret(&self, connection_id: &str, key: &str, secret: &str) -> Result<(), String> { - sqlx::query( - "INSERT OR REPLACE INTO connection_secrets (connection_id, key, secret) \ - VALUES (?, ?, ?)", - ) - .bind(connection_id) - .bind(key) - .bind(secret) - .execute(&self.db) + let connection_id = connection_id.to_string(); + let key = key.to_string(); + let secret = secret.to_string(); + self.with_conn(move |conn| { + conn.execute( + "INSERT OR REPLACE INTO connection_secrets (connection_id, key, secret) VALUES (?, ?, ?)", + params![connection_id, key, secret], + ) + .map(|_| ()) + .map_err(|e| e.to_string()) + }) .await - .map_err(|e| e.to_string())?; - Ok(()) } pub async fn delete_secret(&self, connection_id: &str, key: &str) -> Result<(), String> { - sqlx::query("DELETE FROM connection_secrets WHERE connection_id = ? AND key = ?") - .bind(connection_id) - .bind(key) - .execute(&self.db) - .await - .map_err(|e| e.to_string())?; - Ok(()) + let connection_id = connection_id.to_string(); + let key = key.to_string(); + self.with_conn(move |conn| { + conn.execute( + "DELETE FROM connection_secrets WHERE connection_id = ?1 AND key = ?2", + params![connection_id, key], + ) + .map(|_| ()) + .map_err(|e| e.to_string()) + }) + .await } } -// --------------------------------------------------------------------------- // Layout -// --------------------------------------------------------------------------- impl Storage { pub async fn save_sidebar_layout(&self, layout: &serde_json::Value) -> Result<(), String> { let json = serde_json::to_string(layout).map_err(|e| e.to_string())?; - sqlx::query("INSERT OR REPLACE INTO sidebar_layout (id, layout_json) VALUES (1, ?)") - .bind(&json) - .execute(&self.db) - .await - .map_err(|e| e.to_string())?; - Ok(()) + self.with_conn(move |conn| { + conn.execute("INSERT OR REPLACE INTO sidebar_layout (id, layout_json) VALUES (1, ?1)", [json]) + .map(|_| ()) + .map_err(|e| e.to_string()) + }) + .await } pub async fn load_sidebar_layout(&self) -> Result, String> { - let row: Option<(String,)> = sqlx::query_as("SELECT layout_json FROM sidebar_layout WHERE id = 1") - .fetch_optional(&self.db) - .await - .map_err(|e| e.to_string())?; - match row { - Some((json,)) => serde_json::from_str(&json).map(Some).map_err(|e| e.to_string()), - None => Ok(None), - } + let json: Option = self + .with_conn(|conn| { + conn.query_row("SELECT layout_json FROM sidebar_layout WHERE id = 1", [], |row| row.get(0)) + .optional() + .map_err(|e| e.to_string()) + }) + .await?; + json.map(|value| serde_json::from_str(&value).map_err(|e| e.to_string())).transpose() } } -// --------------------------------------------------------------------------- // Schema cache -// --------------------------------------------------------------------------- impl Storage { pub async fn save_schema_cache(&self, cache_key: &str, payload: &serde_json::Value) -> Result<(), String> { + let cache_key = cache_key.to_string(); let json = serde_json::to_string(payload).map_err(|e| e.to_string())?; - sqlx::query( - "INSERT OR REPLACE INTO schema_cache (cache_key, payload_json, updated_at) \ - VALUES (?, ?, datetime('now'))", - ) - .bind(cache_key) - .bind(&json) - .execute(&self.db) + self.with_conn(move |conn| { + conn.execute( + "INSERT OR REPLACE INTO schema_cache (cache_key, payload_json, updated_at) \ + VALUES (?1, ?2, datetime('now'))", + params![cache_key, json], + ) + .map(|_| ()) + .map_err(|e| e.to_string()) + }) .await - .map_err(|e| e.to_string())?; - Ok(()) } pub async fn load_schema_cache(&self, cache_key: &str) -> Result, String> { - let row: Option<(String,)> = sqlx::query_as("SELECT payload_json FROM schema_cache WHERE cache_key = ?") - .bind(cache_key) - .fetch_optional(&self.db) - .await - .map_err(|e| e.to_string())?; - match row { - Some((json,)) => serde_json::from_str(&json).map(Some).map_err(|e| e.to_string()), - None => Ok(None), - } + let cache_key = cache_key.to_string(); + let json: Option = self + .with_conn(move |conn| { + conn.query_row("SELECT payload_json FROM schema_cache WHERE cache_key = ?1", [cache_key], |row| { + row.get(0) + }) + .optional() + .map_err(|e| e.to_string()) + }) + .await?; + json.map(|value| serde_json::from_str(&value).map_err(|e| e.to_string())).transpose() } pub async fn delete_schema_cache_prefix(&self, prefix: &str) -> Result<(), String> { - sqlx::query("DELETE FROM schema_cache WHERE cache_key = ? OR substr(cache_key, 1, ?) = ?") - .bind(prefix) - .bind(prefix.len() as i64) - .bind(prefix) - .execute(&self.db) - .await - .map_err(|e| e.to_string())?; - Ok(()) + let prefix = prefix.to_string(); + let prefix_len = prefix.len() as i64; + self.with_conn(move |conn| { + conn.execute( + "DELETE FROM schema_cache WHERE cache_key = ?1 OR substr(cache_key, 1, ?2) = ?3", + params![prefix.clone(), prefix_len, prefix], + ) + .map(|_| ()) + .map_err(|e| e.to_string()) + }) + .await } } -// --------------------------------------------------------------------------- // JSON migration -// --------------------------------------------------------------------------- impl Storage { pub async fn migrate_from_json(&self, data_dir: &Path) -> Result<(), String> { @@ -812,12 +773,16 @@ impl Storage { let configs: Vec = serde_json::from_str(&json).unwrap_or_default(); for config in &configs { let config_json = serde_json::to_string(config).map_err(|e| e.to_string())?; - sqlx::query("INSERT OR IGNORE INTO connections (id, config_json) VALUES (?, ?)") - .bind(&config.id) - .bind(&config_json) - .execute(&self.db) - .await - .map_err(|e| e.to_string())?; + let id = config.id.clone(); + self.with_conn(move |conn| { + conn.execute( + "INSERT OR IGNORE INTO connections (id, config_json) VALUES (?1, ?2)", + params![id, config_json], + ) + .map(|_| ()) + .map_err(|e| e.to_string()) + }) + .await?; } let _ = tokio::fs::rename(&path, data_dir.join("connections.json.bak")).await; Ok(()) @@ -829,21 +794,22 @@ impl Storage { return Ok(()); } let json = tokio::fs::read_to_string(&path).await.map_err(|e| e.to_string())?; - let secrets: std::collections::HashMap = serde_json::from_str(&json).unwrap_or_default(); + let secrets: HashMap = serde_json::from_str(&json).unwrap_or_default(); for (key, secret) in &secrets { - // key format: "connection:{id}:{field}" let parts: Vec<&str> = key.splitn(3, ':').collect(); if parts.len() == 3 && parts[0] == "connection" { - sqlx::query( - "INSERT OR IGNORE INTO connection_secrets \ - (connection_id, key, secret) VALUES (?, ?, ?)", - ) - .bind(parts[1]) - .bind(parts[2]) - .bind(secret) - .execute(&self.db) - .await - .map_err(|e| e.to_string())?; + let connection_id = parts[1].to_string(); + let field = parts[2].to_string(); + let secret = secret.clone(); + self.with_conn(move |conn| { + conn.execute( + "INSERT OR IGNORE INTO connection_secrets (connection_id, key, secret) VALUES (?1, ?2, ?3)", + params![connection_id, field, secret], + ) + .map(|_| ()) + .map_err(|e| e.to_string()) + }) + .await?; } } let _ = tokio::fs::rename(&path, data_dir.join("secrets.json.bak")).await; @@ -858,31 +824,7 @@ impl Storage { let json = tokio::fs::read_to_string(&path).await.map_err(|e| e.to_string())?; let entries: Vec = serde_json::from_str(&json).unwrap_or_default(); for entry in &entries { - sqlx::query( - "INSERT OR IGNORE INTO history \ - (id, connection_name, database, sql_text, executed_at, \ - execution_time_ms, success, error, activity_kind, connection_id, operation, target, \ - affected_rows, rollback_sql, details_json) \ - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", - ) - .bind(&entry.id) - .bind(&entry.connection_name) - .bind(&entry.database) - .bind(&entry.sql) - .bind(&entry.executed_at) - .bind(entry.execution_time_ms as i64) - .bind(entry.success) - .bind(&entry.error) - .bind(&entry.activity_kind) - .bind(&entry.connection_id) - .bind(&entry.operation) - .bind(&entry.target) - .bind(entry.affected_rows) - .bind(&entry.rollback_sql) - .bind(&entry.details_json) - .execute(&self.db) - .await - .map_err(|e| e.to_string())?; + self.save_history_entry(entry).await?; } let _ = tokio::fs::rename(&path, data_dir.join("query_history.json.bak")).await; Ok(()) @@ -894,15 +836,18 @@ impl Storage { return Ok(()); } let json = tokio::fs::read_to_string(&path).await.map_err(|e| e.to_string())?; - // Only migrate if the table is empty - let count: (i64,) = - sqlx::query_as("SELECT COUNT(*) FROM ai_config").fetch_one(&self.db).await.map_err(|e| e.to_string())?; - if count.0 == 0 { - sqlx::query("INSERT OR IGNORE INTO ai_config (id, config_json) VALUES (1, ?)") - .bind(&json) - .execute(&self.db) - .await - .map_err(|e| e.to_string())?; + let count: i64 = self + .with_conn(|conn| { + conn.query_row("SELECT COUNT(*) FROM ai_config", [], |row| row.get(0)).map_err(|e| e.to_string()) + }) + .await?; + if count == 0 { + self.with_conn(move |conn| { + conn.execute("INSERT OR IGNORE INTO ai_config (id, config_json) VALUES (1, ?1)", [json]) + .map(|_| ()) + .map_err(|e| e.to_string()) + }) + .await?; } let _ = tokio::fs::rename(&path, data_dir.join("ai_config.json.bak")).await; Ok(()) @@ -916,23 +861,27 @@ impl Storage { let json = tokio::fs::read_to_string(&path).await.map_err(|e| e.to_string())?; let conversations: Vec = serde_json::from_str(&json).unwrap_or_default(); for conv in &conversations { + let conv = conv.clone(); let messages_json = serde_json::to_string(&conv.messages).map_err(|e| e.to_string())?; - sqlx::query( - "INSERT OR IGNORE INTO ai_conversations \ - (id, title, connection_name, database, messages_json, \ - created_at, updated_at) \ - VALUES (?, ?, ?, ?, ?, ?, ?)", - ) - .bind(&conv.id) - .bind(&conv.title) - .bind(&conv.connection_name) - .bind(&conv.database) - .bind(&messages_json) - .bind(&conv.created_at) - .bind(&conv.updated_at) - .execute(&self.db) - .await - .map_err(|e| e.to_string())?; + self.with_conn(move |conn| { + conn.execute( + "INSERT OR IGNORE INTO ai_conversations \ + (id, title, connection_name, database, messages_json, created_at, updated_at) \ + VALUES (?, ?, ?, ?, ?, ?, ?)", + params![ + conv.id, + conv.title, + conv.connection_name, + conv.database, + messages_json, + conv.created_at, + conv.updated_at + ], + ) + .map(|_| ()) + .map_err(|e| e.to_string()) + }) + .await?; } let _ = tokio::fs::rename(&path, data_dir.join("ai_conversations.json.bak")).await; Ok(()) @@ -944,54 +893,47 @@ impl Storage { return Ok(()); } let json = tokio::fs::read_to_string(&path).await.map_err(|e| e.to_string())?; - let count: (i64,) = sqlx::query_as("SELECT COUNT(*) FROM sidebar_layout") - .fetch_one(&self.db) - .await - .map_err(|e| e.to_string())?; - if count.0 == 0 { - sqlx::query("INSERT OR IGNORE INTO sidebar_layout (id, layout_json) VALUES (1, ?)") - .bind(&json) - .execute(&self.db) - .await - .map_err(|e| e.to_string())?; + let count: i64 = self + .with_conn(|conn| { + conn.query_row("SELECT COUNT(*) FROM sidebar_layout", [], |row| row.get(0)).map_err(|e| e.to_string()) + }) + .await?; + if count == 0 { + self.with_conn(move |conn| { + conn.execute("INSERT OR IGNORE INTO sidebar_layout (id, layout_json) VALUES (1, ?1)", [json]) + .map(|_| ()) + .map_err(|e| e.to_string()) + }) + .await?; } let _ = tokio::fs::rename(&path, data_dir.join("sidebar_layout.json.bak")).await; Ok(()) } } -// --------------------------------------------------------------------------- -// Helpers -// --------------------------------------------------------------------------- - -async fn persist_secret_in_tx( - tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, +fn persist_secret_in_tx( + tx: &rusqlite::Transaction<'_>, connection_id: &str, key: &str, secret: &str, ) -> Result<(), String> { if secret.is_empty() { - sqlx::query("DELETE FROM connection_secrets WHERE connection_id = ? AND key = ?") - .bind(connection_id) - .bind(key) - .execute(&mut **tx) - .await + tx.execute("DELETE FROM connection_secrets WHERE connection_id = ?1 AND key = ?2", params![connection_id, key]) .map_err(|e| e.to_string())?; } else { - sqlx::query( - "INSERT OR REPLACE INTO connection_secrets \ - (connection_id, key, secret) VALUES (?, ?, ?)", + tx.execute( + "INSERT OR REPLACE INTO connection_secrets (connection_id, key, secret) VALUES (?, ?, ?)", + params![connection_id, key, secret], ) - .bind(connection_id) - .bind(key) - .bind(secret) - .execute(&mut **tx) - .await .map_err(|e| e.to_string())?; } Ok(()) } +fn map_from_sql_err(err: serde_json::Error) -> rusqlite::Error { + rusqlite::Error::FromSqlConversionFailure(0, rusqlite::types::Type::Text, Box::new(err)) +} + #[cfg(test)] mod tests { use super::{DesktopSettings, Storage}; diff --git a/src-tauri/Cargo.toml b/src-tauri/Cargo.toml index b00369d2e..584d51550 100644 --- a/src-tauri/Cargo.toml +++ b/src-tauri/Cargo.toml @@ -21,7 +21,7 @@ serde_json = "1.0" log = "0.4" tauri = { version = "2.10.3", features = ["tray-icon"] } tauri-plugin-log = "2" -sqlx = { version = "0.8", features = ["runtime-tokio", "tls-native-tls", "mysql", "postgres", "sqlite", "json", "chrono", "uuid", "rust_decimal"] } +sqlx = { version = "0.8", features = ["runtime-tokio", "tls-native-tls", "mysql", "postgres", "json", "chrono", "uuid", "rust_decimal"] } rust_decimal = { version = "1", features = ["serde"] } tokio = { version = "1", features = ["full"] } uuid = { version = "1", features = ["v4", "serde"] } diff --git a/src-tauri/src/commands/connection.rs b/src-tauri/src/commands/connection.rs index 2d25e62d7..e50b203ef 100644 --- a/src-tauri/src/commands/connection.rs +++ b/src-tauri/src/commands/connection.rs @@ -244,10 +244,7 @@ pub async fn test_connection(state: State<'_, Arc>, config: Connection Err(e) => Err(e), }, DatabaseType::Sqlite => match db::sqlite::connect_path(&expand_tilde(&config.host)).await { - Ok(pool) => { - pool.close().await; - Ok("Connection successful".to_string()) - } + Ok(_) => Ok("Connection successful".to_string()), Err(e) => Err(e), }, DatabaseType::Redis => db::redis_driver::connect(&url).await.map(|_| "Connection successful".to_string()), @@ -439,7 +436,7 @@ pub async fn disconnect_db(state: State<'_, Arc>, connection_id: Strin let _ = p.disconnect().await; } PoolKind::Postgres(p) => p.close().await, - PoolKind::Sqlite(p) => p.close().await, + PoolKind::Sqlite(_) => {} PoolKind::Redis(_) => {} PoolKind::DuckDb(_) => {} PoolKind::MongoDb(_) => {}