diff --git a/crates/dbx-core/src/agent_connection.rs b/crates/dbx-core/src/agent_connection.rs index ed191afd4b..c710bec846 100644 --- a/crates/dbx-core/src/agent_connection.rs +++ b/crates/dbx-core/src/agent_connection.rs @@ -1,6 +1,6 @@ use percent_encoding::{percent_decode_str, utf8_percent_encode, NON_ALPHANUMERIC}; -use crate::models::connection::{default_redis_key_separator, ConnectionConfig, DatabaseType}; +use crate::models::connection::{ConnectionConfig, DatabaseType}; use crate::path_utils::expand_tilde; pub fn agent_connect_params(config: &ConnectionConfig, host: &str, port: u16, database: &str) -> serde_json::Value { @@ -392,6 +392,7 @@ mod tests { use super::*; use crate::models::connection::{ default_connect_timeout_secs, default_idle_timeout_secs, default_query_timeout_secs, + default_redis_key_separator, }; fn config(db_type: DatabaseType, database: Option<&str>) -> ConnectionConfig { diff --git a/crates/dbx-core/src/connection.rs b/crates/dbx-core/src/connection.rs index 37a46f62d2..c234d586ba 100644 --- a/crates/dbx-core/src/connection.rs +++ b/crates/dbx-core/src/connection.rs @@ -17,10 +17,8 @@ use crate::db; use crate::db::agent_driver::AgentMethod; use crate::db::proxy_tunnel::ProxyTunnelManager; use crate::db::ssh_tunnel::TunnelManager; -use crate::external; use crate::models::connection::{ - default_redis_key_separator, parse_jdbc_host_port, parse_mongo_first_host, rewrite_jdbc_url_host, ConnectionConfig, - DatabaseType, + parse_jdbc_host_port, parse_mongo_first_host, rewrite_jdbc_url_host, ConnectionConfig, DatabaseType, }; use crate::path_utils::expand_tilde; use crate::plugins::{PluginDriverSession, PluginRegistry, PluginRuntimeEnv}; @@ -1107,30 +1105,23 @@ fn native_postgres_url_config(config: &ConnectionConfig) -> Option bool { - #[cfg(feature = "duckdb-bundled")] - { - let left = expand_tilde(left); - let right = expand_tilde(right); + let left = expand_tilde(left); + let right = expand_tilde(right); - if db::duckdb_driver::is_memory_database_path(&left) || db::duckdb_driver::is_memory_database_path(&right) { - return left.trim().eq_ignore_ascii_case(right.trim()); - } - - if let (Ok(left_path), Ok(right_path)) = (std::fs::canonicalize(&left), std::fs::canonicalize(&right)) { - return left_path == right_path; - } - - if cfg!(windows) { - left.eq_ignore_ascii_case(&right) - } else { - left == right - } + if db::duckdb_driver::is_memory_database_path(&left) || db::duckdb_driver::is_memory_database_path(&right) { + return left.trim().eq_ignore_ascii_case(right.trim()); } - #[cfg(not(feature = "duckdb-bundled"))] - { - let _ = (left, right); - false + + if let (Ok(left_path), Ok(right_path)) = (std::fs::canonicalize(&left), std::fs::canonicalize(&right)) { + return left_path == right_path; + } + + if cfg!(windows) { + left.eq_ignore_ascii_case(&right) + } else { + left == right } } @@ -2004,6 +1995,7 @@ mod tests { let _ = std::fs::remove_dir_all(dir); } + #[cfg(feature = "duckdb-bundled")] #[tokio::test] async fn duckdb_client_session_reuses_base_pool_to_avoid_file_locks() { let (state, dir) = test_app_state().await; diff --git a/crates/dbx-core/src/db/turso_driver.rs b/crates/dbx-core/src/db/turso_driver.rs index 1f4f9acd13..c7ccf7940d 100644 --- a/crates/dbx-core/src/db/turso_driver.rs +++ b/crates/dbx-core/src/db/turso_driver.rs @@ -633,7 +633,6 @@ mod tests { assert_eq!(value_by_column(&columns, &row, "notnull").as_deref(), Some("1")); } - #[test] #[test] fn split_sql_statements_splits_on_semicolons() { assert_eq!(split_sql_statements("SELECT 1"), vec!["SELECT 1"]); diff --git a/crates/dbx-core/src/query.rs b/crates/dbx-core/src/query.rs index dade9bfdab..33747e6443 100644 --- a/crates/dbx-core/src/query.rs +++ b/crates/dbx-core/src/query.rs @@ -1,3 +1,4 @@ +#[cfg(feature = "duckdb-bundled")] use chrono::{DateTime, Duration as ChronoDuration, NaiveDate, NaiveDateTime, NaiveTime, Utc}; #[cfg(feature = "duckdb-bundled")] use duckdb::types::{TimeUnit, Value, ValueRef}; @@ -9,8 +10,10 @@ use tokio_util::sync::CancellationToken; use crate::connection::{AppState, PoolKind}; use crate::db; -use crate::models::connection::{default_redis_key_separator, DatabaseType}; -use crate::sql::{split_sql_batches, split_sql_statements, starts_with_duckdb_result_sql_keyword}; +use crate::models::connection::DatabaseType; +#[cfg(feature = "duckdb-bundled")] +use crate::sql::starts_with_duckdb_result_sql_keyword; +use crate::sql::{split_sql_batches, split_sql_statements}; pub const QUERY_TIMEOUT: Duration = Duration::from_secs(30); pub const MAX_ROWS: usize = 10000; @@ -304,6 +307,7 @@ fn duckdb_time_unit_to_nanos(unit: TimeUnit, value: i64) -> Option { } } +#[cfg(feature = "duckdb-bundled")] fn format_naive_datetime(value: NaiveDateTime) -> String { if value.and_utc().timestamp_subsec_nanos() == 0 { value.format("%Y-%m-%d %H:%M:%S").to_string() @@ -312,6 +316,7 @@ fn format_naive_datetime(value: NaiveDateTime) -> String { } } +#[cfg(feature = "duckdb-bundled")] fn format_temporal_without_empty_fraction(value: String) -> String { if !value.contains('.') { return value; @@ -616,7 +621,7 @@ pub async fn do_execute( options: QueryExecutionOptions, ) -> Result { let query_timeout = resolve_query_timeout(options.timeout_secs); - let (duckdb_attached_names, conn_name_if_readonly) = { + let (_duckdb_attached_names, conn_name_if_readonly) = { let configs = state.configs.read().await; let config = crate::connection::config_for_pool_key(pool_key, &configs); let attached = config @@ -644,7 +649,7 @@ pub async fn do_execute( } let sql = sql.to_string(); let database = database.map(str::to_string); - let attached_names = duckdb_attached_names; + let attached_names = _duckdb_attached_names; let max_rows = options.max_rows; drop(connections); wait_for_query_opt(cancel_token, query_timeout, async move { @@ -1615,7 +1620,7 @@ async fn exec_tx_none_inner( #[cfg(test)] mod tests { use super::*; - use crate::models::connection::{ConnectionConfig, DatabaseType}; + use crate::models::connection::{default_redis_key_separator, ConnectionConfig, DatabaseType}; #[tokio::test] async fn wait_for_query_returns_cancelled_when_token_is_cancelled() { diff --git a/crates/dbx-core/src/schema.rs b/crates/dbx-core/src/schema.rs index b1d8d292e3..35c6b09643 100644 --- a/crates/dbx-core/src/schema.rs +++ b/crates/dbx-core/src/schema.rs @@ -1,12 +1,30 @@ -use crate::connection::{connection_url_for_endpoint, database_connection_config, AppState, MysqlMode, PoolKind}; +use crate::connection::{connection_url_for_endpoint, database_connection_config, AppState, PoolKind}; use crate::db; -use crate::models::connection::{default_redis_key_separator, ConnectionConfig, DatabaseType}; +use crate::models::connection::{ConnectionConfig, DatabaseType}; use crate::query::{agent_execute_query_params, QueryExecutionOptions}; -use std::collections::HashSet; use std::future::Future; use std::sync::Arc; use std::time::Duration; +#[cfg(feature = "duckdb-bundled")] +mod duckdb_metadata; +mod normalization; +mod providers; + +#[cfg(feature = "duckdb-bundled")] +use self::duckdb_metadata::duckdb_attached_database_names; +#[cfg(feature = "duckdb-bundled")] +pub use self::duckdb_metadata::{ + duckdb_attach_database, duckdb_list_databases, duckdb_list_databases_with_attached, duckdb_list_schemas, + duckdb_list_schemas_with_attached, duckdb_primary_catalog, duckdb_query_columns, duckdb_query_columns_in_database, + duckdb_query_columns_in_database_with_attached, duckdb_query_tables, duckdb_query_tables_in_database, + duckdb_query_tables_in_database_with_attached, +}; +use self::normalization::{ + deduplicate_column_infos, filter_completion_objects, filter_objects_by_types, filter_table_infos_for_config, + filter_yashandb_recyclebin_objects, +}; + macro_rules! extract_pool { ($connections:expr, $key:expr, $variant:ident) => { $connections.get($key).and_then(|v| match v { @@ -16,16 +34,6 @@ macro_rules! extract_pool { }; } -macro_rules! dispatch_mysql { - ($p:expr, $mode:expr, $mysql:path, $ob:path $(, $arg:expr)*) => { - if *$mode == MysqlMode::OceanBaseOracle { - $ob($p $(, $arg)*).await - } else { - $mysql($p $(, $arg)*).await - } - }; -} - macro_rules! try_sqlserver { ($connections:expr, $pool_key:expr, $method:ident $(, $arg:expr)*) => { if let Some(client) = extract_pool!(&$connections, $pool_key, SqlServer) { @@ -46,221 +54,6 @@ macro_rules! try_agent { }; } -#[cfg(feature = "duckdb-bundled")] -pub fn duckdb_query_tables(con: &duckdb::Connection) -> Result, String> { - duckdb_query_tables_in_database(con, "main", "main") -} - -#[cfg(feature = "duckdb-bundled")] -pub fn duckdb_query_tables_in_database( - con: &duckdb::Connection, - database: &str, - schema: &str, -) -> Result, String> { - duckdb_query_tables_in_database_with_attached(con, database, schema, &[]) -} - -#[cfg(feature = "duckdb-bundled")] -pub fn duckdb_query_tables_in_database_with_attached( - con: &duckdb::Connection, - database: &str, - schema: &str, - attached_names: &[String], -) -> Result, String> { - let database = duckdb_catalog_name(con, database, attached_names)?; - let mut stmt = con.prepare( - "SELECT table_name, table_type FROM information_schema.tables WHERE table_catalog = ? AND table_schema = ? ORDER BY table_name" - ).map_err(|e| e.to_string())?; - let rows = stmt - .query_map((database.as_str(), schema), |row| { - Ok(db::TableInfo { - name: row.get::<_, String>(0)?, - table_type: row.get::<_, String>(1)?, - comment: None, - parent_schema: None, - parent_name: None, - }) - }) - .map_err(|e| e.to_string())?; - Ok(rows.filter_map(|r| r.ok()).collect()) -} - -#[cfg(feature = "duckdb-bundled")] -pub fn duckdb_attach_database(con: &duckdb::Connection, name: &str, path: &str) -> Result<(), String> { - let name = name.trim(); - let path = path.trim(); - if name.is_empty() || path.is_empty() { - return Err("DuckDB attached database name and path are required".to_string()); - } - let sql = format!("ATTACH {} AS {}", duckdb_quote_string(path), duckdb_quote_ident(name)); - con.execute_batch(&sql).map_err(|e| e.to_string()) -} - -#[cfg(feature = "duckdb-bundled")] -pub fn duckdb_list_databases(con: &duckdb::Connection) -> Result, String> { - duckdb_list_databases_with_attached(con, &[]) -} - -#[cfg(feature = "duckdb-bundled")] -pub fn duckdb_list_databases_with_attached( - con: &duckdb::Connection, - attached_names: &[String], -) -> Result, String> { - let primary = duckdb_primary_catalog(con, attached_names)?; - let mut stmt = con.prepare("SHOW DATABASES").map_err(|e| e.to_string())?; - let rows = stmt - .query_map([], |row| { - let name = row.get::<_, String>(0)?; - Ok(db::DatabaseInfo { name: if name == primary { "main".to_string() } else { name } }) - }) - .map_err(|e| e.to_string())?; - Ok(rows.filter_map(|row| row.ok()).collect()) -} - -#[cfg(feature = "duckdb-bundled")] -pub fn duckdb_list_schemas(con: &duckdb::Connection, database: &str) -> Result, String> { - duckdb_list_schemas_with_attached(con, database, &[]) -} - -#[cfg(feature = "duckdb-bundled")] -pub fn duckdb_list_schemas_with_attached( - con: &duckdb::Connection, - database: &str, - attached_names: &[String], -) -> Result, String> { - let database = duckdb_catalog_name(con, database, attached_names)?; - let mut stmt = con - .prepare( - "SELECT schema_name FROM information_schema.schemata WHERE catalog_name = ? AND schema_name NOT IN ('information_schema', 'pg_catalog') ORDER BY schema_name", - ) - .map_err(|e| e.to_string())?; - let rows = stmt.query_map([database.as_str()], |row| row.get::<_, String>(0)).map_err(|e| e.to_string())?; - Ok(rows.filter_map(|r| r.ok()).collect()) -} - -#[cfg(feature = "duckdb-bundled")] -fn duckdb_catalog_name(con: &duckdb::Connection, database: &str, attached_names: &[String]) -> Result { - if database.trim().is_empty() || database == "main" { - return duckdb_primary_catalog(con, attached_names); - } - Ok(database.to_string()) -} - -#[cfg(feature = "duckdb-bundled")] -pub fn duckdb_primary_catalog(con: &duckdb::Connection, attached_names: &[String]) -> Result { - if attached_names.is_empty() { - return duckdb_current_database(con); - } - let attached: std::collections::HashSet = attached_names.iter().map(|name| name.to_lowercase()).collect(); - let mut stmt = con.prepare("SHOW DATABASES").map_err(|e| e.to_string())?; - let rows = stmt.query_map([], |row| row.get::<_, String>(0)).map_err(|e| e.to_string())?; - for row in rows { - let name = row.map_err(|e| e.to_string())?; - if !attached.contains(&name.to_lowercase()) { - return Ok(name); - } - } - duckdb_current_database(con) -} - -#[cfg(feature = "duckdb-bundled")] -fn duckdb_current_database(con: &duckdb::Connection) -> Result { - con.query_row("SELECT current_database()", [], |row| row.get::<_, String>(0)).map_err(|e| e.to_string()) -} - -#[cfg(feature = "duckdb-bundled")] -fn duckdb_quote_ident(value: &str) -> String { - format!("\"{}\"", value.replace('"', "\"\"")) -} - -#[cfg(feature = "duckdb-bundled")] -fn duckdb_quote_string(value: &str) -> String { - format!("'{}'", value.replace('\'', "''")) -} - -#[cfg(feature = "duckdb-bundled")] -pub fn duckdb_query_columns(con: &duckdb::Connection, table: &str) -> Result, String> { - duckdb_query_columns_in_database(con, "main", "main", table) -} - -#[cfg(feature = "duckdb-bundled")] -pub fn duckdb_query_columns_in_database( - con: &duckdb::Connection, - database: &str, - schema: &str, - table: &str, -) -> Result, String> { - duckdb_query_columns_in_database_with_attached(con, database, schema, table, &[]) -} - -#[cfg(feature = "duckdb-bundled")] -pub fn duckdb_query_columns_in_database_with_attached( - con: &duckdb::Connection, - database: &str, - schema: &str, - table: &str, - attached_names: &[String], -) -> Result, String> { - let database = duckdb_catalog_name(con, database, attached_names)?; - let mut pk_stmt = con - .prepare( - "SELECT kcu.column_name - FROM information_schema.table_constraints tc - JOIN information_schema.key_column_usage kcu - ON tc.constraint_name = kcu.constraint_name - AND tc.table_schema = kcu.table_schema - AND tc.table_name = kcu.table_name - WHERE tc.constraint_type = 'PRIMARY KEY' - AND tc.table_catalog = ? - AND tc.table_schema = ? - AND tc.table_name = ? - ORDER BY kcu.ordinal_position", - ) - .map_err(|e| e.to_string())?; - let pk_rows = pk_stmt - .query_map((database.as_str(), schema, table), |row| row.get::<_, String>(0)) - .map_err(|e| e.to_string())?; - let primary_keys: std::collections::HashSet = pk_rows.filter_map(|r| r.ok()).collect(); - - let mut stmt = con - .prepare( - "SELECT column_name, data_type, is_nullable, column_default - FROM information_schema.columns - WHERE table_catalog = ? AND table_schema = ? AND table_name = ? - ORDER BY ordinal_position", - ) - .map_err(|e| e.to_string())?; - let rows = stmt - .query_map((database.as_str(), schema, table), |row| { - let name = row.get::<_, String>(0)?; - Ok(db::ColumnInfo { - is_primary_key: primary_keys.contains(&name), - name, - data_type: row.get::<_, String>(1)?, - is_nullable: row.get::<_, String>(2).unwrap_or_default() == "YES", - column_default: row.get::<_, Option>(3)?, - extra: None, - comment: None, - numeric_precision: None, - numeric_scale: None, - character_maximum_length: None, - }) - }) - .map_err(|e| e.to_string())?; - Ok(rows.filter_map(|r| r.ok()).collect()) -} - -#[cfg(feature = "duckdb-bundled")] -async fn duckdb_attached_database_names(state: &AppState, connection_id: &str) -> Vec { - state - .configs - .read() - .await - .get(connection_id) - .map(|config| config.attached_databases.iter().map(|database| database.name.clone()).collect()) - .unwrap_or_default() -} - fn clickhouse_metadata_database<'a>(database: &'a str, schema: &'a str) -> &'a str { if database.is_empty() { schema @@ -317,22 +110,13 @@ async fn list_databases_once(state: &AppState, connection_id: &str) -> Result { - db::mysql::list_databases_show(p).await - } - PoolKind::Mysql(p, mode) => dispatch_mysql!(p, mode, db::mysql::list_databases, db::ob_oracle::list_databases), - PoolKind::Postgres(p) => db::postgres::list_databases(p).await, - PoolKind::Sqlite(p) => db::sqlite::list_databases(p).await, - PoolKind::Rqlite(client) => db::rqlite_driver::list_databases(client).await, - PoolKind::Turso(client) => db::turso_driver::list_databases(client).await, - #[cfg(feature = "duckdb-bundled")] - PoolKind::DuckDb(con) => { - let con = con.lock().map_err(|e| e.to_string())?; - duckdb_list_databases_with_attached(&con, &duckdb_attached_names) - } - _ => Ok(vec![]), + #[cfg(feature = "duckdb-bundled")] + if let PoolKind::DuckDb(con) = pool { + let con = con.lock().map_err(|e| e.to_string())?; + return duckdb_list_databases_with_attached(&con, &duckdb_attached_names); } + + providers::native::list_databases(pool, db_config.as_ref()).await } pub async fn list_schemas_core(state: &AppState, connection_id: &str, database: &str) -> Result, String> { @@ -402,16 +186,14 @@ async fn list_schemas_once(state: &AppState, connection_id: &str, database: &str let connections = state.connections.read().await; let pool = connections.get(&pool_key).ok_or("Pool not found")?; - match pool { - PoolKind::Postgres(p) => db::postgres::list_schemas(p).await, - #[cfg(feature = "duckdb-bundled")] - PoolKind::DuckDb(con) => { - let duckdb_attached_names = duckdb_attached_database_names(state, connection_id).await; - let con = con.lock().map_err(|e| e.to_string())?; - duckdb_list_schemas_with_attached(&con, database, &duckdb_attached_names) - } - _ => Ok(vec![]), + #[cfg(feature = "duckdb-bundled")] + if let PoolKind::DuckDb(con) = pool { + let duckdb_attached_names = duckdb_attached_database_names(state, connection_id).await; + let con = con.lock().map_err(|e| e.to_string())?; + return duckdb_list_schemas_with_attached(&con, database, &duckdb_attached_names); } + + providers::native::list_schemas(pool).await } pub async fn list_tables_core( @@ -536,126 +318,18 @@ async fn list_tables_once( let connections = state.connections.read().await; let pool = connections.get(&pool_key).ok_or("Pool not found")?; - match pool { - PoolKind::Mysql(p, _) if db_config.as_ref().is_some_and(is_doris_family_config) => { - db::mysql::list_tables_show(p, database) - .await - .map(|tables| filter_table_infos_for_config(tables, filter, limit, db_config.as_ref())) - } - PoolKind::Mysql(p, mode) => { - dispatch_mysql!(p, mode, db::mysql::list_tables, db::ob_oracle::list_tables, schema) - .map(|tables| filter_table_infos_for_config(tables, filter, limit, db_config.as_ref())) - } - PoolKind::Postgres(p) => db::postgres::list_tables(p, schema) - .await - .map(|tables| filter_table_infos_for_config(tables, filter, limit, db_config.as_ref())), - PoolKind::Sqlite(p) => db::sqlite::list_tables(p, schema) - .await - .map(|tables| filter_table_infos_for_config(tables, filter, limit, db_config.as_ref())), - PoolKind::Rqlite(client) => db::rqlite_driver::list_tables(client, schema) - .await - .map(|tables| filter_table_infos_for_config(tables, filter, limit, db_config.as_ref())), - PoolKind::Turso(client) => db::turso_driver::list_tables(client, schema) - .await - .map(|tables| filter_table_infos_for_config(tables, filter, limit, db_config.as_ref())), - PoolKind::MongoDb(client) => db::mongo_driver::list_collections(client, database) - .await - .map(|names| collection_names_to_tables(names, "COLLECTION")) - .map(|tables| filter_table_infos_for_config(tables, filter, limit, db_config.as_ref())), - PoolKind::Elasticsearch(client) => db::elasticsearch_driver::list_indices(client) - .await - .map(|names| collection_names_to_tables(names, "INDEX")) - .map(|tables| filter_table_infos_for_config(tables, filter, limit, db_config.as_ref())), - _ => Ok(vec![]), - } -} - -fn collection_names_to_tables(names: Vec, table_type: &str) -> Vec { - names - .into_iter() - .map(|name| db::TableInfo { - name, - table_type: table_type.to_string(), - comment: None, - parent_schema: None, - parent_name: None, - }) - .collect() -} - -fn filter_table_infos(tables: Vec, filter: Option<&str>, limit: Option) -> Vec { - let filter = filter.unwrap_or("").to_lowercase(); - let limit = limit.unwrap_or(usize::MAX); - tables - .into_iter() - .filter(|table| filter.is_empty() || table.name.to_lowercase().contains(&filter)) - .take(limit) - .collect() -} - -fn filter_table_infos_for_config( - tables: Vec, - filter: Option<&str>, - limit: Option, - config: Option<&ConnectionConfig>, -) -> Vec { - filter_table_infos(filter_yashandb_recyclebin_tables(tables, config), filter, limit) -} - -fn filter_yashandb_recyclebin_tables( - tables: Vec, - config: Option<&ConnectionConfig>, -) -> Vec { - if !is_yashandb_config(config) { - return tables; - } - tables.into_iter().filter(|table| !is_recyclebin_object_name(&table.name)).collect() -} - -fn filter_yashandb_recyclebin_objects( - objects: Vec, - config: Option<&ConnectionConfig>, -) -> Vec { - if !is_yashandb_config(config) { - return objects; - } - objects.into_iter().filter(|object| !is_recyclebin_object_name(&object.name)).collect() -} - -fn is_yashandb_config(config: Option<&ConnectionConfig>) -> bool { - config.is_some_and(|config| config.db_type == DatabaseType::Yashandb) -} - -fn is_recyclebin_object_name(name: &str) -> bool { - name.to_ascii_uppercase().starts_with("BIN$") + providers::native::list_tables(pool, db_config.as_ref(), database, schema) + .await + .map(|tables| filter_table_infos_for_config(tables, filter, limit, db_config.as_ref())) } #[cfg(test)] mod tests { - use super::{ - clickhouse_metadata_database, deduplicate_column_infos, filter_objects_by_types, - filter_yashandb_recyclebin_objects, filter_yashandb_recyclebin_tables, - is_agent_postgres_metadata_fallback_config, - }; + use super::{clickhouse_metadata_database, is_agent_postgres_metadata_fallback_config}; #[cfg(feature = "duckdb-bundled")] use super::{duckdb_attach_database, duckdb_list_databases, duckdb_query_tables_in_database}; use crate::models::connection::{default_redis_key_separator, ConnectionConfig, DatabaseType}; - fn test_column(name: &str, comment: Option<&str>, is_primary_key: bool) -> super::db::ColumnInfo { - super::db::ColumnInfo { - name: name.to_string(), - data_type: "VARCHAR".to_string(), - is_nullable: true, - column_default: None, - is_primary_key, - extra: None, - comment: comment.map(|value| value.to_string()), - numeric_precision: None, - numeric_scale: None, - character_maximum_length: None, - } - } - fn test_connection_config(db_type: DatabaseType) -> ConnectionConfig { ConnectionConfig { id: "test".to_string(), @@ -747,23 +421,6 @@ mod tests { assert_eq!(clickhouse_metadata_database("default", "testdb"), "default"); } - #[test] - fn deduplicates_columns_and_preserves_later_comment() { - let columns = deduplicate_column_infos(vec![ - test_column("ID", None, false), - test_column("ID", Some("源主键"), true), - test_column("TFBH", Some(""), false), - test_column("TFBH", Some("台账编号"), false), - ]); - - assert_eq!(columns.len(), 2); - assert_eq!(columns[0].name, "ID"); - assert_eq!(columns[0].comment.as_deref(), Some("源主键")); - assert!(columns[0].is_primary_key); - assert_eq!(columns[1].name, "TFBH"); - assert_eq!(columns[1].comment.as_deref(), Some("台账编号")); - } - #[test] fn postgres_like_agent_metadata_fallback_targets_pg_compatible_agents() { assert!(is_agent_postgres_metadata_fallback_config(&test_connection_config(DatabaseType::Kingbase))); @@ -772,107 +429,6 @@ mod tests { assert!(!is_agent_postgres_metadata_fallback_config(&test_connection_config(DatabaseType::Postgres))); assert!(!is_agent_postgres_metadata_fallback_config(&test_connection_config(DatabaseType::Mysql))); } - - #[test] - fn filters_list_objects_by_normalized_object_types() { - let objects = vec![ - super::db::ObjectInfo { - name: "orders".to_string(), - object_type: "BASE TABLE".to_string(), - schema: None, - comment: None, - created_at: None, - updated_at: None, - parent_schema: None, - parent_name: None, - }, - super::db::ObjectInfo { - name: "active_orders".to_string(), - object_type: "MATERIALIZED VIEW".to_string(), - schema: None, - comment: None, - created_at: None, - updated_at: None, - parent_schema: None, - parent_name: None, - }, - super::db::ObjectInfo { - name: "payroll".to_string(), - object_type: "PACKAGE BODY".to_string(), - schema: None, - comment: None, - created_at: None, - updated_at: None, - parent_schema: None, - parent_name: None, - }, - ]; - - let filtered = filter_objects_by_types(objects, Some(&["VIEW".to_string(), "PACKAGE_BODY".to_string()])); - - assert_eq!( - filtered.iter().map(|object| object.name.as_str()).collect::>(), - ["active_orders", "payroll"] - ); - } - - #[test] - fn filters_yashandb_recyclebin_tables() { - let tables = vec![ - super::db::TableInfo { - name: "USERS".to_string(), - table_type: "TABLE".to_string(), - comment: None, - parent_schema: None, - parent_name: None, - }, - super::db::TableInfo { - name: "BIN$abc123==$0".to_string(), - table_type: "TABLE".to_string(), - comment: None, - parent_schema: None, - parent_name: None, - }, - ]; - - let filtered = - filter_yashandb_recyclebin_tables(tables.clone(), Some(&test_connection_config(DatabaseType::Yashandb))); - let oracle = filter_yashandb_recyclebin_tables(tables, Some(&test_connection_config(DatabaseType::Oracle))); - - assert_eq!(filtered.iter().map(|table| table.name.as_str()).collect::>(), ["USERS"]); - assert_eq!(oracle.len(), 2); - } - - #[test] - fn filters_yashandb_recyclebin_objects() { - let objects = vec![ - super::db::ObjectInfo { - name: "ORDERS".to_string(), - object_type: "TABLE".to_string(), - schema: Some("HR".to_string()), - comment: None, - created_at: None, - updated_at: None, - parent_schema: None, - parent_name: None, - }, - super::db::ObjectInfo { - name: "bin$deleted".to_string(), - object_type: "TABLE".to_string(), - schema: Some("HR".to_string()), - comment: None, - created_at: None, - updated_at: None, - parent_schema: None, - parent_name: None, - }, - ]; - - let filtered = - filter_yashandb_recyclebin_objects(objects, Some(&test_connection_config(DatabaseType::Yashandb))); - - assert_eq!(filtered.iter().map(|object| object.name.as_str()).collect::>(), ["ORDERS"]); - } } pub async fn list_objects_core( @@ -1007,65 +563,24 @@ async fn list_objects_once_unfiltered( let connections = state.connections.read().await; let pool = connections.get(&pool_key).ok_or("Pool not found")?; - match pool { - PoolKind::Mysql(p, mode) => { - // Note: mysql and ob_oracle take different second args (database vs schema) - if *mode == MysqlMode::OceanBaseOracle { - db::ob_oracle::list_objects(p, schema).await - } else if db_config.as_ref().is_some_and(is_doris_family_config) { - db::mysql::list_table_objects_show(p, database).await - } else { - db::mysql::list_objects(p, database).await - } - } - PoolKind::Postgres(p) => db::postgres::list_objects(p, schema).await, - _ => { - drop(connections); - Ok(list_tables_core(state, connection_id, database, schema, None, None) - .await? - .into_iter() - .map(|table| db::ObjectInfo { - name: table.name, - object_type: table.table_type, - schema: if schema.is_empty() { None } else { Some(schema.to_string()) }, - comment: table.comment, - created_at: None, - updated_at: None, - parent_schema: table.parent_schema, - parent_name: table.parent_name, - }) - .collect()) - } - } -} - -fn filter_objects_by_types(objects: Vec, object_types: Option<&[String]>) -> Vec { - let Some(object_types) = object_types else { - return objects; - }; - if object_types.is_empty() { - return objects; - } - let wanted: HashSet = - object_types.iter().map(|object_type| normalize_object_info_type(object_type)).collect(); - objects.into_iter().filter(|object| wanted.contains(&normalize_object_info_type(&object.object_type))).collect() -} - -fn normalize_object_info_type(object_type: &str) -> String { - let value = object_type.to_ascii_uppercase().replace(' ', "_"); - if value.contains("PACKAGE_BODY") { - "PACKAGE_BODY".to_string() - } else if value.contains("PACKAGE") { - "PACKAGE".to_string() - } else if value.contains("VIEW") { - "VIEW".to_string() - } else if value.contains("PROC") { - "PROCEDURE".to_string() - } else if value.contains("FUNC") { - "FUNCTION".to_string() - } else { - "TABLE".to_string() + if let Some(objects) = providers::native::list_objects(pool, db_config.as_ref(), database, schema).await? { + return Ok(objects); } + drop(connections); + Ok(list_tables_core(state, connection_id, database, schema, None, None) + .await? + .into_iter() + .map(|table| db::ObjectInfo { + name: table.name, + object_type: table.table_type, + schema: if schema.is_empty() { None } else { Some(schema.to_string()) }, + comment: table.comment, + created_at: None, + updated_at: None, + parent_schema: table.parent_schema, + parent_name: table.parent_name, + }) + .collect()) } async fn list_completion_objects_once( @@ -1146,33 +661,17 @@ async fn list_completion_objects_once( } let pool = connections.get(&pool_key).ok_or("Pool not found")?; - match pool { - PoolKind::Mysql(p, mode) if *mode != MysqlMode::OceanBaseOracle => { - db::mysql::list_completion_objects(p, database).await - } - PoolKind::Mysql(p, mode) if *mode == MysqlMode::OceanBaseOracle => { - db::ob_oracle::list_objects(p, schema).await.map(filter_completion_objects) - } - PoolKind::Postgres(p) => db::postgres::list_objects(p, schema).await.map(filter_completion_objects), - PoolKind::SqlServer(_) => { + match providers::native::list_completion_objects(pool, database, schema).await? { + Some(objects) => Ok(filter_completion_objects(objects)), + None if matches!(pool, PoolKind::SqlServer(_)) => { drop(connections); let objects = list_objects_once(state, connection_id, database, schema, None).await?; Ok(filter_completion_objects(objects)) } - _ => Ok(Vec::new()), + None => Ok(Vec::new()), } } -fn filter_completion_objects(objects: Vec) -> Vec { - objects - .into_iter() - .filter(|object| { - let object_type = object.object_type.to_ascii_uppercase(); - object_type.contains("PROCEDURE") || object_type.contains("FUNCTION") || object_type.contains("TRIGGER") - }) - .collect() -} - fn is_agent_postgres_metadata_fallback_config(config: &ConnectionConfig) -> bool { matches!(config.db_type, DatabaseType::Kingbase | DatabaseType::Highgo | DatabaseType::Vastbase) } @@ -1334,70 +833,9 @@ pub async fn get_columns_core( let connections = state.connections.read().await; let pool = connections.get(&pool_key).ok_or("Pool not found")?; - match pool { - PoolKind::Mysql(p, _) if db_config.as_ref().is_some_and(is_doris_family_config) => { - db::mysql::get_columns_show(p, database, table).await.map(deduplicate_column_infos) - } - PoolKind::Mysql(p, mode) => { - dispatch_mysql!(p, mode, db::mysql::get_columns, db::ob_oracle::get_columns, database, table) - .map(deduplicate_column_infos) - } - PoolKind::Postgres(p) => db::postgres::get_columns(p, schema, table).await.map(deduplicate_column_infos), - PoolKind::Sqlite(p) => db::sqlite::get_columns(p, schema, table).await.map(deduplicate_column_infos), - PoolKind::Rqlite(client) => { - db::rqlite_driver::get_columns(client, schema, table).await.map(deduplicate_column_infos) - } - PoolKind::Turso(client) => { - db::turso_driver::get_columns(client, schema, table).await.map(deduplicate_column_infos) - } - PoolKind::Elasticsearch(client) => { - db::elasticsearch_driver::get_columns(client, table).await.map(deduplicate_column_infos) - } - _ => Ok(vec![]), - } -} - -fn deduplicate_column_infos(columns: Vec) -> Vec { - let mut result: Vec = Vec::with_capacity(columns.len()); - for column in columns { - if let Some(existing) = result.iter_mut().find(|existing| existing.name == column.name) { - existing.is_primary_key |= column.is_primary_key; - existing.is_nullable &= column.is_nullable; - merge_optional_string(&mut existing.column_default, column.column_default); - merge_optional_string(&mut existing.extra, column.extra); - merge_optional_string(&mut existing.comment, column.comment); - if existing.numeric_precision.is_none() { - existing.numeric_precision = column.numeric_precision; - } - if existing.numeric_scale.is_none() { - existing.numeric_scale = column.numeric_scale; - } - if existing.character_maximum_length.is_none() { - existing.character_maximum_length = column.character_maximum_length; - } - if existing.data_type.trim().is_empty() && !column.data_type.trim().is_empty() { - existing.data_type = column.data_type; - } - } else { - result.push(column); - } - } - result -} - -fn merge_optional_string(target: &mut Option, candidate: Option) { - let Some(candidate) = candidate else { - return; - }; - if candidate.trim().is_empty() { - if target.is_none() { - *target = Some(candidate); - } - return; - } - if target.as_ref().is_none_or(|value| value.trim().is_empty()) { - *target = Some(candidate); - } + providers::native::get_columns(pool, db_config.as_ref(), database, schema, table) + .await + .map(deduplicate_column_infos) } pub async fn list_indexes_core( @@ -1418,17 +856,7 @@ pub async fn list_indexes_core( let connections = state.connections.read().await; let pool = connections.get(&pool_key).ok_or("Pool not found")?; - match pool { - PoolKind::Mysql(p, mode) => { - dispatch_mysql!(p, mode, db::mysql::list_indexes, db::ob_oracle::list_indexes, schema, table) - } - PoolKind::Postgres(p) => db::postgres::list_indexes(p, schema, table).await, - PoolKind::Sqlite(p) => db::sqlite::list_indexes(p, schema, table).await, - PoolKind::Rqlite(client) => db::rqlite_driver::list_indexes(client, schema, table).await, - PoolKind::Turso(client) => db::turso_driver::list_indexes(client, schema, table).await, - PoolKind::MongoDb(client) => db::mongo_driver::list_indexes(client, database, table).await, - _ => Ok(vec![]), - } + providers::native::list_indexes(pool, database, schema, table).await } pub async fn list_foreign_keys_core( @@ -1449,16 +877,7 @@ pub async fn list_foreign_keys_core( let connections = state.connections.read().await; let pool = connections.get(&pool_key).ok_or("Pool not found")?; - match pool { - PoolKind::Mysql(p, mode) => { - dispatch_mysql!(p, mode, db::mysql::list_foreign_keys, db::ob_oracle::list_foreign_keys, schema, table) - } - PoolKind::Postgres(p) => db::postgres::list_foreign_keys(p, schema, table).await, - PoolKind::Sqlite(p) => db::sqlite::list_foreign_keys(p, schema, table).await, - PoolKind::Rqlite(client) => db::rqlite_driver::list_foreign_keys(client, schema, table).await, - PoolKind::Turso(client) => db::turso_driver::list_foreign_keys(client, schema, table).await, - _ => Ok(vec![]), - } + providers::native::list_foreign_keys(pool, schema, table).await } pub async fn list_triggers_core( @@ -1479,16 +898,7 @@ pub async fn list_triggers_core( let connections = state.connections.read().await; let pool = connections.get(&pool_key).ok_or("Pool not found")?; - match pool { - PoolKind::Mysql(p, mode) => { - dispatch_mysql!(p, mode, db::mysql::list_triggers, db::ob_oracle::list_triggers, schema, table) - } - PoolKind::Postgres(p) => db::postgres::list_triggers(p, schema, table).await, - PoolKind::Sqlite(p) => db::sqlite::list_triggers(p, schema, table).await, - PoolKind::Rqlite(client) => db::rqlite_driver::list_triggers(client, schema, table).await, - PoolKind::Turso(client) => db::turso_driver::list_triggers(client, schema, table).await, - _ => Ok(vec![]), - } + providers::native::list_triggers(pool, schema, table).await } pub async fn get_table_ddl_core( @@ -1545,36 +955,13 @@ pub async fn get_table_ddl_core( let pool = connections.get(&pool_key).ok_or("Pool not found")?; let db_config = connection_config(state, connection_id).await; - match pool { - PoolKind::Mysql(p, _) => mysql_ddl(p, table).await, - PoolKind::Postgres(p) if db_config.as_ref().is_some_and(is_opengauss_family_config) => { - match opengauss_table_ddl(p, schema, table).await { - Ok(ddl) => Ok(ddl), - Err(_) => pg_ddl(p, schema, table).await, - } - } - PoolKind::Postgres(p) => pg_ddl(p, schema, table).await, - PoolKind::Sqlite(p) => sqlite_ddl(p, table).await, - PoolKind::Rqlite(client) => db::rqlite_driver::table_ddl(client, table).await, - PoolKind::Turso(client) => db::turso_driver::table_ddl(client, table).await, - _ => Err("DDL not supported for this database type".to_string()), - } + providers::native::table_ddl(pool, db_config.as_ref(), schema, table).await } async fn connection_config(state: &AppState, connection_id: &str) -> Option { state.configs.read().await.get(connection_id).cloned() } -fn is_opengauss_family_config(config: &ConnectionConfig) -> bool { - matches!(config.db_type, DatabaseType::OpenGauss | DatabaseType::Gaussdb) - || matches!(config.driver_profile.as_deref(), Some("opengauss" | "gaussdb")) -} - -fn is_doris_family_config(config: &ConnectionConfig) -> bool { - matches!(config.db_type, DatabaseType::Doris | DatabaseType::StarRocks) - || matches!(config.driver_profile.as_deref(), Some("doris" | "selectdb" | "starrocks")) -} - fn sql_string(value: &str) -> String { format!("'{}'", value.replace('\'', "''")) } @@ -1804,28 +1191,11 @@ pub async fn get_object_source_core( return Ok(result); } } else { - match connections.get(&pool_key).ok_or("Pool not found")? { - PoolKind::Mysql(pool, _) => mysql_object_source(pool, name, &object_type).await?, - PoolKind::Postgres(pool) => postgres_object_source(pool, schema, name, &object_type).await?, - PoolKind::Sqlite(pool) => first_string_cell( - db::sqlite::execute_query(pool, &sqlite_object_source_sql(name, &object_type)).await?, - )?, - PoolKind::Rqlite(client) => { - return db::rqlite_driver::object_source(client, name, &object_type).await; - } - PoolKind::Turso(client) => { - return db::turso_driver::object_source(client, name, &object_type).await; - } - PoolKind::ClickHouse(client) if matches!(object_type, db::ObjectSourceKind::View) => { - let result = db::clickhouse_driver::execute_query( - client, - database, - &format!("SHOW CREATE TABLE {}", mysql_ident(name)), - ) - .await?; - first_string_cell(result)? - } - _ => return Err("Object source is not supported for this database type".to_string()), + let pool = connections.get(&pool_key).ok_or("Pool not found")?; + if let Some(source) = providers::native::object_source(pool, database, schema, name, &object_type).await? { + source + } else { + return Err("Object source is not supported for this database type".to_string()); } } }; diff --git a/crates/dbx-core/src/schema/duckdb_metadata.rs b/crates/dbx-core/src/schema/duckdb_metadata.rs new file mode 100644 index 0000000000..030de20cf1 --- /dev/null +++ b/crates/dbx-core/src/schema/duckdb_metadata.rs @@ -0,0 +1,202 @@ +use crate::connection::AppState; +use crate::db; + +pub fn duckdb_query_tables(con: &duckdb::Connection) -> Result, String> { + duckdb_query_tables_in_database(con, "main", "main") +} + +pub fn duckdb_query_tables_in_database( + con: &duckdb::Connection, + database: &str, + schema: &str, +) -> Result, String> { + duckdb_query_tables_in_database_with_attached(con, database, schema, &[]) +} + +pub fn duckdb_query_tables_in_database_with_attached( + con: &duckdb::Connection, + database: &str, + schema: &str, + attached_names: &[String], +) -> Result, String> { + let database = duckdb_catalog_name(con, database, attached_names)?; + let mut stmt = con + .prepare( + "SELECT table_name, table_type FROM information_schema.tables WHERE table_catalog = ? AND table_schema = ? ORDER BY table_name", + ) + .map_err(|e| e.to_string())?; + let rows = stmt + .query_map((database.as_str(), schema), |row| { + Ok(db::TableInfo { + name: row.get::<_, String>(0)?, + table_type: row.get::<_, String>(1)?, + comment: None, + parent_schema: None, + parent_name: None, + }) + }) + .map_err(|e| e.to_string())?; + Ok(rows.filter_map(|r| r.ok()).collect()) +} + +pub fn duckdb_attach_database(con: &duckdb::Connection, name: &str, path: &str) -> Result<(), String> { + let name = name.trim(); + let path = path.trim(); + if name.is_empty() || path.is_empty() { + return Err("DuckDB attached database name and path are required".to_string()); + } + let sql = format!("ATTACH {} AS {}", duckdb_quote_string(path), duckdb_quote_ident(name)); + con.execute_batch(&sql).map_err(|e| e.to_string()) +} + +pub fn duckdb_list_databases(con: &duckdb::Connection) -> Result, String> { + duckdb_list_databases_with_attached(con, &[]) +} + +pub fn duckdb_list_databases_with_attached( + con: &duckdb::Connection, + attached_names: &[String], +) -> Result, String> { + let primary = duckdb_primary_catalog(con, attached_names)?; + let mut stmt = con.prepare("SHOW DATABASES").map_err(|e| e.to_string())?; + let rows = stmt + .query_map([], |row| { + let name = row.get::<_, String>(0)?; + Ok(db::DatabaseInfo { name: if name == primary { "main".to_string() } else { name } }) + }) + .map_err(|e| e.to_string())?; + Ok(rows.filter_map(|row| row.ok()).collect()) +} + +pub fn duckdb_list_schemas(con: &duckdb::Connection, database: &str) -> Result, String> { + duckdb_list_schemas_with_attached(con, database, &[]) +} + +pub fn duckdb_list_schemas_with_attached( + con: &duckdb::Connection, + database: &str, + attached_names: &[String], +) -> Result, String> { + let database = duckdb_catalog_name(con, database, attached_names)?; + let mut stmt = con + .prepare( + "SELECT schema_name FROM information_schema.schemata WHERE catalog_name = ? AND schema_name NOT IN ('information_schema', 'pg_catalog') ORDER BY schema_name", + ) + .map_err(|e| e.to_string())?; + let rows = stmt.query_map([database.as_str()], |row| row.get::<_, String>(0)).map_err(|e| e.to_string())?; + Ok(rows.filter_map(|r| r.ok()).collect()) +} + +fn duckdb_catalog_name(con: &duckdb::Connection, database: &str, attached_names: &[String]) -> Result { + if database.trim().is_empty() || database == "main" { + return duckdb_primary_catalog(con, attached_names); + } + Ok(database.to_string()) +} + +pub fn duckdb_primary_catalog(con: &duckdb::Connection, attached_names: &[String]) -> Result { + if attached_names.is_empty() { + return duckdb_current_database(con); + } + let attached: std::collections::HashSet = attached_names.iter().map(|name| name.to_lowercase()).collect(); + let mut stmt = con.prepare("SHOW DATABASES").map_err(|e| e.to_string())?; + let rows = stmt.query_map([], |row| row.get::<_, String>(0)).map_err(|e| e.to_string())?; + for row in rows { + let name = row.map_err(|e| e.to_string())?; + if !attached.contains(&name.to_lowercase()) { + return Ok(name); + } + } + duckdb_current_database(con) +} + +fn duckdb_current_database(con: &duckdb::Connection) -> Result { + con.query_row("SELECT current_database()", [], |row| row.get::<_, String>(0)).map_err(|e| e.to_string()) +} + +fn duckdb_quote_ident(value: &str) -> String { + format!("\"{}\"", value.replace('"', "\"\"")) +} + +fn duckdb_quote_string(value: &str) -> String { + format!("'{}'", value.replace('\'', "''")) +} + +pub fn duckdb_query_columns(con: &duckdb::Connection, table: &str) -> Result, String> { + duckdb_query_columns_in_database(con, "main", "main", table) +} + +pub fn duckdb_query_columns_in_database( + con: &duckdb::Connection, + database: &str, + schema: &str, + table: &str, +) -> Result, String> { + duckdb_query_columns_in_database_with_attached(con, database, schema, table, &[]) +} + +pub fn duckdb_query_columns_in_database_with_attached( + con: &duckdb::Connection, + database: &str, + schema: &str, + table: &str, + attached_names: &[String], +) -> Result, String> { + let database = duckdb_catalog_name(con, database, attached_names)?; + let mut pk_stmt = con + .prepare( + "SELECT kcu.column_name + FROM information_schema.table_constraints tc + JOIN information_schema.key_column_usage kcu + ON tc.constraint_name = kcu.constraint_name + AND tc.table_schema = kcu.table_schema + AND tc.table_name = kcu.table_name + WHERE tc.constraint_type = 'PRIMARY KEY' + AND tc.table_catalog = ? + AND tc.table_schema = ? + AND tc.table_name = ? + ORDER BY kcu.ordinal_position", + ) + .map_err(|e| e.to_string())?; + let pk_rows = pk_stmt + .query_map((database.as_str(), schema, table), |row| row.get::<_, String>(0)) + .map_err(|e| e.to_string())?; + let primary_keys: std::collections::HashSet = pk_rows.filter_map(|r| r.ok()).collect(); + + let mut stmt = con + .prepare( + "SELECT column_name, data_type, is_nullable, column_default + FROM information_schema.columns + WHERE table_catalog = ? AND table_schema = ? AND table_name = ? + ORDER BY ordinal_position", + ) + .map_err(|e| e.to_string())?; + let rows = stmt + .query_map((database.as_str(), schema, table), |row| { + let name = row.get::<_, String>(0)?; + Ok(db::ColumnInfo { + is_primary_key: primary_keys.contains(&name), + name, + data_type: row.get::<_, String>(1)?, + is_nullable: row.get::<_, String>(2).unwrap_or_default() == "YES", + column_default: row.get::<_, Option>(3)?, + extra: None, + comment: None, + numeric_precision: None, + numeric_scale: None, + character_maximum_length: None, + }) + }) + .map_err(|e| e.to_string())?; + Ok(rows.filter_map(|r| r.ok()).collect()) +} + +pub(super) async fn duckdb_attached_database_names(state: &AppState, connection_id: &str) -> Vec { + state + .configs + .read() + .await + .get(connection_id) + .map(|config| config.attached_databases.iter().map(|database| database.name.clone()).collect()) + .unwrap_or_default() +} diff --git a/crates/dbx-core/src/schema/normalization.rs b/crates/dbx-core/src/schema/normalization.rs new file mode 100644 index 0000000000..cafb91ff3d --- /dev/null +++ b/crates/dbx-core/src/schema/normalization.rs @@ -0,0 +1,322 @@ +use crate::db; +use crate::models::connection::{ConnectionConfig, DatabaseType}; +use std::collections::HashSet; + +pub(super) fn filter_table_infos( + tables: Vec, + filter: Option<&str>, + limit: Option, +) -> Vec { + let filter = filter.unwrap_or("").to_lowercase(); + let limit = limit.unwrap_or(usize::MAX); + tables + .into_iter() + .filter(|table| filter.is_empty() || table.name.to_lowercase().contains(&filter)) + .take(limit) + .collect() +} + +pub(super) fn filter_table_infos_for_config( + tables: Vec, + filter: Option<&str>, + limit: Option, + config: Option<&ConnectionConfig>, +) -> Vec { + filter_table_infos(filter_yashandb_recyclebin_tables(tables, config), filter, limit) +} + +pub(super) fn filter_yashandb_recyclebin_tables( + tables: Vec, + config: Option<&ConnectionConfig>, +) -> Vec { + if !is_yashandb_config(config) { + return tables; + } + tables.into_iter().filter(|table| !is_recyclebin_object_name(&table.name)).collect() +} + +pub(super) fn filter_yashandb_recyclebin_objects( + objects: Vec, + config: Option<&ConnectionConfig>, +) -> Vec { + if !is_yashandb_config(config) { + return objects; + } + objects.into_iter().filter(|object| !is_recyclebin_object_name(&object.name)).collect() +} + +fn is_yashandb_config(config: Option<&ConnectionConfig>) -> bool { + config.is_some_and(|config| config.db_type == DatabaseType::Yashandb) +} + +fn is_recyclebin_object_name(name: &str) -> bool { + name.to_ascii_uppercase().starts_with("BIN$") +} + +pub(super) fn filter_objects_by_types( + objects: Vec, + object_types: Option<&[String]>, +) -> Vec { + let Some(object_types) = object_types else { + return objects; + }; + if object_types.is_empty() { + return objects; + } + let wanted: HashSet = + object_types.iter().map(|object_type| normalize_object_info_type(object_type)).collect(); + objects.into_iter().filter(|object| wanted.contains(&normalize_object_info_type(&object.object_type))).collect() +} + +fn normalize_object_info_type(object_type: &str) -> String { + let value = object_type.to_ascii_uppercase().replace(' ', "_"); + if value.contains("PACKAGE_BODY") { + "PACKAGE_BODY".to_string() + } else if value.contains("PACKAGE") { + "PACKAGE".to_string() + } else if value.contains("VIEW") { + "VIEW".to_string() + } else if value.contains("PROC") { + "PROCEDURE".to_string() + } else if value.contains("FUNC") { + "FUNCTION".to_string() + } else { + "TABLE".to_string() + } +} + +pub(super) fn filter_completion_objects(objects: Vec) -> Vec { + objects + .into_iter() + .filter(|object| { + let object_type = object.object_type.to_ascii_uppercase(); + object_type.contains("PROCEDURE") || object_type.contains("FUNCTION") || object_type.contains("TRIGGER") + }) + .collect() +} + +pub(super) fn deduplicate_column_infos(columns: Vec) -> Vec { + let mut result: Vec = Vec::with_capacity(columns.len()); + for column in columns { + if let Some(existing) = result.iter_mut().find(|existing| existing.name == column.name) { + existing.is_primary_key |= column.is_primary_key; + existing.is_nullable &= column.is_nullable; + merge_optional_string(&mut existing.column_default, column.column_default); + merge_optional_string(&mut existing.extra, column.extra); + merge_optional_string(&mut existing.comment, column.comment); + if existing.numeric_precision.is_none() { + existing.numeric_precision = column.numeric_precision; + } + if existing.numeric_scale.is_none() { + existing.numeric_scale = column.numeric_scale; + } + if existing.character_maximum_length.is_none() { + existing.character_maximum_length = column.character_maximum_length; + } + if existing.data_type.trim().is_empty() && !column.data_type.trim().is_empty() { + existing.data_type = column.data_type; + } + } else { + result.push(column); + } + } + result +} + +fn merge_optional_string(target: &mut Option, candidate: Option) { + let Some(candidate) = candidate else { + return; + }; + if candidate.trim().is_empty() { + if target.is_none() { + *target = Some(candidate); + } + return; + } + if target.as_ref().is_none_or(|value| value.trim().is_empty()) { + *target = Some(candidate); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::models::connection::{default_redis_key_separator, ConnectionConfig}; + + fn test_connection_config(db_type: DatabaseType) -> ConnectionConfig { + ConnectionConfig { + id: "test".to_string(), + name: "test".to_string(), + db_type, + driver_profile: None, + driver_label: None, + url_params: None, + host: "127.0.0.1".to_string(), + port: 5432, + username: "user".to_string(), + password: "secret".to_string(), + database: Some("demo".to_string()), + visible_databases: None, + attached_databases: Vec::new(), + color: None, + transport_layers: Vec::new(), + connect_timeout_secs: 5, + query_timeout_secs: 30, + idle_timeout_secs: 60, + ssl: false, + ca_cert_path: String::new(), + client_cert_path: String::new(), + client_key_path: String::new(), + sysdba: false, + oracle_connection_type: None, + connection_string: None, + redis_connection_mode: None, + redis_sentinel_master: String::new(), + redis_sentinel_nodes: String::new(), + redis_sentinel_username: String::new(), + redis_sentinel_password: String::new(), + redis_sentinel_tls: false, + redis_cluster_nodes: String::new(), + redis_key_separator: default_redis_key_separator(), + etcd_endpoints: String::new(), + external_config: None, + jdbc_driver_class: None, + jdbc_driver_paths: Vec::new(), + one_time: false, + read_only: false, + } + } + + fn test_column(name: &str, comment: Option<&str>, is_primary_key: bool) -> db::ColumnInfo { + db::ColumnInfo { + name: name.to_string(), + data_type: "VARCHAR".to_string(), + is_nullable: true, + column_default: None, + is_primary_key, + extra: None, + comment: comment.map(|value| value.to_string()), + numeric_precision: None, + numeric_scale: None, + character_maximum_length: None, + } + } + + #[test] + fn deduplicates_columns_and_preserves_later_comment() { + let columns = deduplicate_column_infos(vec![ + test_column("ID", None, false), + test_column("ID", Some("source primary key"), true), + test_column("TFBH", Some(""), false), + test_column("TFBH", Some("ledger number"), false), + ]); + + assert_eq!(columns.len(), 2); + assert_eq!(columns[0].name, "ID"); + assert_eq!(columns[0].comment.as_deref(), Some("source primary key")); + assert!(columns[0].is_primary_key); + assert_eq!(columns[1].name, "TFBH"); + assert_eq!(columns[1].comment.as_deref(), Some("ledger number")); + } + + #[test] + fn filters_list_objects_by_normalized_object_types() { + let objects = vec![ + db::ObjectInfo { + name: "orders".to_string(), + object_type: "BASE TABLE".to_string(), + schema: None, + comment: None, + created_at: None, + updated_at: None, + parent_schema: None, + parent_name: None, + }, + db::ObjectInfo { + name: "active_orders".to_string(), + object_type: "MATERIALIZED VIEW".to_string(), + schema: None, + comment: None, + created_at: None, + updated_at: None, + parent_schema: None, + parent_name: None, + }, + db::ObjectInfo { + name: "payroll".to_string(), + object_type: "PACKAGE BODY".to_string(), + schema: None, + comment: None, + created_at: None, + updated_at: None, + parent_schema: None, + parent_name: None, + }, + ]; + + let filtered = filter_objects_by_types(objects, Some(&["VIEW".to_string(), "PACKAGE_BODY".to_string()])); + + assert_eq!( + filtered.iter().map(|object| object.name.as_str()).collect::>(), + ["active_orders", "payroll"] + ); + } + + #[test] + fn filters_yashandb_recyclebin_tables() { + let tables = vec![ + db::TableInfo { + name: "USERS".to_string(), + table_type: "TABLE".to_string(), + comment: None, + parent_schema: None, + parent_name: None, + }, + db::TableInfo { + name: "BIN$abc123==$0".to_string(), + table_type: "TABLE".to_string(), + comment: None, + parent_schema: None, + parent_name: None, + }, + ]; + + let filtered = + filter_yashandb_recyclebin_tables(tables.clone(), Some(&test_connection_config(DatabaseType::Yashandb))); + let oracle = filter_yashandb_recyclebin_tables(tables, Some(&test_connection_config(DatabaseType::Oracle))); + + assert_eq!(filtered.iter().map(|table| table.name.as_str()).collect::>(), ["USERS"]); + assert_eq!(oracle.len(), 2); + } + + #[test] + fn filters_yashandb_recyclebin_objects() { + let objects = vec![ + db::ObjectInfo { + name: "ORDERS".to_string(), + object_type: "TABLE".to_string(), + schema: Some("HR".to_string()), + comment: None, + created_at: None, + updated_at: None, + parent_schema: None, + parent_name: None, + }, + db::ObjectInfo { + name: "bin$deleted".to_string(), + object_type: "TABLE".to_string(), + schema: Some("HR".to_string()), + comment: None, + created_at: None, + updated_at: None, + parent_schema: None, + parent_name: None, + }, + ]; + + let filtered = + filter_yashandb_recyclebin_objects(objects, Some(&test_connection_config(DatabaseType::Yashandb))); + + assert_eq!(filtered.iter().map(|object| object.name.as_str()).collect::>(), ["ORDERS"]); + } +} diff --git a/crates/dbx-core/src/schema/providers/mod.rs b/crates/dbx-core/src/schema/providers/mod.rs new file mode 100644 index 0000000000..c08432d777 --- /dev/null +++ b/crates/dbx-core/src/schema/providers/mod.rs @@ -0,0 +1 @@ +pub(super) mod native; diff --git a/crates/dbx-core/src/schema/providers/native.rs b/crates/dbx-core/src/schema/providers/native.rs new file mode 100644 index 0000000000..a302207ac8 --- /dev/null +++ b/crates/dbx-core/src/schema/providers/native.rs @@ -0,0 +1,250 @@ +use crate::connection::{MysqlMode, PoolKind}; +use crate::db; +use crate::models::connection::{ConnectionConfig, DatabaseType}; + +pub(in crate::schema) async fn list_databases( + pool: &PoolKind, + config: Option<&ConnectionConfig>, +) -> Result, String> { + match pool { + PoolKind::Mysql(p, _) if config.is_some_and(is_doris_family_config) => db::mysql::list_databases_show(p).await, + PoolKind::Mysql(p, mode) if *mode == MysqlMode::OceanBaseOracle => db::ob_oracle::list_databases(p).await, + PoolKind::Mysql(p, _) => db::mysql::list_databases(p).await, + PoolKind::Postgres(p) => db::postgres::list_databases(p).await, + PoolKind::Sqlite(p) => db::sqlite::list_databases(p).await, + PoolKind::Rqlite(client) => db::rqlite_driver::list_databases(client).await, + PoolKind::Turso(client) => db::turso_driver::list_databases(client).await, + _ => Ok(vec![]), + } +} + +pub(in crate::schema) async fn list_schemas(pool: &PoolKind) -> Result, String> { + match pool { + PoolKind::Postgres(p) => db::postgres::list_schemas(p).await, + _ => Ok(vec![]), + } +} + +pub(in crate::schema) async fn list_tables( + pool: &PoolKind, + config: Option<&ConnectionConfig>, + database: &str, + schema: &str, +) -> Result, String> { + match pool { + PoolKind::Mysql(p, _) if config.is_some_and(is_doris_family_config) => { + db::mysql::list_tables_show(p, database).await + } + PoolKind::Mysql(p, mode) if *mode == MysqlMode::OceanBaseOracle => db::ob_oracle::list_tables(p, schema).await, + PoolKind::Mysql(p, _) => db::mysql::list_tables(p, schema).await, + PoolKind::Postgres(p) => db::postgres::list_tables(p, schema).await, + PoolKind::Sqlite(p) => db::sqlite::list_tables(p, schema).await, + PoolKind::Rqlite(client) => db::rqlite_driver::list_tables(client, schema).await, + PoolKind::Turso(client) => db::turso_driver::list_tables(client, schema).await, + PoolKind::MongoDb(client) => db::mongo_driver::list_collections(client, database) + .await + .map(|names| collection_names_to_tables(names, "COLLECTION")), + PoolKind::Elasticsearch(client) => { + db::elasticsearch_driver::list_indices(client).await.map(|names| collection_names_to_tables(names, "INDEX")) + } + _ => Ok(vec![]), + } +} + +pub(in crate::schema) async fn list_objects( + pool: &PoolKind, + config: Option<&ConnectionConfig>, + database: &str, + schema: &str, +) -> Result>, String> { + match pool { + PoolKind::Mysql(p, mode) if *mode == MysqlMode::OceanBaseOracle => { + db::ob_oracle::list_objects(p, schema).await.map(Some) + } + PoolKind::Mysql(p, _) if config.is_some_and(is_doris_family_config) => { + db::mysql::list_table_objects_show(p, database).await.map(Some) + } + PoolKind::Mysql(p, _) => db::mysql::list_objects(p, database).await.map(Some), + PoolKind::Postgres(p) => db::postgres::list_objects(p, schema).await.map(Some), + _ => Ok(None), + } +} + +pub(in crate::schema) async fn list_completion_objects( + pool: &PoolKind, + database: &str, + schema: &str, +) -> Result>, String> { + match pool { + PoolKind::Mysql(p, mode) if *mode != MysqlMode::OceanBaseOracle => { + db::mysql::list_completion_objects(p, database).await.map(Some) + } + PoolKind::Mysql(p, mode) if *mode == MysqlMode::OceanBaseOracle => { + db::ob_oracle::list_objects(p, schema).await.map(Some) + } + PoolKind::Postgres(p) => db::postgres::list_objects(p, schema).await.map(Some), + _ => Ok(None), + } +} + +pub(in crate::schema) async fn get_columns( + pool: &PoolKind, + config: Option<&ConnectionConfig>, + database: &str, + schema: &str, + table: &str, +) -> Result, String> { + match pool { + PoolKind::Mysql(p, _) if config.is_some_and(is_doris_family_config) => { + db::mysql::get_columns_show(p, database, table).await + } + PoolKind::Mysql(p, mode) if *mode == MysqlMode::OceanBaseOracle => { + db::ob_oracle::get_columns(p, database, table).await + } + PoolKind::Mysql(p, _) => db::mysql::get_columns(p, database, table).await, + PoolKind::Postgres(p) => db::postgres::get_columns(p, schema, table).await, + PoolKind::Sqlite(p) => db::sqlite::get_columns(p, schema, table).await, + PoolKind::Rqlite(client) => db::rqlite_driver::get_columns(client, schema, table).await, + PoolKind::Turso(client) => db::turso_driver::get_columns(client, schema, table).await, + PoolKind::Elasticsearch(client) => db::elasticsearch_driver::get_columns(client, table).await, + _ => Ok(vec![]), + } +} + +pub(in crate::schema) async fn list_indexes( + pool: &PoolKind, + database: &str, + schema: &str, + table: &str, +) -> Result, String> { + match pool { + PoolKind::Mysql(p, mode) if *mode == MysqlMode::OceanBaseOracle => { + db::ob_oracle::list_indexes(p, schema, table).await + } + PoolKind::Mysql(p, _) => db::mysql::list_indexes(p, schema, table).await, + PoolKind::Postgres(p) => db::postgres::list_indexes(p, schema, table).await, + PoolKind::Sqlite(p) => db::sqlite::list_indexes(p, schema, table).await, + PoolKind::Rqlite(client) => db::rqlite_driver::list_indexes(client, schema, table).await, + PoolKind::Turso(client) => db::turso_driver::list_indexes(client, schema, table).await, + PoolKind::MongoDb(client) => db::mongo_driver::list_indexes(client, database, table).await, + _ => Ok(vec![]), + } +} + +pub(in crate::schema) async fn list_foreign_keys( + pool: &PoolKind, + schema: &str, + table: &str, +) -> Result, String> { + match pool { + PoolKind::Mysql(p, mode) if *mode == MysqlMode::OceanBaseOracle => { + db::ob_oracle::list_foreign_keys(p, schema, table).await + } + PoolKind::Mysql(p, _) => db::mysql::list_foreign_keys(p, schema, table).await, + PoolKind::Postgres(p) => db::postgres::list_foreign_keys(p, schema, table).await, + PoolKind::Sqlite(p) => db::sqlite::list_foreign_keys(p, schema, table).await, + PoolKind::Rqlite(client) => db::rqlite_driver::list_foreign_keys(client, schema, table).await, + PoolKind::Turso(client) => db::turso_driver::list_foreign_keys(client, schema, table).await, + _ => Ok(vec![]), + } +} + +pub(in crate::schema) async fn list_triggers( + pool: &PoolKind, + schema: &str, + table: &str, +) -> Result, String> { + match pool { + PoolKind::Mysql(p, mode) if *mode == MysqlMode::OceanBaseOracle => { + db::ob_oracle::list_triggers(p, schema, table).await + } + PoolKind::Mysql(p, _) => db::mysql::list_triggers(p, schema, table).await, + PoolKind::Postgres(p) => db::postgres::list_triggers(p, schema, table).await, + PoolKind::Sqlite(p) => db::sqlite::list_triggers(p, schema, table).await, + PoolKind::Rqlite(client) => db::rqlite_driver::list_triggers(client, schema, table).await, + PoolKind::Turso(client) => db::turso_driver::list_triggers(client, schema, table).await, + _ => Ok(vec![]), + } +} + +pub(in crate::schema) async fn table_ddl( + pool: &PoolKind, + config: Option<&ConnectionConfig>, + schema: &str, + table: &str, +) -> Result { + match pool { + PoolKind::Mysql(p, _) => super::super::mysql_ddl(p, table).await, + PoolKind::Postgres(p) if config.is_some_and(is_opengauss_family_config) => { + match super::super::opengauss_table_ddl(p, schema, table).await { + Ok(ddl) => Ok(ddl), + Err(_) => super::super::pg_ddl(p, schema, table).await, + } + } + PoolKind::Postgres(p) => super::super::pg_ddl(p, schema, table).await, + PoolKind::Sqlite(p) => super::super::sqlite_ddl(p, table).await, + PoolKind::Rqlite(client) => db::rqlite_driver::table_ddl(client, table).await, + PoolKind::Turso(client) => db::turso_driver::table_ddl(client, table).await, + _ => Err("DDL not supported for this database type".to_string()), + } +} + +pub(in crate::schema) async fn object_source( + pool: &PoolKind, + database: &str, + schema: &str, + name: &str, + object_type: &db::ObjectSourceKind, +) -> Result, String> { + match pool { + PoolKind::Mysql(pool, _) => super::super::mysql_object_source(pool, name, object_type).await.map(Some), + PoolKind::Postgres(pool) => { + super::super::postgres_object_source(pool, schema, name, object_type).await.map(Some) + } + PoolKind::Sqlite(pool) => { + let source = super::super::first_string_cell( + db::sqlite::execute_query(pool, &super::super::sqlite_object_source_sql(name, object_type)).await?, + )?; + Ok(Some(source)) + } + PoolKind::Rqlite(client) => { + db::rqlite_driver::object_source(client, name, object_type).await.map(|source| Some(source.source)) + } + PoolKind::Turso(client) => { + db::turso_driver::object_source(client, name, object_type).await.map(|source| Some(source.source)) + } + PoolKind::ClickHouse(client) if matches!(object_type, db::ObjectSourceKind::View) => { + let result = db::clickhouse_driver::execute_query( + client, + database, + &format!("SHOW CREATE TABLE {}", super::super::mysql_ident(name)), + ) + .await?; + super::super::first_string_cell(result).map(Some) + } + _ => Ok(None), + } +} + +fn collection_names_to_tables(names: Vec, table_type: &str) -> Vec { + names + .into_iter() + .map(|name| db::TableInfo { + name, + table_type: table_type.to_string(), + comment: None, + parent_schema: None, + parent_name: None, + }) + .collect() +} + +fn is_opengauss_family_config(config: &ConnectionConfig) -> bool { + matches!(config.db_type, DatabaseType::OpenGauss | DatabaseType::Gaussdb) + || matches!(config.driver_profile.as_deref(), Some("opengauss" | "gaussdb")) +} + +fn is_doris_family_config(config: &ConnectionConfig) -> bool { + matches!(config.db_type, DatabaseType::Doris | DatabaseType::StarRocks) + || matches!(config.driver_profile.as_deref(), Some("doris" | "selectdb" | "starrocks")) +} diff --git a/crates/dbx-core/src/table_structure_sql.rs b/crates/dbx-core/src/table_structure_sql.rs index 39ea8f9173..a94f76b20d 100644 --- a/crates/dbx-core/src/table_structure_sql.rs +++ b/crates/dbx-core/src/table_structure_sql.rs @@ -2184,6 +2184,7 @@ mod tests { ); } + #[cfg(feature = "duckdb-bundled")] #[test] fn builds_duckdb_create_table_statements() { let mut name = column("name"); diff --git a/crates/dbx-core/src/transfer.rs b/crates/dbx-core/src/transfer.rs index ed4509a5cb..2c0494817b 100644 --- a/crates/dbx-core/src/transfer.rs +++ b/crates/dbx-core/src/transfer.rs @@ -6,9 +6,10 @@ use tokio::sync::RwLock; use crate::connection::{AppState, PoolKind}; use crate::db; use crate::db::mongo_driver::MongoDocumentResult; -use crate::models::connection::{default_redis_key_separator, DatabaseType}; +use crate::models::connection::DatabaseType; use crate::object_source_sql::{build_executable_object_source_statements, EditableObjectSourceSqlInput}; use crate::query::{agent_execute_query_params, QueryExecutionOptions}; +#[cfg(feature = "duckdb-bundled")] use crate::sql::starts_with_executable_sql_keyword; static CANCELLED: std::sync::LazyLock>> = @@ -3126,9 +3127,14 @@ where #[cfg(test)] mod tests { use super::*; + #[cfg(feature = "duckdb-bundled")] use crate::connection::{AppState, PoolKind}; + #[cfg(feature = "duckdb-bundled")] + use crate::models::connection::default_redis_key_separator; + #[cfg(feature = "duckdb-bundled")] use crate::storage::Storage; use serde_json::json; + #[cfg(feature = "duckdb-bundled")] use std::sync::Arc; #[cfg(feature = "duckdb-bundled")]