From e2236502b4a002e96e22f594929d4d071b5292ea Mon Sep 17 00:00:00 2001 From: Joyce Babu Date: Wed, 29 May 2024 16:12:07 +0530 Subject: [PATCH] chore: add proper logging --- Cargo.lock | 27 ++++++++++++++++ Cargo.toml | 2 ++ src/main.rs | 90 ++++++++++++++++++++++++++--------------------------- 3 files changed, 73 insertions(+), 46 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index cd3a19d..ad75718 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -131,8 +131,10 @@ dependencies = [ "clap", "ctrlc", "dirs", + "env_logger", "indicatif", "libz-sys", + "log", "openssl", "regex", "rpassword", @@ -171,6 +173,19 @@ version = "0.3.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a357d28ed41a50f9c765dbfe56cbc04a64e53e5fc58ba79fbc34c10ef3df831f" +[[package]] +name = "env_logger" +version = "0.9.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a12e6657c4c97ebab115a42dcee77225f7f482cdd841cf7088c657a42e9e00e7" +dependencies = [ + "atty", + "humantime", + "log", + "regex", + "termcolor", +] + [[package]] name = "equivalent" version = "1.0.1" @@ -240,6 +255,12 @@ dependencies = [ "libc", ] +[[package]] +name = "humantime" +version = "2.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a3a5bfb195931eeb336b2a7b4d761daec841b97f947d34394601737a7bba5e4" + [[package]] name = "indexmap" version = "1.9.3" @@ -342,6 +363,12 @@ dependencies = [ "scopeguard", ] +[[package]] +name = "log" +version = "0.4.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "90ed8c1e510134f979dbc4f070f87d4313098b704861a105fe34231c70a3901c" + [[package]] name = "memchr" version = "2.6.4" diff --git a/Cargo.toml b/Cargo.toml index aa276a0..9756c2a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -21,6 +21,8 @@ serde_json = "1.0.107" serde = { version = "1.0", features = ["derive"] } dirs = "5.0.1" libz-sys = { version = "1.1.12", default-features = false, features = ["libc"] } +log = "0.4" +env_logger = "0.9" [target.'cfg(all(target_os = "linux", target_env = "musl"))'.dependencies] openssl = { version = "0.10", features = ["vendored"] } diff --git a/src/main.rs b/src/main.rs index e381ec4..0d1e106 100644 --- a/src/main.rs +++ b/src/main.rs @@ -7,13 +7,15 @@ extern crate ssh2; use bzip2::read::BzDecoder; use clap::{Arg, Command as App}; +use env_logger::Env; +use log::{error, warn, info, debug}; use regex::Regex; use serde::{Deserialize, Serialize}; use ssh2::Session; use std::collections::HashMap; use std::error::Error; use std::fs::{self, File}; -use std::io::{copy, BufReader, BufWriter, Read}; +use std::io::{copy, BufReader, BufWriter, Read, Write}; use std::net::TcpStream; use std::path::Path; use std::process; @@ -44,13 +46,12 @@ impl Config { let ssh_host = cap .get(2) .unwrap_or_else(|| { - eprintln!("Failed to parse ssh host"); + error!("Failed to parse ssh host"); process::exit(1); }) .as_str(); let mut config: Config = match config_file.exists() { - // let config_path = ; true => Config::parse(config_file.to_str().unwrap(), ssh_host), false => Config { ssh_host: ssh_host.to_string(), @@ -62,7 +63,7 @@ impl Config { match cap.get(1) { Some(val) => config.ssh_user = val.as_str().to_string(), None => { - eprintln!("SSH username cannot be empty"); + error!("SSH username cannot be empty"); process::exit(1); } } @@ -78,7 +79,7 @@ impl Config { match Config::parse_config_file(path, host) { Ok(config) => config, Err(err) => { - eprintln!("Failed to parse config file - {}", err); + error!("Failed to parse config file - {}", err); Config { ssh_host: host.to_string(), ..Default::default() @@ -100,6 +101,13 @@ impl Config { } fn main() -> Result<(), Box> { + // Initialize the logger with the default environment variable "RUST_LOG" + env_logger::Builder::from_env(Env::default().default_filter_or("info")) + .format(|buf, record| { + writeln!(buf, "[{}] {}", record.level(), record.args()) + }) + .init(); + let matches = App::new("Database Importer") .about("Import database from remote server") .arg( @@ -156,7 +164,7 @@ fn main() -> Result<(), Box> { let host = match matches.value_of("host") { Some(n) => n, None => { - eprintln!("Target host cannot be empty"); + error!("Target host cannot be empty"); process::exit(1); } }; @@ -166,7 +174,7 @@ fn main() -> Result<(), Box> { let tables = match matches.values_of("table") { Some(n) => n, None => { - eprintln!("Target tables cannot be empty"); + error!("Target tables cannot be empty"); process::exit(1); } }; @@ -182,7 +190,7 @@ fn main() -> Result<(), Box> { None => { // User option was not specified. // Use ssh username, if available - eprintln!("Using SSH username as database username"); + info!("Using SSH username as database username"); config.ssh_user.to_owned() } }, @@ -195,18 +203,18 @@ fn main() -> Result<(), Box> { None => { // User option was not specified. // Use ssh username, if available - eprintln!("Inferring database name as {}_db", db_user); + info!("Inferring database name as {}_db", db_user); format!("{}_db", db_user) } }, }; - let db_pass = match matches.value_of("db_pass") { + let db_pass = match matches.value_of("dbpass") { Some(val) => val.to_owned(), None => match config.db_pass { Some(val) => val, None => { - eprintln!("Database password cannot be empty"); + error!("Database password cannot be empty"); process::exit(2); } }, @@ -217,7 +225,7 @@ fn main() -> Result<(), Box> { // Connect to the remote server let ssh_host_port = format!("{}:{}", config.ssh_host, config.ssh_port); let tcp = TcpStream::connect(&ssh_host_port).unwrap_or_else(|err| { - eprintln!("{}", err); + error!("{}", err); process::exit(1); }); @@ -226,24 +234,21 @@ fn main() -> Result<(), Box> { sess.handshake().expect("SSH handshake failed"); // Try to authenticate with the first identity in the agent. - eprint!("Attept to authenticate with ssh-agent..."); + info!("Attempt to authenticate with ssh-agent..."); let _ = sess.userauth_agent(&config.ssh_user); // Make sure we succeeded if !sess.authenticated() { - eprintln!("FAILED"); - - eprintln!("Falling back to password login"); - eprint!("Enter password: "); - + warn!("SSH-agent authentication failed. Falling back to password login."); + info!("Enter password: "); let ssh_pass = rpassword::read_password().unwrap(); sess.userauth_password(&config.ssh_user, &ssh_pass) .unwrap_or_else(|_| { - eprintln!("Failed to authenticate to remote server"); + error!("Failed to authenticate to remote server"); process::exit(1); }); } else { - eprintln!("OK"); + info!("SSH-agent authentication succeeded."); } let mut remote_temp_file = String::new(); @@ -254,13 +259,6 @@ fn main() -> Result<(), Box> { channel.read_to_string(&mut remote_temp_file).unwrap(); let remote_temp_file = remote_temp_file.trim(); - // ctrlc::set_handler(move || { - // // Handle early termination - // let status = sess.sftp().and_then(|sftp| { - // sftp.unlink(Path::new(remote_temp_file)) - // }); - // }).expect("Error setting Ctrl-C handler"); - let pass_arg = format!("-p{}", &db_pass); let mut v = vec!["mysqldump", "-u", &db_user, &pass_arg, &db_name]; @@ -279,33 +277,29 @@ fn main() -> Result<(), Box> { remote_temp_file ); - eprint!("Exporting database on target server..."); + info!("Exporting database on target server..."); let mut channel = sess.channel_session().ok().unwrap(); let exit_status = channel .exec(&arg) .and_then(|_| channel.close()) - // .and_then(|_| channel.wait_close()) .and_then(|_| channel.exit_status()); match exit_status { - Ok(0) => (), + Ok(0) => info!("Database export succeeded."), _ => { - eprintln!("ERR"); - eprintln!("Failed to export database"); + error!("Failed to export database"); process::exit(4); } } - eprintln!("OK"); - let (remote_file, stat) = sess .scp_recv(Path::new(&remote_temp_file)) .unwrap_or_else(|err| { - eprintln!("Failed to download file - {}", err); + error!("Failed to download file - {}", err); process::exit(2); }); - eprintln!("Exported file size: {}", stat.size()); + info!("Exported file size: {}", stat.size()); let temp_file = Builder::new() .suffix(".sql.bz2") @@ -325,13 +319,13 @@ fn main() -> Result<(), Box> { match copy(&mut remote_file, &mut target) { Ok(_) => (), Err(err) => { - eprintln!("Failed to download exported database dump - {}", err); + error!("Failed to download exported database dump - {}", err); process::exit(3); } } + debug!("Database dump downloaded to {:?}", path); progressbar.finish_and_clear(); - println!("Downloading database dump...OK"); let mut channel = sess.channel_session().ok().unwrap(); let arg = format!("rm -f {}", remote_temp_file); @@ -341,8 +335,8 @@ fn main() -> Result<(), Box> { .and_then(|_| channel.exit_status()); match exit_status { - Ok(0) => eprintln!("Removed temporary file from remote filesystem...OK"), - _ => eprintln!("Failed to delete temporary file from remote filesystem"), + Ok(0) => info!("Removed temporary file from remote filesystem."), + _ => warn!("Failed to delete temporary file from remote filesystem."), } // TODO: Delete temporary file from remote filesystem @@ -356,36 +350,40 @@ fn main() -> Result<(), Box> { .spawn() .expect("Failed to spawn mysql client"); + debug!("Spawning mysql client"); if let Some(stdin) = &mut cmd.stdin { let mut writer = BufWriter::new(stdin); + debug!("Copying uncompressed dump to mysql client stdin"); // Decompress and import to local mysql database match copy(&mut Box::new(BzDecoder::new(f)), &mut writer) { - Ok(_) => { - eprintln!("Import completed successfully") - } + Ok(_) => info!("Database import completed successfully."), Err(err) => { - eprintln!("Failed to import database dump - {}", err); + error!("Failed to import database dump - {}", err); process::exit(5); } }; } + debug!("Waiting for import to complete"); + let result = match cmd.try_wait() { Ok(Some(status)) => status.code(), Ok(None) => cmd.wait().unwrap().code(), Err(e) => { - eprintln!("error attempting to wait: {}", e); + error!("Error attempting to wait: {}", e); process::exit(5); } }; + debug!("Removing temporary file {:?}", path); fs::remove_file(path)?; match result { Some(0) => (), - n => eprintln!("Import failed with exit code {:?}", n), + n => error!("Import failed with exit code {:?}", n), }; Ok(()) } +