optionally read probes from a file with the client

Author: matthieugouel

src/client/handler.rs

@@ -1,13 +1,24 @@
 use anyhow::Result;
+use caracat::models::Probe;
 use log::trace;
 use std::io::{stdin, BufRead};
+use std::path::PathBuf;
 
 use crate::auth::{KafkaAuth, SaslAuth};
 use crate::client::producer::produce;
 use crate::config::AppConfig;
 use crate::probe::decode_probe;
 
-pub async fn handle(config: &AppConfig, agents: &str) -> Result<()> {
+fn read_probes<R: BufRead>(buf_reader: R) -> Result<Vec<Probe>> {
+    let mut probes = Vec::new();
+    for line in buf_reader.lines() {
+        let probe = line?;
+        probes.push(decode_probe(&probe)?);
+    }
+    Ok(probes)
+}
+
+pub async fn handle(config: &AppConfig, agents: &str, probes_file: Option<PathBuf>) -> Result<()> {
     trace!("Client handler");
     trace!("{:?}", config);
 
@@ -26,16 +37,24 @@ pub async fn handle(config: &AppConfig, agents: &str) -> Result<()> {
         }
     };
 
-    // Get probes from stdin
-    let mut probes = Vec::new();
-    for line in stdin().lock().lines() {
-        let probe = line?;
-        probes.push(decode_probe(&probe)?);
-    }
+    // Read probes from file or stdin
+    let probes = match probes_file {
+        Some(probes_file) => {
+            let file = std::fs::File::open(probes_file)?;
+            let buf_reader = std::io::BufReader::new(file);
+            read_probes(buf_reader)?
+        }
+        None => {
+            let stdin = stdin();
+            let buf_reader = stdin.lock();
+            read_probes(buf_reader)?
+        }
+    };
 
-    // Split agents
+    // Split the agents
     let agents = agents.split(',').collect::<Vec<&str>>();
 
+    // Produce Kafka messages
     produce(config, auth, agents, probes).await;
 
     Ok(())

src/main.rs

@@ -12,6 +12,7 @@ use clap_verbosity_flag::{InfoLevel, Verbosity};
 use env_logger::Builder;
 use log::error;
 use std::io::{stdin, IsTerminal, Write};
+use std::path::PathBuf;
 
 use crate::config::app_config;
 
@@ -41,6 +42,10 @@ enum Command {
         /// Agent IDs (comma separated)
         #[arg(index = 1)]
         agents: String,
+
+        /// Probes file (read stdin if not provided)
+        #[arg(short, long)]
+        probes_file: Option<PathBuf>,
     },
 }
 
@@ -79,14 +84,18 @@ async fn main() -> Result<()> {
                 Err(e) => error!("Error: {}", e),
             }
         }
-        Command::Client { config, agents } => {
+        Command::Client {
+            config,
+            agents,
+            probes_file,
+        } => {
             if stdin().is_terminal() {
                 App::command().print_help().unwrap();
                 ::std::process::exit(2);
             }
 
             let app_config = app_config(&config);
-            match client::handle(&app_config, &agents).await {
+            match client::handle(&app_config, &agents, probes_file).await {
                 Ok(_) => (),
                 Err(e) => error!("Error: {}", e),
             }