github/dbx
4
0
Fork 0
mirror of https://github.com/t8y2/dbx.git synced 2026-10-23 11:36:21 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor(schema): split metadata providers

Browse Source
This commit is contained in:
t8y2
2026-06-10 19:10:53 +08:00
parent 27aef6059d
commit 33ab7aeae6
11 changed files with 882 additions and 733 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/dbx-core/src/agent_connection.rs
+2 -1
View File
@@ -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 {
crates/dbx-core/src/connection.rs
+16 -24
View File
@@ -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<ConnectionCon
}
}
#[cfg(feature = "duckdb-bundled")]
fn duckdb_paths_match(left: &str, right: &str) -> 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;
crates/dbx-core/src/db/turso_driver.rs
-1
View File
@@ -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"]);
crates/dbx-core/src/query.rs
+10 -5
View File
@@ -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<i64> {
}
}
#[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<db::QueryResult, String> {
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() {
crates/dbx-core/src/schema.rs
+71 -701
View File
@@ -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<Vec<db::TableInfo>, 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<Vec<db::TableInfo>, 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<Vec<db::TableInfo>, 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<Vec<db::DatabaseInfo>, 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<Vec<db::DatabaseInfo>, 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<Vec<String>, 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<Vec<String>, 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<String, String> {
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<String, String> {
if attached_names.is_empty() {
return duckdb_current_database(con);
}
let attached: std::collections::HashSet<String> = 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<String, String> {
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<Vec<db::ColumnInfo>, 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<Vec<db::ColumnInfo>, 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<Vec<db::ColumnInfo>, 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<String> = 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<String>>(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<String> {
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<Ve
let connections = state.connections.read().await;
let pool = connections.get(connection_id).ok_or("Connection not found")?;
match pool {
PoolKind::Mysql(p, _) if db_config.as_ref().is_some_and(is_doris_family_config) => {
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<Vec<String>, 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<String>, table_type: &str) -> Vec<db::TableInfo> {
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<db::TableInfo>, filter: Option<&str>, limit: Option<usize>) -> Vec<db::TableInfo> {
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<db::TableInfo>,
filter: Option<&str>,
limit: Option<usize>,
config: Option<&ConnectionConfig>,
) -> Vec<db::TableInfo> {
filter_table_infos(filter_yashandb_recyclebin_tables(tables, config), filter, limit)
}
fn filter_yashandb_recyclebin_tables(
tables: Vec<db::TableInfo>,
config: Option<&ConnectionConfig>,
) -> Vec<db::TableInfo> {
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<db::ObjectInfo>,
config: Option<&ConnectionConfig>,
) -> Vec<db::ObjectInfo> {
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::<Vec<_>>(),
["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::<Vec<_>>(), ["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::<Vec<_>>(), ["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<db::ObjectInfo>, object_types: Option<&[String]>) -> Vec<db::ObjectInfo> {
let Some(object_types) = object_types else {
return objects;
};
if object_types.is_empty() {
return objects;
}
let wanted: HashSet<String> =
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<db::ObjectInfo>) -> Vec<db::ObjectInfo> {
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<db::ColumnInfo>) -> Vec<db::ColumnInfo> {
let mut result: Vec<db::ColumnInfo> = 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<String>, candidate: Option<String>) {
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<ConnectionConfig> {
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());
}
}
};
crates/dbx-core/src/schema/duckdb_metadata.rs
+202
View File
@@ -0,0 +1,202 @@
use crate::connection::AppState;
use crate::db;
pub fn duckdb_query_tables(con: &duckdb::Connection) -> Result<Vec<db::TableInfo>, String> {
duckdb_query_tables_in_database(con, "main", "main")
}
pub fn duckdb_query_tables_in_database(
con: &duckdb::Connection,
database: &str,
schema: &str,
) -> Result<Vec<db::TableInfo>, 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<Vec<db::TableInfo>, 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<Vec<db::DatabaseInfo>, String> {
duckdb_list_databases_with_attached(con, &[])
}
pub fn duckdb_list_databases_with_attached(
con: &duckdb::Connection,
attached_names: &[String],
) -> Result<Vec<db::DatabaseInfo>, 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<Vec<String>, String> {
duckdb_list_schemas_with_attached(con, database, &[])
}
pub fn duckdb_list_schemas_with_attached(
con: &duckdb::Connection,
database: &str,
attached_names: &[String],
) -> Result<Vec<String>, 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<String, String> {
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<String, String> {
if attached_names.is_empty() {
return duckdb_current_database(con);
}
let attached: std::collections::HashSet<String> = 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<String, String> {
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<Vec<db::ColumnInfo>, 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<Vec<db::ColumnInfo>, 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<Vec<db::ColumnInfo>, 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<String> = 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<String>>(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<String> {
state
.configs
.read()
.await
.get(connection_id)
.map(|config| config.attached_databases.iter().map(|database| database.name.clone()).collect())
.unwrap_or_default()
}
crates/dbx-core/src/schema/normalization.rs
+322
View File
@@ -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<db::TableInfo>,
filter: Option<&str>,
limit: Option<usize>,
) -> Vec<db::TableInfo> {
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<db::TableInfo>,
filter: Option<&str>,
limit: Option<usize>,
config: Option<&ConnectionConfig>,
) -> Vec<db::TableInfo> {
filter_table_infos(filter_yashandb_recyclebin_tables(tables, config), filter, limit)
}
pub(super) fn filter_yashandb_recyclebin_tables(
tables: Vec<db::TableInfo>,
config: Option<&ConnectionConfig>,
) -> Vec<db::TableInfo> {
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<db::ObjectInfo>,
config: Option<&ConnectionConfig>,
) -> Vec<db::ObjectInfo> {
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<db::ObjectInfo>,
object_types: Option<&[String]>,
) -> Vec<db::ObjectInfo> {
let Some(object_types) = object_types else {
return objects;
};
if object_types.is_empty() {
return objects;
}
let wanted: HashSet<String> =
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<db::ObjectInfo>) -> Vec<db::ObjectInfo> {
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<db::ColumnInfo>) -> Vec<db::ColumnInfo> {
let mut result: Vec<db::ColumnInfo> = 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<String>, candidate: Option<String>) {
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::<Vec<_>>(),
["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::<Vec<_>>(), ["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::<Vec<_>>(), ["ORDERS"]);
}
}
crates/dbx-core/src/schema/providers/mod.rs
+1
View File
@@ -0,0 +1 @@
pub(super) mod native;
crates/dbx-core/src/schema/providers/native.rs
+250
View File
@@ -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<Vec<db::DatabaseInfo>, 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<Vec<String>, 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<Vec<db::TableInfo>, 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<Option<Vec<db::ObjectInfo>>, 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<Option<Vec<db::ObjectInfo>>, 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<Vec<db::ColumnInfo>, 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<Vec<db::IndexInfo>, 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<Vec<db::ForeignKeyInfo>, 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<Vec<db::TriggerInfo>, 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<String, String> {
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<Option<String>, 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<String>, table_type: &str) -> Vec<db::TableInfo> {
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"))
}
crates/dbx-core/src/table_structure_sql.rs
+1
View File
@@ -2184,6 +2184,7 @@ mod tests {
);
}
#[cfg(feature = "duckdb-bundled")]
#[test]
fn builds_duckdb_create_table_statements() {
let mut name = column("name");
crates/dbx-core/src/transfer.rs
+7 -1
View File
@@ -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<RwLock<HashSet<String>>> =
@@ -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")]