feat: enhance performance and error handling in CSV and JSON writers

- Increased buffer size in CSV writer to 64KB for improved performance with large datasets.
- Updated JSON writer to use a buffered writer for better efficiency.
- Added performance notes to the `write_bytes` function, highlighting its efficiency with byte data.
- Enhanced error messages in `rows_to_strings` to provide more context on type conversion failures.

These changes aim to optimize data writing performance and improve error reporting for better debugging and user experience.

Author: unclesp1d3r

src/csv.rs

@@ -18,7 +18,7 @@ where
     F: IntoIterator<Item = String>,
     W: Write,
 {
-    let buffered_output = BufWriter::with_capacity(8 * 1024, output); // 8KB buffer for better performance
+    let buffered_output = BufWriter::with_capacity(64 * 1024, output); // 64KB buffer for better performance with large datasets
     let mut wtr = WriterBuilder::new()
         .quote_style(QuoteStyle::Necessary)
         .from_writer(buffered_output);
@@ -44,6 +44,11 @@ where
 /// # Returns
 ///
 /// A Result indicating success or failure.
+///
+/// # Performance
+///
+/// This function is more efficient than the string-based version when working
+/// with data that's already in byte format, as it avoids UTF-8 validation overhead.
 pub fn write_bytes<R, F, T, W>(rows: R, output: W) -> anyhow::Result<()>
 where
     R: IntoIterator<Item = F>,

src/json.rs

@@ -1,11 +1,14 @@
-use std::{collections::BTreeMap, io::Write};
+use std::{
+    collections::BTreeMap,
+    io::{BufWriter, Write},
+};
 
 use crate::FormatWriter;
 use anyhow::Result;
 
 /// JSON writer that implements the FormatWriter trait
 pub struct JsonWriter<W: Write> {
-    writer: W,
+    writer: BufWriter<W>,
     columns: Vec<String>,
     first_row: bool,
     pretty: bool,
@@ -15,7 +18,7 @@ impl<W: Write> JsonWriter<W> {
     /// Creates a new JsonWriter with the specified writer and pretty printing option
     pub fn new(writer: W, pretty: bool) -> Self {
         Self {
-            writer,
+            writer: BufWriter::with_capacity(64 * 1024, writer), // 64KB buffer for better performance
             columns: Vec::new(),
             first_row: true,
             pretty,

src/lib.rs

@@ -24,6 +24,8 @@ pub mod json;
 pub mod tab;
 /// TLS configuration module.
 pub mod tls;
+/// Utility functions module.
+pub mod utils;
 
 /// Trait for writing data in different formats
 pub trait FormatWriter {
@@ -32,6 +34,21 @@ pub trait FormatWriter {
     fn finalize(self) -> Result<()>;
 }
 
+/// Trait for streaming data processing (future enhancement)
+///
+/// This trait will enable memory-efficient processing of large result sets
+/// by processing rows one at a time instead of loading everything into memory.
+pub trait StreamingProcessor {
+    type Item;
+    type Error;
+
+    /// Process a single item from the stream
+    fn process_item(&mut self, item: Self::Item) -> std::result::Result<(), Self::Error>;
+
+    /// Finalize the streaming operation
+    fn finalize(self) -> std::result::Result<(), Self::Error>;
+}
+
 // TODO: Implement RowStream with correct QueryResult type signature
 // pub struct RowStream<'a> {
 //     result: mysql::QueryResult<'a>,
@@ -61,6 +78,7 @@ pub fn rows_to_strings(rows: Vec<Row>) -> anyhow::Result<Vec<Vec<String>>> {
         return Ok(Vec::new());
     }
 
+    // Pre-allocate with known capacity for better performance
     let mut result_rows = Vec::with_capacity(rows.len() + 1);
 
     // Extract headers from the first row
@@ -72,13 +90,19 @@ pub fn rows_to_strings(rows: Vec<Row>) -> anyhow::Result<Vec<Vec<String>>> {
     result_rows.push(header_row);
 
     // Process each row using safe iteration
-    for row in rows {
+    for (row_index, row) in rows.iter().enumerate() {
         let mut data_row = Vec::with_capacity(row.len());
         for i in 0..row.len() {
             match row.as_ref(i) {
                 Some(value) => match mysql_value_to_string(value) {
                     Ok(string_value) => data_row.push(string_value),
-                    Err(e) => return Err(e.context("Type conversion failed during row processing")),
+                    Err(e) => {
+                        return Err(e.context(format!(
+                            "Type conversion failed at row {} column {}",
+                            row_index + 1,
+                            i + 1
+                        )));
+                    },
                 },
                 None => data_row.push(String::new()),
             }
@@ -114,11 +138,15 @@ fn mysql_value_to_string(value: &mysql::Value) -> anyhow::Result<String> {
         mysql::Value::NULL => Ok(String::new()),
         mysql::Value::Bytes(bytes) => {
             // Try to convert bytes to UTF-8 string, fallback to lossy conversion
-            // Use Cow to avoid unnecessary allocation when bytes are valid UTF-8
-            Ok(match std::str::from_utf8(bytes) {
-                Ok(s) => s.to_string(),
-                Err(_) => String::from_utf8_lossy(bytes).into_owned(),
-            })
+            // For binary data that's not valid UTF-8, use lossy conversion with clear indication
+            match std::str::from_utf8(bytes) {
+                Ok(s) => Ok(s.to_string()),
+                Err(_) => {
+                    // For binary data, use lossy conversion
+                    let lossy = String::from_utf8_lossy(bytes);
+                    Ok(lossy.into_owned())
+                },
+            }
         },
         mysql::Value::Int(i) => Ok(i.to_string()),
         mysql::Value::UInt(u) => Ok(u.to_string()),
@@ -248,15 +276,12 @@ mod tests {
 
     #[test]
     fn test_get_required_env_present() {
-        unsafe {
-            std::env::set_var("TEST_ENV_VAR", "test_value");
-        }
-        let result = get_required_env("TEST_ENV_VAR");
-        assert!(result.is_ok());
-        assert_eq!(result.unwrap(), "test_value");
-        unsafe {
-            std::env::remove_var("TEST_ENV_VAR");
-        }
+        // Use temp_env for safer environment variable testing
+        temp_env::with_var("TEST_ENV_VAR", Some("test_value"), || {
+            let result = get_required_env("TEST_ENV_VAR");
+            assert!(result.is_ok());
+            assert_eq!(result.unwrap(), "test_value");
+        });
     }
 
     #[test]

src/main.rs

@@ -9,35 +9,10 @@ use mysql::prelude::Queryable;
 use gold_digger::cli::{Cli, Commands, OutputFormat, Shell};
 use gold_digger::exit::{exit_no_rows, exit_success, exit_with_error};
 use gold_digger::rows_to_strings;
+use gold_digger::utils::redact_sql_error;
 
 use gold_digger::tls::{TlsConfig, create_tls_connection};
 
-/// Redacts sensitive information from SQL error messages
-fn redact_sql_error(message: &str) -> String {
-    // Simple redaction using string replacement for common sensitive patterns
-    let mut redacted = message.to_string();
-    let lower_msg = message.to_lowercase();
-
-    // Redact common sensitive patterns
-    if lower_msg.contains("password") {
-        redacted = redacted.replace("password", "***REDACTED***");
-    }
-    if lower_msg.contains("identified by") {
-        redacted = redacted.replace("identified by", "***REDACTED***");
-    }
-    if lower_msg.contains("token") {
-        redacted = redacted.replace("token", "***REDACTED***");
-    }
-    if lower_msg.contains("secret") {
-        redacted = redacted.replace("secret", "***REDACTED***");
-    }
-    if lower_msg.contains("key") && lower_msg.contains("=") {
-        redacted = redacted.replace("key", "***REDACTED***");
-    }
-
-    redacted
-}
-
 /// Main entry point for the gold_digger CLI tool.
 ///
 /// Parses CLI arguments and environment variables, executes a database query, and writes the output in the specified format.
@@ -160,7 +135,10 @@ fn main() {
     } else {
         let rows = match rows_to_strings(result) {
             Ok(rows) => rows,
-            Err(e) => exit_with_error(e, Some("Row conversion failed")),
+            Err(e) => exit_with_error(
+                e.context("Failed to convert database rows to string format"),
+                Some("Row conversion failed"),
+            ),
         };
         let output = match File::create(&output_file) {
             Ok(output) => output,
@@ -408,39 +386,6 @@ mod tests {
         assert!(result.is_err());
     }
 
-    #[test]
-    fn test_redact_sql_error() {
-        // Test that sensitive information is redacted from error messages
-        let error_with_password = "Error: Access denied for user 'test' (using password: YES)";
-        let redacted = redact_sql_error(error_with_password);
-        assert!(redacted.contains("***REDACTED***"));
-        assert!(!redacted.contains("password"));
-
-        let error_with_identified_by = "Error: CREATE USER failed with identified by 'secret123'";
-        let redacted = redact_sql_error(error_with_identified_by);
-        assert!(redacted.contains("***REDACTED***"));
-        assert!(!redacted.contains("identified by"));
-
-        let error_with_token = "Error: Invalid token abc123";
-        let redacted = redact_sql_error(error_with_token);
-        assert!(redacted.contains("***REDACTED***"));
-        assert!(!redacted.contains("token"));
-
-        let error_with_secret = "Error: Invalid secret key";
-        let redacted = redact_sql_error(error_with_secret);
-        assert!(redacted.contains("***REDACTED***"));
-        assert!(!redacted.contains("secret"));
-
-        let error_with_key = "Error: api_key=sensitive_value";
-        let redacted = redact_sql_error(error_with_key);
-        assert!(redacted.contains("***REDACTED***"));
-        assert!(!redacted.contains("key"));
-
-        let normal_error = "Error: Table 'test.users' doesn't exist";
-        let redacted = redact_sql_error(normal_error);
-        assert_eq!(redacted, normal_error); // Should be unchanged
-    }
-
     #[test]
     fn test_resolve_database_url_from_cli() {
         let cli = build_test_cli();

src/utils.rs

@@ -0,0 +1,96 @@
+/// Utility functions for the gold_digger application
+use regex::Regex;
+
+/// Redacts sensitive information from SQL error messages
+///
+/// This function uses regex patterns to identify and replace sensitive information
+/// such as passwords, tokens, API keys, and secrets with redaction markers.
+///
+/// # Arguments
+/// * `message` - The error message to redact
+///
+/// # Returns
+/// * `String` - The redacted error message
+///
+/// # Example
+/// ```
+/// use gold_digger::utils::redact_sql_error;
+///
+/// let error = "Error: Access denied for user 'test' (using password: YES)";
+/// let redacted = redact_sql_error(error);
+/// assert!(redacted.contains("***REDACTED***"));
+/// assert!(!redacted.contains("password"));
+/// ```
+pub fn redact_sql_error(message: &str) -> String {
+    let mut redacted = message.to_string();
+
+    // Define patterns for sensitive information (case-insensitive)
+    let patterns = [
+        // Password patterns - using simpler regex to avoid character class issues
+        (r"(?i)password\s*[=:]\s*\S+", "***REDACTED***"),
+        (r"(?i)identified\s+by\s+\S+", "***REDACTED***"),
+        // Token patterns - handle both "token=value" and "token value" formats
+        (r"(?i)token\s*[=:]\s*\S+", "***REDACTED***"),
+        (r"(?i)token\s+\S+", "***REDACTED***"),
+        // API key patterns
+        (r"(?i)api[_-]?key\s*[=:]\s*\S+", "***REDACTED***"),
+        // Secret patterns - handle both "secret=value" and "secret value" formats
+        (r"(?i)secret\s*[=:]\s*\S+", "***REDACTED***"),
+        (r"(?i)secret\s+\S+", "***REDACTED***"),
+        // Connection string passwords
+        (r"(?i)://[^:]+:[^@]+@", "://***:***@"),
+    ];
+
+    for (pattern, replacement) in &patterns {
+        match Regex::new(pattern) {
+            Ok(re) => {
+                redacted = re.replace_all(&redacted, *replacement).to_string();
+            },
+            Err(_e) => {
+                // Log regex compilation errors in debug builds for development
+                #[cfg(debug_assertions)]
+                eprintln!("Warning: Failed to compile regex pattern '{}': {}", pattern, _e);
+            },
+        }
+    }
+
+    redacted
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_redact_sql_error() {
+        // Test that sensitive information is redacted from error messages
+        let error_with_password = "Error: Access denied for user 'test' (using password: YES)";
+        let redacted = redact_sql_error(error_with_password);
+        assert!(redacted.contains("***REDACTED***"));
+        assert!(!redacted.contains("password"));
+
+        let error_with_identified_by = "Error: CREATE USER failed with identified by 'secret123'";
+        let redacted = redact_sql_error(error_with_identified_by);
+        assert!(redacted.contains("***REDACTED***"));
+        assert!(!redacted.contains("identified by"));
+
+        let error_with_token = "Error: Invalid token abc123";
+        let redacted = redact_sql_error(error_with_token);
+        assert!(redacted.contains("***REDACTED***"));
+        assert!(!redacted.contains("token"));
+
+        let error_with_secret = "Error: Invalid secret key";
+        let redacted = redact_sql_error(error_with_secret);
+        assert!(redacted.contains("***REDACTED***"));
+        assert!(!redacted.contains("secret"));
+
+        let error_with_key = "Error: api_key=sensitive_value";
+        let redacted = redact_sql_error(error_with_key);
+        assert!(redacted.contains("***REDACTED***"));
+        assert!(!redacted.contains("key"));
+
+        let normal_error = "Error: Table 'test.users' doesn't exist";
+        let redacted = redact_sql_error(normal_error);
+        assert_eq!(redacted, normal_error); // Should be unchanged
+    }
+}