Add cmd line parser (#512)

This commit adds the clap library and configures the necessary args to
parse from the command line,  expanding the current option of a single
file and adding support for environment variables.

Signed-off-by: Sebastian Webber <sebastian@swebber.me>

Cargo.toml

@@ -46,6 +46,7 @@ trust-dns-resolver = "0.22.0"
 tokio-test = "0.4.2"
 serde_json = "1"
 itertools = "0.10"
+clap = { version = "4.3.1", features = ["derive", "env"] }
 
 [target.'cfg(not(target_env = "msvc"))'.dependencies]
 jemallocator = "0.5.0"

src/admin.rs

@@ -1,7 +1,7 @@
 use crate::pool::BanReason;
 use crate::stats::pool::PoolStats;
 use bytes::{Buf, BufMut, BytesMut};
-use log::{debug, error, info, trace};
+use log::{error, info, trace};
 use nix::sys::signal::{self, Signal};
 use nix::unistd::Pid;
 use std::collections::HashMap;

src/cmd_args.rs

@@ -0,0 +1,17 @@
+use clap::Parser;
+use log::LevelFilter;
+
+/// PgCat: Nextgen PostgreSQL Pooler
+#[derive(Parser, Debug)]
+#[command(author, version, about, long_about = None)]
+pub(crate) struct Args {
+    #[arg(default_value_t = String::from("pgcat.toml"), env)]
+    pub config_file: String,
+
+    #[arg(short, long, default_value_t = LevelFilter::Info, env)]
+    pub log_level: log::LevelFilter,
+}
+
+pub(crate) fn parse() -> Args {
+    return Args::parse();
+}

src/main.rs

@@ -68,8 +68,11 @@ use pgcat::pool::{ClientServerMap, ConnectionPool};
 use pgcat::prometheus::start_metric_server;
 use pgcat::stats::{Collector, Reporter, REPORTER};
 
+mod cmd_args;
+
 fn main() -> Result<(), Box<dyn std::error::Error>> {
-    pgcat::multi_logger::MultiLogger::init().unwrap();
+    let args = cmd_args::parse();
+    pgcat::multi_logger::MultiLogger::init(args.log_level).unwrap();
 
     info!("Welcome to PgCat! Meow. (Version {})", VERSION);
 
@@ -78,20 +81,12 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
         std::process::exit(exitcode::CONFIG);
     }
 
-    let args = std::env::args().collect::<Vec<String>>();
-
-    let config_file = if args.len() == 2 {
-        args[1].to_string()
-    } else {
-        String::from("pgcat.toml")
-    };
-
     // Create a transient runtime for loading the config for the first time.
     {
         let runtime = Builder::new_multi_thread().worker_threads(1).build()?;
 
         runtime.block_on(async {
-            match pgcat::config::parse(&config_file).await {
+            match pgcat::config::parse(args.config_file.as_str()).await {
                 Ok(_) => (),
                 Err(err) => {
                     error!("Config parse error: {:?}", err);

src/multi_logger.rs

@@ -1,4 +1,4 @@
-use log::{Level, Log, Metadata, Record, SetLoggerError};
+use log::{Level, LevelFilter, Log, Metadata, Record, SetLoggerError};
 
 // This is a special kind of logger that allows sending logs to different
 // targets depending on the log level.
@@ -25,8 +25,11 @@ pub struct MultiLogger {
 }
 
 impl MultiLogger {
-    fn new() -> Self {
-        let stderr_logger = env_logger::builder().format_timestamp_micros().build();
+    fn new(filter: LevelFilter) -> Self {
+        let stderr_logger = env_logger::builder()
+            .filter(None, filter)
+            .format_timestamp_micros()
+            .build();
         let stdout_logger = env_logger::Builder::from_env("STDOUT_LOG")
             .format_timestamp_micros()
             .target(env_logger::Target::Stdout)
@@ -38,8 +41,8 @@ impl MultiLogger {
         }
     }
 
-    pub fn init() -> Result<(), SetLoggerError> {
-        let logger = Self::new();
+    pub fn init(filter: LevelFilter) -> Result<(), SetLoggerError> {
+        let logger = Self::new(filter);
 
         log::set_max_level(logger.stderr_logger.filter());
         log::set_boxed_logger(Box::new(logger))
@@ -75,6 +78,6 @@ mod test {
 
     #[test]
     fn test_init() {
-        MultiLogger::init().unwrap();
+        MultiLogger::init(LevelFilter::Error).unwrap();
     }
 }