Add dynamic completion for users and databases

src/server/session_handler.rs

@@ -18,15 +18,16 @@ use crate::{
         authorization::check_authorization,
         sql::{
             database_operations::{
-                create_databases, drop_databases, list_all_databases_for_user, list_databases,
+                complete_database_name, create_databases, drop_databases,
+                list_all_databases_for_user, list_databases,
             },
             database_privilege_operations::{
                 apply_privilege_diffs, get_all_database_privileges, get_databases_privilege_data,
             },
             user_operations::{
-                create_database_users, drop_database_users, list_all_database_users_for_unix_user,
-                list_database_users, lock_database_users, set_password_for_database_user,
-                unlock_database_users,
+                complete_user_name, create_database_users, drop_database_users,
+                list_all_database_users_for_unix_user, list_database_users, lock_database_users,
+                set_password_for_database_user, unlock_database_users,
             },
         },
     },
@@ -171,6 +172,33 @@ async fn session_handler_with_db_connection(
                 let result = check_authorization(dbs_or_users, unix_user).await;
                 Response::CheckAuthorization(result)
             }
+            Request::CompleteDatabaseName(partial_database_name) => {
+                // TODO: more correct validation here
+                if !partial_database_name
+                    .chars()
+                    .all(|c| c.is_ascii_alphanumeric() || c == '_' || c == '-')
+                {
+                    Response::CompleteDatabaseName(vec![])
+                } else {
+                    let result =
+                        complete_database_name(partial_database_name, unix_user, db_connection)
+                            .await;
+                    Response::CompleteDatabaseName(result)
+                }
+            }
+            Request::CompleteUserName(partial_user_name) => {
+                // TODO: more correct validation here
+                if !partial_user_name
+                    .chars()
+                    .all(|c| c.is_ascii_alphanumeric() || c == '_' || c == '-')
+                {
+                    Response::CompleteUserName(vec![])
+                } else {
+                    let result =
+                        complete_user_name(partial_user_name, unix_user, db_connection).await;
+                    Response::CompleteUserName(result)
+                }
+            }
             Request::CreateDatabases(databases_names) => {
                 let result = create_databases(databases_names, unix_user, db_connection).await;
                 Response::CreateDatabases(result)

src/server/sql/database_operations.rs

@@ -5,6 +5,7 @@ use sqlx::prelude::*;
 
 use serde::{Deserialize, Serialize};
 
+use crate::core::protocol::CompleteDatabaseNameResponse;
 use crate::core::types::MySQLDatabase;
 use crate::{
     core::{
@@ -43,6 +44,45 @@ pub(super) async fn unsafe_database_exists(
     Ok(result?.is_some())
 }
 
+pub async fn complete_database_name(
+    database_prefix: String,
+    unix_user: &UnixUser,
+    connection: &mut MySqlConnection,
+) -> CompleteDatabaseNameResponse {
+    let result = sqlx::query(
+        r#"
+          SELECT `SCHEMA_NAME` AS `database`
+          FROM `information_schema`.`SCHEMATA`
+          WHERE `SCHEMA_NAME` NOT IN ('information_schema', 'performance_schema', 'mysql', 'sys')
+            AND `SCHEMA_NAME` REGEXP ?
+            AND `SCHEMA_NAME` LIKE ?
+        "#,
+    )
+    .bind(create_user_group_matching_regex(unix_user))
+    .bind(format!("{}%", database_prefix))
+    .fetch_all(connection)
+    .await;
+
+    match result {
+        Ok(rows) => rows
+            .into_iter()
+            .filter_map(|row| {
+                let database: String = row.try_get("database").ok()?;
+                Some(database.into())
+            })
+            .collect(),
+        Err(err) => {
+            tracing::error!(
+                "Failed to complete database name for prefix '{}' and user '{}': {:?}",
+                database_prefix,
+                unix_user.username,
+                err
+            );
+            vec![]
+        }
+    }
+}
+
 pub async fn create_databases(
     database_names: Vec<MySQLDatabase>,
     unix_user: &UnixUser,

src/server/sql/user_operations.rs

@@ -51,6 +51,44 @@ async fn unsafe_user_exists(
     result
 }
 
+pub async fn complete_user_name(
+    user_prefix: String,
+    unix_user: &UnixUser,
+    connection: &mut MySqlConnection,
+) -> Vec<MySQLUser> {
+    let result = sqlx::query(
+        r#"
+          SELECT `User` AS `user`
+          FROM `mysql`.`user`
+          WHERE `User` REGEXP ?
+            AND `User` LIKE ?
+        "#,
+    )
+    .bind(create_user_group_matching_regex(unix_user))
+    .bind(format!("{}%", user_prefix))
+    .fetch_all(connection)
+    .await;
+
+    match result {
+        Ok(rows) => rows
+            .into_iter()
+            .filter_map(|row| {
+                let user: String = try_get_with_binary_fallback(&row, "user").ok()?;
+                Some(user.into())
+            })
+            .collect(),
+        Err(err) => {
+            tracing::error!(
+                "Failed to complete user name for prefix '{}' and user '{}': {:?}",
+                user_prefix,
+                unix_user.username,
+                err
+            );
+            vec![]
+        }
+    }
+}
+
 pub async fn create_database_users(
     db_users: Vec<MySQLUser>,
     unix_user: &UnixUser,