t8y2
b77c02c829
Fix SqlStatementSplitter.push_current_statement to preserve leading comments with their following SQL statement instead of stripping them.
1271 lines
48 KiB
Rust
1271 lines
48 KiB
Rust
use std::collections::HashMap;
|
|
use std::sync::Arc;
|
|
|
|
use crate::connection::{AppState, MysqlMode, PoolKind};
|
|
use crate::db;
|
|
use crate::models::connection::DatabaseType;
|
|
|
|
pub fn duckdb_query_tables(con: &duckdb::Connection) -> Result<Vec<db::TableInfo>, String> {
|
|
duckdb_query_tables_in_database(con, "main", "main")
|
|
}
|
|
|
|
pub fn duckdb_query_tables_in_database(
|
|
con: &duckdb::Connection,
|
|
database: &str,
|
|
schema: &str,
|
|
) -> Result<Vec<db::TableInfo>, String> {
|
|
duckdb_query_tables_in_database_with_attached(con, database, schema, &[])
|
|
}
|
|
|
|
pub fn duckdb_query_tables_in_database_with_attached(
|
|
con: &duckdb::Connection,
|
|
database: &str,
|
|
schema: &str,
|
|
attached_names: &[String],
|
|
) -> Result<Vec<db::TableInfo>, String> {
|
|
let database = duckdb_catalog_name(con, database, attached_names)?;
|
|
let mut stmt = con.prepare(
|
|
"SELECT table_name, table_type FROM information_schema.tables WHERE table_catalog = ? AND table_schema = ? ORDER BY table_name"
|
|
).map_err(|e| e.to_string())?;
|
|
let rows = stmt
|
|
.query_map((database.as_str(), schema), |row| {
|
|
Ok(db::TableInfo { name: row.get::<_, String>(0)?, table_type: row.get::<_, String>(1)?, comment: None })
|
|
})
|
|
.map_err(|e| e.to_string())?;
|
|
Ok(rows.filter_map(|r| r.ok()).collect())
|
|
}
|
|
|
|
pub fn duckdb_attach_database(con: &duckdb::Connection, name: &str, path: &str) -> Result<(), String> {
|
|
let name = name.trim();
|
|
let path = path.trim();
|
|
if name.is_empty() || path.is_empty() {
|
|
return Err("DuckDB attached database name and path are required".to_string());
|
|
}
|
|
let sql = format!("ATTACH {} AS {}", duckdb_quote_string(path), duckdb_quote_ident(name));
|
|
con.execute_batch(&sql).map_err(|e| e.to_string())
|
|
}
|
|
|
|
pub fn duckdb_list_databases(con: &duckdb::Connection) -> Result<Vec<db::DatabaseInfo>, String> {
|
|
duckdb_list_databases_with_attached(con, &[])
|
|
}
|
|
|
|
pub fn duckdb_list_databases_with_attached(
|
|
con: &duckdb::Connection,
|
|
attached_names: &[String],
|
|
) -> Result<Vec<db::DatabaseInfo>, String> {
|
|
let primary = duckdb_primary_catalog(con, attached_names)?;
|
|
let mut stmt = con.prepare("SHOW DATABASES").map_err(|e| e.to_string())?;
|
|
let rows = stmt
|
|
.query_map([], |row| {
|
|
let name = row.get::<_, String>(0)?;
|
|
Ok(db::DatabaseInfo { name: if name == primary { "main".to_string() } else { name } })
|
|
})
|
|
.map_err(|e| e.to_string())?;
|
|
Ok(rows.filter_map(|row| row.ok()).collect())
|
|
}
|
|
|
|
pub fn duckdb_list_schemas(con: &duckdb::Connection, database: &str) -> Result<Vec<String>, String> {
|
|
duckdb_list_schemas_with_attached(con, database, &[])
|
|
}
|
|
|
|
pub fn duckdb_list_schemas_with_attached(
|
|
con: &duckdb::Connection,
|
|
database: &str,
|
|
attached_names: &[String],
|
|
) -> Result<Vec<String>, String> {
|
|
let database = duckdb_catalog_name(con, database, attached_names)?;
|
|
let mut stmt = con
|
|
.prepare(
|
|
"SELECT schema_name FROM information_schema.schemata WHERE catalog_name = ? AND schema_name NOT IN ('information_schema', 'pg_catalog') ORDER BY schema_name",
|
|
)
|
|
.map_err(|e| e.to_string())?;
|
|
let rows = stmt.query_map([database.as_str()], |row| row.get::<_, String>(0)).map_err(|e| e.to_string())?;
|
|
Ok(rows.filter_map(|r| r.ok()).collect())
|
|
}
|
|
|
|
fn duckdb_catalog_name(con: &duckdb::Connection, database: &str, attached_names: &[String]) -> Result<String, String> {
|
|
if database.trim().is_empty() || database == "main" {
|
|
return duckdb_primary_catalog(con, attached_names);
|
|
}
|
|
Ok(database.to_string())
|
|
}
|
|
|
|
pub fn duckdb_primary_catalog(con: &duckdb::Connection, attached_names: &[String]) -> Result<String, String> {
|
|
if attached_names.is_empty() {
|
|
return duckdb_current_database(con);
|
|
}
|
|
let attached: std::collections::HashSet<String> = attached_names.iter().map(|name| name.to_lowercase()).collect();
|
|
let mut stmt = con.prepare("SHOW DATABASES").map_err(|e| e.to_string())?;
|
|
let rows = stmt.query_map([], |row| row.get::<_, String>(0)).map_err(|e| e.to_string())?;
|
|
for row in rows {
|
|
let name = row.map_err(|e| e.to_string())?;
|
|
if !attached.contains(&name.to_lowercase()) {
|
|
return Ok(name);
|
|
}
|
|
}
|
|
duckdb_current_database(con)
|
|
}
|
|
|
|
fn duckdb_current_database(con: &duckdb::Connection) -> Result<String, String> {
|
|
con.query_row("SELECT current_database()", [], |row| row.get::<_, String>(0)).map_err(|e| e.to_string())
|
|
}
|
|
|
|
fn duckdb_quote_ident(value: &str) -> String {
|
|
format!("\"{}\"", value.replace('"', "\"\""))
|
|
}
|
|
|
|
fn duckdb_quote_string(value: &str) -> String {
|
|
format!("'{}'", value.replace('\'', "''"))
|
|
}
|
|
|
|
pub fn duckdb_query_columns(con: &duckdb::Connection, table: &str) -> Result<Vec<db::ColumnInfo>, String> {
|
|
duckdb_query_columns_in_database(con, "main", "main", table)
|
|
}
|
|
|
|
pub fn duckdb_query_columns_in_database(
|
|
con: &duckdb::Connection,
|
|
database: &str,
|
|
schema: &str,
|
|
table: &str,
|
|
) -> Result<Vec<db::ColumnInfo>, String> {
|
|
duckdb_query_columns_in_database_with_attached(con, database, schema, table, &[])
|
|
}
|
|
|
|
pub fn duckdb_query_columns_in_database_with_attached(
|
|
con: &duckdb::Connection,
|
|
database: &str,
|
|
schema: &str,
|
|
table: &str,
|
|
attached_names: &[String],
|
|
) -> Result<Vec<db::ColumnInfo>, String> {
|
|
let database = duckdb_catalog_name(con, database, attached_names)?;
|
|
let mut pk_stmt = con
|
|
.prepare(
|
|
"SELECT kcu.column_name
|
|
FROM information_schema.table_constraints tc
|
|
JOIN information_schema.key_column_usage kcu
|
|
ON tc.constraint_name = kcu.constraint_name
|
|
AND tc.table_schema = kcu.table_schema
|
|
AND tc.table_name = kcu.table_name
|
|
WHERE tc.constraint_type = 'PRIMARY KEY'
|
|
AND tc.table_catalog = ?
|
|
AND tc.table_schema = ?
|
|
AND tc.table_name = ?
|
|
ORDER BY kcu.ordinal_position",
|
|
)
|
|
.map_err(|e| e.to_string())?;
|
|
let pk_rows = pk_stmt
|
|
.query_map((database.as_str(), schema, table), |row| row.get::<_, String>(0))
|
|
.map_err(|e| e.to_string())?;
|
|
let primary_keys: std::collections::HashSet<String> = pk_rows.filter_map(|r| r.ok()).collect();
|
|
|
|
let mut stmt = con
|
|
.prepare(
|
|
"SELECT column_name, data_type, is_nullable, column_default
|
|
FROM information_schema.columns
|
|
WHERE table_catalog = ? AND table_schema = ? AND table_name = ?
|
|
ORDER BY ordinal_position",
|
|
)
|
|
.map_err(|e| e.to_string())?;
|
|
let rows = stmt
|
|
.query_map((database.as_str(), schema, table), |row| {
|
|
let name = row.get::<_, String>(0)?;
|
|
Ok(db::ColumnInfo {
|
|
is_primary_key: primary_keys.contains(&name),
|
|
name,
|
|
data_type: row.get::<_, String>(1)?,
|
|
is_nullable: row.get::<_, String>(2).unwrap_or_default() == "YES",
|
|
column_default: row.get::<_, Option<String>>(3)?,
|
|
extra: None,
|
|
comment: None,
|
|
numeric_precision: None,
|
|
numeric_scale: None,
|
|
character_maximum_length: None,
|
|
})
|
|
})
|
|
.map_err(|e| e.to_string())?;
|
|
Ok(rows.filter_map(|r| r.ok()).collect())
|
|
}
|
|
|
|
pub fn extract_duckdb(
|
|
connections: &HashMap<String, PoolKind>,
|
|
key: &str,
|
|
) -> Option<Arc<std::sync::Mutex<duckdb::Connection>>> {
|
|
match connections.get(key)? {
|
|
PoolKind::DuckDb(con) => Some(con.clone()),
|
|
_ => None,
|
|
}
|
|
}
|
|
|
|
pub fn extract_external(
|
|
connections: &HashMap<String, PoolKind>,
|
|
key: &str,
|
|
) -> Option<Arc<crate::external::ExternalPool>> {
|
|
match connections.get(key)? {
|
|
PoolKind::ExternalTabular(pool) => Some(pool.clone()),
|
|
_ => None,
|
|
}
|
|
}
|
|
|
|
pub fn extract_sqlserver(
|
|
connections: &HashMap<String, PoolKind>,
|
|
key: &str,
|
|
) -> Option<Arc<tokio::sync::Mutex<db::sqlserver::SqlServerClient>>> {
|
|
match connections.get(key)? {
|
|
PoolKind::SqlServer(client) => Some(client.clone()),
|
|
_ => None,
|
|
}
|
|
}
|
|
|
|
pub fn extract_clickhouse(
|
|
connections: &HashMap<String, PoolKind>,
|
|
key: &str,
|
|
) -> Option<db::clickhouse_driver::ChClient> {
|
|
match connections.get(key)? {
|
|
PoolKind::ClickHouse(client) => Some(client.clone()),
|
|
_ => None,
|
|
}
|
|
}
|
|
|
|
pub fn extract_agent(
|
|
connections: &HashMap<String, PoolKind>,
|
|
key: &str,
|
|
) -> Option<Arc<tokio::sync::Mutex<db::agent_driver::AgentDriverClient>>> {
|
|
match connections.get(key)? {
|
|
PoolKind::Agent(client) => Some(client.clone()),
|
|
_ => None,
|
|
}
|
|
}
|
|
|
|
async fn duckdb_attached_database_names(state: &AppState, connection_id: &str) -> Vec<String> {
|
|
state
|
|
.configs
|
|
.read()
|
|
.await
|
|
.get(connection_id)
|
|
.map(|config| config.attached_databases.iter().map(|database| database.name.clone()).collect())
|
|
.unwrap_or_default()
|
|
}
|
|
|
|
fn clickhouse_metadata_database<'a>(database: &'a str, schema: &'a str) -> &'a str {
|
|
if database.is_empty() {
|
|
schema
|
|
} else {
|
|
database
|
|
}
|
|
}
|
|
|
|
pub async fn list_databases_core(state: &AppState, connection_id: &str) -> Result<Vec<db::DatabaseInfo>, String> {
|
|
log::info!("[list_databases] connection_id={connection_id}");
|
|
{
|
|
let connections = state.connections.read().await;
|
|
if extract_external(&connections, connection_id).is_some() {
|
|
return Ok(vec![db::DatabaseInfo { name: "main".to_string() }]);
|
|
}
|
|
if let Some(PoolKind::ExternalDriver { config, session, .. }) = connections.get(connection_id) {
|
|
let config = config.clone();
|
|
let session = session.clone();
|
|
drop(connections);
|
|
return session
|
|
.invoke::<Vec<db::DatabaseInfo>>("listDatabases", serde_json::json!({ "connection": config.as_ref() }))
|
|
.await;
|
|
}
|
|
if let Some(client) = extract_clickhouse(&connections, connection_id) {
|
|
drop(connections);
|
|
return db::clickhouse_driver::list_databases(&client).await;
|
|
}
|
|
if let Some(client) = extract_sqlserver(&connections, connection_id) {
|
|
drop(connections);
|
|
let mut client = client.lock().await;
|
|
return db::sqlserver::list_databases(&mut client).await;
|
|
}
|
|
if let Some(client) = extract_agent(&connections, connection_id) {
|
|
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().await;
|
|
}
|
|
}
|
|
|
|
let duckdb_attached_names = duckdb_attached_database_names(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, mode) => {
|
|
if *mode == MysqlMode::OceanBaseOracle {
|
|
db::ob_oracle::list_databases(p).await
|
|
} else {
|
|
db::mysql::list_databases(p).await
|
|
}
|
|
}
|
|
PoolKind::Postgres(p) => db::postgres::list_databases(p).await,
|
|
PoolKind::Sqlite(p) => db::sqlite::list_databases(p).await,
|
|
PoolKind::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<Vec<String>, String> {
|
|
let pool_key = state.get_or_create_pool(connection_id, Some(database)).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::<Vec<String>>(
|
|
"listSchemas",
|
|
serde_json::json!({ "connection": config.as_ref(), "database": database }),
|
|
)
|
|
.await;
|
|
}
|
|
if let Some(client) = extract_sqlserver(&connections, &pool_key) {
|
|
drop(connections);
|
|
let mut client = client.lock().await;
|
|
return db::sqlserver::list_schemas(&mut client).await;
|
|
}
|
|
if let Some(client) = extract_agent(&connections, &pool_key) {
|
|
drop(connections);
|
|
let mut client = client.lock().await;
|
|
return client.list_schemas(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_schemas(p).await,
|
|
PoolKind::DuckDb(con) => {
|
|
let duckdb_attached_names = duckdb_attached_database_names(state, connection_id).await;
|
|
let con = con.lock().map_err(|e| e.to_string())?;
|
|
duckdb_list_schemas_with_attached(&con, database, &duckdb_attached_names)
|
|
}
|
|
_ => Ok(vec![]),
|
|
}
|
|
}
|
|
|
|
pub async fn list_tables_core(
|
|
state: &AppState,
|
|
connection_id: &str,
|
|
database: &str,
|
|
schema: &str,
|
|
filter: Option<&str>,
|
|
limit: Option<usize>,
|
|
) -> Result<Vec<db::TableInfo>, String> {
|
|
let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?;
|
|
let duckdb_attached_names = duckdb_attached_database_names(state, connection_id).await;
|
|
|
|
{
|
|
let connections = state.connections.read().await;
|
|
if let Some(ext_pool) = extract_external(&connections, &pool_key) {
|
|
drop(connections);
|
|
let cache = ext_pool.cache.clone();
|
|
return tokio::task::spawn_blocking(move || {
|
|
let con = cache.lock().map_err(|e| e.to_string())?;
|
|
duckdb_query_tables(&con)
|
|
})
|
|
.await
|
|
.map_err(|e| e.to_string())?;
|
|
}
|
|
if let Some(PoolKind::ExternalDriver { config, session, .. }) = connections.get(&pool_key) {
|
|
let config = config.clone();
|
|
let session = session.clone();
|
|
drop(connections);
|
|
return session
|
|
.invoke::<Vec<db::TableInfo>>(
|
|
"listTables",
|
|
serde_json::json!({ "connection": config.as_ref(), "database": database, "schema": schema }),
|
|
)
|
|
.await;
|
|
}
|
|
if let Some(con) = extract_duckdb(&connections, &pool_key) {
|
|
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);
|
|
}
|
|
if let Some(client) = extract_clickhouse(&connections, &pool_key) {
|
|
drop(connections);
|
|
return db::clickhouse_driver::list_tables(&client, clickhouse_metadata_database(database, schema)).await;
|
|
}
|
|
if let Some(client) = extract_sqlserver(&connections, &pool_key) {
|
|
drop(connections);
|
|
let mut client = client.lock().await;
|
|
return db::sqlserver::list_tables(&mut client, schema, filter, limit).await;
|
|
}
|
|
if let Some(client) = extract_agent(&connections, &pool_key) {
|
|
drop(connections);
|
|
let mut client = client.lock().await;
|
|
return client.list_tables(schema).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_tables(p, schema).await.map(|tables| filter_table_infos(tables, filter, limit))
|
|
} else {
|
|
db::mysql::list_tables(p, schema).await.map(|tables| filter_table_infos(tables, filter, limit))
|
|
}
|
|
}
|
|
PoolKind::Postgres(p) => {
|
|
db::postgres::list_tables(p, schema).await.map(|tables| filter_table_infos(tables, filter, limit))
|
|
}
|
|
PoolKind::Sqlite(p) => {
|
|
db::sqlite::list_tables(p, schema).await.map(|tables| filter_table_infos(tables, filter, limit))
|
|
}
|
|
_ => Ok(vec![]),
|
|
}
|
|
}
|
|
|
|
fn filter_table_infos(tables: Vec<db::TableInfo>, filter: Option<&str>, limit: Option<usize>) -> Vec<db::TableInfo> {
|
|
let filter = filter.unwrap_or("").to_lowercase();
|
|
let limit = limit.unwrap_or(usize::MAX);
|
|
tables
|
|
.into_iter()
|
|
.filter(|table| filter.is_empty() || table.name.to_lowercase().contains(&filter))
|
|
.take(limit)
|
|
.collect()
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::{
|
|
clickhouse_metadata_database, duckdb_attach_database, duckdb_list_databases, duckdb_query_tables_in_database,
|
|
};
|
|
|
|
#[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);
|
|
}
|
|
|
|
#[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);
|
|
}
|
|
|
|
#[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");
|
|
}
|
|
}
|
|
|
|
pub async fn list_objects_core(
|
|
state: &AppState,
|
|
connection_id: &str,
|
|
database: &str,
|
|
schema: &str,
|
|
) -> Result<Vec<db::ObjectInfo>, String> {
|
|
let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?;
|
|
|
|
{
|
|
let connections = state.connections.read().await;
|
|
if let Some(ext_pool) = extract_external(&connections, &pool_key) {
|
|
drop(connections);
|
|
let cache = ext_pool.cache.clone();
|
|
return tokio::task::spawn_blocking(move || {
|
|
let con = cache.lock().map_err(|e| e.to_string())?;
|
|
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,
|
|
})
|
|
.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);
|
|
return session
|
|
.invoke::<Vec<db::ObjectInfo>>(
|
|
"listObjects",
|
|
serde_json::json!({ "connection": config.as_ref(), "database": database, "schema": schema }),
|
|
)
|
|
.await;
|
|
}
|
|
if let Some(client) = extract_sqlserver(&connections, &pool_key) {
|
|
drop(connections);
|
|
let mut client = client.lock().await;
|
|
return db::sqlserver::list_objects(&mut client, schema).await;
|
|
}
|
|
if let Some(client) = extract_agent(&connections, &pool_key) {
|
|
drop(connections);
|
|
let mut client = client.lock().await;
|
|
return client.list_objects(schema).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_objects(p, schema).await
|
|
} else {
|
|
db::mysql::list_objects(p, database).await
|
|
}
|
|
}
|
|
PoolKind::Postgres(p) => db::postgres::list_objects(p, schema).await,
|
|
_ => {
|
|
drop(connections);
|
|
Ok(list_tables_core(state, connection_id, database, schema, None, None)
|
|
.await?
|
|
.into_iter()
|
|
.map(|table| db::ObjectInfo {
|
|
name: table.name,
|
|
object_type: table.table_type,
|
|
schema: if schema.is_empty() { None } else { Some(schema.to_string()) },
|
|
comment: table.comment,
|
|
created_at: None,
|
|
updated_at: None,
|
|
})
|
|
.collect())
|
|
}
|
|
}
|
|
}
|
|
|
|
pub async fn get_columns_core(
|
|
state: &AppState,
|
|
connection_id: &str,
|
|
database: &str,
|
|
schema: &str,
|
|
table: &str,
|
|
) -> Result<Vec<db::ColumnInfo>, String> {
|
|
let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?;
|
|
let duckdb_attached_names = duckdb_attached_database_names(state, connection_id).await;
|
|
|
|
{
|
|
let connections = state.connections.read().await;
|
|
if let Some(ext_pool) = extract_external(&connections, &pool_key) {
|
|
drop(connections);
|
|
let cache = ext_pool.cache.clone();
|
|
let table = table.to_string();
|
|
return tokio::task::spawn_blocking(move || {
|
|
let con = cache.lock().map_err(|e| e.to_string())?;
|
|
duckdb_query_columns(&con, &table)
|
|
})
|
|
.await
|
|
.map_err(|e| e.to_string())?;
|
|
}
|
|
if let Some(PoolKind::ExternalDriver { config, session, .. }) = connections.get(&pool_key) {
|
|
let config = config.clone();
|
|
let session = session.clone();
|
|
drop(connections);
|
|
return session
|
|
.invoke::<Vec<db::ColumnInfo>>(
|
|
"getColumns",
|
|
serde_json::json!({
|
|
"connection": config.as_ref(),
|
|
"database": database,
|
|
"schema": schema,
|
|
"table": table,
|
|
}),
|
|
)
|
|
.await;
|
|
}
|
|
if let Some(con) = extract_duckdb(&connections, &pool_key) {
|
|
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_clickhouse(&connections, &pool_key) {
|
|
drop(connections);
|
|
return db::clickhouse_driver::get_columns(&client, clickhouse_metadata_database(database, schema), table)
|
|
.await;
|
|
}
|
|
if let Some(client) = extract_sqlserver(&connections, &pool_key) {
|
|
drop(connections);
|
|
let mut client = client.lock().await;
|
|
return db::sqlserver::get_columns(&mut client, schema, table).await;
|
|
}
|
|
if let Some(client) = extract_agent(&connections, &pool_key) {
|
|
drop(connections);
|
|
let mut client = client.lock().await;
|
|
return client.get_columns(schema, table).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::get_columns(p, database, table).await
|
|
} else {
|
|
db::mysql::get_columns(p, database, table).await
|
|
}
|
|
}
|
|
PoolKind::Postgres(p) => db::postgres::get_columns(p, schema, table).await,
|
|
PoolKind::Sqlite(p) => db::sqlite::get_columns(p, schema, table).await,
|
|
_ => Ok(vec![]),
|
|
}
|
|
}
|
|
|
|
pub async fn list_indexes_core(
|
|
state: &AppState,
|
|
connection_id: &str,
|
|
database: &str,
|
|
schema: &str,
|
|
table: &str,
|
|
) -> Result<Vec<db::IndexInfo>, String> {
|
|
let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?;
|
|
|
|
{
|
|
let connections = state.connections.read().await;
|
|
if let Some(client) = extract_sqlserver(&connections, &pool_key) {
|
|
drop(connections);
|
|
let mut client = client.lock().await;
|
|
return db::sqlserver::list_indexes(&mut client, schema, table).await;
|
|
}
|
|
if let Some(client) = extract_agent(&connections, &pool_key) {
|
|
drop(connections);
|
|
let mut client = client.lock().await;
|
|
return client.list_indexes(schema, table).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_indexes(p, schema, table).await
|
|
} else {
|
|
db::mysql::list_indexes(p, schema, table).await
|
|
}
|
|
}
|
|
PoolKind::Postgres(p) => db::postgres::list_indexes(p, schema, table).await,
|
|
PoolKind::Sqlite(p) => db::sqlite::list_indexes(p, schema, table).await,
|
|
_ => Ok(vec![]),
|
|
}
|
|
}
|
|
|
|
pub async fn list_foreign_keys_core(
|
|
state: &AppState,
|
|
connection_id: &str,
|
|
database: &str,
|
|
schema: &str,
|
|
table: &str,
|
|
) -> Result<Vec<db::ForeignKeyInfo>, String> {
|
|
let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?;
|
|
|
|
{
|
|
let connections = state.connections.read().await;
|
|
if let Some(client) = extract_sqlserver(&connections, &pool_key) {
|
|
drop(connections);
|
|
let mut client = client.lock().await;
|
|
return db::sqlserver::list_foreign_keys(&mut client, schema, table).await;
|
|
}
|
|
if let Some(client) = extract_agent(&connections, &pool_key) {
|
|
drop(connections);
|
|
let mut client = client.lock().await;
|
|
return client.list_foreign_keys(schema, table).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, 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,
|
|
_ => Ok(vec![]),
|
|
}
|
|
}
|
|
|
|
pub async fn list_triggers_core(
|
|
state: &AppState,
|
|
connection_id: &str,
|
|
database: &str,
|
|
schema: &str,
|
|
table: &str,
|
|
) -> Result<Vec<db::TriggerInfo>, String> {
|
|
let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?;
|
|
|
|
{
|
|
let connections = state.connections.read().await;
|
|
if let Some(client) = extract_sqlserver(&connections, &pool_key) {
|
|
drop(connections);
|
|
let mut client = client.lock().await;
|
|
return db::sqlserver::list_triggers(&mut client, schema, table).await;
|
|
}
|
|
if let Some(client) = extract_agent(&connections, &pool_key) {
|
|
drop(connections);
|
|
let mut client = client.lock().await;
|
|
return client.list_triggers(schema, table).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, 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,
|
|
_ => Ok(vec![]),
|
|
}
|
|
}
|
|
|
|
pub async fn get_table_ddl_core(
|
|
state: &AppState,
|
|
connection_id: &str,
|
|
database: &str,
|
|
schema: &str,
|
|
table: &str,
|
|
) -> Result<String, String> {
|
|
let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?;
|
|
|
|
{
|
|
let connections = state.connections.read().await;
|
|
if let Some(con) = extract_duckdb(&connections, &pool_key) {
|
|
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_clickhouse(&connections, &pool_key) {
|
|
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_sqlserver(&connections, &pool_key) {
|
|
drop(connections);
|
|
let mut client = client.lock().await;
|
|
return build_sqlserver_ddl(&mut client, schema, table).await;
|
|
}
|
|
if let Some(client) = extract_agent(&connections, &pool_key) {
|
|
drop(connections);
|
|
let mut client = client.lock().await;
|
|
return client.get_table_ddl(schema, table).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) => pg_ddl(p, schema, table).await,
|
|
PoolKind::Sqlite(p) => sqlite_ddl(p, table).await,
|
|
_ => Err("DDL not supported for this database type".to_string()),
|
|
}
|
|
}
|
|
|
|
fn sql_string(value: &str) -> String {
|
|
format!("'{}'", value.replace('\'', "''"))
|
|
}
|
|
|
|
fn pg_ident(value: &str) -> String {
|
|
format!("\"{}\"", 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 => "view",
|
|
db::ObjectSourceKind::Procedure | db::ObjectSourceKind::Function => "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'",
|
|
}
|
|
}
|
|
|
|
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 {
|
|
match kind {
|
|
db::ObjectSourceKind::View => {
|
|
format!("SELECT pg_get_viewdef('{}.{}'::regclass, true)", pg_ident(schema), pg_ident(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
|
|
)
|
|
}
|
|
}
|
|
}
|
|
|
|
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::Procedure => "PROCEDURE",
|
|
db::ObjectSourceKind::Function => "FUNCTION",
|
|
};
|
|
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)),
|
|
}
|
|
}
|
|
|
|
fn first_string_cell(result: db::QueryResult) -> Result<String, String> {
|
|
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<String, String> {
|
|
use mysql_async::prelude::*;
|
|
let sql = mysql_object_source_sql(name, kind);
|
|
let mut conn = pool.get_conn().await.map_err(|e| e.to_string())?;
|
|
let result = conn.query_iter(&sql).await.map_err(|e| e.to_string())?;
|
|
let rows: Vec<mysql_async::Row> = 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::<String, usize>(index)
|
|
.and_then(|result| result.ok())
|
|
.or_else(|| {
|
|
row.get_opt::<Vec<u8>, 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<db::ObjectSource, String> {
|
|
let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?;
|
|
let source = {
|
|
let connections = state.connections.read().await;
|
|
if let Some(client) = extract_sqlserver(&connections, &pool_key) {
|
|
drop(connections);
|
|
let mut client = client.lock().await;
|
|
first_string_cell(
|
|
db::sqlserver::execute_query(&mut client, &sqlserver_object_source_sql(schema, name, &object_type))
|
|
.await?,
|
|
)?
|
|
} else if let Some(client) = extract_agent(&connections, &pool_key) {
|
|
drop(connections);
|
|
let mut client = client.lock().await;
|
|
let result: db::ObjectSource = client.get_object_source(schema, name, &object_type).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) => first_string_cell(
|
|
db::postgres::execute_query(pool, &postgres_object_source_sql(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::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,
|
|
})
|
|
}
|
|
|
|
#[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() {
|
|
assert_eq!(
|
|
postgres_object_source_sql("public", "active_users", &ObjectSourceKind::View),
|
|
"SELECT pg_get_viewdef('\"public\".\"active_users\"'::regclass, true)"
|
|
);
|
|
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_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"
|
|
);
|
|
}
|
|
}
|
|
|
|
#[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';"));
|
|
}
|
|
}
|
|
|
|
pub async fn mysql_ddl(pool: &db::mysql::MySqlPool, table: &str) -> Result<String, String> {
|
|
use mysql_async::prelude::*;
|
|
let sql = format!("SHOW CREATE TABLE `{}`", table.replace('`', "``"));
|
|
let mut conn = pool.get_conn().await.map_err(|e| e.to_string())?;
|
|
let result = conn.query_iter(&sql).await.map_err(|e| e.to_string())?;
|
|
let rows: Vec<mysql_async::Row> = result.collect_and_drop().await.map_err(|e| e.to_string())?;
|
|
let row = rows.first().ok_or("DDL not found")?;
|
|
row.get_opt::<String, usize>(1)
|
|
.and_then(|result| result.ok())
|
|
.or_else(|| {
|
|
row.get_opt::<Vec<u8>, 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<String, String> {
|
|
let pool = pool.clone();
|
|
let table = table.to_string();
|
|
tokio::task::spawn_blocking(move || {
|
|
pool.with_connection(|conn| {
|
|
conn.query_row("SELECT sql FROM sqlite_master WHERE type='table' AND name=?1", [table], |row| {
|
|
row.get::<_, String>(0)
|
|
})
|
|
.map_err(|e| e.to_string())
|
|
})
|
|
})
|
|
.await
|
|
.map_err(|e| e.to_string())?
|
|
}
|
|
|
|
pub async fn pg_ddl(pool: &sqlx::postgres::PgPool, schema: &str, table: &str) -> Result<String, String> {
|
|
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))
|
|
}
|
|
|
|
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<String> = columns
|
|
.iter()
|
|
.map(|c| {
|
|
let mut line = format!(" {} {}", pg_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| pg_ident(k)).collect::<Vec<_>>().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::<Vec<_>>().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::<Vec<_>>().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<String, String> {
|
|
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?;
|
|
|
|
let mut ddl = format!("CREATE TABLE [{schema}].[{table}] (\n");
|
|
let col_lines: Vec<String> = columns
|
|
.iter()
|
|
.map(|c| {
|
|
let mut line = format!(" [{}] {}", 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| format!("[{k}]")).collect::<Vec<_>>().join(", ")
|
|
));
|
|
}
|
|
for fk in &fkeys {
|
|
ddl.push_str(&format!(
|
|
",\n CONSTRAINT [{}] FOREIGN KEY ([{}]) REFERENCES [{}]([{}])",
|
|
fk.name, fk.column, fk.ref_table, fk.ref_column
|
|
));
|
|
}
|
|
ddl.push_str("\n);\n");
|
|
|
|
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| format!("[{c}]")).collect::<Vec<_>>().join(", ");
|
|
let include = idx
|
|
.included_columns
|
|
.as_deref()
|
|
.filter(|c| !c.is_empty())
|
|
.map(|cols| format!(" INCLUDE ({})", cols.iter().map(|c| format!("[{c}]")).collect::<Vec<_>>().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 [{schema}].[{table}] ({cols}){include}{filter};",
|
|
idx.name
|
|
));
|
|
}
|
|
Ok(ddl)
|
|
}
|