diff --git a/crates/dbx-core/src/db/mysql.rs b/crates/dbx-core/src/db/mysql.rs index ddf328031b..6d7c402f4d 100644 --- a/crates/dbx-core/src/db/mysql.rs +++ b/crates/dbx-core/src/db/mysql.rs @@ -6,7 +6,9 @@ use sqlx::{Column, Executor, Row, TypeInfo, ValueRef}; use std::time::{Duration, Instant}; use crate::sql::starts_with_executable_sql_keyword; -use crate::types::{ColumnInfo, DatabaseInfo, ForeignKeyInfo, IndexInfo, QueryResult, TableInfo, TriggerInfo}; +use crate::types::{ + ColumnInfo, DatabaseInfo, ForeignKeyInfo, IndexInfo, ObjectInfo, QueryResult, TableInfo, TriggerInfo, +}; fn quote_value(s: &str) -> String { format!("'{}'", s.replace('\\', "\\\\").replace('\'', "\\'")) @@ -244,6 +246,39 @@ pub async fn list_tables(pool: &MySqlPool, database: &str) -> Result String { + format!( + "SELECT TABLE_NAME AS object_name, \ + CASE WHEN TABLE_TYPE = 'VIEW' THEN 'VIEW' ELSE 'TABLE' END AS object_type, \ + TABLE_COMMENT AS object_comment, \ + CASE WHEN TABLE_TYPE = 'VIEW' THEN 1 ELSE 0 END AS sort_order \ + FROM information_schema.TABLES \ + WHERE TABLE_SCHEMA = {db} \ + UNION ALL \ + SELECT ROUTINE_NAME AS object_name, ROUTINE_TYPE AS object_type, NULL AS object_comment, \ + CASE WHEN ROUTINE_TYPE = 'PROCEDURE' THEN 2 ELSE 3 END AS sort_order \ + FROM information_schema.ROUTINES \ + WHERE ROUTINE_SCHEMA = {db} AND ROUTINE_TYPE IN ('PROCEDURE', 'FUNCTION') \ + ORDER BY sort_order, object_name", + db = quote_value(database), + ) +} + +pub async fn list_objects(pool: &MySqlPool, database: &str) -> Result, String> { + let sql = list_objects_sql(database); + let rows: Vec = sqlx::raw_sql(&sql).fetch_all(pool).await.map_err(|e| e.to_string())?; + + Ok(rows + .iter() + .map(|row| ObjectInfo { + name: get_str_by_name(row, "object_name"), + object_type: get_str_by_name(row, "object_type"), + schema: Some(database.to_string()), + comment: get_opt_str(row, "object_comment").filter(|s| !s.is_empty()), + }) + .collect()) +} + pub async fn get_columns(pool: &MySqlPool, database: &str, table: &str) -> Result, String> { let sql = format!( "SELECT c.COLUMN_NAME, c.COLUMN_TYPE, c.IS_NULLABLE, c.COLUMN_DEFAULT, c.EXTRA, c.COLUMN_COMMENT, \ @@ -471,4 +506,14 @@ mod tests { assert_eq!(numeric_metadata_u64_to_i32(Some(i32::MAX as u64 + 1)), None); assert_eq!(numeric_metadata_u64_to_i32(None), None); } + + #[test] + fn mysql_list_objects_sql_includes_routines() { + let sql = list_objects_sql("app"); + + assert!(sql.contains("information_schema.TABLES")); + assert!(sql.contains("information_schema.ROUTINES")); + assert!(sql.contains("'PROCEDURE'")); + assert!(sql.contains("'FUNCTION'")); + } } diff --git a/crates/dbx-core/src/db/ob_oracle.rs b/crates/dbx-core/src/db/ob_oracle.rs index 340298f6f9..5aaff5ad16 100644 --- a/crates/dbx-core/src/db/ob_oracle.rs +++ b/crates/dbx-core/src/db/ob_oracle.rs @@ -1,7 +1,7 @@ use sqlx::mysql::{MySqlPool, MySqlRow}; use sqlx::Row; -use crate::types::{ColumnInfo, DatabaseInfo, ForeignKeyInfo, IndexInfo, TableInfo, TriggerInfo}; +use crate::types::{ColumnInfo, DatabaseInfo, ForeignKeyInfo, IndexInfo, ObjectInfo, TableInfo, TriggerInfo}; fn quote_value(s: &str) -> String { format!("'{}'", s.replace('\\', "\\\\").replace('\'', "\\'")) @@ -52,6 +52,37 @@ pub async fn list_tables(pool: &MySqlPool, schema: &str) -> Result String { + format!( + "SELECT TABLE_NAME AS OBJECT_NAME, 'TABLE' AS OBJECT_TYPE, 0 AS SORT_ORDER \ + FROM ALL_TABLES WHERE OWNER = {s} \ + UNION ALL \ + SELECT VIEW_NAME AS OBJECT_NAME, 'VIEW' AS OBJECT_TYPE, 1 AS SORT_ORDER \ + FROM ALL_VIEWS WHERE OWNER = {s} \ + UNION ALL \ + SELECT OBJECT_NAME, OBJECT_TYPE, CASE WHEN OBJECT_TYPE = 'PROCEDURE' THEN 2 ELSE 3 END AS SORT_ORDER \ + FROM ALL_PROCEDURES \ + WHERE OWNER = {s} AND OBJECT_TYPE IN ('PROCEDURE', 'FUNCTION') AND PROCEDURE_NAME IS NULL \ + ORDER BY SORT_ORDER, OBJECT_NAME", + s = quote_value(schema), + ) +} + +pub async fn list_objects(pool: &MySqlPool, schema: &str) -> Result, String> { + let sql = list_objects_sql(schema); + let rows: Vec = sqlx::raw_sql(&sql).fetch_all(pool).await.map_err(|e| e.to_string())?; + + Ok(rows + .iter() + .map(|row| ObjectInfo { + name: get_str(row, 0), + object_type: get_str(row, 1), + schema: Some(schema.to_string()), + comment: None, + }) + .collect()) +} + pub async fn get_columns(pool: &MySqlPool, schema: &str, table: &str) -> Result, String> { let sql = format!( "SELECT c.COLUMN_NAME, c.DATA_TYPE, c.NULLABLE, c.DATA_DEFAULT, \ @@ -198,3 +229,19 @@ pub async fn list_triggers(pool: &MySqlPool, schema: &str, table: &str) -> Resul }) .collect()) } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn ob_oracle_list_objects_sql_includes_routines() { + let sql = list_objects_sql("DLJPM"); + + assert!(sql.contains("ALL_TABLES")); + assert!(sql.contains("ALL_VIEWS")); + assert!(sql.contains("ALL_PROCEDURES")); + assert!(sql.contains("'PROCEDURE'")); + assert!(sql.contains("'FUNCTION'")); + } +} diff --git a/crates/dbx-core/src/db/postgres.rs b/crates/dbx-core/src/db/postgres.rs index bb5b19f2a7..cdac598543 100644 --- a/crates/dbx-core/src/db/postgres.rs +++ b/crates/dbx-core/src/db/postgres.rs @@ -8,7 +8,9 @@ use std::time::{Duration, Instant}; use super::file_validator::validate_file_path; use crate::sql::starts_with_executable_sql_keyword; -use crate::types::{ColumnInfo, DatabaseInfo, ForeignKeyInfo, IndexInfo, QueryResult, TableInfo, TriggerInfo}; +use crate::types::{ + ColumnInfo, DatabaseInfo, ForeignKeyInfo, IndexInfo, ObjectInfo, QueryResult, TableInfo, TriggerInfo, +}; fn pg_temporal_to_json_value(row: &PgRow, idx: usize) -> Option { if let Ok(v) = row.try_get::, _>(idx) { @@ -226,6 +228,44 @@ pub async fn list_tables(pool: &PgPool, schema: &str) -> Result, .collect()) } +fn list_objects_sql() -> &'static str { + "SELECT c.relname AS object_name, \ + CASE c.relkind \ + WHEN 'v' THEN 'VIEW' \ + WHEN 'm' THEN 'VIEW' \ + ELSE 'TABLE' \ + END AS object_type, \ + obj_description(c.oid) AS object_comment, \ + CASE c.relkind WHEN 'v' THEN 1 WHEN 'm' THEN 1 ELSE 0 END AS sort_order \ + FROM pg_catalog.pg_class c \ + JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \ + WHERE n.nspname = $1 AND c.relkind IN ('r','v','m','f','p') \ + UNION ALL \ + SELECT p.proname AS object_name, \ + CASE p.prokind WHEN 'p' THEN 'PROCEDURE' ELSE 'FUNCTION' END AS object_type, \ + obj_description(p.oid) AS object_comment, \ + CASE p.prokind WHEN 'p' THEN 2 ELSE 3 END AS sort_order \ + FROM pg_catalog.pg_proc p \ + JOIN pg_catalog.pg_namespace n ON n.oid = p.pronamespace \ + WHERE n.nspname = $1 AND p.prokind IN ('p','f') \ + ORDER BY sort_order, object_name" +} + +pub async fn list_objects(pool: &PgPool, schema: &str) -> Result, String> { + let rows: Vec = + sqlx::query(list_objects_sql()).bind(schema).fetch_all(pool).await.map_err(|e| e.to_string())?; + + Ok(rows + .iter() + .map(|row| ObjectInfo { + name: row.get::("object_name"), + object_type: row.get::("object_type"), + schema: Some(schema.to_string()), + comment: row.get::, _>("object_comment").filter(|s| !s.is_empty()), + }) + .collect()) +} + pub async fn list_schemas(pool: &PgPool) -> Result, String> { let rows: Vec = sqlx::query( "SELECT n.nspname AS schema_name FROM pg_catalog.pg_namespace n \ @@ -520,3 +560,18 @@ pub async fn list_triggers(pool: &PgPool, schema: &str, table: &str) -> Result { + 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, + }) + .collect()) + } + } } pub async fn get_columns_core(