Let mysql connection shut down gracefully

src/cli/database_command.rs

@@ -1,7 +1,7 @@
 use anyhow::Context;
 use clap::Parser;
 use prettytable::{Cell, Row, Table};
-use sqlx::MySqlConnection;
+use sqlx::{Connection, MySqlConnection};
 
 use crate::core::{self, database_operations::DatabasePrivileges};
 
@@ -124,19 +124,23 @@ struct DatabaseEditPermArgs {
     yes: bool,
 }
 
-pub async fn handle_command(args: DatabaseArgs, conn: MySqlConnection) -> anyhow::Result<()> {
+pub async fn handle_command(args: DatabaseArgs, mut conn: MySqlConnection) -> anyhow::Result<()> {
-    match args.subcmd {
+    let result = match args.subcmd {
-        DatabaseCommand::Create(args) => create_databases(args, conn).await,
+        DatabaseCommand::Create(args) => create_databases(args, &mut conn).await,
-        DatabaseCommand::Drop(args) => drop_databases(args, conn).await,
+        DatabaseCommand::Drop(args) => drop_databases(args, &mut conn).await,
-        DatabaseCommand::List(args) => list_databases(args, conn).await,
+        DatabaseCommand::List(args) => list_databases(args, &mut conn).await,
-        DatabaseCommand::ShowPerm(args) => show_databases(args, conn).await,
+        DatabaseCommand::ShowPerm(args) => show_databases(args, &mut conn).await,
-        DatabaseCommand::EditPerm(args) => edit_permissions(args, conn).await,
+        DatabaseCommand::EditPerm(args) => edit_permissions(args, &mut conn).await,
-    }
+    };
+
+    conn.close().await?;
+
+    result
 }
 
 async fn create_databases(
     args: DatabaseCreateArgs,
-    mut conn: MySqlConnection,
+    conn: &mut MySqlConnection,
 ) -> anyhow::Result<()> {
     if args.name.is_empty() {
         anyhow::bail!("No database names provided");
@@ -144,7 +148,7 @@ async fn create_databases(
 
     for name in args.name {
         // TODO: This can be optimized by fetching all the database privileges in one query.
-        if let Err(e) = core::database_operations::create_database(&name, &mut conn).await {
+        if let Err(e) = core::database_operations::create_database(&name, conn).await {
             eprintln!("Failed to create database '{}': {}", name, e);
             eprintln!("Skipping...");
         }
@@ -153,14 +157,14 @@ async fn create_databases(
     Ok(())
 }
 
-async fn drop_databases(args: DatabaseDropArgs, mut conn: MySqlConnection) -> anyhow::Result<()> {
+async fn drop_databases(args: DatabaseDropArgs, conn: &mut MySqlConnection) -> anyhow::Result<()> {
     if args.name.is_empty() {
         anyhow::bail!("No database names provided");
     }
 
     for name in args.name {
         // TODO: This can be optimized by fetching all the database privileges in one query.
-        if let Err(e) = core::database_operations::drop_database(&name, &mut conn).await {
+        if let Err(e) = core::database_operations::drop_database(&name, conn).await {
             eprintln!("Failed to drop database '{}': {}", name, e);
             eprintln!("Skipping...");
         }
@@ -169,8 +173,8 @@ async fn drop_databases(args: DatabaseDropArgs, mut conn: MySqlConnection) -> an
     Ok(())
 }
 
-async fn list_databases(args: DatabaseListArgs, mut conn: MySqlConnection) -> anyhow::Result<()> {
+async fn list_databases(args: DatabaseListArgs, conn: &mut MySqlConnection) -> anyhow::Result<()> {
-    let databases = core::database_operations::get_database_list(&mut conn).await?;
+    let databases = core::database_operations::get_database_list(conn).await?;
 
     if databases.is_empty() {
         println!("No databases to show.");
@@ -188,14 +192,17 @@ async fn list_databases(args: DatabaseListArgs, mut conn: MySqlConnection) -> an
     Ok(())
 }
 
-async fn show_databases(args: DatabaseShowPermArgs, mut conn: MySqlConnection) -> anyhow::Result<()> {
+async fn show_databases(
+    args: DatabaseShowPermArgs,
+    conn: &mut MySqlConnection,
+) -> anyhow::Result<()> {
     let database_users_to_show = if args.name.is_empty() {
-        core::database_operations::get_all_database_privileges(&mut conn).await?
+        core::database_operations::get_all_database_privileges(conn).await?
     } else {
         // TODO: This can be optimized by fetching all the database privileges in one query.
         let mut result = Vec::with_capacity(args.name.len());
         for name in args.name {
-            match core::database_operations::get_database_privileges(&name, &mut conn).await {
+            match core::database_operations::get_database_privileges(&name, conn).await {
                 Ok(db) => result.extend(db),
                 Err(e) => {
                     eprintln!("Failed to show database '{}': {}", name, e);
@@ -305,11 +312,14 @@ fn parse_permission_table_cli_arg(arg: &str) -> anyhow::Result<DatabasePrivilege
     Ok(result)
 }
 
-async fn edit_permissions(args: DatabaseEditPermArgs, mut conn: MySqlConnection) -> anyhow::Result<()> {
+async fn edit_permissions(
+    args: DatabaseEditPermArgs,
+    conn: &mut MySqlConnection,
+) -> anyhow::Result<()> {
     let _data = if let Some(name) = &args.name {
-        core::database_operations::get_database_privileges(name, &mut conn).await?
+        core::database_operations::get_database_privileges(name, conn).await?
     } else {
-        core::database_operations::get_all_database_privileges(&mut conn).await?
+        core::database_operations::get_all_database_privileges(conn).await?
     };
 
     if !args.text {

src/cli/user_command.rs

@@ -2,7 +2,7 @@ use std::vec;
 
 use anyhow::Context;
 use clap::Parser;
-use sqlx::MySqlConnection;
+use sqlx::{Connection, MySqlConnection};
 
 use crate::core::user_operations::validate_ownership_of_user_name;
 
@@ -57,24 +57,28 @@ struct UserShowArgs {
     username: Vec<String>,
 }
 
-pub async fn handle_command(args: UserArgs, conn: MySqlConnection) -> anyhow::Result<()> {
+pub async fn handle_command(args: UserArgs, mut conn: MySqlConnection) -> anyhow::Result<()> {
-    match args.subcmd {
+    let result = match args.subcmd {
-        UserCommand::Create(args) => create_users(args, conn).await,
+        UserCommand::Create(args) => create_users(args, &mut conn).await,
-        UserCommand::Drop(args) => drop_users(args, conn).await,
+        UserCommand::Drop(args) => drop_users(args, &mut conn).await,
-        UserCommand::Passwd(args) => change_password_for_user(args, conn).await,
+        UserCommand::Passwd(args) => change_password_for_user(args, &mut conn).await,
-        UserCommand::Show(args) => show_users(args, conn).await,
+        UserCommand::Show(args) => show_users(args, &mut conn).await,
-    }
+    };
+
+    conn.close().await?;
+
+    result
 }
 
 // TODO: provide a better error message when the user already exists
-async fn create_users(args: UserCreateArgs, mut conn: MySqlConnection) -> anyhow::Result<()> {
+async fn create_users(args: UserCreateArgs, conn: &mut MySqlConnection) -> anyhow::Result<()> {
     if args.username.is_empty() {
         anyhow::bail!("No usernames provided");
     }
 
     for username in args.username {
         if let Err(e) =
-            crate::core::user_operations::create_database_user(&username, &mut conn).await
+            crate::core::user_operations::create_database_user(&username, conn).await
         {
             eprintln!("{}", e);
             eprintln!("Skipping...");
@@ -84,14 +88,14 @@ async fn create_users(args: UserCreateArgs, mut conn: MySqlConnection) -> anyhow
 }
 
 // TODO: provide a better error message when the user does not exist
-async fn drop_users(args: UserDeleteArgs, mut conn: MySqlConnection) -> anyhow::Result<()> {
+async fn drop_users(args: UserDeleteArgs, conn: &mut MySqlConnection) -> anyhow::Result<()> {
     if args.username.is_empty() {
         anyhow::bail!("No usernames provided");
     }
 
     for username in args.username {
         if let Err(e) =
-            crate::core::user_operations::delete_database_user(&username, &mut conn).await
+            crate::core::user_operations::delete_database_user(&username, conn).await
         {
             eprintln!("{}", e);
             eprintln!("Skipping...");
@@ -102,7 +106,7 @@ async fn drop_users(args: UserDeleteArgs, mut conn: MySqlConnection) -> anyhow::
 
 async fn change_password_for_user(
     args: UserPasswdArgs,
-    mut conn: MySqlConnection,
+    conn: &mut MySqlConnection,
 ) -> anyhow::Result<()> {
     // NOTE: although this also is checked in `set_password_for_database_user`, we check it here
     //       to provide a more natural order of error messages.
@@ -131,18 +135,18 @@ async fn change_password_for_user(
     crate::core::user_operations::set_password_for_database_user(
         &args.username,
         &password,
-        &mut conn,
+        conn,
     )
     .await?;
 
     Ok(())
 }
 
-async fn show_users(args: UserShowArgs, mut conn: MySqlConnection) -> anyhow::Result<()> {
+async fn show_users(args: UserShowArgs, conn: &mut MySqlConnection) -> anyhow::Result<()> {
     let user = crate::core::common::get_current_unix_user()?;
 
     let users = if args.username.is_empty() {
-        crate::core::user_operations::get_all_database_users_for_user(&user, &mut conn).await?
+        crate::core::user_operations::get_all_database_users_for_user(&user, conn).await?
     } else {
         let mut result = vec![];
         for username in args.username {
@@ -153,7 +157,7 @@ async fn show_users(args: UserShowArgs, mut conn: MySqlConnection) -> anyhow::Re
             }
 
             let user =
-                crate::core::user_operations::get_database_user_for_user(&username, &mut conn)
+                crate::core::user_operations::get_database_user_for_user(&username, conn)
                     .await?;
             if let Some(user) = user {
                 result.push(user);