Let mysql connection shut down gracefully

This commit is contained in:
Oystein Kristoffer Tveit 2024-04-26 00:29:39 +02:00
parent 561241d589
commit 3deeeb45c2
Signed by: oysteikt
GPG Key ID: 9F2F7D8250F35146
2 changed files with 52 additions and 38 deletions

View File

@ -1,7 +1,7 @@
use anyhow::Context; use anyhow::Context;
use clap::Parser; use clap::Parser;
use prettytable::{Cell, Row, Table}; use prettytable::{Cell, Row, Table};
use sqlx::MySqlConnection; use sqlx::{Connection, MySqlConnection};
use crate::core::{self, database_operations::DatabasePrivileges}; use crate::core::{self, database_operations::DatabasePrivileges};
@ -124,19 +124,23 @@ struct DatabaseEditPermArgs {
yes: bool, 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( async fn create_databases(
args: DatabaseCreateArgs, args: DatabaseCreateArgs,
mut conn: MySqlConnection, conn: &mut MySqlConnection,
) -> anyhow::Result<()> { ) -> anyhow::Result<()> {
if args.name.is_empty() { if args.name.is_empty() {
anyhow::bail!("No database names provided"); anyhow::bail!("No database names provided");
@ -144,7 +148,7 @@ async fn create_databases(
for name in args.name { for name in args.name {
// TODO: This can be optimized by fetching all the database privileges in one query. // 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!("Failed to create database '{}': {}", name, e);
eprintln!("Skipping..."); eprintln!("Skipping...");
} }
@ -153,14 +157,14 @@ async fn create_databases(
Ok(()) 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() { if args.name.is_empty() {
anyhow::bail!("No database names provided"); anyhow::bail!("No database names provided");
} }
for name in args.name { for name in args.name {
// TODO: This can be optimized by fetching all the database privileges in one query. // 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!("Failed to drop database '{}': {}", name, e);
eprintln!("Skipping..."); eprintln!("Skipping...");
} }
@ -169,8 +173,8 @@ async fn drop_databases(args: DatabaseDropArgs, mut conn: MySqlConnection) -> an
Ok(()) 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() { if databases.is_empty() {
println!("No databases to show."); println!("No databases to show.");
@ -188,14 +192,17 @@ async fn list_databases(args: DatabaseListArgs, mut conn: MySqlConnection) -> an
Ok(()) 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() { 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 { } else {
// TODO: This can be optimized by fetching all the database privileges in one query. // TODO: This can be optimized by fetching all the database privileges in one query.
let mut result = Vec::with_capacity(args.name.len()); let mut result = Vec::with_capacity(args.name.len());
for name in args.name { 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), Ok(db) => result.extend(db),
Err(e) => { Err(e) => {
eprintln!("Failed to show database '{}': {}", name, 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) 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 { 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 { } else {
core::database_operations::get_all_database_privileges(&mut conn).await? core::database_operations::get_all_database_privileges(conn).await?
}; };
if !args.text { if !args.text {

View File

@ -2,7 +2,7 @@ use std::vec;
use anyhow::Context; use anyhow::Context;
use clap::Parser; use clap::Parser;
use sqlx::MySqlConnection; use sqlx::{Connection, MySqlConnection};
use crate::core::user_operations::validate_ownership_of_user_name; use crate::core::user_operations::validate_ownership_of_user_name;
@ -57,24 +57,28 @@ struct UserShowArgs {
username: Vec<String>, 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 // 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() { if args.username.is_empty() {
anyhow::bail!("No usernames provided"); anyhow::bail!("No usernames provided");
} }
for username in args.username { for username in args.username {
if let Err(e) = 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!("{}", e);
eprintln!("Skipping..."); 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 // 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() { if args.username.is_empty() {
anyhow::bail!("No usernames provided"); anyhow::bail!("No usernames provided");
} }
for username in args.username { for username in args.username {
if let Err(e) = 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!("{}", e);
eprintln!("Skipping..."); eprintln!("Skipping...");
@ -102,7 +106,7 @@ async fn drop_users(args: UserDeleteArgs, mut conn: MySqlConnection) -> anyhow::
async fn change_password_for_user( async fn change_password_for_user(
args: UserPasswdArgs, args: UserPasswdArgs,
mut conn: MySqlConnection, conn: &mut MySqlConnection,
) -> anyhow::Result<()> { ) -> anyhow::Result<()> {
// NOTE: although this also is checked in `set_password_for_database_user`, we check it here // 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. // 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( crate::core::user_operations::set_password_for_database_user(
&args.username, &args.username,
&password, &password,
&mut conn, conn,
) )
.await?; .await?;
Ok(()) 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 user = crate::core::common::get_current_unix_user()?;
let users = if args.username.is_empty() { 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 { } else {
let mut result = vec![]; let mut result = vec![];
for username in args.username { for username in args.username {
@ -153,7 +157,7 @@ async fn show_users(args: UserShowArgs, mut conn: MySqlConnection) -> anyhow::Re
} }
let user = 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?; .await?;
if let Some(user) = user { if let Some(user) = user {
result.push(user); result.push(user);