Add dynamic completion for users and databases

src/core/bootstrap.rs

@@ -159,7 +159,7 @@ fn connect_to_external_server(
 /// Drop privileges to the real user and group of the process.
 /// If the process is not running with elevated privileges, this function
 /// is a no-op.
-fn drop_privs() -> anyhow::Result<()> {
+pub fn drop_privs() -> anyhow::Result<()> {
     tracing::debug!("Dropping privileges");
     let real_uid = nix::unistd::getuid();
     let real_gid = nix::unistd::getgid();

src/core/completion.rs

@@ -0,0 +1,5 @@
+mod mysql_database_completer;
+mod mysql_user_completer;
+
+pub use mysql_database_completer::*;
+pub use mysql_user_completer::*;

src/core/completion/mysql_database_completer.rs

@@ -0,0 +1,73 @@
+use clap_complete::CompletionCandidate;
+use clap_verbosity_flag::Verbosity;
+use futures_util::SinkExt;
+use tokio::net::UnixStream as TokioUnixStream;
+use tokio_stream::StreamExt;
+
+use crate::{
+    client::commands::erroneous_server_response,
+    core::{
+        bootstrap::bootstrap_server_connection_and_drop_privileges,
+        protocol::{Request, Response, create_client_to_server_message_stream},
+    },
+};
+
+pub fn mysql_database_completer(current: &std::ffi::OsStr) -> Vec<CompletionCandidate> {
+    match tokio::runtime::Builder::new_current_thread()
+        .enable_all()
+        .build()
+    {
+        Ok(runtime) => match runtime.block_on(mysql_database_completer_(current)) {
+            Ok(completions) => completions,
+            Err(err) => {
+                eprintln!("Error getting MySQL database completions: {}", err);
+                Vec::new()
+            }
+        },
+        Err(err) => {
+            eprintln!("Error starting Tokio runtime: {}", err);
+            Vec::new()
+        }
+    }
+}
+
+/// Connect to the server to get MySQL database completions.
+async fn mysql_database_completer_(
+    current: &std::ffi::OsStr,
+) -> anyhow::Result<Vec<CompletionCandidate>> {
+    let server_connection =
+        bootstrap_server_connection_and_drop_privileges(None, None, Verbosity::new(0, 1))?;
+
+    let tokio_socket = TokioUnixStream::from_std(server_connection)?;
+    let mut server_connection = create_client_to_server_message_stream(tokio_socket);
+
+    while let Some(Ok(message)) = server_connection.next().await {
+        match message {
+            Response::Error(err) => {
+                anyhow::bail!("{}", err);
+            }
+            Response::Ready => break,
+            message => {
+                eprintln!("Unexpected message from server: {:?}", message);
+            }
+        }
+    }
+
+    let message = Request::CompleteDatabaseName(current.to_string_lossy().to_string());
+
+    if let Err(err) = server_connection.send(message).await {
+        server_connection.close().await.ok();
+        anyhow::bail!(anyhow::Error::from(err).context("Failed to communicate with server"));
+    }
+
+    let result = match server_connection.next().await {
+        Some(Ok(Response::CompleteDatabaseName(suggestions))) => suggestions,
+        response => return erroneous_server_response(response).map(|_| vec![]),
+    };
+
+    server_connection.send(Request::Exit).await?;
+
+    let result = result.into_iter().map(CompletionCandidate::new).collect();
+
+    Ok(result)
+}

src/core/completion/mysql_user_completer.rs

@@ -0,0 +1,73 @@
+use clap_complete::CompletionCandidate;
+use clap_verbosity_flag::Verbosity;
+use futures_util::SinkExt;
+use tokio::net::UnixStream as TokioUnixStream;
+use tokio_stream::StreamExt;
+
+use crate::{
+    client::commands::erroneous_server_response,
+    core::{
+        bootstrap::bootstrap_server_connection_and_drop_privileges,
+        protocol::{Request, Response, create_client_to_server_message_stream},
+    },
+};
+
+pub fn mysql_user_completer(current: &std::ffi::OsStr) -> Vec<CompletionCandidate> {
+    match tokio::runtime::Builder::new_current_thread()
+        .enable_all()
+        .build()
+    {
+        Ok(runtime) => match runtime.block_on(mysql_user_completer_(current)) {
+            Ok(completions) => completions,
+            Err(err) => {
+                eprintln!("Error getting MySQL user completions: {}", err);
+                Vec::new()
+            }
+        },
+        Err(err) => {
+            eprintln!("Error starting Tokio runtime: {}", err);
+            Vec::new()
+        }
+    }
+}
+
+/// Connect to the server to get MySQL user completions.
+async fn mysql_user_completer_(
+    current: &std::ffi::OsStr,
+) -> anyhow::Result<Vec<CompletionCandidate>> {
+    let server_connection =
+        bootstrap_server_connection_and_drop_privileges(None, None, Verbosity::new(0, 1))?;
+
+    let tokio_socket = TokioUnixStream::from_std(server_connection)?;
+    let mut server_connection = create_client_to_server_message_stream(tokio_socket);
+
+    while let Some(Ok(message)) = server_connection.next().await {
+        match message {
+            Response::Error(err) => {
+                anyhow::bail!("{}", err);
+            }
+            Response::Ready => break,
+            message => {
+                eprintln!("Unexpected message from server: {:?}", message);
+            }
+        }
+    }
+
+    let message = Request::CompleteUserName(current.to_string_lossy().to_string());
+
+    if let Err(err) = server_connection.send(message).await {
+        server_connection.close().await.ok();
+        anyhow::bail!(anyhow::Error::from(err).context("Failed to communicate with server"));
+    }
+
+    let result = match server_connection.next().await {
+        Some(Ok(Response::CompleteUserName(suggestions))) => suggestions,
+        response => return erroneous_server_response(response).map(|_| vec![]),
+    };
+
+    server_connection.send(Request::Exit).await?;
+
+    let result = result.into_iter().map(CompletionCandidate::new).collect();
+
+    Ok(result)
+}

src/core/protocol/commands.rs

@@ -1,4 +1,6 @@
 mod check_authorization;
+mod complete_database_name;
+mod complete_user_name;
 mod create_databases;
 mod create_users;
 mod drop_databases;
@@ -15,6 +17,8 @@ mod passwd_user;
 mod unlock_users;
 
 pub use check_authorization::*;
+pub use complete_database_name::*;
+pub use complete_user_name::*;
 pub use create_databases::*;
 pub use create_users::*;
 pub use drop_databases::*;
@@ -64,6 +68,9 @@ pub fn create_client_to_server_message_stream(socket: UnixStream) -> ClientToSer
 pub enum Request {
     CheckAuthorization(CheckAuthorizationRequest),
 
+    CompleteDatabaseName(CompleteDatabaseNameRequest),
+    CompleteUserName(CompleteUserNameRequest),
+
     CreateDatabases(CreateDatabasesRequest),
     DropDatabases(DropDatabasesRequest),
     ListDatabases(ListDatabasesRequest),
@@ -88,6 +95,9 @@ pub enum Request {
 pub enum Response {
     CheckAuthorization(CheckAuthorizationResponse),
 
+    CompleteDatabaseName(CompleteDatabaseNameResponse),
+    CompleteUserName(CompleteUserNameResponse),
+
     // Specific data for specific commands
     CreateDatabases(CreateDatabasesResponse),
     DropDatabases(DropDatabasesResponse),

src/core/protocol/commands/complete_database_name.rs

@@ -0,0 +1,5 @@
+use crate::core::types::MySQLDatabase;
+
+pub type CompleteDatabaseNameRequest = String;
+
+pub type CompleteDatabaseNameResponse = Vec<MySQLDatabase>;

src/core/protocol/commands/complete_user_name.rs

@@ -0,0 +1,5 @@
+use crate::core::types::MySQLUser;
+
+pub type CompleteUserNameRequest = String;
+
+pub type CompleteUserNameResponse = Vec<MySQLUser>;

src/core/types.rs

@@ -1,4 +1,5 @@
 use std::{
+    ffi::OsString,
     fmt,
     ops::{Deref, DerefMut},
     str::FromStr,
@@ -49,6 +50,12 @@ impl From<String> for MySQLUser {
     }
 }
 
+impl From<MySQLUser> for OsString {
+    fn from(val: MySQLUser) -> Self {
+        val.0.into()
+    }
+}
+
 #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize, Default)]
 pub struct MySQLDatabase(String);
 
@@ -92,6 +99,12 @@ impl From<String> for MySQLDatabase {
     }
 }
 
+impl From<MySQLDatabase> for OsString {
+    fn from(val: MySQLDatabase) -> Self {
+        val.0.into()
+    }
+}
+
 #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
 pub enum DbOrUser {
     Database(MySQLDatabase),