diff --git a/Cargo.lock b/Cargo.lock index 4d2687bf28..03288939d7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1678,6 +1678,7 @@ name = "dbx-core" version = "0.1.0" dependencies = [ "anyhow", + "async-trait", "base64 0.22.1", "calamine", "chrono", diff --git a/crates/dbx-core/Cargo.toml b/crates/dbx-core/Cargo.toml index 9107c2e89f..4f8f4696a0 100644 --- a/crates/dbx-core/Cargo.toml +++ b/crates/dbx-core/Cargo.toml @@ -30,3 +30,4 @@ calamine = "0.30.1" odbc-api = "25" rust-gaussdb = { git = "https://github.com/t8y2/rust-gaussdb", branch = "main" } base64 = "0.22" +async-trait = "0.1" diff --git a/crates/dbx-core/src/connection.rs b/crates/dbx-core/src/connection.rs index 1d5257d5fc..3cb68e8fec 100644 --- a/crates/dbx-core/src/connection.rs +++ b/crates/dbx-core/src/connection.rs @@ -6,6 +6,7 @@ use tokio::sync::Mutex; use crate::db; use crate::db::ssh_tunnel::TunnelManager; +use crate::external; use crate::models::connection::{parse_mongo_first_host, ConnectionConfig, DatabaseType}; use crate::plugins::{PluginDriverSession, PluginRegistry}; use crate::query_cancel::RunningQueries; @@ -43,6 +44,7 @@ pub enum PoolKind { Elasticsearch(db::elasticsearch_driver::EsClient), Dameng(Arc>), Gaussdb(Arc>), + ExternalTabular(Arc), ExternalDriver { driver_id: String, config: ConnectionConfig, session: Arc }, } @@ -437,6 +439,7 @@ mod tests { ssl: false, sysdba: false, connection_string: None, + external_config: None, jdbc_driver_class: None, jdbc_driver_paths: Vec::new(), } diff --git a/crates/dbx-core/src/connection_secrets.rs b/crates/dbx-core/src/connection_secrets.rs index 1622805080..2bf3c1a68d 100644 --- a/crates/dbx-core/src/connection_secrets.rs +++ b/crates/dbx-core/src/connection_secrets.rs @@ -289,6 +289,7 @@ mod tests { ssl: false, sysdba: false, connection_string: None, + external_config: None, jdbc_driver_class: None, jdbc_driver_paths: Vec::new(), } diff --git a/crates/dbx-core/src/external/duckdb_cache.rs b/crates/dbx-core/src/external/duckdb_cache.rs new file mode 100644 index 0000000000..50bfc9bb14 --- /dev/null +++ b/crates/dbx-core/src/external/duckdb_cache.rs @@ -0,0 +1,171 @@ +use std::sync::Arc; + +use super::types::{CacheState, ExternalTableRef, ExternalTableSnapshot}; + +/// Loads an external table snapshot into an in-memory DuckDB connection. +pub fn load_snapshot_to_duckdb(con: &duckdb::Connection, snapshot: &ExternalTableSnapshot) -> Result<(), String> { + let table_name = sanitize_table_name(&snapshot.table_ref.table_name); + + con.execute(&format!("DROP TABLE IF EXISTS \"{}\"", table_name), []) + .map_err(|e| format!("Failed to drop table: {e}"))?; + + let col_defs: Vec = snapshot + .columns + .iter() + .map(|c| { + format!( + "\"{}\" {}{}", + c.name.replace('"', "\"\""), + &c.duckdb_type, + if c.is_nullable { "" } else { " NOT NULL" } + ) + }) + .collect(); + + let pk_cols: Vec<&str> = snapshot.columns.iter().filter(|c| c.is_primary_key).map(|c| c.name.as_str()).collect(); + + let mut create_sql = format!("CREATE TABLE \"{}\" ({}", table_name, col_defs.join(", ")); + if !pk_cols.is_empty() { + create_sql.push_str(&format!( + ", PRIMARY KEY ({})", + pk_cols.iter().map(|c| format!("\"{}\"", c.replace('"', "\"\""))).collect::>().join(", ") + )); + } + create_sql.push(')'); + + con.execute(&create_sql, []).map_err(|e| format!("Failed to create table: {e}"))?; + + if snapshot.rows.is_empty() { + return Ok(()); + } + + let placeholders: Vec<&str> = (0..snapshot.columns.len()).map(|_| "?").collect(); + let insert_sql = format!("INSERT INTO \"{}\" VALUES ({})", table_name, placeholders.join(", ")); + let mut stmt = con.prepare(&insert_sql).map_err(|e| format!("Failed to prepare insert: {e}"))?; + + for row in &snapshot.rows { + let params: Vec> = row + .iter() + .enumerate() + .map(|(i, val)| { + let col_type = snapshot.columns.get(i).map(|c| c.duckdb_type.as_str()).unwrap_or("VARCHAR"); + json_to_duckdb_param(val, col_type) + }) + .collect(); + + let param_refs: Vec<&dyn duckdb::ToSql> = params.iter().map(|p| p.as_ref()).collect(); + stmt.execute(param_refs.as_slice()).map_err(|e| format!("Failed to insert row: {e}"))?; + } + + Ok(()) +} + +fn json_to_duckdb_param(val: &serde_json::Value, col_type: &str) -> Box { + match val { + serde_json::Value::Null => Box::new(None::), + serde_json::Value::Bool(b) => Box::new(*b), + serde_json::Value::Number(n) => { + if col_type.contains("INT") { + Box::new(n.as_i64().unwrap_or(0)) + } else if col_type.contains("DOUBLE") || col_type.contains("FLOAT") || col_type.contains("DECIMAL") { + Box::new(n.as_f64().unwrap_or(0.0)) + } else { + Box::new(n.to_string()) + } + } + serde_json::Value::String(s) => Box::new(s.clone()), + _ => Box::new(val.to_string()), + } +} + +/// Sanitize a source table name into a stable DuckDB table identifier. +pub fn sanitize_table_name(name: &str) -> String { + let sanitized: String = name.chars().map(|c| if c.is_alphanumeric() || c == '_' { c } else { '_' }).collect(); + + if sanitized.is_empty() { + "_unnamed_".to_string() + } else if sanitized.chars().next().is_some_and(|c| c.is_ascii_digit()) { + format!("_{sanitized}") + } else { + sanitized + } +} + +/// Manages a DuckDB cache for one external source. +pub struct ExternalPool { + pub source: Arc, + pub cache: Arc>, + pub cache_state: std::sync::Mutex, + pub table_map: std::sync::Mutex>, +} + +impl ExternalPool { + pub fn new( + source: Arc, + cache: Arc>, + ) -> Self { + Self { + source, + cache, + cache_state: std::sync::Mutex::new(CacheState::Empty), + table_map: std::sync::Mutex::new(std::collections::HashMap::new()), + } + } + + pub async fn refresh_cache(&self) -> Result<(), String> { + { + let mut state = self.cache_state.lock().map_err(|e| e.to_string())?; + *state = CacheState::Loading; + } + + let tables = self.source.list_tables().await?; + let mut new_table_map = std::collections::HashMap::new(); + + for table_ref in &tables { + let snapshot = self.source.load_table(table_ref).await?; + let mut sanitized_name = sanitize_table_name(&table_ref.table_name); + + if new_table_map.contains_key(&sanitized_name) { + let base = sanitized_name.clone(); + let mut suffix = 2; + loop { + sanitized_name = format!("{base}_{suffix}"); + if !new_table_map.contains_key(&sanitized_name) { + break; + } + suffix += 1; + } + } + + let cache = self.cache.clone(); + tokio::task::spawn_blocking(move || { + let con = cache.lock().map_err(|e| e.to_string())?; + load_snapshot_to_duckdb(&con, &snapshot) + }) + .await + .map_err(|e| e.to_string())??; + + new_table_map.insert(sanitized_name, table_ref.clone()); + } + + { + let mut map = self.table_map.lock().map_err(|e| e.to_string())?; + *map = new_table_map; + } + { + let mut state = self.cache_state.lock().map_err(|e| e.to_string())?; + *state = CacheState::Fresh; + } + + Ok(()) + } +} + +impl std::fmt::Debug for ExternalPool { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ExternalPool") + .field("source", &self.source.display_name()) + .field("cache_state", &self.cache_state) + .finish() + } +} diff --git a/crates/dbx-core/src/external/mod.rs b/crates/dbx-core/src/external/mod.rs new file mode 100644 index 0000000000..1fc33bef5c --- /dev/null +++ b/crates/dbx-core/src/external/mod.rs @@ -0,0 +1,7 @@ +pub mod duckdb_cache; +pub mod traits; +pub mod types; + +pub use duckdb_cache::ExternalPool; +pub use traits::ExternalTabularSource; +pub use types::*; diff --git a/crates/dbx-core/src/external/traits.rs b/crates/dbx-core/src/external/traits.rs new file mode 100644 index 0000000000..6fdafa33f1 --- /dev/null +++ b/crates/dbx-core/src/external/traits.rs @@ -0,0 +1,21 @@ +use async_trait::async_trait; + +use super::types::{ExternalCapabilities, ExternalColumnDef, ExternalTableRef, ExternalTableSnapshot}; + +/// Trait for external tabular data sources. +#[async_trait] +pub trait ExternalTabularSource: Send + Sync + std::fmt::Debug { + fn capabilities(&self) -> ExternalCapabilities; + + async fn list_tables(&self) -> Result, String>; + + async fn get_columns(&self, table: &ExternalTableRef) -> Result, String>; + + async fn load_table(&self, table: &ExternalTableRef) -> Result; + + async fn source_version(&self, table: &ExternalTableRef) -> Result; + + async fn test_connection(&self) -> Result; + + fn display_name(&self) -> String; +} diff --git a/crates/dbx-core/src/external/types.rs b/crates/dbx-core/src/external/types.rs new file mode 100644 index 0000000000..75161ae4e8 --- /dev/null +++ b/crates/dbx-core/src/external/types.rs @@ -0,0 +1,57 @@ +use serde::{Deserialize, Serialize}; + +/// Reference to a specific table within an external source. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash)] +pub struct ExternalTableRef { + pub source_id: String, + pub table_name: String, + pub display_name: String, +} + +/// Column definition for an external table snapshot. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ExternalColumnDef { + pub name: String, + pub duckdb_type: String, + pub is_nullable: bool, + pub is_primary_key: bool, + pub comment: Option, +} + +/// A full table snapshot loaded from an external source. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ExternalTableSnapshot { + pub table_ref: ExternalTableRef, + pub columns: Vec, + pub rows: Vec>, + pub source_version: String, +} + +/// Capability flags for an external data source. +#[derive(Debug, Clone, Serialize, Deserialize, Default)] +pub struct ExternalCapabilities { + pub can_read: bool, + pub can_write: bool, + pub can_append: bool, + pub can_delete_rows: bool, + pub supports_multiple_tables: bool, + pub supports_refresh: bool, + pub supports_file_watch: bool, + pub supports_schema_detection: bool, +} + +/// Cache state tracking for external source snapshots. +#[derive(Debug, Clone, PartialEq)] +pub enum CacheState { + Empty, + Fresh, + Stale, + Loading, + Error(String), +} + +impl Default for CacheState { + fn default() -> Self { + Self::Empty + } +} diff --git a/crates/dbx-core/src/lib.rs b/crates/dbx-core/src/lib.rs index 30697b8bd2..cd34fdf92e 100644 --- a/crates/dbx-core/src/lib.rs +++ b/crates/dbx-core/src/lib.rs @@ -2,6 +2,7 @@ pub mod ai; pub mod connection; pub mod connection_secrets; pub mod db; +pub mod external; pub mod history; pub mod models; pub mod mongo_ops; diff --git a/crates/dbx-core/src/models/connection.rs b/crates/dbx-core/src/models/connection.rs index ea938b9efb..5315673246 100644 --- a/crates/dbx-core/src/models/connection.rs +++ b/crates/dbx-core/src/models/connection.rs @@ -43,6 +43,9 @@ pub struct ConnectionConfig { pub sysdba: bool, #[serde(default)] pub connection_string: Option, + /// Typed configuration for external tabular sources. + #[serde(default)] + pub external_config: Option, #[serde(default)] pub jdbc_driver_class: Option, #[serde(default)] @@ -57,7 +60,7 @@ pub fn default_ssh_connect_timeout_secs() -> u64 { 5 } -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq)] #[serde(rename_all = "lowercase")] pub enum DatabaseType { Mysql, @@ -397,6 +400,7 @@ mod tests { ssl: false, sysdba: false, connection_string: None, + external_config: None, jdbc_driver_class: None, jdbc_driver_paths: Vec::new(), } diff --git a/crates/dbx-core/src/query.rs b/crates/dbx-core/src/query.rs index 54573dbc54..e6c1672ed1 100644 --- a/crates/dbx-core/src/query.rs +++ b/crates/dbx-core/src/query.rs @@ -237,6 +237,22 @@ pub async fn do_execute( .await .map(truncate_result) } + PoolKind::ExternalTabular(ext_pool) => { + if !starts_with_executable_sql_keyword(sql, &["SELECT", "WITH", "SHOW", "DESCRIBE", "EXPLAIN", "PRAGMA"]) { + return Err("External data sources are read-only. Only SELECT queries are supported.".to_string()); + } + let con = ext_pool.cache.clone(); + let sql = sql.to_string(); + drop(connections); + wait_for_query(cancel_token, async move { + let task = tokio::task::spawn_blocking(move || { + let con = con.lock().map_err(|e| e.to_string())?; + duckdb_execute(&con, &sql) + }); + task.await.map_err(|e| e.to_string())? + }) + .await + } PoolKind::ExternalDriver { config, session, .. } => { let config = config.clone(); let session = session.clone(); @@ -414,6 +430,7 @@ pub async fn execute_statements_in_transaction( | PoolKind::MongoDb(_) | PoolKind::Oracle(_) | PoolKind::Elasticsearch(_) + | PoolKind::ExternalTabular(_) | PoolKind::ExternalDriver { .. } => TxPath::None, }) }; diff --git a/crates/dbx-core/src/schema.rs b/crates/dbx-core/src/schema.rs index f86b104ae4..45693f3e1f 100644 --- a/crates/dbx-core/src/schema.rs +++ b/crates/dbx-core/src/schema.rs @@ -72,6 +72,16 @@ pub fn extract_duckdb( } } +pub fn extract_external( + connections: &HashMap, + key: &str, +) -> Option> { + match connections.get(key)? { + PoolKind::ExternalTabular(pool) => Some(pool.clone()), + _ => None, + } +} + pub fn extract_sqlserver( connections: &HashMap, key: &str, @@ -122,6 +132,9 @@ pub fn extract_gaussdb( pub async fn list_databases_core(state: &AppState, connection_id: &str) -> Result, String> { { let connections = state.connections.lock().await; + if extract_external(&connections, connection_id).is_some() { + return Ok(vec![db::DatabaseInfo { name: "main".to_string() }]); + } if let Some(PoolKind::ExternalDriver { config, session, .. }) = connections.get(connection_id) { let config = config.clone(); let session = session.clone(); @@ -230,6 +243,16 @@ pub async fn list_tables_core( { let connections = state.connections.lock().await; + if let Some(ext_pool) = extract_external(&connections, &pool_key) { + drop(connections); + let cache = ext_pool.cache.clone(); + return tokio::task::spawn_blocking(move || { + let con = cache.lock().map_err(|e| e.to_string())?; + duckdb_query_tables(&con) + }) + .await + .map_err(|e| e.to_string())?; + } if let Some(PoolKind::ExternalDriver { config, session, .. }) = connections.get(&pool_key) { let config = config.clone(); let session = session.clone(); @@ -301,6 +324,17 @@ pub async fn get_columns_core( { let connections = state.connections.lock().await; + if let Some(ext_pool) = extract_external(&connections, &pool_key) { + drop(connections); + let cache = ext_pool.cache.clone(); + let table = table.to_string(); + return tokio::task::spawn_blocking(move || { + let con = cache.lock().map_err(|e| e.to_string())?; + duckdb_query_columns(&con, &table) + }) + .await + .map_err(|e| e.to_string())?; + } if let Some(PoolKind::ExternalDriver { config, session, .. }) = connections.get(&pool_key) { let config = config.clone(); let session = session.clone(); diff --git a/crates/dbx-core/src/transfer.rs b/crates/dbx-core/src/transfer.rs index b981d40783..297a3a55af 100644 --- a/crates/dbx-core/src/transfer.rs +++ b/crates/dbx-core/src/transfer.rs @@ -569,6 +569,17 @@ pub async fn execute_on_pool(state: &AppState, pool_key: &str, sql: &str) -> Res .await .map_err(|e| e.to_string())? } + PoolKind::ExternalTabular(ext_pool) => { + let con = ext_pool.cache.clone(); + let sql = sql.to_string(); + drop(connections); + tokio::task::spawn_blocking(move || { + let con = con.lock().map_err(|e| e.to_string())?; + crate::query::duckdb_execute(&con, &sql) + }) + .await + .map_err(|e| e.to_string())? + } _ => Err("Unsupported database type for transfer".to_string()), } } @@ -603,6 +614,18 @@ pub async fn get_columns_for_transfer( .map_err(|e| e.to_string())?; } + if let Some(PoolKind::ExternalTabular(ext_pool)) = connections.get(pool_key) { + let con = ext_pool.cache.clone(); + drop(connections); + let table = table.to_string(); + return tokio::task::spawn_blocking(move || { + let con = con.lock().map_err(|e| e.to_string())?; + crate::schema::duckdb_query_columns(&con, &table) + }) + .await + .map_err(|e| e.to_string())?; + } + if let Some(PoolKind::ClickHouse(client)) = connections.get(pool_key) { let client = client.clone(); let database = database.to_string(); diff --git a/src-tauri/src/commands/connection.rs b/src-tauri/src/commands/connection.rs index 756c9c45a7..f98221f0f8 100644 --- a/src-tauri/src/commands/connection.rs +++ b/src-tauri/src/commands/connection.rs @@ -260,6 +260,7 @@ pub async fn disconnect_db(state: State<'_, Arc>, connection_id: Strin PoolKind::Elasticsearch(_) => {} PoolKind::Dameng(_) => {} PoolKind::Gaussdb(_) => {} + PoolKind::ExternalTabular(_) => {} PoolKind::ExternalDriver { .. } => {} } }