From 772942eed02cea053de46ad2c365ef342bda03fb Mon Sep 17 00:00:00 2001 From: h7x4 Date: Sun, 31 May 2026 05:09:52 +0900 Subject: [PATCH] WIP: show-db: limit output to 5 tables/users by default --- src/client/commands/edit_privs.rs | 4 +- src/client/commands/show_db.rs | 21 +- .../mysql_dbadm.rs | 6 +- src/core/protocol/commands.rs | 12 +- src/core/protocol/commands/list_databases.rs | 32 +- src/server/session_handler.rs | 6 +- src/server/sql/database_operations.rs | 277 +++++++++++------- 7 files changed, 231 insertions(+), 127 deletions(-) diff --git a/src/client/commands/edit_privs.rs b/src/client/commands/edit_privs.rs index 7b87d61..d6e5193 100644 --- a/src/client/commands/edit_privs.rs +++ b/src/client/commands/edit_privs.rs @@ -23,7 +23,7 @@ use crate::{ parse_privilege_data_from_editor_content, reduce_privilege_diffs, }, protocol::{ - ClientToServerMessageStream, ListDatabasesError, ListUsersError, + ClientToServerMessageStream, ListDatabasesError, ListDatabasesRequest, ListUsersError, ModifyDatabasePrivilegesError, Request, Response, print_modify_database_privileges_output_status, request_validation::ValidationError, }, @@ -132,7 +132,7 @@ async fn databases_exist( .map(|diff| diff.get_database_name().clone()) .collect(); - let message = Request::ListDatabases(Some(database_list)); + let message = Request::ListDatabases(ListDatabasesRequest::new(Some(database_list), false)); server_connection.send(message).await?; let result = match server_connection.next().await { diff --git a/src/client/commands/show_db.rs b/src/client/commands/show_db.rs index aa13c26..69be193 100644 --- a/src/client/commands/show_db.rs +++ b/src/client/commands/show_db.rs @@ -8,8 +8,8 @@ use crate::{ core::{ completion::mysql_database_completer, protocol::{ - ClientToServerMessageStream, ListDatabasesError, Request, Response, - print_list_databases_output_status, print_list_databases_output_status_json, + ClientToServerMessageStream, ListDatabasesError, ListDatabasesRequest, Request, + Response, print_list_databases_output_status, print_list_databases_output_status_json, request_validation::ValidationError, }, types::MySQLDatabase, @@ -27,6 +27,10 @@ pub struct ShowDbArgs { #[arg(short, long)] json: bool, + /// Show all tables and users for each database + #[arg(short = 'a', long)] + all: bool, + /// Show sizes in bytes instead of human-readable format #[arg(short, long)] bytes: bool, @@ -36,11 +40,14 @@ pub async fn show_databases( args: ShowDbArgs, mut server_connection: ClientToServerMessageStream, ) -> anyhow::Result<()> { - let message = if args.name.is_empty() { - Request::ListDatabases(None) - } else { - Request::ListDatabases(Some(args.name.clone())) - }; + let message = Request::ListDatabases(ListDatabasesRequest::new( + if args.name.is_empty() { + None + } else { + Some(args.name.clone()) + }, + args.all || args.json, + )); server_connection.send(message).await?; diff --git a/src/client/mysql_admutils_compatibility/mysql_dbadm.rs b/src/client/mysql_admutils_compatibility/mysql_dbadm.rs index d1cee05..6cc9265 100644 --- a/src/client/mysql_admutils_compatibility/mysql_dbadm.rs +++ b/src/client/mysql_admutils_compatibility/mysql_dbadm.rs @@ -22,8 +22,8 @@ use crate::{ completion::{mysql_database_completer, prefix_completer}, database_privileges::DatabasePrivilegeRow, protocol::{ - ClientToServerMessageStream, ListPrivilegesError, Request, Response, - create_client_to_server_message_stream, + ClientToServerMessageStream, ListDatabasesRequest, ListPrivilegesError, Request, + Response, create_client_to_server_message_stream, }, types::MySQLDatabase, }, @@ -285,7 +285,7 @@ async fn show_databases( args.name.iter().map(trim_db_name_to_32_chars).collect(); let message = if database_names.is_empty() { - let message = Request::ListDatabases(None); + let message = Request::ListDatabases(ListDatabasesRequest::new(None, false)); server_connection.send(message).await?; let response = server_connection.next().await; let databases = match response { diff --git a/src/core/protocol/commands.rs b/src/core/protocol/commands.rs index b76da3f..f44f459 100644 --- a/src/core/protocol/commands.rs +++ b/src/core/protocol/commands.rs @@ -142,7 +142,8 @@ impl Request { Request::ListDatabases(req) => format!( "{}{}", self.command_name(), - req.as_ref() + req.names + .as_ref() .map_or("".to_string(), |r| format!("({})", r.len())) ), Request::ListPrivileges(req) => format!( @@ -206,9 +207,12 @@ impl Request { Request::CompleteUserName(_) => Default::default(), Request::CreateDatabases(databases) => databases.iter().cloned().collect(), Request::DropDatabases(databases) => databases.iter().cloned().collect(), - Request::ListDatabases(databases) => { - databases.clone().unwrap_or_default().into_iter().collect() - } + Request::ListDatabases(request) => request + .names + .clone() + .unwrap_or_default() + .into_iter() + .collect(), Request::ListPrivileges(databases) => { databases.clone().unwrap_or_default().into_iter().collect() } diff --git a/src/core/protocol/commands/list_databases.rs b/src/core/protocol/commands/list_databases.rs index 3d33dcf..cc8cde0 100644 --- a/src/core/protocol/commands/list_databases.rs +++ b/src/core/protocol/commands/list_databases.rs @@ -14,7 +14,21 @@ use crate::{ server::sql::database_operations::DatabaseRow, }; -pub type ListDatabasesRequest = Option>; +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct ListDatabasesRequest { + pub names: Option>, + #[serde(default)] + pub include_all_tables_and_users: bool, +} + +impl ListDatabasesRequest { + pub fn new(names: Option>, include_all_tables_and_users: bool) -> Self { + Self { + names, + include_all_tables_and_users, + } + } +} pub type ListDatabasesResponse = BTreeMap>; @@ -144,7 +158,7 @@ mod tests { #[test] fn test_serialize_deserialize_request() { - let request = Some(vec!["db1".into(), "db2".into()]); + let request = ListDatabasesRequest::new(Some(vec!["db1".into(), "db2".into()]), true); let json = serde_json::to_string_pretty(&request).unwrap(); println!("Serialized request:\n{}", json); @@ -152,6 +166,20 @@ mod tests { assert_eq!(request, deserialized); } + #[test] + fn test_deserialize_request_without_include_all_tables_and_users_defaults_to_false() { + let json = serde_json::json!({ + "names": ["db1", "db2"] + }) + .to_string(); + + let deserialized: ListDatabasesRequest = serde_json::from_str(&json).unwrap(); + assert_eq!( + deserialized, + ListDatabasesRequest::new(Some(vec!["db1".into(), "db2".into()]), false) + ); + } + #[test] fn test_serialize_deserialize_response() { let response: ListDatabasesResponse = vec![ diff --git a/src/server/session_handler.rs b/src/server/session_handler.rs index b683fe6..6b48ca6 100644 --- a/src/server/session_handler.rs +++ b/src/server/session_handler.rs @@ -332,14 +332,15 @@ async fn handle_request( .await; Response::DropDatabases(result) } - Request::ListDatabases(ref database_names) => { - if let Some(database_names) = database_names { + Request::ListDatabases(ref request) => { + if let Some(database_names) = &request.names { let result = list_databases( database_names, unix_user, db_connection, db_is_mariadb, group_denylist, + request.include_all_tables_and_users, ) .await; Response::ListDatabases(result) @@ -349,6 +350,7 @@ async fn handle_request( db_connection, db_is_mariadb, group_denylist, + request.include_all_tables_and_users, ) .await; Response::ListAllDatabases(result) diff --git a/src/server/sql/database_operations.rs b/src/server/sql/database_operations.rs index 1a7c381..f80c9d2 100644 --- a/src/server/sql/database_operations.rs +++ b/src/server/sql/database_operations.rs @@ -24,6 +24,8 @@ use crate::{ server::{common::create_user_group_matching_regex, sql::quote_identifier}, }; +const MAX_SHOW_DB_RELATED_ITEMS: usize = 5; + // NOTE: this function is unsafe because it does no input validation. pub(super) async fn unsafe_database_exists( database_name: &str, @@ -248,12 +250,84 @@ impl FromRow<'_, sqlx::mysql::MySqlRow> for DatabaseRow { } } +fn list_database_query(include_all_tables_and_users: bool) -> AssertSqlSafe { + let limit_clause = if include_all_tables_and_users { + "".to_string() + } else { + format!(" LIMIT {}", MAX_SHOW_DB_RELATED_ITEMS) + }; + + AssertSqlSafe(format!( + r" + SELECT + CAST(s.SCHEMA_NAME AS CHAR(64)) AS `database`, + t.tables, + u.users, + s.DEFAULT_COLLATION_NAME AS `collation`, + s.DEFAULT_CHARACTER_SET_NAME AS `character_set`, + CAST(COALESCE(sz.size_bytes, 0) AS UNSIGNED) AS size_bytes + FROM information_schema.SCHEMATA s + + LEFT JOIN ( + SELECT + x.TABLE_SCHEMA, + GROUP_CONCAT(x.TABLE_NAME ORDER BY x.TABLE_NAME SEPARATOR ',') AS tables + FROM ( + SELECT + TABLE_SCHEMA, + TABLE_NAME + FROM information_schema.TABLES + WHERE TABLE_SCHEMA = ? + ORDER BY TABLE_NAME{limit_clause} + ) x + GROUP BY x.TABLE_SCHEMA + ) t + ON t.TABLE_SCHEMA = s.SCHEMA_NAME + + LEFT JOIN ( + SELECT + x.DB, + GROUP_CONCAT(DISTINCT x.User ORDER BY x.User SEPARATOR ',') AS users + FROM ( + SELECT + DB, + User + FROM mysql.db + WHERE DB = ? + ORDER BY User{limit_clause} + ) x + GROUP BY x.DB + ) u + ON u.DB = s.SCHEMA_NAME + + LEFT JOIN ( + SELECT + TABLE_SCHEMA, + SUM(DATA_LENGTH + INDEX_LENGTH) AS size_bytes + FROM information_schema.TABLES + WHERE TABLE_SCHEMA = ? + GROUP BY TABLE_SCHEMA + ) sz + ON sz.TABLE_SCHEMA = s.SCHEMA_NAME + + WHERE s.SCHEMA_NAME REGEXP ? + AND s.SCHEMA_NAME NOT IN ( + 'information_schema', + 'performance_schema', + 'mysql', + 'sys' + ) + " + )) +} + pub async fn list_databases( database_names: &[MySQLDatabase], unix_user: &UnixUser, connection: &mut MySqlConnection, _db_is_mariadb: bool, group_denylist: &GroupDenylist, + include_all_tables_and_users: bool, ) -> ListDatabasesResponse { let mut results = BTreeMap::new(); @@ -269,58 +343,19 @@ pub async fn list_databases( continue; } - let result = sqlx::query_as::<_, DatabaseRow>( - r" - SELECT - CAST(s.SCHEMA_NAME AS CHAR(64)) AS `database`, - t.tables, - u.users, - s.DEFAULT_COLLATION_NAME AS `collation`, - s.DEFAULT_CHARACTER_SET_NAME AS `character_set`, - CAST(COALESCE(t.size_bytes, 0) AS UNSIGNED) AS `size_bytes` - FROM information_schema.SCHEMATA s + let query = list_database_query(include_all_tables_and_users); - LEFT JOIN ( - SELECT - TABLE_SCHEMA, - GROUP_CONCAT( - DISTINCT CAST(TABLE_NAME AS CHAR(64)) - ORDER BY TABLE_NAME - SEPARATOR ',' - ) AS tables, - SUM(DATA_LENGTH + INDEX_LENGTH) AS size_bytes - FROM information_schema.TABLES - WHERE TABLE_SCHEMA = ? - GROUP BY TABLE_SCHEMA - ) t - ON t.TABLE_SCHEMA = s.SCHEMA_NAME - - LEFT JOIN ( - SELECT - DB, - GROUP_CONCAT( - DISTINCT CAST(User AS CHAR(64)) - ORDER BY User - SEPARATOR ',' - ) AS users - FROM mysql.db - WHERE DB = ? - GROUP BY DB - ) u - ON u.DB = s.SCHEMA_NAME - - WHERE s.SCHEMA_NAME = ?; - ", - ) - .bind(database_name.to_string()) - .bind(database_name.to_string()) - .bind(database_name.to_string()) - .fetch_optional(&mut *connection) - .await - .map_err(|err| ListDatabasesError::MySqlError(err.to_string())) - .and_then(|database| { - database.map_or_else(|| Err(ListDatabasesError::DatabaseDoesNotExist), Ok) - }); + let result = sqlx::query_as::<_, DatabaseRow>(query) + .bind(database_name.to_string()) + .bind(database_name.to_string()) + .bind(database_name.to_string()) + .bind(database_name.to_string()) + .fetch_optional(&mut *connection) + .await + .map_err(|err| ListDatabasesError::MySqlError(err.to_string())) + .and_then(|database| { + database.map_or_else(|| Err(ListDatabasesError::DatabaseDoesNotExist), Ok) + }); if let Err(err) = &result { tracing::error!("Failed to list database '{}': {:?}", &database_name, err); @@ -334,69 +369,97 @@ pub async fn list_databases( results } +fn list_all_databases_for_user_query(include_all_tables_and_users: bool) -> AssertSqlSafe { + let limit_clause = if include_all_tables_and_users { + "".to_string() + } else { + format!(" LIMIT {}", MAX_SHOW_DB_RELATED_ITEMS) + }; + + AssertSqlSafe(format!( + r" + SELECT + CAST(s.SCHEMA_NAME AS CHAR(64)) AS `database`, + t.tables, + u.users, + s.DEFAULT_COLLATION_NAME AS `collation`, + s.DEFAULT_CHARACTER_SET_NAME AS `character_set`, + CAST(COALESCE(sz.size_bytes, 0) AS UNSIGNED) AS size_bytes + FROM information_schema.SCHEMATA s + + LEFT JOIN ( + SELECT + x.TABLE_SCHEMA, + GROUP_CONCAT(x.TABLE_NAME ORDER BY x.TABLE_NAME SEPARATOR ',') AS tables + FROM ( + SELECT + TABLE_SCHEMA, + TABLE_NAME + FROM information_schema.TABLES + WHERE TABLE_SCHEMA REGEXP ? + ORDER BY TABLE_NAME{limit_clause} + ) x + GROUP BY x.TABLE_SCHEMA + ) t + ON t.TABLE_SCHEMA = s.SCHEMA_NAME + + LEFT JOIN ( + SELECT + x.DB, + GROUP_CONCAT(DISTINCT x.User ORDER BY x.User SEPARATOR ',') AS users + FROM ( + SELECT + DB, + User + FROM mysql.db + WHERE DB REGEXP ? + ORDER BY User{limit_clause} + ) x + GROUP BY x.DB + ) u + ON u.DB = s.SCHEMA_NAME + + LEFT JOIN ( + SELECT + TABLE_SCHEMA, + SUM(DATA_LENGTH + INDEX_LENGTH) AS size_bytes + FROM information_schema.TABLES + WHERE TABLE_SCHEMA REGEXP ? + GROUP BY TABLE_SCHEMA + ) sz + ON sz.TABLE_SCHEMA = s.SCHEMA_NAME + + WHERE s.SCHEMA_NAME REGEXP ? + AND s.SCHEMA_NAME NOT IN ( + 'information_schema', + 'performance_schema', + 'mysql', + 'sys' + ) + + ORDER BY s.SCHEMA_NAME + " + )) +} + pub async fn list_all_databases_for_user( unix_user: &UnixUser, connection: &mut MySqlConnection, _db_is_mariadb: bool, group_denylist: &GroupDenylist, + include_all_tables_and_users: bool, ) -> ListAllDatabasesResponse { - let result = sqlx::query_as::<_, DatabaseRow>( - r" - SELECT - CAST(s.SCHEMA_NAME AS CHAR(64)) AS `database`, - t.tables, - u.users, - s.DEFAULT_COLLATION_NAME AS collation, - s.DEFAULT_CHARACTER_SET_NAME AS character_set, - CAST(COALESCE(t.size_bytes, 0) AS UNSIGNED) AS size_bytes - FROM information_schema.SCHEMATA s + let query = list_all_databases_for_user_query(include_all_tables_and_users); + let user_group_regex = create_user_group_matching_regex(unix_user, group_denylist); - LEFT JOIN ( - SELECT - TABLE_SCHEMA, - GROUP_CONCAT( - DISTINCT CAST(TABLE_NAME AS CHAR(64)) - ORDER BY TABLE_NAME - SEPARATOR ',' - ) AS tables, - SUM(DATA_LENGTH + INDEX_LENGTH) AS size_bytes - FROM information_schema.TABLES - WHERE TABLE_SCHEMA REGEXP ? - GROUP BY TABLE_SCHEMA - ) t - ON t.TABLE_SCHEMA = s.SCHEMA_NAME - - LEFT JOIN ( - SELECT - DB, - GROUP_CONCAT( - DISTINCT CAST(User AS CHAR(64)) - ORDER BY User - SEPARATOR ',' - ) AS users - FROM mysql.db - WHERE DB REGEXP ? - GROUP BY DB - ) u - ON u.DB = s.SCHEMA_NAME - - WHERE s.SCHEMA_NAME REGEXP ? - AND s.SCHEMA_NAME NOT IN ( - 'information_schema', - 'performance_schema', - 'mysql', - 'sys' - ) - - ORDER BY s.SCHEMA_NAME - ", - ) - .bind(create_user_group_matching_regex(unix_user, group_denylist)) - .bind(create_user_group_matching_regex(unix_user, group_denylist)) - .bind(create_user_group_matching_regex(unix_user, group_denylist)) - .fetch_all(connection) - .await - .map_err(|err| ListAllDatabasesError::MySqlError(err.to_string())); + let result = sqlx::query_as::<_, DatabaseRow>(query) + .bind(&user_group_regex) + .bind(&user_group_regex) + .bind(&user_group_regex) + .bind(&user_group_regex) + .fetch_all(connection) + .await + .map_err(|err| ListAllDatabasesError::MySqlError(err.to_string())); // TODO: should we assert that the users are also owned by the unix_user from the request?