use crate::connection::{connection_url_for_endpoint, database_connection_config, AppState, MysqlMode, PoolKind}; use crate::db; use crate::models::connection::{ConnectionConfig, DatabaseType}; use crate::query::{agent_execute_query_params, should_discard_pool_after_error, QueryExecutionOptions}; use std::future::Future; use std::sync::Arc; use std::time::{Duration, Instant}; macro_rules! extract_pool { ($connections:expr, $key:expr, $variant:ident) => { $connections.get($key).and_then(|v| match v { PoolKind::$variant(val) => Some(val.clone()), _ => None, }) }; } 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) { drop($connections); let mut client = client.lock().await; return db::sqlserver::$method(&mut client $(, $arg)*).await; } }; } #[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")] pub fn duckdb_completion_assistant_search( con: &duckdb::Connection, request: &db::CompletionAssistantRequest, attached_names: &[String], ) -> Result { let limit = request.max_results.unwrap_or(100).clamp(1, 1000); let kinds = if request.object_kinds.is_empty() { vec![db::CompletionAssistantObjectKind::Table, db::CompletionAssistantObjectKind::View] } else { request.object_kinds.clone() }; let mut candidates = Vec::new(); if kinds.iter().any(|kind| matches!(kind, db::CompletionAssistantObjectKind::Schema)) { candidates.extend(duckdb_completion_schemas(con, request, attached_names, limit)?); if candidates.len() >= limit { return Ok(db::CompletionAssistantResponse { candidates, incomplete: true, fallback_used: false }); } } if kinds.iter().any(db::CompletionAssistantObjectKind::is_table_like) { candidates.extend(duckdb_completion_tables(con, request, &kinds, attached_names, limit - candidates.len())?); if candidates.len() >= limit { return Ok(db::CompletionAssistantResponse { candidates, incomplete: true, fallback_used: false }); } } if kinds.iter().any(|kind| matches!(kind, db::CompletionAssistantObjectKind::Column)) { candidates.extend(duckdb_completion_columns(con, request, attached_names, limit - candidates.len())?); if candidates.len() >= limit { return Ok(db::CompletionAssistantResponse { candidates, incomplete: true, fallback_used: false }); } } Ok(db::CompletionAssistantResponse { candidates, incomplete: false, fallback_used: false }) } #[cfg(feature = "duckdb-bundled")] fn duckdb_completion_schemas( con: &duckdb::Connection, request: &db::CompletionAssistantRequest, attached_names: &[String], limit: usize, ) -> Result, String> { if limit == 0 { return Ok(Vec::new()); } let database = duckdb_catalog_name(con, &request.database, attached_names)?; let pattern = duckdb_completion_like_pattern(request); let mut stmt = con .prepare( "SELECT schema_name FROM information_schema.schemata WHERE catalog_name = ? AND schema_name NOT IN ('information_schema', 'pg_catalog') AND lower(schema_name) LIKE lower(?) ESCAPE '\\' ORDER BY schema_name LIMIT ?", ) .map_err(|e| e.to_string())?; let rows = stmt .query_map((database.as_str(), pattern.as_str(), limit as i64), |row| { let schema = row.get::<_, String>(0)?; Ok(db::CompletionAssistantCandidate { name: schema.clone(), kind: db::CompletionAssistantCandidateKind::Schema, database: Some(request.database.clone()), schema: Some(schema), parent_schema: None, parent_name: None, comment: None, data_type: None, }) }) .map_err(|e| e.to_string())?; Ok(rows.filter_map(|row| row.ok()).collect()) } #[cfg(feature = "duckdb-bundled")] fn duckdb_completion_tables( con: &duckdb::Connection, request: &db::CompletionAssistantRequest, kinds: &[db::CompletionAssistantObjectKind], attached_names: &[String], limit: usize, ) -> Result, String> { if limit == 0 { return Ok(Vec::new()); } let database = duckdb_catalog_name(con, &request.database, attached_names)?; let schema = request.parent_schema.as_deref().or(request.schema.as_deref()).unwrap_or("main"); let include_tables = kinds.iter().any(|kind| matches!(kind, db::CompletionAssistantObjectKind::Table)); let include_views = kinds.iter().any(|kind| matches!(kind, db::CompletionAssistantObjectKind::View)); let pattern = duckdb_completion_like_pattern(request); let mut stmt = con .prepare( "SELECT table_name, table_type FROM information_schema.tables WHERE table_catalog = ? AND table_schema = ? AND ((? AND table_type = 'BASE TABLE') OR (? AND table_type = 'VIEW')) AND lower(table_name) LIKE lower(?) ESCAPE '\\' ORDER BY table_name LIMIT ?", ) .map_err(|e| e.to_string())?; let rows = stmt .query_map((database.as_str(), schema, include_tables, include_views, pattern.as_str(), limit as i64), |row| { let table_type = row.get::<_, String>(1)?; Ok(db::CompletionAssistantCandidate { name: row.get(0)?, kind: if table_type.eq_ignore_ascii_case("VIEW") { db::CompletionAssistantCandidateKind::View } else { db::CompletionAssistantCandidateKind::Table }, database: Some(request.database.clone()), schema: Some(schema.to_string()), parent_schema: None, parent_name: None, comment: None, data_type: None, }) }) .map_err(|e| e.to_string())?; Ok(rows.filter_map(|row| row.ok()).collect()) } #[cfg(feature = "duckdb-bundled")] fn duckdb_completion_columns( con: &duckdb::Connection, request: &db::CompletionAssistantRequest, attached_names: &[String], limit: usize, ) -> Result, String> { if limit == 0 { return Ok(Vec::new()); } let Some(table) = request.parent_name.as_deref().filter(|table| !table.trim().is_empty()) else { return Ok(Vec::new()); }; let database = duckdb_catalog_name(con, &request.database, attached_names)?; let schema = request.parent_schema.as_deref().or(request.schema.as_deref()).unwrap_or("main"); let pattern = duckdb_completion_like_pattern(request); let mut stmt = con .prepare( "SELECT column_name, data_type FROM information_schema.columns WHERE table_catalog = ? AND table_schema = ? AND table_name = ? AND lower(column_name) LIKE lower(?) ESCAPE '\\' ORDER BY ordinal_position LIMIT ?", ) .map_err(|e| e.to_string())?; let rows = stmt .query_map((database.as_str(), schema, table, pattern.as_str(), limit as i64), |row| { Ok(db::CompletionAssistantCandidate { name: row.get(0)?, kind: db::CompletionAssistantCandidateKind::Column, database: Some(request.database.clone()), schema: Some(schema.to_string()), parent_schema: Some(schema.to_string()), parent_name: Some(table.to_string()), comment: None, data_type: Some(row.get(1)?), }) }) .map_err(|e| e.to_string())?; Ok(rows.filter_map(|row| row.ok()).collect()) } #[cfg(feature = "duckdb-bundled")] fn duckdb_completion_like_pattern(request: &db::CompletionAssistantRequest) -> String { let mask = request.mask.trim().trim_matches('%'); let escaped = mask.replace('\\', "\\\\").replace('%', "\\%").replace('_', "\\_"); match request.match_mode.as_ref().unwrap_or(&db::CompletionAssistantMatchMode::Prefix) { db::CompletionAssistantMatchMode::Prefix => format!("{escaped}%"), db::CompletionAssistantMatchMode::Contains => format!("%{escaped}%"), } } #[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 } else { database } } fn agent_metadata_timeout(config: Option<&ConnectionConfig>) -> Option { let Some(config) = config else { return Some(Duration::from_secs(60)); }; match config.effective_query_timeout_secs() { 0 => None, seconds => Some(Duration::from_secs(seconds.max(60))), } } pub async fn list_databases_core(state: &AppState, connection_id: &str) -> Result, String> { retry_metadata_connection(state, connection_id, None, || list_databases_once(state, connection_id)).await } pub async fn list_sqlserver_linked_servers_core( state: &AppState, connection_id: &str, ) -> Result, String> { let connections = state.connections.read().await; if let Some(client) = extract_pool!(&connections, connection_id, SqlServer) { drop(connections); let mut client = client.lock().await; return db::sqlserver::list_linked_servers(&mut client).await; } Ok(vec![]) } pub async fn list_sqlserver_linked_server_catalogs_core( state: &AppState, connection_id: &str, server: &str, ) -> Result, String> { let connections = state.connections.read().await; if let Some(client) = extract_pool!(&connections, connection_id, SqlServer) { drop(connections); let mut client = client.lock().await; return db::sqlserver::list_linked_server_catalogs(&mut client, server).await; } Ok(vec![]) } pub async fn list_sqlserver_linked_server_schemas_core( state: &AppState, connection_id: &str, server: &str, catalog: &str, ) -> Result, String> { let connections = state.connections.read().await; if let Some(client) = extract_pool!(&connections, connection_id, SqlServer) { drop(connections); let mut client = client.lock().await; return db::sqlserver::list_linked_server_schemas(&mut client, server, catalog).await; } Ok(vec![]) } pub async fn list_sqlserver_linked_server_tables_core( state: &AppState, connection_id: &str, server: &str, catalog: &str, schema: &str, filter: Option<&str>, limit: Option, offset: Option, ) -> Result, String> { let connections = state.connections.read().await; if let Some(client) = extract_pool!(&connections, connection_id, SqlServer) { drop(connections); let mut client = client.lock().await; return db::sqlserver::list_linked_server_tables(&mut client, server, catalog, schema, filter, limit, offset) .await; } Ok(vec![]) } async fn list_databases_once(state: &AppState, connection_id: &str) -> Result, String> { log::info!("[list_databases] connection_id={connection_id}"); let db_config = connection_config(state, connection_id).await; { let connections = state.connections.read().await; #[cfg(feature = "duckdb-bundled")] if extract_pool!(&connections, connection_id, ExternalTabular).is_some() { return Ok(vec![db::DatabaseInfo { name: "main".to_string() }]); } if let Some(PoolKind::ExternalDriver { config, session, .. }) = connections.get(connection_id) { let config = config.clone(); let session = session.clone(); drop(connections); return session .invoke::>("listDatabases", serde_json::json!({ "connection": config.as_ref() })) .await; } if let Some(client) = extract_pool!(&connections, connection_id, ClickHouse) { drop(connections); return db::clickhouse_driver::list_databases(&client).await; } if let Some(client) = extract_pool!(&connections, connection_id, InfluxDb) { drop(connections); return db::influxdb_driver::list_databases(&client).await; } try_sqlserver!(connections, connection_id, list_databases); if let Some(client) = extract_pool!(&connections, connection_id, Agent) { let is_mongo = state.configs.read().await.get(connection_id).is_some_and(|c| c.db_type == DatabaseType::MongoDb); if is_mongo { drop(connections); let dbs = crate::mongo_ops::mongo_list_databases_core(state, connection_id).await?; return Ok(dbs.into_iter().map(|name| db::DatabaseInfo { name }).collect()); } drop(connections); let mut client = client.lock().await; return client.list_databases(agent_metadata_timeout(db_config.as_ref())).await; } } #[cfg(feature = "duckdb-bundled")] let duckdb_attached_names = duckdb_attached_database_names(state, connection_id).await; let db_config = connection_config(state, connection_id).await; 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 .map(|databases| filter_mysql_system_databases_for_config(databases, db_config.as_ref())) } 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, #[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![]), } } pub async fn list_schemas_core(state: &AppState, connection_id: &str, database: &str) -> Result, String> { list_schemas_core_with_visible_filter(state, connection_id, database, false).await } pub async fn list_schemas_core_with_visible_filter( state: &AppState, connection_id: &str, database: &str, apply_visible_filter: bool, ) -> Result, String> { retry_metadata_connection(state, connection_id, Some(database), || { list_schemas_once(state, connection_id, database, apply_visible_filter) }) .await } pub async fn list_schema_infos_core( state: &AppState, connection_id: &str, database: &str, ) -> Result, String> { retry_metadata_connection(state, connection_id, Some(database), || { list_schema_infos_once(state, connection_id, database) }) .await } async fn list_schema_infos_once( state: &AppState, connection_id: &str, database: &str, ) -> Result, String> { let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?; { let connections = state.connections.read().await; if let Some(PoolKind::Postgres(pool)) = connections.get(&pool_key) { return db::postgres::list_schema_infos(pool).await; } } let schemas = list_schemas_once(state, connection_id, database, false).await?; Ok(schemas.into_iter().map(|name| db::SchemaInfo { name, comment: None }).collect()) } pub async fn list_data_types_core( state: &AppState, connection_id: &str, database: &str, ) -> Result, String> { retry_metadata_connection(state, connection_id, Some(database), || async { let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?; let db_config = connection_config(state, connection_id).await; let connections = state.connections.read().await; if let Some(PoolKind::ExternalDriver { config, session, .. }) = connections.get(&pool_key) { let config = config.clone(); let session = session.clone(); drop(connections); return session .invoke::>( "listDataTypes", serde_json::json!({ "connection": config.as_ref(), "database": database }), ) .await .map(deduplicate_data_type_names); } if let Some(client) = extract_pool!(&connections, &pool_key, Agent) { drop(connections); let mut client = client.lock().await; return client .list_data_types::>(database, agent_metadata_timeout(db_config.as_ref())) .await .map(deduplicate_data_type_names); } Ok(Vec::new()) }) .await } fn deduplicate_data_type_names(names: Vec) -> Vec { let mut seen = std::collections::HashSet::new(); let mut result = Vec::new(); for name in names { let trimmed = name.trim(); if trimmed.is_empty() { continue; } let key = trimmed.to_ascii_lowercase(); if seen.insert(key) { result.push(trimmed.to_string()); } } result } async fn list_schemas_once( state: &AppState, connection_id: &str, database: &str, apply_visible_filter: bool, ) -> Result, String> { let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?; let db_config = connection_config(state, connection_id).await; let visible_schema_filter = visible_schema_filter(db_config.as_ref(), database, apply_visible_filter); { let connections = state.connections.read().await; if let Some(PoolKind::ExternalDriver { config, session, .. }) = connections.get(&pool_key) { let config = config.clone(); let session = session.clone(); drop(connections); return session .invoke::>( "listSchemas", serde_json::json!({ "connection": config.as_ref(), "database": database }), ) .await; } try_sqlserver!(connections, &pool_key, list_schemas); if let Some(client) = extract_pool!(&connections, &pool_key, Agent) { let fallback_config = db_config.clone(); drop(connections); let mut client = client.lock().await; match client .list_schemas_filtered::>( database, visible_schema_filter.as_deref(), agent_metadata_timeout(db_config.as_ref()), ) .await { Ok(schemas) if !schemas.is_empty() => { return Ok(filter_visible_schema_names(schemas, visible_schema_filter.as_deref())) } Ok(schemas) => { if let Some(config) = fallback_config.as_ref() { match native_postgres_metadata_pool(state, connection_id, database, config).await { Ok(Some(pool)) => { return db::postgres::list_schemas(&pool).await.map(|schemas| { filter_visible_schema_names(schemas, visible_schema_filter.as_deref()) }) } Ok(None) => { return Ok(filter_visible_schema_names(schemas, visible_schema_filter.as_deref())) } Err(error) => { log::warn!( "[schema][agent:list_schemas:fallback-failed] connection_id={} database={} error={}", connection_id, database, error ); } } } return Ok(filter_visible_schema_names(schemas, visible_schema_filter.as_deref())); } Err(agent_error) => { if let Some(config) = fallback_config.as_ref() { if let Some(pool) = native_postgres_metadata_pool(state, connection_id, database, config).await? { return db::postgres::list_schemas(&pool) .await .map(|schemas| filter_visible_schema_names(schemas, visible_schema_filter.as_deref())) .map_err(|fallback_error| { format!( "{agent_error}\n\nNative PostgreSQL metadata fallback failed: {fallback_error}" ) }); } } return Err(agent_error); } } } } let connections = state.connections.read().await; let pool = connections.get(&pool_key).ok_or("Pool not found")?; match pool { PoolKind::Mysql(p, mode) if *mode == MysqlMode::OceanBaseOracle => db::ob_oracle::list_schemas(p) .await .map(|schemas| filter_visible_schema_names(schemas, visible_schema_filter.as_deref())), PoolKind::Postgres(p) => db::postgres::list_schemas(p) .await .map(|schemas| filter_visible_schema_names(schemas, visible_schema_filter.as_deref())), #[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) .map(|schemas| filter_visible_schema_names(schemas, visible_schema_filter.as_deref())) } _ => Ok(vec![]), } } fn visible_schema_filter( config: Option<&ConnectionConfig>, database: &str, apply_visible_filter: bool, ) -> Option> { if !apply_visible_filter { return None; } config?.visible_schemas.as_ref()?.get(database).cloned() } fn filter_visible_schema_names(schemas: Vec, visible: Option<&[String]>) -> Vec { let Some(visible) = visible else { return schemas; }; let visible: std::collections::HashSet<&str> = visible.iter().map(String::as_str).collect(); schemas.into_iter().filter(|schema| visible.contains(schema.as_str())).collect() } pub async fn list_tables_core( state: &AppState, connection_id: &str, database: &str, schema: &str, filter: Option<&str>, limit: Option, offset: Option, object_types: Option<&[String]>, ) -> Result, String> { retry_metadata_connection(state, connection_id, Some(database), || { list_tables_once(state, connection_id, database, schema, filter, limit, offset, object_types) }) .await } /// List vector database collections, returning structured info (name, id, dimension). /// Only works for PoolKind::VectorDb connections; returns an error for other types. pub async fn list_vector_collections_core( state: &AppState, connection_id: &str, database: &str, ) -> Result, String> { let pool_key = state.get_or_create_pool(connection_id, if database.is_empty() { None } else { Some(database) }).await?; let client = { let connections = state.connections.read().await; match connections.get(&pool_key) { Some(PoolKind::VectorDb(client)) => client.clone(), _ => return Err("Not a vector database connection".to_string()), } }; db::vector_driver::list_collections_with_db(&client, database).await } pub async fn get_table_comment_core( state: &AppState, connection_id: &str, database: &str, schema: &str, table: &str, ) -> Result, String> { if crate::sql_dialect::parse_sqlserver_linked_schema_ref(schema).is_some() { return Err("Table comments are not available for linked server tables".to_string()); } retry_metadata_connection(state, connection_id, Some(database), || async { let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?; let db_config = connection_config(state, connection_id).await; { let connections = state.connections.read().await; try_sqlserver!(connections, &pool_key, get_table_comment, schema, table); if let Some(client) = extract_pool!(&connections, &pool_key, Agent) { if db_config.as_ref().is_some_and(|config| { matches!(config.db_type, DatabaseType::Oracle | DatabaseType::OceanbaseOracle) }) { let sql = oracle_table_comment_sql(schema, table); let timeout = agent_metadata_timeout(db_config.as_ref()); drop(connections); let mut client = client.lock().await; let result = client .execute_query_with_timeout::( agent_execute_query_params( &sql, Some(database), Some(schema), QueryExecutionOptions { max_rows: Some(1), fetch_size: None, page_size: None, result_session_id: None, client_session_id: None, timeout_secs: None, execution_id: None, }, ), timeout, ) .await?; return oracle_table_comment_from_query_result(result); } } } let connections = state.connections.read().await; let pool = connections.get(&pool_key).ok_or("Pool not found")?; match pool { PoolKind::Mysql(p, mode) if *mode != MysqlMode::OceanBaseOracle && !db_config.as_ref().is_some_and(is_doris_family_config) && !db_config.as_ref().is_some_and(is_manticoresearch_config) => { db::mysql::get_table_comment(p, schema, table).await } PoolKind::Postgres(p) if !db_config.as_ref().is_some_and(is_questdb_config) => { db::postgres::get_table_comment(p, schema, table).await } _ => Err("Table comment lookup is not supported for this connection".to_string()), } }) .await } fn oracle_table_comment_sql(schema: &str, table: &str) -> String { format!( "SELECT COMMENTS FROM ALL_TAB_COMMENTS WHERE OWNER = {} AND TABLE_NAME = {} AND TABLE_TYPE IN ('TABLE', 'VIEW')", sql_string(schema), sql_string(table), ) } fn oracle_table_comment_from_query_result(result: db::QueryResult) -> Result, String> { Ok(result .rows .first() .and_then(|row| row.iter().find_map(|value| value.as_str().map(str::to_string))) .filter(|value| !value.trim().is_empty())) } async fn list_tables_once( state: &AppState, connection_id: &str, database: &str, schema: &str, filter: Option<&str>, limit: Option, offset: Option, object_types: Option<&[String]>, ) -> Result, String> { let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?; #[cfg(feature = "duckdb-bundled")] let duckdb_attached_names = duckdb_attached_database_names(state, connection_id).await; let db_config = connection_config(state, connection_id).await; { let connections = state.connections.read().await; #[cfg(feature = "duckdb-bundled")] if let Some(ext_pool) = extract_pool!(&connections, &pool_key, ExternalTabular) { drop(connections); let cache = ext_pool.cache.clone(); return tokio::task::spawn_blocking(move || { let con = cache.lock().map_err(|e| e.to_string())?; duckdb_query_tables(&con) }) .await .map_err(|e| e.to_string())? .map(|tables| filter_table_infos(tables, filter, limit, offset, object_types)); } if let Some(PoolKind::ExternalDriver { config, session, .. }) = connections.get(&pool_key) { let config = config.clone(); let session = session.clone(); drop(connections); if uses_presto_like_information_schema_tables(&config.db_type) { return external_driver_presto_like_tables(session, config.as_ref(), database, schema) .await .map(|tables| filter_table_infos(tables, filter, limit, offset, object_types)); } return session .invoke::>( "listTables", serde_json::json!({ "connection": config.as_ref(), "database": database, "schema": schema }), ) .await .map(|tables| filter_table_infos(tables, filter, limit, offset, object_types)); } #[cfg(feature = "duckdb-bundled")] if let Some(con) = extract_pool!(&connections, &pool_key, DuckDb) { drop(connections); let con = con.lock().map_err(|e| e.to_string())?; return duckdb_query_tables_in_database_with_attached(&con, database, schema, &duckdb_attached_names) .map(|tables| filter_table_infos(tables, filter, limit, offset, object_types)); } if let Some(client) = extract_pool!(&connections, &pool_key, ClickHouse) { drop(connections); return db::clickhouse_driver::list_tables(&client, clickhouse_metadata_database(database, schema)) .await .map(|tables| filter_table_infos(tables, filter, limit, offset, object_types)); } if let Some(client) = extract_pool!(&connections, &pool_key, InfluxDb) { drop(connections); return db::influxdb_driver::list_tables(&client, database) .await .map(|tables| filter_table_infos(tables, filter, limit, offset, object_types)); } if let Some(linked) = crate::sql_dialect::parse_sqlserver_linked_schema_ref(schema) { if let Some(client) = extract_pool!(&connections, &pool_key, SqlServer) { drop(connections); let mut client = client.lock().await; return db::sqlserver::list_linked_server_tables( &mut client, &linked.server, &linked.catalog, &linked.schema, filter, None, None, ) .await .map(|tables| filter_table_infos(tables, filter, limit, offset, object_types)); } } if object_types.is_some() { if let Some(client) = extract_pool!(&connections, &pool_key, SqlServer) { drop(connections); let mut client = client.lock().await; return db::sqlserver::list_tables(&mut client, schema, filter, None, None) .await .map(|tables| filter_table_infos(tables, filter, limit, offset, object_types)); } } try_sqlserver!(connections, &pool_key, list_tables, schema, filter, limit, offset); if let Some(client) = extract_pool!(&connections, &pool_key, Agent) { let fallback_config = db_config.clone(); drop(connections); let mut client = client.lock().await; match client .list_tables_filtered::>( database, schema, object_types, agent_metadata_timeout(db_config.as_ref()), ) .await { Ok(tables) if !tables.is_empty() => { return Ok(filter_table_infos(tables, filter, limit, offset, object_types)) } Ok(tables) => { if let Some(config) = fallback_config.as_ref() { match native_postgres_metadata_pool(state, connection_id, database, config).await { Ok(Some(pool)) => { return if object_types.is_some() { db::postgres::list_tables_filtered(&pool, schema, filter, None, None) .await .map(|tables| filter_table_infos(tables, filter, limit, offset, object_types)) } else { db::postgres::list_tables_filtered(&pool, schema, filter, limit, offset).await }; } Ok(None) => return Ok(filter_table_infos(tables, filter, limit, offset, object_types)), Err(error) => { log::warn!( "[schema][agent:list_tables:fallback-failed] connection_id={} database={} schema={} error={}", connection_id, database, schema, error ); } } } return Ok(filter_table_infos(tables, filter, limit, offset, object_types)); } Err(agent_error) => { if let Some(config) = fallback_config.as_ref() { if let Some(pool) = native_postgres_metadata_pool(state, connection_id, database, config).await? { let result = if object_types.is_some() { db::postgres::list_tables_filtered(&pool, schema, filter, None, None) .await .map(|tables| filter_table_infos(tables, filter, limit, offset, object_types)) } else { db::postgres::list_tables_filtered(&pool, schema, filter, limit, offset).await }; return result.map_err(|fallback_error| { format!("{agent_error}\n\nNative PostgreSQL metadata fallback failed: {fallback_error}") }); } } return Err(agent_error); } } } } 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(tables, filter, limit, offset, object_types)) } PoolKind::Mysql(p, mode) => { if *mode == MysqlMode::OceanBaseOracle { let tables = db::ob_oracle::list_tables(p, schema).await?; Ok(filter_table_infos(tables, filter, limit, offset, object_types)) } else { db::mysql::list_tables_filtered( p, mysql_table_metadata_catalog(database, schema), filter, limit, offset, object_types, ) .await .map(|tables| filter_table_infos(tables, None, None, None, object_types)) } } PoolKind::Postgres(p) if db_config.as_ref().is_some_and(is_questdb_config) => { db::questdb::list_tables(p, schema) .await .map(|tables| filter_table_infos(tables, filter, limit, offset, object_types)) } PoolKind::Postgres(p) => { if object_types.is_some() { db::postgres::list_tables_filtered(p, schema, filter, None, None) .await .map(|tables| filter_table_infos(tables, filter, limit, offset, object_types)) } else { db::postgres::list_tables_filtered(p, schema, filter, limit, offset).await } } PoolKind::Sqlite(p) => db::sqlite::list_tables(p, schema) .await .map(|tables| filter_table_infos(tables, filter, limit, offset, object_types)), PoolKind::Rqlite(client) => db::rqlite_driver::list_tables(client, schema) .await .map(|tables| filter_table_infos(tables, filter, limit, offset, object_types)), PoolKind::MongoDb(client) => db::mongo_driver::list_collections(client, database) .await .map(|names| collection_names_to_tables(names, "COLLECTION")) .map(|tables| filter_table_infos(tables, filter, limit, offset, object_types)), PoolKind::Elasticsearch(client) => db::elasticsearch_driver::list_indices(client) .await .map(|names| collection_names_to_tables(names, "INDEX")) .map(|tables| filter_table_infos(tables, filter, limit, offset, object_types)), PoolKind::VectorDb(client) => db::vector_driver::list_collections(client) .await .map(|infos| collection_names_to_tables(infos.into_iter().map(|i| i.name).collect(), "COLLECTION")) .map(|tables| filter_table_infos(tables, filter, limit, offset, object_types)), _ => 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, offset: Option, object_types: Option<&[String]>, ) -> Vec { let filter = filter.unwrap_or(""); let limit = limit.unwrap_or(usize::MAX); let offset = offset.unwrap_or(0); tables .into_iter() .filter(|table| crate::sql::contains_or_fuzzy_match(&table.name, filter)) .filter(|table| table_info_matches_object_types(table, object_types)) .skip(offset) .take(limit) .collect() } fn table_info_matches_object_types(table: &db::TableInfo, object_types: Option<&[String]>) -> bool { let Some(object_types) = object_types else { return true; }; if object_types.is_empty() { return true; } let table_type = normalize_table_info_object_type(&table.table_type); object_types.iter().any(|object_type| normalize_table_info_object_type(object_type) == table_type) } fn normalize_table_info_object_type(value: &str) -> String { let upper = value.to_ascii_uppercase().replace(' ', "_"); if upper.contains("MATERIALIZED") && upper.contains("VIEW") { return "MATERIALIZED_VIEW".to_string(); } if upper.contains("VIEW") { return "VIEW".to_string(); } if upper.contains("COLLECTION") { return "COLLECTION".to_string(); } if upper.contains("INDEX") { return "INDEX".to_string(); } "TABLE".to_string() } fn uses_presto_like_information_schema_tables(db_type: &DatabaseType) -> bool { matches!(db_type, DatabaseType::PrestoSql | DatabaseType::Trino) } async fn external_driver_presto_like_tables( session: Arc, config: &ConnectionConfig, database: &str, schema: &str, ) -> Result, String> { let result: db::QueryResult = session .invoke( "executeQuery", serde_json::json!({ "connection": config, "database": database, "schema": schema, "sql": presto_like_information_schema_tables_sql(database, schema), "maxRows": 100000, "fetchSize": 1000, "timeoutSecs": 60 }), ) .await?; Ok(presto_like_tables_from_query_result(&result)) } async fn external_driver_presto_like_objects( session: Arc, config: &ConnectionConfig, database: &str, schema: &str, ) -> Result, String> { let tables = external_driver_presto_like_tables(session, config, database, schema).await?; Ok(tables .into_iter() .map(|table| db::ObjectInfo { name: table.name, object_type: table.table_type, schema: 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 presto_like_information_schema_tables_sql(database: &str, schema: &str) -> String { let source = if database.trim().is_empty() { "information_schema.tables".to_string() } else { format!("{}.information_schema.tables", quote_presto_like_identifier(database)) }; format!( "SELECT table_name, CASE table_type WHEN 'BASE TABLE' THEN 'TABLE' ELSE table_type END AS table_type \ FROM {source} \ WHERE table_schema = {} AND table_type IN ('BASE TABLE', 'VIEW') \ ORDER BY table_type, table_name", sql_string_literal(schema) ) } fn presto_like_tables_from_query_result(result: &db::QueryResult) -> Vec { result .rows .iter() .filter_map(|row| { let name = query_result_cell_string(row, 0)?; if name.trim().is_empty() { return None; } Some(db::TableInfo { name, table_type: normalize_information_schema_table_type( query_result_cell_string(row, 1).as_deref().unwrap_or("TABLE"), ), comment: None, parent_schema: None, parent_name: None, }) }) .collect() } fn query_result_cell_string(row: &[serde_json::Value], index: usize) -> Option { let value = row.get(index)?; if value.is_null() { return None; } value.as_str().map(ToString::to_string).or_else(|| Some(value.to_string())) } fn normalize_information_schema_table_type(table_type: &str) -> String { match table_type.trim().to_ascii_uppercase().replace(' ', "_").as_str() { "BASE_TABLE" => "TABLE".to_string(), "VIEW" => "VIEW".to_string(), "MATERIALIZED_VIEW" => "MATERIALIZED_VIEW".to_string(), _ => table_type.to_string(), } } fn mysql_table_metadata_catalog<'a>(database: &'a str, _schema: &str) -> &'a str { database } fn quote_presto_like_identifier(identifier: &str) -> String { format!("\"{}\"", identifier.replace('"', "\"\"")) } fn sql_string_literal(value: &str) -> String { format!("'{}'", value.replace('\'', "''")) } #[cfg(test)] mod tests { use super::db; use super::{ clickhouse_metadata_database, deduplicate_column_infos, filter_mysql_system_databases_for_config, filter_table_infos, filter_visible_schema_names, is_agent_postgres_metadata_fallback_config, is_retryable_metadata_error, mysql_table_metadata_catalog, normalize_information_schema_table_type, oracle_table_comment_from_query_result, oracle_table_comment_sql, presto_like_information_schema_tables_sql, presto_like_tables_from_query_result, visible_schema_filter, }; #[cfg(feature = "duckdb-bundled")] use super::{ duckdb_attach_database, duckdb_completion_assistant_search, duckdb_list_databases, duckdb_query_tables_in_database, }; use crate::models::connection::{ConnectionConfig, DatabaseType}; use std::collections::HashMap; 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(), 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, visible_schemas: None, attached_databases: Vec::new(), color: None, transport_layers: Vec::new(), connect_timeout_secs: 5, query_timeout_secs: 30, idle_timeout_secs: 60, keepalive_interval_secs: 0, 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: crate::models::connection::default_redis_key_separator(), redis_scan_page_size: None, etcd_endpoints: String::new(), gbase_server: String::new(), informix_server: String::new(), external_config: None, jdbc_driver_class: None, jdbc_driver_paths: Vec::new(), one_time: false, read_only: false, } } #[test] fn mysql_table_child_metadata_uses_database_not_schema() { assert_eq!(mysql_table_metadata_catalog("app_db", ""), "app_db"); assert_eq!(mysql_table_metadata_catalog("app_db", "public"), "app_db"); } #[test] fn metadata_retry_recovers_missing_pool_only_as_transient_state() { assert!(is_retryable_metadata_error("Pool not found")); assert!(is_retryable_metadata_error("connection reset by peer")); assert!(!is_retryable_metadata_error("Unknown column 'email' in 'field list'")); assert!(!is_retryable_metadata_error("Access denied for user")); } #[test] fn visible_schema_filter_only_applies_when_requested() { let mut config = test_connection_config(DatabaseType::Oracle); config.visible_schemas = Some(HashMap::from([("ORCLPDB1".to_string(), vec!["APP".to_string(), "REPORTING".to_string()])])); assert_eq!(visible_schema_filter(Some(&config), "ORCLPDB1", false), None); assert_eq!( visible_schema_filter(Some(&config), "ORCLPDB1", true), Some(vec!["APP".to_string(), "REPORTING".to_string()]) ); assert_eq!(visible_schema_filter(Some(&config), "OTHER", true), None); } #[test] fn filter_visible_schema_names_preserves_database_order() { let schemas = vec!["APP".to_string(), "SYS".to_string(), "REPORTING".to_string()]; let visible = vec!["REPORTING".to_string(), "APP".to_string()]; assert_eq!(filter_visible_schema_names(schemas, Some(&visible)), vec!["APP", "REPORTING"]); } fn test_table_info(name: &str) -> super::db::TableInfo { super::db::TableInfo { name: name.to_string(), table_type: "BASE TABLE".to_string(), comment: None, parent_schema: None, parent_name: None, } } fn test_database_info(name: &str) -> super::db::DatabaseInfo { super::db::DatabaseInfo { name: name.to_string() } } #[test] fn manticoresearch_database_list_filters_mysql_system_databases() { let databases = vec![ test_database_info("Manticore"), test_database_info("information_schema"), test_database_info("mysql"), test_database_info("performance_schema"), test_database_info("sys"), ]; let config = test_connection_config(DatabaseType::ManticoreSearch); let filtered = filter_mysql_system_databases_for_config(databases, Some(&config)); assert_eq!(filtered.into_iter().map(|database| database.name).collect::>(), vec!["Manticore"]); } #[test] fn manticoresearch_show_metadata_uses_unqualified_table_names() { let config = test_connection_config(DatabaseType::ManticoreSearch); assert_eq!(super::mysql_show_metadata_database_for_config(Some(&config), "Manticore"), ""); } #[test] fn doris_show_metadata_keeps_database_qualifier() { let config = test_connection_config(DatabaseType::Doris); assert_eq!(super::mysql_show_metadata_database_for_config(Some(&config), "analytics"), "analytics"); } #[test] fn doris_database_list_keeps_system_databases() { let databases = vec![test_database_info("information_schema"), test_database_info("analytics")]; let config = test_connection_config(DatabaseType::Doris); let filtered = filter_mysql_system_databases_for_config(databases, Some(&config)); assert_eq!( filtered.into_iter().map(|database| database.name).collect::>(), vec!["information_schema", "analytics"] ); } #[test] fn filter_table_infos_applies_filter_offset_and_limit() { let tables = vec![ test_table_info("alpha"), test_table_info("audit_log"), test_table_info("audit_record"), test_table_info("users"), ]; let filtered = filter_table_infos(tables, Some("audit"), Some(1), Some(1), None); assert_eq!(filtered.len(), 1); assert_eq!(filtered[0].name, "audit_record"); } #[test] fn filter_table_infos_matches_fuzzy_subsequences() { let tables = vec![test_table_info("system_user"), test_table_info("user_order"), test_table_info("alpha")]; let system_user = filter_table_infos(tables.clone(), Some("sysu"), None, None, None); assert_eq!(system_user.into_iter().map(|table| table.name).collect::>(), vec!["system_user"]); let user_order = filter_table_infos(tables, Some("uo"), None, None, None); assert_eq!(user_order.into_iter().map(|table| table.name).collect::>(), vec!["user_order"]); } #[test] fn filter_table_infos_skips_fuzzy_for_single_character_filters() { let tables = vec![test_table_info("orders"), test_table_info("user_order")]; let filtered = filter_table_infos(tables, Some("u"), None, None, None); assert_eq!(filtered.into_iter().map(|table| table.name).collect::>(), vec!["user_order"]); } #[test] fn filter_table_infos_keeps_special_filter_characters_literal() { let tables = vec![test_table_info("user_%"), test_table_info("user_account"), test_table_info("userXpercent")]; let filtered = filter_table_infos(tables, Some("user_%"), None, None, None); assert_eq!(filtered.into_iter().map(|table| table.name).collect::>(), vec!["user_%"]); } #[test] fn filter_table_infos_filters_object_type_before_offset_and_limit() { let tables = vec![ test_table_info("orders"), super::db::TableInfo { name: "active_orders".to_string(), table_type: "VIEW".to_string(), comment: None, parent_schema: None, parent_name: None, }, test_table_info("users"), super::db::TableInfo { name: "active_users".to_string(), table_type: "VIEW".to_string(), comment: None, parent_schema: None, parent_name: None, }, ]; let object_types = vec!["VIEW".to_string()]; let filtered = filter_table_infos(tables, None, Some(1), Some(1), Some(&object_types)); assert_eq!(filtered.len(), 1); assert_eq!(filtered[0].name, "active_users"); } #[test] fn presto_like_information_schema_sql_uses_catalog_and_schema_without_system_jdbc() { let sql = presto_like_information_schema_tables_sql("hive", "sales_analytics"); assert_eq!( sql, "SELECT table_name, CASE table_type WHEN 'BASE TABLE' THEN 'TABLE' ELSE table_type END AS table_type FROM \"hive\".information_schema.tables WHERE table_schema = 'sales_analytics' AND table_type IN ('BASE TABLE', 'VIEW') ORDER BY table_type, table_name" ); assert!(!sql.contains("system.jdbc.tables")); } #[test] fn presto_like_information_schema_sql_escapes_identifiers_and_literals() { let sql = presto_like_information_schema_tables_sql("hi\"ve", "sales'analytics"); assert!(sql.contains("\"hi\"\"ve\".information_schema.tables")); assert!(sql.contains("table_schema = 'sales''analytics'")); } #[test] fn presto_like_tables_from_query_result_normalizes_base_table_type() { let result = super::db::QueryResult { columns: vec!["table_name".to_string(), "table_type".to_string()], column_types: vec![], column_sortables: vec![], rows: vec![ vec![serde_json::json!("daily_revenue"), serde_json::json!("BASE TABLE")], vec![serde_json::json!("revenue_view"), serde_json::json!("VIEW")], ], affected_rows: 0, execution_time_ms: 1, truncated: false, session_id: None, has_more: false, }; let tables = presto_like_tables_from_query_result(&result); assert_eq!(tables[0].name, "daily_revenue"); assert_eq!(tables[0].table_type, "TABLE"); assert_eq!(tables[1].name, "revenue_view"); assert_eq!(tables[1].table_type, "VIEW"); assert_eq!(normalize_information_schema_table_type("MATERIALIZED VIEW"), "MATERIALIZED_VIEW"); } #[cfg(feature = "duckdb-bundled")] #[test] fn duckdb_list_databases_includes_attached_database() { let unique = uuid::Uuid::new_v4(); let path = std::env::temp_dir().join(format!("dbx-attached-{unique}.duckdb")); let _ = std::fs::remove_file(&path); let con = duckdb::Connection::open_in_memory().unwrap(); duckdb_attach_database(&con, "analytics", path.to_str().unwrap()).unwrap(); let databases = duckdb_list_databases(&con).unwrap(); assert!(databases.iter().any(|database| database.name == "main")); assert!(databases.iter().any(|database| database.name == "analytics")); let _ = std::fs::remove_file(path); } #[cfg(feature = "duckdb-bundled")] #[test] fn duckdb_query_tables_filters_by_attached_database() { let unique = uuid::Uuid::new_v4(); let path = std::env::temp_dir().join(format!("dbx-attached-tables-{unique}.duckdb")); let _ = std::fs::remove_file(&path); let con = duckdb::Connection::open_in_memory().unwrap(); con.execute_batch("CREATE TABLE main_table(id INTEGER);").unwrap(); duckdb_attach_database(&con, "analytics", path.to_str().unwrap()).unwrap(); con.execute_batch("CREATE TABLE analytics.attached_table(id INTEGER);").unwrap(); let main_tables = duckdb_query_tables_in_database(&con, "main", "main").unwrap(); let attached_tables = duckdb_query_tables_in_database(&con, "analytics", "main").unwrap(); assert!(main_tables.iter().any(|table| table.name == "main_table")); assert!(!main_tables.iter().any(|table| table.name == "attached_table")); assert!(attached_tables.iter().any(|table| table.name == "attached_table")); assert!(!attached_tables.iter().any(|table| table.name == "main_table")); let _ = std::fs::remove_file(path); } #[cfg(feature = "duckdb-bundled")] #[test] fn duckdb_completion_assistant_searches_catalog_metadata_with_limit() { let con = duckdb::Connection::open_in_memory().unwrap(); con.execute_batch( "CREATE TABLE account(id INTEGER, display_name VARCHAR); CREATE VIEW account_view AS SELECT id FROM account;", ) .unwrap(); let request = db::CompletionAssistantRequest { connection_id: "c1".to_string(), database: "main".to_string(), schema: Some("main".to_string()), object_kinds: vec![db::CompletionAssistantObjectKind::Table, db::CompletionAssistantObjectKind::View], mask: "account".to_string(), case_sensitive: false, global_search: false, max_results: Some(1), search_in_comments: false, search_in_definitions: false, parent_schema: Some("main".to_string()), parent_name: None, match_mode: Some(db::CompletionAssistantMatchMode::Prefix), }; let tables = duckdb_completion_assistant_search(&con, &request, &[]).unwrap(); assert_eq!(tables.candidates.len(), 1); assert!(tables.incomplete); assert!(!tables.fallback_used); assert_eq!(tables.candidates[0].name, "account"); let columns = duckdb_completion_assistant_search( &con, &db::CompletionAssistantRequest { object_kinds: vec![db::CompletionAssistantObjectKind::Column], mask: "name".to_string(), max_results: Some(10), parent_name: Some("account".to_string()), match_mode: Some(db::CompletionAssistantMatchMode::Contains), ..request }, &[], ) .unwrap(); assert_eq!(columns.candidates.len(), 1); assert_eq!(columns.candidates[0].name, "display_name"); } #[test] fn detects_unsupported_agent_completion_assistant_errors() { assert!(super::is_agent_completion_assistant_unsupported( "Agent RPC error (-1): Unknown method: completion_assistant_search_v1" )); assert!(super::is_agent_completion_assistant_unsupported( "Agent RPC error (-1): Completion assistant search is not supported by this agent" )); assert!(!super::is_agent_completion_assistant_unsupported("Agent RPC error (-1): Connection failed")); } #[test] fn clickhouse_metadata_uses_schema_when_database_is_empty() { assert_eq!(clickhouse_metadata_database("", "testdb"), "testdb"); assert_eq!(clickhouse_metadata_database("testdb", ""), "testdb"); 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))); assert!(is_agent_postgres_metadata_fallback_config(&test_connection_config(DatabaseType::Highgo))); assert!(is_agent_postgres_metadata_fallback_config(&test_connection_config(DatabaseType::Vastbase))); 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 agent_metadata_timeout_defaults_to_sixty_seconds_and_honors_longer_config() { assert_eq!(super::agent_metadata_timeout(None), Some(std::time::Duration::from_secs(60))); let mut config = test_connection_config(DatabaseType::Oracle); assert_eq!(super::agent_metadata_timeout(Some(&config)), Some(std::time::Duration::from_secs(60))); config.query_timeout_secs = 120; assert_eq!(super::agent_metadata_timeout(Some(&config)), Some(std::time::Duration::from_secs(120))); config.query_timeout_secs = 0; assert_eq!(super::agent_metadata_timeout(Some(&config)), None); } #[test] fn oracle_table_comment_sql_targets_single_table_and_escapes_literals() { let sql = oracle_table_comment_sql("APP'S", "USER'S"); assert!(sql.contains("ALL_TAB_COMMENTS")); assert!(sql.contains("OWNER = 'APP''S'")); assert!(sql.contains("TABLE_NAME = 'USER''S'")); assert!(sql.contains("TABLE_TYPE IN ('TABLE', 'VIEW')")); assert!(!sql.contains("ALL_OBJECTS")); } #[test] fn oracle_table_comment_from_query_result_returns_optional_non_blank_comment() { let result = db::QueryResult { columns: vec!["COMMENTS".to_string()], column_types: Vec::new(), column_sortables: Vec::new(), rows: vec![vec![serde_json::json!("Customer table")]], affected_rows: 0, execution_time_ms: 0, truncated: false, session_id: None, has_more: false, }; assert_eq!(oracle_table_comment_from_query_result(result).unwrap().as_deref(), Some("Customer table")); let empty = db::QueryResult { columns: vec!["COMMENTS".to_string()], column_types: Vec::new(), column_sortables: Vec::new(), rows: vec![vec![serde_json::json!(" ")]], affected_rows: 0, execution_time_ms: 0, truncated: false, session_id: None, has_more: false, }; assert_eq!(oracle_table_comment_from_query_result(empty).unwrap(), None); } } pub async fn list_objects_core( state: &AppState, connection_id: &str, database: &str, schema: &str, ) -> Result, String> { retry_metadata_connection(state, connection_id, Some(database), || { list_objects_once(state, connection_id, database, schema) }) .await } pub async fn list_object_statistics_core( state: &AppState, connection_id: &str, database: &str, schema: &str, ) -> Result, String> { retry_metadata_connection(state, connection_id, Some(database), || { list_object_statistics_once(state, connection_id, database, schema) }) .await } pub async fn list_completion_objects_core( state: &AppState, connection_id: &str, database: &str, schema: &str, ) -> Result, String> { retry_metadata_connection(state, connection_id, Some(database), || { list_completion_objects_once(state, connection_id, database, schema) }) .await } pub async fn completion_assistant_search_core( state: &AppState, request: db::CompletionAssistantRequest, ) -> Result { let started_at = Instant::now(); let request_summary = format!( "connection_id={} database={} schema={:?} kinds={:?} mask={} limit={:?}", request.connection_id, request.database, request.schema, request.object_kinds, request.mask, request.max_results ); retry_metadata_connection(state, &request.connection_id, Some(&request.database), || async { let pool_key = state.get_or_create_pool(&request.connection_id, Some(&request.database)).await?; log::debug!("[schema][completion_assistant:start] {request_summary}"); { let connections = state.connections.read().await; try_sqlserver!(connections, &pool_key, completion_assistant_search, &request); } { let connections = state.connections.read().await; if let Some(pool) = connections.get(&pool_key).and_then(|pool| match pool { PoolKind::Sqlite(pool) => Some(pool.clone()), _ => None, }) { drop(connections); return db::sqlite::completion_assistant_search(&pool, &request).await; } } #[cfg(feature = "duckdb-bundled")] { let duckdb_attached_names = duckdb_attached_database_names(state, &request.connection_id).await; let connections = state.connections.read().await; if let Some(con) = extract_pool!(&connections, &pool_key, DuckDb) { drop(connections); let con = con.lock().map_err(|e| e.to_string())?; return duckdb_completion_assistant_search(&con, &request, &duckdb_attached_names); } } { let connections = state.connections.read().await; if let Some(pool) = connections.get(&pool_key).and_then(|pool| match pool { PoolKind::Postgres(pool) => Some(pool.clone()), _ => None, }) { drop(connections); return db::postgres::completion_assistant_search(&pool, &request).await; } } { let connections = state.connections.read().await; if let Some(pool) = connections.get(&pool_key).and_then(|pool| match pool { PoolKind::Mysql(pool, mode) if *mode != MysqlMode::OceanBaseOracle => Some(pool.clone()), _ => None, }) { drop(connections); return db::mysql::completion_assistant_search(&pool, &request).await; } } { let connections = state.connections.read().await; if let Some(client) = extract_pool!(&connections, &pool_key, Agent) { let db_config = connection_config(state, &request.connection_id).await; drop(connections); let mut client = client.lock().await; match client .completion_assistant_search::( &request, agent_metadata_timeout(db_config.as_ref()), ) .await { Ok(mut response) => { response.fallback_used = false; return Ok(response); } Err(error) if is_agent_completion_assistant_unsupported(&error) => { log::debug!( "[schema][completion_assistant:agent-fallback] {} reason={}", request_summary, error ); } Err(error) => return Err(error), } } } let response = completion_assistant_fallback_core(state, &request).await; if let Ok(response) = &response { log::debug!( "[schema][completion_assistant:done] {} elapsed_ms={} candidates={} fallback_used={}", request_summary, started_at.elapsed().as_millis(), response.candidates.len(), response.fallback_used ); } response }) .await } fn is_agent_completion_assistant_unsupported(error: &str) -> bool { error.contains("Unknown method: completion_assistant_search_v1") || error.contains("Method not found: completion_assistant_search_v1") || error.contains("method not found: completion_assistant_search_v1") || error.contains("Completion assistant search is not supported") } async fn completion_assistant_fallback_core( state: &AppState, request: &db::CompletionAssistantRequest, ) -> Result { let limit = request.max_results.unwrap_or(100).clamp(1, 1000); let kinds = if request.object_kinds.is_empty() { vec![db::CompletionAssistantObjectKind::Table, db::CompletionAssistantObjectKind::View] } else { request.object_kinds.clone() }; let mut candidates = Vec::new(); let schema = request.parent_schema.as_deref().or(request.schema.as_deref()).unwrap_or(""); let filter = request.mask.trim().trim_matches('%'); if kinds.iter().any(|kind| matches!(kind, db::CompletionAssistantObjectKind::Schema)) { let schemas = list_schemas_core(state, &request.connection_id, &request.database).await?; for schema_name in schemas { if completion_name_matches(&schema_name, filter, request.match_mode.as_ref()) { candidates.push(db::CompletionAssistantCandidate { name: schema_name.clone(), kind: db::CompletionAssistantCandidateKind::Schema, database: Some(request.database.clone()), schema: Some(schema_name), parent_schema: None, parent_name: None, comment: None, data_type: None, }); } if candidates.len() >= limit { return Ok(db::CompletionAssistantResponse { candidates, incomplete: true, fallback_used: true }); } } } if kinds.iter().any(db::CompletionAssistantObjectKind::is_table_like) { let object_types = completion_table_object_types(&kinds); let tables = list_tables_core( state, &request.connection_id, &request.database, schema, if filter.is_empty() { None } else { Some(filter) }, Some(limit), None, object_types.as_deref(), ) .await?; for table in tables { let kind = if table.table_type.to_uppercase().contains("VIEW") { db::CompletionAssistantCandidateKind::View } else { db::CompletionAssistantCandidateKind::Table }; candidates.push(db::CompletionAssistantCandidate { name: table.name, kind, database: Some(request.database.clone()), schema: if schema.is_empty() { None } else { Some(schema.to_string()) }, parent_schema: table.parent_schema, parent_name: table.parent_name, comment: table.comment, data_type: None, }); if candidates.len() >= limit { return Ok(db::CompletionAssistantResponse { candidates, incomplete: true, fallback_used: true }); } } } if kinds.iter().any(|kind| matches!(kind, db::CompletionAssistantObjectKind::Column)) { if let Some(table) = request.parent_name.as_deref().filter(|table| !table.trim().is_empty()) { let columns = get_columns_core(state, &request.connection_id, &request.database, schema, table).await?; for column in columns { if completion_name_matches(&column.name, filter, request.match_mode.as_ref()) { candidates.push(db::CompletionAssistantCandidate { name: column.name, kind: db::CompletionAssistantCandidateKind::Column, database: Some(request.database.clone()), schema: if schema.is_empty() { None } else { Some(schema.to_string()) }, parent_schema: if schema.is_empty() { None } else { Some(schema.to_string()) }, parent_name: Some(table.to_string()), comment: column.comment, data_type: Some(column.data_type), }); } if candidates.len() >= limit { return Ok(db::CompletionAssistantResponse { candidates, incomplete: true, fallback_used: true }); } } } } Ok(db::CompletionAssistantResponse { candidates, incomplete: false, fallback_used: true }) } fn completion_table_object_types(kinds: &[db::CompletionAssistantObjectKind]) -> Option> { let mut object_types = Vec::new(); if kinds.iter().any(|kind| matches!(kind, db::CompletionAssistantObjectKind::Table)) { object_types.push("table".to_string()); } if kinds.iter().any(|kind| matches!(kind, db::CompletionAssistantObjectKind::View)) { object_types.push("view".to_string()); } if object_types.is_empty() { None } else { Some(object_types) } } fn completion_name_matches(name: &str, filter: &str, mode: Option<&db::CompletionAssistantMatchMode>) -> bool { if filter.is_empty() { return true; } let name = name.to_lowercase(); let filter = filter.to_lowercase(); match mode.unwrap_or(&db::CompletionAssistantMatchMode::Prefix) { db::CompletionAssistantMatchMode::Prefix => name.starts_with(&filter), db::CompletionAssistantMatchMode::Contains => name.contains(&filter), } } async fn list_object_statistics_once( state: &AppState, connection_id: &str, database: &str, schema: &str, ) -> Result, String> { let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?; let db_config = connection_config(state, connection_id).await; let connections = state.connections.read().await; try_sqlserver!(connections, &pool_key, list_object_statistics, schema); let pool = connections.get(&pool_key).ok_or("Pool not found")?; match pool { PoolKind::Mysql(p, mode) => { if *mode == MysqlMode::OceanBaseOracle || db_config.as_ref().is_some_and(is_manticoresearch_config) { Ok(vec![]) } else { db::mysql::list_object_statistics(p, database).await } } PoolKind::Postgres(p) if db_config.as_ref().is_some_and(is_questdb_config) => Ok(vec![]), PoolKind::Postgres(p) => db::postgres::list_object_statistics(p, schema).await, PoolKind::ClickHouse(client) => { db::clickhouse_driver::list_object_statistics(client, clickhouse_metadata_database(database, schema)).await } _ => Ok(vec![]), } } async fn list_objects_once( state: &AppState, connection_id: &str, database: &str, schema: &str, ) -> Result, String> { let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?; let db_config = connection_config(state, connection_id).await; { let connections = state.connections.read().await; #[cfg(feature = "duckdb-bundled")] if let Some(ext_pool) = extract_pool!(&connections, &pool_key, ExternalTabular) { drop(connections); let cache = ext_pool.cache.clone(); return tokio::task::spawn_blocking(move || { let con = cache.lock().map_err(|e| e.to_string())?; Ok(duckdb_query_tables(&con)? .into_iter() .map(|table| db::ObjectInfo { name: table.name, object_type: table.table_type, schema: None, comment: table.comment, created_at: None, updated_at: None, parent_schema: table.parent_schema, parent_name: table.parent_name, }) .collect()) }) .await .map_err(|e| e.to_string())?; } if let Some(PoolKind::ExternalDriver { config, session, .. }) = connections.get(&pool_key) { let config = config.clone(); let session = session.clone(); drop(connections); if uses_presto_like_information_schema_tables(&config.db_type) { return external_driver_presto_like_objects(session, config.as_ref(), database, schema).await; } return session .invoke::>( "listObjects", serde_json::json!({ "connection": config.as_ref(), "database": database, "schema": schema }), ) .await; } try_sqlserver!(connections, &pool_key, list_objects, schema); if let Some(client) = extract_pool!(&connections, &pool_key, Agent) { let is_oracle = db_config.as_ref().is_some_and(|config| config.db_type == DatabaseType::Oracle); let fallback_config = db_config.clone(); drop(connections); if is_oracle { return oracle_agent_list_objects(client, database, schema, agent_metadata_timeout(db_config.as_ref())) .await; } let mut client = client.lock().await; match client .list_objects::>(database, schema, agent_metadata_timeout(db_config.as_ref())) .await { Ok(objects) if !objects.is_empty() => return Ok(objects), Ok(objects) => { if let Some(config) = fallback_config.as_ref() { match native_postgres_metadata_pool(state, connection_id, database, config).await { Ok(Some(pool)) => return db::postgres::list_objects(&pool, schema).await, Ok(None) => return Ok(objects), Err(error) => { log::warn!( "[schema][agent:list_objects:fallback-failed] connection_id={} database={} schema={} error={}", connection_id, database, schema, error ); } } } return Ok(objects); } Err(agent_error) => { if let Some(config) = fallback_config.as_ref() { if let Some(pool) = native_postgres_metadata_pool(state, connection_id, database, config).await? { return db::postgres::list_objects(&pool, schema).await.map_err(|fallback_error| { format!("{agent_error}\n\nNative PostgreSQL metadata fallback failed: {fallback_error}") }); } } return Err(agent_error); } } } } 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_manticoresearch_config) { db::manticoresearch::list_objects(p, database).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) if db_config.as_ref().is_some_and(is_questdb_config) => { db::questdb::list_objects(p, schema).await } PoolKind::Postgres(p) => db::postgres::list_objects(p, schema).await, _ => { drop(connections); Ok(list_tables_core(state, connection_id, database, schema, None, None, 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( state: &AppState, connection_id: &str, database: &str, schema: &str, ) -> Result, String> { let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?; let db_config = connection_config(state, connection_id).await; let connections = state.connections.read().await; if let Some(PoolKind::ExternalDriver { config, session, .. }) = connections.get(&pool_key) { let config = config.clone(); let session = session.clone(); drop(connections); return session .invoke::>( "listObjects", serde_json::json!({ "connection": config.as_ref(), "database": database, "schema": schema }), ) .await .map(filter_completion_objects); } if let Some(client) = extract_pool!(&connections, &pool_key, Agent) { let is_oracle = db_config.as_ref().is_some_and(|config| config.db_type == DatabaseType::Oracle); let fallback_config = db_config.clone(); drop(connections); let objects = if is_oracle { oracle_agent_list_objects(client, database, schema, agent_metadata_timeout(db_config.as_ref())).await? } else { let mut client = client.lock().await; match client .list_objects::>(database, schema, agent_metadata_timeout(db_config.as_ref())) .await { Ok(objects) if !objects.is_empty() => objects, Ok(objects) => { if let Some(config) = fallback_config.as_ref() { match native_postgres_metadata_pool(state, connection_id, database, config).await { Ok(Some(pool)) => { return db::postgres::list_objects(&pool, schema).await.map(filter_completion_objects) } Ok(None) => objects, Err(error) => { log::warn!( "[schema][agent:list_completion_objects:fallback-failed] connection_id={} database={} schema={} error={}", connection_id, database, schema, error ); objects } } } else { objects } } Err(agent_error) => { if let Some(config) = fallback_config.as_ref() { if let Some(pool) = native_postgres_metadata_pool(state, connection_id, database, config).await? { return db::postgres::list_objects(&pool, schema) .await .map(filter_completion_objects) .map_err(|fallback_error| { format!( "{agent_error}\n\nNative PostgreSQL metadata fallback failed: {fallback_error}" ) }); } } return Err(agent_error); } } }; return Ok(filter_completion_objects(objects)); } 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) if db_config.as_ref().is_some_and(is_questdb_config) => { db::questdb::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(_) => { drop(connections); let objects = list_objects_once(state, connection_id, database, schema).await?; Ok(filter_completion_objects(objects)) } _ => 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 { // Kingbase has dedicated agent metadata SQL and may carry JDBC-specific URL // parameters that the native PostgreSQL driver cannot parse. matches!(config.db_type, DatabaseType::Highgo | DatabaseType::Vastbase) } async fn native_postgres_metadata_pool( state: &AppState, connection_id: &str, database: &str, config: &ConnectionConfig, ) -> Result, String> { if !is_agent_postgres_metadata_fallback_config(config) { return Ok(None); } let mut postgres_config = database_connection_config(config, Some(database)); postgres_config.db_type = DatabaseType::Postgres; let (host, port) = state.connection_host_port(connection_id, &postgres_config).await?; let url = connection_url_for_endpoint(&postgres_config, &host, port); let connect_timeout = Duration::from_secs(postgres_config.effective_connect_timeout_secs()); db::postgres::connect(&url, connect_timeout).await.map(Some) } async fn retry_metadata_connection( state: &AppState, connection_id: &str, database: Option<&str>, mut operation: F, ) -> Result where F: FnMut() -> Fut, Fut: Future>, { let result = operation().await; match result { Err(error) if is_retryable_metadata_error(&error) => { state.reconnect_pool(connection_id, database).await?; operation().await } _ => result, } } fn is_retryable_metadata_error(error: &str) -> bool { error == "Pool not found" || crate::query::is_connection_error(error) } pub async fn get_columns_core( state: &AppState, connection_id: &str, database: &str, schema: &str, table: &str, ) -> Result, String> { retry_metadata_connection(state, connection_id, Some(database), || async { let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?; #[cfg(feature = "duckdb-bundled")] let duckdb_attached_names = duckdb_attached_database_names(state, connection_id).await; let db_config = connection_config(state, connection_id).await; { let connections = state.connections.read().await; #[cfg(feature = "duckdb-bundled")] if let Some(ext_pool) = extract_pool!(&connections, &pool_key, ExternalTabular) { drop(connections); let cache = ext_pool.cache.clone(); let table = table.to_string(); return tokio::task::spawn_blocking(move || { let con = cache.lock().map_err(|e| e.to_string())?; duckdb_query_columns(&con, &table) }) .await .map_err(|e| e.to_string())?; } if let Some(PoolKind::ExternalDriver { config, session, .. }) = connections.get(&pool_key) { let config = config.clone(); let session = session.clone(); drop(connections); let columns = session .invoke::>( "getColumns", serde_json::json!({ "connection": config.as_ref(), "database": database, "schema": schema, "table": table, }), ) .await?; return Ok(deduplicate_column_infos(columns)); } #[cfg(feature = "duckdb-bundled")] if let Some(con) = extract_pool!(&connections, &pool_key, DuckDb) { drop(connections); let con = con.lock().map_err(|e| e.to_string())?; return duckdb_query_columns_in_database_with_attached( &con, database, schema, table, &duckdb_attached_names, ); } if let Some(client) = extract_pool!(&connections, &pool_key, ClickHouse) { drop(connections); return db::clickhouse_driver::get_columns(&client, clickhouse_metadata_database(database, schema), table) .await .map(deduplicate_column_infos); } if let Some(client) = extract_pool!(&connections, &pool_key, InfluxDb) { drop(connections); return db::influxdb_driver::get_columns(&client, database, table).await.map(deduplicate_column_infos); } if let Some(linked) = crate::sql_dialect::parse_sqlserver_linked_schema_ref(schema) { if let Some(client) = extract_pool!(&connections, &pool_key, SqlServer) { drop(connections); let mut client = client.lock().await; return db::sqlserver::get_linked_server_columns( &mut client, &linked.server, &linked.catalog, &linked.schema, table, ) .await .map(deduplicate_column_infos); } } try_sqlserver!(connections, &pool_key, get_columns, schema, table); if let Some(client) = extract_pool!(&connections, &pool_key, Agent) { let fallback_config = db_config.clone(); drop(connections); let mut client = client.lock().await; match client .get_columns::>( database, schema, table, agent_metadata_timeout(db_config.as_ref()), ) .await { Ok(columns) if !columns.is_empty() => return Ok(deduplicate_column_infos(columns)), Ok(columns) => { if let Some(config) = fallback_config.as_ref() { match native_postgres_metadata_pool(state, connection_id, database, config).await { Ok(Some(pool)) => { return db::postgres::get_columns(&pool, schema, table) .await .map(deduplicate_column_infos); } Ok(None) => return Ok(deduplicate_column_infos(columns)), Err(error) => { log::warn!( "[schema][agent:get_columns:fallback-failed] connection_id={} database={} schema={} table={} error={}", connection_id, database, schema, table, error ); } } } return Ok(deduplicate_column_infos(columns)); } Err(agent_error) => { if let Some(config) = fallback_config.as_ref() { if let Some(pool) = native_postgres_metadata_pool(state, connection_id, database, config).await? { return db::postgres::get_columns(&pool, schema, table) .await .map(deduplicate_column_infos) .map_err(|fallback_error| { format!( "{agent_error}\n\nNative PostgreSQL metadata fallback failed: {fallback_error}" ) }); } } return Err(agent_error); } } } } 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_manticoresearch_config) => { let metadata_database = mysql_show_metadata_database_for_config(db_config.as_ref(), database); db::manticoresearch::get_columns(p, metadata_database, table).await.map(deduplicate_column_infos) } PoolKind::Mysql(p, _) if db_config.as_ref().is_some_and(is_doris_family_config) => { let metadata_database = mysql_show_metadata_database_for_config(db_config.as_ref(), database); // Doris/StarRocks previously went straight to `SHOW COLUMNS` for // speed (see perf(doris) commit), but `SHOW COLUMNS` reports the // `Key` column as `YES`/`NO` rather than MySQL's `PRI`, so primary // keys were never detected. `get_columns` queries // information_schema.COLUMNS first — where `COLUMN_KEY = 'PRI'` // correctly identifies primary keys (and only real primary keys, // not duplicate-key sort columns) — and falls back to `SHOW COLUMNS` // automatically when information_schema is unavailable. db::mysql::get_columns(p, metadata_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) if db_config.as_ref().is_some_and(is_questdb_config) => { db::questdb::get_columns(p, schema, table).await.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) } _ => Ok(vec![]), } }) .await } 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); } } pub async fn list_indexes_core( state: &AppState, connection_id: &str, database: &str, schema: &str, table: &str, ) -> Result, String> { if crate::sql_dialect::parse_sqlserver_linked_schema_ref(schema).is_some() { return Ok(vec![]); } retry_metadata_connection(state, connection_id, Some(database), || async { let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?; let db_config = connection_config(state, connection_id).await; { let connections = state.connections.read().await; try_sqlserver!(connections, &pool_key, list_indexes, schema, table); if let Some(client) = extract_pool!(&connections, &pool_key, Agent) { drop(connections); let mut client = client.lock().await; return client.list_indexes(database, schema, table, agent_metadata_timeout(db_config.as_ref())).await; } } let connections = state.connections.read().await; let pool = connections.get(&pool_key).ok_or("Pool not found")?; match pool { PoolKind::Mysql(p, mode) => { if db_config.as_ref().is_some_and(is_manticoresearch_config) { return db::manticoresearch::list_indexes(p, table).await; } if *mode == MysqlMode::OceanBaseOracle { db::ob_oracle::list_indexes(p, schema, table).await } else if db_config.as_ref().is_some_and(is_doris_family_config) { db::mysql::list_doris_family_indexes(p, mysql_table_metadata_catalog(database, schema), table).await } else { db::mysql::list_indexes(p, mysql_table_metadata_catalog(database, schema), table).await } } PoolKind::Postgres(p) if db_config.as_ref().is_some_and(is_questdb_config) => { db::questdb::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::MongoDb(client) => db::mongo_driver::list_indexes(client, database, table).await, _ => Ok(vec![]), } }) .await } pub async fn list_foreign_keys_core( state: &AppState, connection_id: &str, database: &str, schema: &str, table: &str, ) -> Result, String> { if crate::sql_dialect::parse_sqlserver_linked_schema_ref(schema).is_some() { return Ok(vec![]); } retry_metadata_connection(state, connection_id, Some(database), || async { let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?; let db_config = connection_config(state, connection_id).await; { let connections = state.connections.read().await; try_sqlserver!(connections, &pool_key, list_foreign_keys, schema, table); if let Some(client) = extract_pool!(&connections, &pool_key, Agent) { drop(connections); let mut client = client.lock().await; return client .list_foreign_keys(database, schema, table, agent_metadata_timeout(db_config.as_ref())) .await; } } let connections = state.connections.read().await; let pool = connections.get(&pool_key).ok_or("Pool not found")?; match pool { PoolKind::Mysql(p, mode) => { if *mode == MysqlMode::OceanBaseOracle { db::ob_oracle::list_foreign_keys(p, schema, table).await } else { db::mysql::list_foreign_keys(p, mysql_table_metadata_catalog(database, 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, _ => Ok(vec![]), } }) .await } pub async fn list_triggers_core( state: &AppState, connection_id: &str, database: &str, schema: &str, table: &str, ) -> Result, String> { if crate::sql_dialect::parse_sqlserver_linked_schema_ref(schema).is_some() { return Ok(vec![]); } retry_metadata_connection(state, connection_id, Some(database), || async { let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?; let db_config = connection_config(state, connection_id).await; { let connections = state.connections.read().await; try_sqlserver!(connections, &pool_key, list_triggers, schema, table); if let Some(client) = extract_pool!(&connections, &pool_key, Agent) { drop(connections); let mut client = client.lock().await; return client.list_triggers(database, schema, table, agent_metadata_timeout(db_config.as_ref())).await; } } let connections = state.connections.read().await; let pool = connections.get(&pool_key).ok_or("Pool not found")?; match pool { PoolKind::Mysql(p, mode) => { if *mode == MysqlMode::OceanBaseOracle { db::ob_oracle::list_triggers(p, schema, table).await } else { db::mysql::list_triggers(p, mysql_table_metadata_catalog(database, 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, _ => Ok(vec![]), } }) .await } pub async fn list_functions_core( state: &AppState, connection_id: &str, database: &str, schema: &str, ) -> Result, String> { retry_metadata_connection(state, connection_id, Some(database), || async { let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?; 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_functions(p, schema).await, _ => Ok(vec![]), } }) .await } pub async fn list_sequences_core( state: &AppState, connection_id: &str, database: &str, schema: &str, with_last_values: bool, ) -> Result, String> { retry_metadata_connection(state, connection_id, Some(database), || async { let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?; 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_sequences(p, schema, with_last_values).await, _ => Ok(vec![]), } }) .await } pub async fn list_rules_core( state: &AppState, connection_id: &str, database: &str, schema: &str, ) -> Result, String> { retry_metadata_connection(state, connection_id, Some(database), || async { let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?; 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_rules(p, schema).await, _ => Ok(vec![]), } }) .await } pub async fn list_owners_core( state: &AppState, connection_id: &str, database: &str, schema: &str, ) -> Result, String> { retry_metadata_connection(state, connection_id, Some(database), || async { let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?; 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_owners(p, schema).await, _ => Ok(vec![]), } }) .await } pub async fn get_table_ddl_core( state: &AppState, connection_id: &str, database: &str, schema: &str, table: &str, object_type: Option, ) -> Result { if crate::sql_dialect::parse_sqlserver_linked_schema_ref(schema).is_some() { return Err("DDL is not supported for SQL Server linked server tables".to_string()); } if matches!(object_type, Some(db::ObjectSourceKind::View)) { let source = get_object_source_core(state, connection_id, database, schema, table, db::ObjectSourceKind::View).await?; let database_type = connection_config(state, connection_id).await.map(|config| config.db_type); return Ok(crate::object_source_sql::build_view_ddl_sql(crate::object_source_sql::BuildViewDdlInput { database_type, schema: if schema.trim().is_empty() { None } else { Some(schema.to_string()) }, name: table.to_string(), source: source.source, })); } if matches!(object_type, Some(db::ObjectSourceKind::MaterializedView)) { let source = get_object_source_core( state, connection_id, database, schema, table, db::ObjectSourceKind::MaterializedView, ) .await?; return Ok(source.source); } let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?; let db_config = connection_config(state, connection_id).await; { let connections = state.connections.read().await; #[cfg(feature = "duckdb-bundled")] if let Some(con) = extract_pool!(&connections, &pool_key, DuckDb) { drop(connections); let tbl = table.replace('\'', "''"); let con = con.lock().map_err(|e| e.to_string())?; let mut stmt = con .prepare(&format!("SELECT sql FROM duckdb_tables() WHERE table_name = '{tbl}'")) .map_err(|e| e.to_string())?; let mut rows = stmt.query([]).map_err(|e| e.to_string())?; if let Some(row) = rows.next().map_err(|e| e.to_string())? { return row.get::<_, String>(0).map_err(|e| e.to_string()); } return Err("Table not found".to_string()); } if let Some(client) = extract_pool!(&connections, &pool_key, ClickHouse) { drop(connections); let clickhouse_database = clickhouse_metadata_database(database, schema); let result = db::clickhouse_driver::execute_query( &client, clickhouse_database, &format!("SHOW CREATE TABLE `{table}`"), ) .await?; return result .rows .first() .and_then(|r| r.first()) .and_then(|v| v.as_str()) .map(|s| s.to_string()) .ok_or_else(|| "Table not found".to_string()); } if let Some(client) = extract_pool!(&connections, &pool_key, SqlServer) { drop(connections); let mut client = client.lock().await; return build_sqlserver_ddl(&mut client, schema, table).await; } if let Some(client) = extract_pool!(&connections, &pool_key, Agent) { drop(connections); if let Some(config) = db_config.as_ref().filter(|config| is_agent_postgres_metadata_fallback_config(config)) { match native_postgres_metadata_pool(state, connection_id, database, config).await { Ok(Some(pool)) => match pg_ddl(&pool, schema, table).await { Ok(ddl) => return Ok(ddl), Err(error) => { log::warn!( "[schema][agent:get_table_ddl:postgres-compatible-native-fallback-failed] connection_id={} database={} schema={} table={} error={}", connection_id, database, schema, table, error ); } }, Ok(None) => {} Err(error) => { log::warn!( "[schema][agent:get_table_ddl:postgres-compatible-native-pool-failed] connection_id={} database={} schema={} table={} error={}", connection_id, database, schema, table, error ); } } } let mut client = client.lock().await; return client.get_table_ddl(database, schema, table, agent_metadata_timeout(db_config.as_ref())).await; } } let connections = state.connections.read().await; let pool = connections.get(&pool_key).ok_or("Pool not found")?; 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) if db_config.as_ref().is_some_and(is_questdb_config) => { match db::questdb::questdb_table_or_view_ddl(p, 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, _ => Err("DDL not supported for this database type".to_string()), } } 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 | DatabaseType::ManticoreSearch) || matches!(config.driver_profile.as_deref(), Some("doris" | "selectdb" | "starrocks" | "manticoresearch")) } fn is_manticoresearch_config(config: &ConnectionConfig) -> bool { matches!(config.db_type, DatabaseType::ManticoreSearch) || matches!(config.driver_profile.as_deref(), Some("manticoresearch")) } fn mysql_show_metadata_database_for_config<'a>(config: Option<&ConnectionConfig>, database: &'a str) -> &'a str { if config.is_some_and(is_manticoresearch_config) { "" } else { database } } fn filter_mysql_system_databases_for_config( databases: Vec, config: Option<&ConnectionConfig>, ) -> Vec { if !config.is_some_and(is_manticoresearch_config) { return databases; } databases.into_iter().filter(|database| !is_mysql_system_database(&database.name)).collect() } fn is_mysql_system_database(name: &str) -> bool { matches!(name.to_ascii_lowercase().as_str(), "information_schema" | "mysql" | "performance_schema" | "sys") } fn is_questdb_config(config: &ConnectionConfig) -> bool { matches!(config.db_type, DatabaseType::Questdb) || matches!(config.driver_profile.as_deref(), Some("questdb")) } fn sql_string(value: &str) -> String { format!("'{}'", value.replace('\'', "''")) } fn pg_ident(value: &str) -> String { format!("\"{}\"", value.replace('"', "\"\"")) } fn sqlserver_ident(value: &str) -> String { format!("[{}]", value.replace(']', "]]")) } fn sqlserver_n_string(value: &str) -> String { format!("N'{}'", value.replace('\'', "''")) } fn mysql_ident(value: &str) -> String { format!("`{}`", value.replace('`', "``")) } fn sqlite_object_type(kind: &db::ObjectSourceKind) -> &'static str { match kind { db::ObjectSourceKind::View | db::ObjectSourceKind::MaterializedView => "view", db::ObjectSourceKind::Procedure | db::ObjectSourceKind::Function | db::ObjectSourceKind::Sequence | db::ObjectSourceKind::Package | db::ObjectSourceKind::PackageBody => "routine", } } fn sqlserver_object_type_filter(kind: &db::ObjectSourceKind) -> &'static str { match kind { db::ObjectSourceKind::View => "'V'", db::ObjectSourceKind::Procedure => "'P'", db::ObjectSourceKind::Function => "'FN','IF','TF','FS','FT'", db::ObjectSourceKind::Sequence | db::ObjectSourceKind::Package | db::ObjectSourceKind::PackageBody | db::ObjectSourceKind::MaterializedView => "''", } } pub fn sqlserver_object_source_sql(schema: &str, name: &str, kind: &db::ObjectSourceKind) -> String { format!( "SELECT m.definition FROM sys.sql_modules m \ JOIN sys.objects o ON o.object_id = m.object_id \ JOIN sys.schemas s ON s.schema_id = o.schema_id \ WHERE s.name = {} AND o.name = {} AND o.type IN ({})", sql_string(schema), sql_string(name), sqlserver_object_type_filter(kind) ) } pub fn postgres_object_source_sql(schema: &str, name: &str, kind: &db::ObjectSourceKind) -> String { postgres_object_source_sql_inner(schema, name, kind, true) } fn postgres_object_source_sql_without_relispopulated(schema: &str, name: &str, kind: &db::ObjectSourceKind) -> String { postgres_object_source_sql_inner(schema, name, kind, false) } fn postgres_function_object_source_sql_without_prokind(schema: &str, name: &str) -> String { format!( "SELECT pg_get_functiondef(p.oid) \ FROM pg_proc p \ JOIN pg_namespace n ON n.oid = p.pronamespace \ WHERE n.nspname = {} AND p.proname = {} AND NOT p.proisagg AND NOT p.proiswindow \ ORDER BY p.oid LIMIT 1", sql_string(schema), sql_string(name) ) } fn postgres_object_source_sql_inner( schema: &str, name: &str, kind: &db::ObjectSourceKind, include_relispopulated: bool, ) -> String { match kind { db::ObjectSourceKind::View | db::ObjectSourceKind::MaterializedView => { let materialized_populated_clause = if include_relispopulated { " || CASE WHEN c.relispopulated THEN ' WITH DATA' ELSE ' WITH NO DATA' END" } else { "" }; let materialized_viewdef = "regexp_replace(pg_get_viewdef(c.oid, 0), ';[[:space:]]*$', '')"; let materialized_source_expr = format!( "CASE WHEN {materialized_viewdef} ~* '^[[:space:]]*CREATE[[:space:]]+(OR[[:space:]]+REPLACE[[:space:]]+)?MATERIALIZED[[:space:]]+VIEW[[:space:]]+' \ THEN {materialized_viewdef} \ ELSE format('CREATE MATERIALIZED VIEW %I.%I AS ', n.nspname, c.relname) || {materialized_viewdef}{materialized_populated_clause} \ END" ); format!( "SELECT CASE WHEN c.relkind = 'm' THEN {} \ ELSE format('CREATE OR REPLACE VIEW %I.%I AS ', n.nspname, c.relname) || pg_get_viewdef(c.oid, 0) \ END \ FROM pg_catalog.pg_class c \ JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \ WHERE n.nspname = {} AND c.relname = {} AND c.relkind IN ('v','m') \ ORDER BY c.oid LIMIT 1", materialized_source_expr, sql_string(schema), sql_string(name) ) } db::ObjectSourceKind::Procedure | db::ObjectSourceKind::Function => { let prokind = if matches!(kind, db::ObjectSourceKind::Procedure) { "p" } else { "f" }; format!( "SELECT pg_get_functiondef(p.oid) \ FROM pg_proc p \ JOIN pg_namespace n ON n.oid = p.pronamespace \ WHERE n.nspname = {} AND p.proname = {} AND p.prokind = '{}' \ ORDER BY p.oid LIMIT 1", sql_string(schema), sql_string(name), prokind ) } db::ObjectSourceKind::Sequence => { format!( "SELECT concat_ws(E'\\n\\n', \ '-- auto-generated definition' || E'\\n' || \ 'create sequence ' || quote_ident(c.relname) || E'\\n' || \ ' as ' || pg_catalog.format_type(s.seqtypid, NULL) || ';', \ 'alter sequence ' || quote_ident(c.relname) || ' owner to ' || quote_ident(pg_get_userbyid(c.relowner)) || ';', \ CASE WHEN owned.relname IS NOT NULL AND a.attname IS NOT NULL \ THEN 'alter sequence ' || quote_ident(c.relname) || ' owned by ' || quote_ident(owned.relname) || '.' || quote_ident(a.attname) || ';' \ END \ ) \ FROM pg_catalog.pg_class c \ JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \ JOIN pg_catalog.pg_sequence s ON s.seqrelid = c.oid \ LEFT JOIN pg_catalog.pg_depend d \ ON d.classid = 'pg_class'::regclass AND d.objid = c.oid AND d.deptype = 'a' \ LEFT JOIN pg_catalog.pg_class owned ON owned.oid = d.refobjid \ LEFT JOIN pg_catalog.pg_attribute a ON a.attrelid = d.refobjid AND a.attnum = d.refobjsubid \ WHERE n.nspname = {} AND c.relname = {} AND c.relkind = 'S' \ ORDER BY c.oid LIMIT 1", sql_string(schema), sql_string(name) ) } db::ObjectSourceKind::Package | db::ObjectSourceKind::PackageBody => "SELECT NULL WHERE FALSE".to_string(), } } pub fn oracle_object_source_sql(schema: &str, name: &str, kind: &db::ObjectSourceKind) -> String { let object_type = match kind { db::ObjectSourceKind::View => "VIEW", db::ObjectSourceKind::MaterializedView => "MATERIALIZED_VIEW", db::ObjectSourceKind::Procedure => "PROCEDURE", db::ObjectSourceKind::Function => "FUNCTION", db::ObjectSourceKind::Sequence => "SEQUENCE", db::ObjectSourceKind::Package => "PACKAGE", db::ObjectSourceKind::PackageBody => "PACKAGE_BODY", }; if schema.trim().is_empty() { format!("SELECT DBMS_METADATA.GET_DDL({}, {}) FROM DUAL", sql_string(object_type), sql_string(name)) } else { format!( "SELECT DBMS_METADATA.GET_DDL({}, {}, {}) FROM DUAL", sql_string(object_type), sql_string(name), sql_string(schema) ) } } pub fn sqlite_object_source_sql(name: &str, kind: &db::ObjectSourceKind) -> String { format!( "SELECT sql FROM sqlite_master WHERE type = {} AND name = {}", sql_string(sqlite_object_type(kind)), sql_string(name) ) } pub fn mysql_object_source_sql(name: &str, kind: &db::ObjectSourceKind) -> String { match kind { db::ObjectSourceKind::View => format!("SHOW CREATE VIEW {}", mysql_ident(name)), db::ObjectSourceKind::Procedure => format!("SHOW CREATE PROCEDURE {}", mysql_ident(name)), db::ObjectSourceKind::Function => format!("SHOW CREATE FUNCTION {}", mysql_ident(name)), db::ObjectSourceKind::Sequence | db::ObjectSourceKind::Package | db::ObjectSourceKind::PackageBody | db::ObjectSourceKind::MaterializedView => String::new(), } } pub fn postgres_view_source_fallback_sql(schema: &str, name: &str) -> String { format!( "SELECT definition \ FROM pg_catalog.pg_views \ WHERE schemaname = {} AND viewname = {} \ LIMIT 1", sql_string(schema), sql_string(name) ) } fn first_string_cell(result: db::QueryResult) -> Result { result .rows .first() .and_then(|row| row.iter().find_map(|value| value.as_str().map(str::to_string))) .filter(|value| !value.trim().is_empty()) .ok_or_else(|| "Object source not found".to_string()) } async fn mysql_object_source( pool: &db::mysql::MySqlPool, name: &str, kind: &db::ObjectSourceKind, ) -> Result { use mysql_async::prelude::*; let sql = mysql_object_source_sql(name, kind); let mut conn = db::mysql::get_conn_with_timeout(pool, db::connection_timeout()).await?; let result = conn.query_iter(&sql).await.map_err(|e| e.to_string())?; let rows: Vec = result.collect_and_drop().await.map_err(|e| e.to_string())?; let row = rows.first().ok_or("Object source not found")?; let index = if matches!(kind, db::ObjectSourceKind::View) { 1 } else { 2 }; row.get_opt::(index) .and_then(|result| result.ok()) .or_else(|| { row.get_opt::, usize>(index) .and_then(|result| result.ok()) .map(|b| String::from_utf8_lossy(&b).to_string()) }) .ok_or_else(|| "Failed to read object source".to_string()) } pub async fn get_object_source_core( state: &AppState, connection_id: &str, database: &str, schema: &str, name: &str, object_type: db::ObjectSourceKind, ) -> Result { retry_metadata_connection(state, connection_id, Some(database), || { get_object_source_once(state, connection_id, database, schema, name, object_type.clone()) }) .await } async fn get_object_source_once( state: &AppState, connection_id: &str, database: &str, schema: &str, name: &str, object_type: db::ObjectSourceKind, ) -> Result { let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?; let db_config = connection_config(state, connection_id).await; let source = { let connections = state.connections.read().await; if let Some(PoolKind::ExternalDriver { config, session, .. }) = connections.get(&pool_key) { let config = config.clone(); let session = session.clone(); drop(connections); let result: db::ObjectSource = session .invoke( "getObjectSource", serde_json::json!({ "connection": config.as_ref(), "database": database, "schema": schema, "name": name, "object_type": &object_type, }), ) .await?; return Ok(result); } if let Some(client) = extract_pool!(&connections, &pool_key, SqlServer) { drop(connections); let mut client = client.lock().await; let result = db::sqlserver::execute_query(&mut client, &sqlserver_object_source_sql(schema, name, &object_type)) .await; drop(client); if matches!(result.as_ref(), Err(err) if should_discard_pool_after_error(Some(DatabaseType::SqlServer), err)) { state.remove_pool_by_key(&pool_key).await; } first_string_cell(result?)? } else if let Some(client) = extract_pool!(&connections, &pool_key, Agent) { drop(connections); if db_config.as_ref().is_some_and(|config| config.db_type == DatabaseType::Oracle) && matches!(object_type, db::ObjectSourceKind::Package | db::ObjectSourceKind::PackageBody) { oracle_agent_object_source( client, database, schema, name, &object_type, agent_metadata_timeout(db_config.as_ref()), ) .await? } else { let mut client = client.lock().await; let result: db::ObjectSource = client .get_object_source(database, schema, name, &object_type, agent_metadata_timeout(db_config.as_ref())) .await?; 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) if db_config.as_ref().is_some_and(is_questdb_config) => { // only view db::questdb::questdb_object_source(pool, name).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::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()), } } }; Ok(db::ObjectSource { name: name.to_string(), object_type, schema: if schema.is_empty() { None } else { Some(schema.to_string()) }, source, }) } fn oracle_owner_filter(schema: &str) -> String { let schema = schema.trim(); if schema.is_empty() { "USER".to_string() } else { sql_string(&schema.to_uppercase()) } } pub fn oracle_list_objects_sql(schema: &str) -> String { format!( "SELECT object_name, CASE object_type WHEN 'PACKAGE BODY' THEN 'PACKAGE_BODY' ELSE object_type END AS object_type, owner \ FROM all_objects \ WHERE owner = {} AND object_type IN ('TABLE', 'VIEW', 'PROCEDURE', 'FUNCTION', 'PACKAGE', 'PACKAGE BODY') \ ORDER BY CASE object_type WHEN 'TABLE' THEN 0 WHEN 'VIEW' THEN 1 WHEN 'PROCEDURE' THEN 2 WHEN 'FUNCTION' THEN 3 WHEN 'PACKAGE' THEN 4 ELSE 5 END, object_name", oracle_owner_filter(schema) ) } async fn oracle_agent_list_objects( client: Arc>, database: &str, schema: &str, timeout_duration: Option, ) -> Result, String> { let sql = oracle_list_objects_sql(schema); let params = agent_execute_query_params( &sql, if database.is_empty() { None } else { Some(database) }, if schema.is_empty() { None } else { Some(schema) }, QueryExecutionOptions { max_rows: Some(10_000), ..Default::default() }, ); let mut client = client.lock().await; let result: db::QueryResult = client.execute_query_with_timeout(params, timeout_duration).await?; Ok(result .rows .into_iter() .filter_map(|row| { let name = row.first()?.as_str()?.to_string(); let object_type = row.get(1)?.as_str()?.to_string(); let schema = row.get(2).and_then(|value| value.as_str()).map(str::to_string); Some(db::ObjectInfo { name, object_type, schema, comment: None, created_at: None, updated_at: None, parent_schema: None, parent_name: None, }) }) .collect()) } async fn oracle_agent_object_source( client: Arc>, database: &str, schema: &str, name: &str, object_type: &db::ObjectSourceKind, timeout_duration: Option, ) -> Result { let sql = oracle_object_source_sql(schema, name, object_type); let params = agent_execute_query_params( &sql, if database.is_empty() { None } else { Some(database) }, if schema.is_empty() { None } else { Some(schema) }, QueryExecutionOptions { max_rows: Some(1), ..Default::default() }, ); let mut client = client.lock().await; let result: db::QueryResult = client.execute_query_with_timeout(params, timeout_duration).await?; first_string_cell(result) } async fn postgres_object_source( pool: &deadpool_postgres::Pool, schema: &str, name: &str, object_type: &db::ObjectSourceKind, ) -> Result { let sql = postgres_object_source_sql(schema, name, object_type); match db::postgres::execute_query(pool, &sql).await.and_then(first_string_cell) { Ok(source) => Ok(source), Err(primary_err) if postgres_missing_relispopulated_error(&primary_err) && matches!(object_type, db::ObjectSourceKind::View | db::ObjectSourceKind::MaterializedView) => { let fallback_sql = postgres_object_source_sql_without_relispopulated(schema, name, object_type); db::postgres::execute_query(pool, &fallback_sql) .await .and_then(first_string_cell) .map_err(|fallback_err| format!("{primary_err}; relispopulated fallback failed: {fallback_err}")) } Err(primary_err) if postgres_missing_prokind_error(&primary_err) && matches!(object_type, db::ObjectSourceKind::Function) => { let fallback_sql = postgres_function_object_source_sql_without_prokind(schema, name); db::postgres::execute_query(pool, &fallback_sql) .await .and_then(first_string_cell) .map_err(|fallback_err| format!("{primary_err}; prokind fallback failed: {fallback_err}")) } Err(primary_err) if matches!(object_type, db::ObjectSourceKind::View) => { let fallback_sql = postgres_view_source_fallback_sql(schema, name); db::postgres::execute_query(pool, &fallback_sql) .await .and_then(first_string_cell) .map_err(|fallback_err| format!("{primary_err}; fallback failed: {fallback_err}")) } Err(err) => Err(err), } } fn postgres_missing_prokind_error(err: &str) -> bool { let lower = err.to_ascii_lowercase(); lower.contains("does not exist") && (lower.contains("column p.prokind") || lower.contains("column \"p\".\"prokind\"") || lower.contains("column \"prokind\"")) } fn postgres_missing_relispopulated_error(err: &str) -> bool { let lower = err.to_ascii_lowercase(); lower.contains("does not exist") && (lower.contains("column c.relispopulated") || lower.contains("column \"c\".\"relispopulated\"") || lower.contains("column \"relispopulated\"")) } #[cfg(test)] mod object_source_tests { use super::*; use crate::types::ObjectSourceKind; #[test] fn builds_sqlserver_object_source_sql_for_schema_scoped_routines() { assert_eq!( sqlserver_object_source_sql("dbo", "refresh_cache", &ObjectSourceKind::Procedure), "SELECT m.definition FROM sys.sql_modules m JOIN sys.objects o ON o.object_id = m.object_id JOIN sys.schemas s ON s.schema_id = o.schema_id WHERE s.name = 'dbo' AND o.name = 'refresh_cache' AND o.type IN ('P')" ); } #[test] fn builds_postgres_object_source_sql_for_views_and_functions() { let view_sql = postgres_object_source_sql("public", "active_users", &ObjectSourceKind::View); assert!(view_sql.contains("CREATE MATERIALIZED VIEW")); assert!(view_sql.contains("CREATE OR REPLACE VIEW")); assert!(view_sql.contains("CASE WHEN c.relispopulated THEN ' WITH DATA' ELSE ' WITH NO DATA' END")); assert!(view_sql.contains("n.nspname = 'public'")); assert!(view_sql.contains("c.relname = 'active_users'")); assert_eq!( postgres_object_source_sql("public", "recalc_score", &ObjectSourceKind::Function), "SELECT pg_get_functiondef(p.oid) FROM pg_proc p JOIN pg_namespace n ON n.oid = p.pronamespace WHERE n.nspname = 'public' AND p.proname = 'recalc_score' AND p.prokind = 'f' ORDER BY p.oid LIMIT 1" ); } #[test] fn builds_postgres_object_source_sql_without_relispopulated_for_legacy_catalogs() { let sql = postgres_object_source_sql_without_relispopulated( "public", "active_users", &ObjectSourceKind::MaterializedView, ); assert!(sql.contains("CREATE MATERIALIZED VIEW")); assert!(sql.contains("pg_get_viewdef(c.oid, 0)")); assert!(!sql.contains("relispopulated")); } #[test] fn builds_postgres_function_source_sql_without_prokind_for_legacy_catalogs() { let sql = postgres_function_object_source_sql_without_prokind("public", "recalc_score"); assert_eq!( sql, "SELECT pg_get_functiondef(p.oid) FROM pg_proc p JOIN pg_namespace n ON n.oid = p.pronamespace WHERE n.nspname = 'public' AND p.proname = 'recalc_score' AND NOT p.proisagg AND NOT p.proiswindow ORDER BY p.oid LIMIT 1" ); } #[test] fn keeps_legacy_materialized_viewdef_when_it_already_contains_create_statement() { let sql = postgres_object_source_sql("public", "active_users", &ObjectSourceKind::MaterializedView); assert!( sql.contains( "~* '^[[:space:]]*CREATE[[:space:]]+(OR[[:space:]]+REPLACE[[:space:]]+)?MATERIALIZED[[:space:]]+VIEW[[:space:]]+'" ) ); assert!(sql.contains( "THEN regexp_replace(pg_get_viewdef(c.oid, 0), ';[[:space:]]*$', '') ELSE format('CREATE MATERIALIZED VIEW" )); } #[test] fn detects_legacy_postgres_relispopulated_errors() { assert!(postgres_missing_relispopulated_error("ERROR: column c.relispopulated does not exist")); assert!(!postgres_missing_relispopulated_error("ERROR: relation public.relispopulated does not exist")); } #[test] fn detects_legacy_postgres_prokind_errors() { assert!(postgres_missing_prokind_error("ERROR: column p.prokind does not exist")); assert!(postgres_missing_prokind_error("ERROR: column \"p\".\"prokind\" does not exist")); assert!(!postgres_missing_prokind_error("ERROR: relation public.prokind does not exist")); } #[test] fn builds_postgres_view_source_sql_without_regclass_cast() { let sql = postgres_object_source_sql("tenant's schema", "active users", &ObjectSourceKind::View); assert!(!sql.contains("::regclass")); assert!(sql.contains("pg_get_viewdef(c.oid, 0)")); assert!(sql.contains("format('CREATE OR REPLACE VIEW %I.%I AS ', n.nspname, c.relname)")); assert!(sql.contains("n.nspname = 'tenant''s schema'")); assert!(sql.contains("c.relname = 'active users'")); assert!(sql.contains("c.relkind IN ('v','m')")); } #[test] fn builds_postgres_view_source_fallback_sql_from_pg_views() { assert_eq!( postgres_view_source_fallback_sql("tenant's schema", "active users"), "SELECT definition FROM pg_catalog.pg_views WHERE schemaname = 'tenant''s schema' AND viewname = 'active users' LIMIT 1" ); } #[test] fn builds_oracle_object_source_sql_using_metadata_api() { assert_eq!( oracle_object_source_sql("HR", "ACTIVE_USERS", &ObjectSourceKind::View), "SELECT DBMS_METADATA.GET_DDL('VIEW', 'ACTIVE_USERS', 'HR') FROM DUAL" ); assert_eq!( oracle_object_source_sql("HR", "PAYROLL", &ObjectSourceKind::PackageBody), "SELECT DBMS_METADATA.GET_DDL('PACKAGE_BODY', 'PAYROLL', 'HR') FROM DUAL" ); assert_eq!( oracle_object_source_sql("", "PAYROLL", &ObjectSourceKind::Package), "SELECT DBMS_METADATA.GET_DDL('PACKAGE', 'PAYROLL') FROM DUAL" ); } #[test] fn builds_oracle_list_objects_sql_with_packages() { let sql = oracle_list_objects_sql("hr"); assert!(sql.contains("'PACKAGE'")); assert!(sql.contains("'PACKAGE BODY'")); assert!(sql.contains("CASE object_type WHEN 'PACKAGE BODY' THEN 'PACKAGE_BODY'")); assert!(sql.contains("owner = 'HR'")); } } #[cfg(test)] mod ddl_tests { use super::*; fn column(name: &str, data_type: &str) -> db::ColumnInfo { db::ColumnInfo { name: name.to_string(), data_type: data_type.to_string(), is_nullable: true, column_default: None, is_primary_key: false, extra: None, comment: None, numeric_precision: None, numeric_scale: None, character_maximum_length: None, } } #[test] fn postgres_table_ddl_includes_column_comments() { let mut display_name = column("display_name", "text"); display_name.comment = Some("User's display name".to_string()); let columns = vec![display_name]; let ddl = render_postgres_table_ddl("public", "users", &columns, &[], &[]); assert!(ddl.contains("COMMENT ON COLUMN \"public\".\"users\".\"display_name\" IS 'User''s display name';")); } #[test] fn postgres_table_ddl_includes_generated_identity() { let mut id = column("id", "integer"); id.is_nullable = false; id.is_primary_key = true; id.extra = Some("generated by default as identity".to_string()); let ddl = render_postgres_table_ddl("public", "users", &[id], &[], &[]); assert!(ddl.contains("\"id\" integer generated by default as identity NOT NULL"), "ddl: {ddl}"); } #[test] fn sqlserver_table_ddl_includes_column_comments() { let mut display_name = column("display]name", "nvarchar(100)"); display_name.comment = Some("User's display name".to_string()); let columns = vec![display_name]; let ddl = render_sqlserver_table_ddl("dbo", "users", &columns, &[], &[]); assert!(ddl.contains("CREATE TABLE [dbo].[users] (\n [display]]name] nvarchar(100)\n);")); assert!(ddl.contains( "EXEC sys.sp_addextendedproperty @name=N'MS_Description', @value=N'User''s display name', @level0type=N'SCHEMA', @level0name=N'dbo', @level1type=N'TABLE', @level1name=N'users', @level2type=N'COLUMN', @level2name=N'display]name';" )); } #[test] fn opengauss_table_ddl_uses_native_tabledef_function() { assert_eq!( opengauss_table_ddl_sql("tenant's schema", "active users"), "SELECT pg_get_tabledef('\"tenant''s schema\".\"active users\"')" ); } } pub async fn mysql_ddl(pool: &db::mysql::MySqlPool, table: &str) -> Result { use mysql_async::prelude::*; let sql = format!("SHOW CREATE TABLE `{}`", table.replace('`', "``")); // Use the health-checked getter so a stale pooled connection (server closed // it after an idle timeout, NAT/firewall dropped the TCP state, etc.) is // detected and replaced before issuing the query. Without this, the first // DDL request after a period of inactivity could surface a low-level // connection error that a manual refresh would have masked. let mut conn = db::mysql::get_conn_with_health_check(pool).await?; let result = conn.query_iter(&sql).await.map_err(|e| e.to_string())?; let rows: Vec = result.collect_and_drop().await.map_err(|e| e.to_string())?; let row = rows.first().ok_or("DDL not found")?; row.get_opt::(1) .and_then(|result| result.ok()) .or_else(|| { row.get_opt::, usize>(1) .and_then(|result| result.ok()) .map(|b| String::from_utf8_lossy(&b).to_string()) }) .ok_or_else(|| "Failed to read DDL".to_string()) } pub async fn sqlite_ddl(pool: &db::sqlite::SqliteHandle, table: &str) -> Result { let pool = pool.clone(); let table = table.to_string(); tokio::task::spawn_blocking(move || { pool.with_connection(|conn| { conn.query_row("SELECT sql FROM sqlite_master WHERE type='table' AND name=?1", [table], |row| { row.get::<_, String>(0) }) .map_err(|e| e.to_string()) }) }) .await .map_err(|e| e.to_string())? } pub async fn opengauss_table_ddl(pool: &deadpool_postgres::Pool, schema: &str, table: &str) -> Result { first_string_cell(db::postgres::execute_query(pool, &opengauss_table_ddl_sql(schema, table)).await?) } pub fn opengauss_table_ddl_sql(schema: &str, table: &str) -> String { let qualified_name = format!("{}.{}", pg_ident(schema), pg_ident(table)); format!("SELECT pg_get_tabledef({})", sql_string(&qualified_name)) } pub async fn pg_ddl(pool: &deadpool_postgres::Pool, schema: &str, table: &str) -> Result { let (columns, indexes, fkeys) = tokio::try_join!( db::postgres::get_columns(pool, schema, table), db::postgres::list_indexes(pool, schema, table), db::postgres::list_foreign_keys(pool, schema, table), )?; Ok(render_postgres_table_ddl(schema, table, &columns, &indexes, &fkeys)) } pub fn render_postgres_table_ddl( schema: &str, table: &str, columns: &[db::ColumnInfo], indexes: &[db::IndexInfo], fkeys: &[db::ForeignKeyInfo], ) -> String { let table_name = format!("{}.{}", pg_ident(schema), pg_ident(table)); let mut ddl = format!("CREATE TABLE {table_name} (\n"); let col_lines: Vec = columns .iter() .map(|c| { let mut line = format!(" {} {}", pg_ident(&c.name), c.data_type); let generated_clause = c .extra .as_deref() .map(str::trim) .filter(|value| !value.is_empty() && value.to_ascii_lowercase().starts_with("generated ")); if let Some(extra) = generated_clause { line.push_str(&format!(" {extra}")); } if !c.is_nullable { line.push_str(" NOT NULL"); } if generated_clause.is_none() { if let Some(ref def) = c.column_default { line.push_str(&format!(" DEFAULT {def}")); } } line }) .collect(); ddl.push_str(&col_lines.join(",\n")); let pks: Vec<&str> = columns.iter().filter(|c| c.is_primary_key).map(|c| c.name.as_str()).collect(); if !pks.is_empty() { ddl.push_str(&format!(",\n PRIMARY KEY ({})", pks.iter().map(|k| pg_ident(k)).collect::>().join(", "))); } for fk in fkeys { ddl.push_str(&format!( ",\n CONSTRAINT {} FOREIGN KEY ({}) REFERENCES {}({})", pg_ident(&fk.name), pg_ident(&fk.column), pg_ident(&fk.ref_table), pg_ident(&fk.ref_column) )); } ddl.push_str("\n);\n"); for col in columns { if let Some(comment) = col.comment.as_deref().filter(|comment| !comment.is_empty()) { ddl.push_str(&format!( "\nCOMMENT ON COLUMN {table_name}.{} IS {};", pg_ident(&col.name), sql_string(comment) )); } } for idx in indexes { if idx.is_primary { continue; } let unique = if idx.is_unique { "UNIQUE " } else { "" }; let cols = idx.columns.iter().map(|c| pg_ident(c)).collect::>().join(", "); let using = idx.index_type.as_deref().map(|t| format!(" USING {t}")).unwrap_or_default(); let include = idx .included_columns .as_deref() .filter(|c| !c.is_empty()) .map(|cols| format!(" INCLUDE ({})", cols.iter().map(|c| pg_ident(c)).collect::>().join(", "))) .unwrap_or_default(); let filter = idx.filter.as_deref().map(|f| format!(" WHERE {f}")).unwrap_or_default(); ddl.push_str(&format!( "\nCREATE {unique}INDEX {} ON {table_name}{using} ({cols}){include}{filter};", pg_ident(&idx.name) )); if let Some(ref c) = idx.comment { ddl.push_str(&format!( "\nCOMMENT ON INDEX {}.{} IS {};", pg_ident(schema), pg_ident(&idx.name), sql_string(c) )); } } ddl } pub async fn build_sqlserver_ddl( client: &mut db::sqlserver::SqlServerClient, schema: &str, table: &str, ) -> Result { let columns = db::sqlserver::get_columns(client, schema, table).await?; let indexes = db::sqlserver::list_indexes(client, schema, table).await?; let fkeys = db::sqlserver::list_foreign_keys(client, schema, table).await?; Ok(render_sqlserver_table_ddl(schema, table, &columns, &indexes, &fkeys)) } pub fn render_sqlserver_table_ddl( schema: &str, table: &str, columns: &[db::ColumnInfo], indexes: &[db::IndexInfo], fkeys: &[db::ForeignKeyInfo], ) -> String { let table_name = format!("{}.{}", sqlserver_ident(schema), sqlserver_ident(table)); let mut ddl = format!("CREATE TABLE {table_name} (\n"); let col_lines: Vec = columns .iter() .map(|c| { let mut line = format!(" {} {}", sqlserver_ident(&c.name), c.data_type); if !c.is_nullable { line.push_str(" NOT NULL"); } if let Some(ref def) = c.column_default { line.push_str(&format!(" DEFAULT {def}")); } line }) .collect(); ddl.push_str(&col_lines.join(",\n")); let pks: Vec<&str> = columns.iter().filter(|c| c.is_primary_key).map(|c| c.name.as_str()).collect(); if !pks.is_empty() { ddl.push_str(&format!( ",\n PRIMARY KEY ({})", pks.iter().map(|k| sqlserver_ident(k)).collect::>().join(", ") )); } for fk in fkeys { ddl.push_str(&format!( ",\n CONSTRAINT {} FOREIGN KEY ({}) REFERENCES {}({})", sqlserver_ident(&fk.name), sqlserver_ident(&fk.column), sqlserver_ident(&fk.ref_table), sqlserver_ident(&fk.ref_column) )); } ddl.push_str("\n);\n"); for column in columns { if let Some(comment) = column.comment.as_deref().map(str::trim).filter(|comment| !comment.is_empty()) { ddl.push_str(&format!( "\nEXEC sys.sp_addextendedproperty @name=N'MS_Description', @value={}, @level0type=N'SCHEMA', @level0name={}, @level1type=N'TABLE', @level1name={}, @level2type=N'COLUMN', @level2name={};", sqlserver_n_string(comment), sqlserver_n_string(schema), sqlserver_n_string(table), sqlserver_n_string(&column.name) )); } } for idx in indexes { if idx.is_primary { continue; } let unique = if idx.is_unique { "UNIQUE " } else { "" }; let idx_type = idx.index_type.as_deref().map(|t| format!("{t} ")).unwrap_or_default(); let cols = idx.columns.iter().map(|c| sqlserver_ident(c)).collect::>().join(", "); let include = idx .included_columns .as_deref() .filter(|c| !c.is_empty()) .map(|cols| { format!(" INCLUDE ({})", cols.iter().map(|c| sqlserver_ident(c)).collect::>().join(", ")) }) .unwrap_or_default(); let filter = idx.filter.as_deref().map(|f| format!(" WHERE {f}")).unwrap_or_default(); ddl.push_str(&format!( "\nCREATE {unique}{idx_type}INDEX {} ON {table_name} ({cols}){include}{filter};", sqlserver_ident(&idx.name) )); } ddl }