fix(schema): list routines in object browser

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<Vec<TableIn
         .collect())
 }
 
+fn list_objects_sql(database: &str) -> 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<Vec<ObjectInfo>, String> {
+    let sql = list_objects_sql(database);
+    let rows: Vec<MySqlRow> = 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<Vec<ColumnInfo>, 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'"));
+    }
 }

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<Vec<TableInfo
     Ok(rows.iter().map(|row| TableInfo { name: get_str(row, 0), table_type: get_str(row, 1), comment: None }).collect())
 }
 
+fn list_objects_sql(schema: &str) -> 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<Vec<ObjectInfo>, String> {
+    let sql = list_objects_sql(schema);
+    let rows: Vec<MySqlRow> = 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<Vec<ColumnInfo>, 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'"));
+    }
+}

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<serde_json::Value> {
     if let Ok(v) = row.try_get::<DateTime<Utc>, _>(idx) {
@@ -226,6 +228,44 @@ pub async fn list_tables(pool: &PgPool, schema: &str) -> Result<Vec<TableInfo>,
         .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<Vec<ObjectInfo>, String> {
+    let rows: Vec<PgRow> =
+        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::<String, _>("object_name"),
+            object_type: row.get::<String, _>("object_type"),
+            schema: Some(schema.to_string()),
+            comment: row.get::<Option<String>, _>("object_comment").filter(|s| !s.is_empty()),
+        })
+        .collect())
+}
+
 pub async fn list_schemas(pool: &PgPool) -> Result<Vec<String>, String> {
     let rows: Vec<PgRow> = 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<V
         })
         .collect())
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn postgres_list_objects_sql_includes_routines() {
+        let sql = list_objects_sql();
+
+        assert!(sql.contains("pg_catalog.pg_class"));
+        assert!(sql.contains("pg_catalog.pg_proc"));
+        assert!(sql.contains("'PROCEDURE'"));
+        assert!(sql.contains("'FUNCTION'"));
+    }
+}

crates/dbx-core/src/schema.rs

@@ -339,16 +339,32 @@ pub async fn list_objects_core(
         }
     }
 
-    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())
+    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,
+                })
+                .collect())
+        }
+    }
 }
 
 pub async fn get_columns_core(