Optimized to coalesce many small chunks into larger output chunks.

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "pg_stage_rs"
-version = "0.1.2"
+version = "0.1.3"
 edition = "2021"
 description = "PostgreSQL dump anonymizer - streaming obfuscation for plain and custom formats"
 

README.md

@@ -46,6 +46,14 @@ pg_dump -Fp mydb | pg_stage_rs --locale ru --format plain > anonymized.sql
 
 # Delete specific tables by regex
 pg_dump -Fp mydb | pg_stage_rs --delete-table-pattern "^audit_.*" > anonymized.sql
+
+# Verbose mode (show dump metadata)
+pg_dump -Fc mydb | pg_stage_rs --verbose > anonymized.dump
+# Output to stderr:
+#   [INFO] pg_dump format version: 1.16.0
+#   [INFO] Compression: Zlib
+#   [INFO] Database: "mydb"
+#   [INFO] TOC entries: 1234
 ```
 
 ### CLI Options
@@ -55,6 +63,7 @@ pg_dump -Fp mydb | pg_stage_rs --delete-table-pattern "^audit_.*" > anonymized.s
 | `-l, --locale` | `en` | Locale for generated data (`en`, `ru`) |
 | `-d, --delimiter` | `\t` | Column delimiter character |
 | `-f, --format` | auto | Force format: `plain`/`p`, `custom`/`c` |
+| `-v, --verbose` | off | Show dump info: format version, compression, TOC count |
 | `--delete-table-pattern` | -- | Regex pattern for tables to remove (repeatable) |
 
 ## Defining Mutations

src/format/custom/blocks.rs

@@ -114,84 +114,95 @@ impl<'a> BlockProcessor<'a> {
     }
 
     /// Pass through a block without mutation.
+    /// Optimized to coalesce many small chunks into larger output chunks.
     pub fn pass_through_block<R: Read, W: Write>(
         &self,
         reader: &mut R,
         writer: &mut W,
     ) -> Result<()> {
-        loop {
-            let chunk_len = self.dio.read_int(reader)?;
-            self.dio.write_int(writer, chunk_len)?;
+        // Use ChunkReader to efficiently read many small chunks
+        let mut chunk_reader = ChunkReader::new(reader, self.dio);
+        let mut read_buf = vec![0u8; READ_BUF_SIZE];
+        let mut output_buf: Vec<u8> = Vec::with_capacity(OUTPUT_CHUNK_SIZE * 2);
 
-            if chunk_len == 0 {
+        loop {
+            let n = chunk_reader.read(&mut read_buf)?;
+            if n == 0 {
                 break;
             }
 
-            let len = chunk_len.unsigned_abs() as usize;
+            output_buf.extend_from_slice(&read_buf[..n]);
 
-            if len > MAX_CHUNK_SIZE {
-                return Err(PgStageError::InvalidFormat(format!(
-                    "Chunk size {} exceeds maximum {}, stream may be corrupted",
-                    len, MAX_CHUNK_SIZE
-                )));
+            // Flush when buffer is large enough
+            if output_buf.len() >= OUTPUT_CHUNK_SIZE {
+                for chunk in output_buf.chunks(OUTPUT_CHUNK_SIZE) {
+                    self.dio.write_int(writer, chunk.len() as i32)?;
+                    writer.write_all(chunk)?;
+                }
+                output_buf.clear();
             }
+        }
 
-            let mut buf = vec![0u8; len];
-            reader.read_exact(&mut buf)?;
-            writer.write_all(&buf)?;
+        // Write remaining data
+        if !output_buf.is_empty() {
+            self.dio.write_int(writer, output_buf.len() as i32)?;
+            writer.write_all(&output_buf)?;
         }
+
+        // Terminator
+        self.dio.write_int(writer, 0)?;
         Ok(())
     }
 
     /// Streaming processing for uncompressed blocks.
-    /// Reads chunks one at a time, processes complete lines, writes output immediately.
+    /// Uses ChunkReader to buffer small chunks efficiently (critical for -Z0 dumps
+    /// which can have millions of tiny chunks).
     fn process_block_uncompressed<R: Read, W: Write>(
         &mut self,
         reader: &mut R,
         writer: &mut W,
     ) -> Result<()> {
+        // Use ChunkReader to efficiently buffer many small chunks into large reads
+        let mut chunk_reader = ChunkReader::new(reader, self.dio);
+
+        let mut read_buf = vec![0u8; READ_BUF_SIZE];
         let mut line_tail: Vec<u8> = Vec::new();
         let mut output_buf: Vec<u8> = Vec::with_capacity(OUTPUT_CHUNK_SIZE * 2);
 
         loop {
-            let chunk_len = self.dio.read_int(reader)?;
-            if chunk_len == 0 {
+            // Read large buffer from ChunkReader (it handles chunk boundaries internally)
+            let n = chunk_reader.read(&mut read_buf)?;
+            if n == 0 {
                 break;
             }
 
-            let len = chunk_len.unsigned_abs() as usize;
-            if len > MAX_CHUNK_SIZE {
-                return Err(PgStageError::InvalidFormat(format!(
-                    "Chunk size {} exceeds maximum {}",
-                    len, MAX_CHUNK_SIZE
-                )));
-            }
-
-            let mut buf = vec![0u8; len];
-            reader.read_exact(&mut buf)?;
-
-            // Prepend leftover from previous chunk
+            // Combine with leftover from previous iteration
             let data = if line_tail.is_empty() {
-                buf
+                &read_buf[..n]
             } else {
-                let mut combined = std::mem::take(&mut line_tail);
-                combined.extend_from_slice(&buf);
-                combined
+                line_tail.extend_from_slice(&read_buf[..n]);
+                line_tail.as_slice()
             };
 
             // Split into complete lines + tail
-            match memrchr(b'\n', &data) {
+            match memrchr(b'\n', data) {
                 Some(last_nl) => {
-                    line_tail = data[last_nl + 1..].to_vec();
-                    self.process_complete_lines(&data[..=last_nl], &mut output_buf);
+                    let complete = &data[..=last_nl];
+                    let tail = &data[last_nl + 1..];
+
+                    self.process_complete_lines(complete, &mut output_buf);
+                    line_tail = tail.to_vec();
 
                     // Flush output when large enough
                     if output_buf.len() >= OUTPUT_CHUNK_SIZE {
                         self.flush_uncompressed(writer, &mut output_buf)?;
                     }
                 }
                 None => {
-                    line_tail = data;
+                    if line_tail.is_empty() {
+                        line_tail = read_buf[..n].to_vec();
+                    }
+                    // else: data already IS line_tail (extended above), keep as-is
                 }
             }
         }

src/format/custom/header.rs

@@ -38,6 +38,7 @@ pub fn parse_header<R: Read, W: Write>(
     reader: &mut R,
     writer: &mut W,
     initial_bytes: &[u8],
+    verbose: bool,
 ) -> Result<Header> {
     // Write initial bytes (the magic we already consumed for detection)
     writer.write_all(initial_bytes)?;
@@ -72,8 +73,9 @@ pub fn parse_header<R: Read, W: Write>(
     let vrev = DumpIO::read_byte(reader)?;
     writer.write_all(&[vrev])?;
 
-    #[cfg(debug_assertions)]
-    eprintln!("pg_dump version: {}.{}.{}", vmaj, vmin, vrev);
+    if verbose {
+        eprintln!("[INFO] pg_dump format version: {}.{}.{}", vmaj, vmin, vrev);
+    }
 
     // custom.py validation: < 1.12 or > 1.16 is unsupported
     if vmaj < 1 || (vmaj == 1 && vmin < 12) {
@@ -97,8 +99,9 @@ pub fn parse_header<R: Read, W: Write>(
     let offset_size = DumpIO::read_byte(reader)? as usize;
     writer.write_all(&[offset_size as u8])?;
 
-    #[cfg(debug_assertions)]
-    eprintln!("int_size={}, offset_size={}", int_size, offset_size);
+    if verbose {
+        eprintln!("[INFO] int_size={}, offset_size={}", int_size, offset_size);
+    }
 
     // Validate sizes
     if int_size == 0 || int_size > 8 || offset_size == 0 || offset_size > 8 {
@@ -162,8 +165,9 @@ pub fn parse_header<R: Read, W: Write>(
         }
     };
 
-    #[cfg(debug_assertions)]
-    eprintln!("Compression: {:?}", compression);
+    if verbose {
+        eprintln!("[INFO] Compression: {:?}", compression);
+    }
 
     // Timestamp: custom.py reads 7 integers (sec, min, hour, mday, mon, year, isdst)
     // The 7th integer is ignored in Python (_isdst), but must be read/written to maintain sync.
@@ -172,19 +176,22 @@ pub fn parse_header<R: Read, W: Write>(
     }
 
     // Database name (string)
-    let _db_name = dio.read_string_bypass(reader, writer)?;
-    #[cfg(debug_assertions)]
-    eprintln!("Database: {:?}", _db_name);
+    let db_name = dio.read_string_bypass(reader, writer)?;
+    if verbose {
+        eprintln!("[INFO] Database: {:?}", db_name.as_deref().unwrap_or(""));
+    }
 
     // Server version (string)
-    let _server_ver = dio.read_string_bypass(reader, writer)?;
-    #[cfg(debug_assertions)]
-    eprintln!("Server version: {:?}", _server_ver);
+    let server_ver = dio.read_string_bypass(reader, writer)?;
+    if verbose {
+        eprintln!("[INFO] Server version: {:?}", server_ver.as_deref().unwrap_or(""));
+    }
 
     // Dump version (string)
-    let _dump_ver = dio.read_string_bypass(reader, writer)?;
-    #[cfg(debug_assertions)]
-    eprintln!("pg_dump version string: {:?}", _dump_ver);
+    let dump_ver = dio.read_string_bypass(reader, writer)?;
+    if verbose {
+        eprintln!("[INFO] pg_dump version: {:?}", dump_ver.as_deref().unwrap_or(""));
+    }
 
     Ok(Header {
         vmaj,

src/format/custom/mod.rs

@@ -15,11 +15,18 @@ use crate::processor::DataProcessor;
 /// Handler for PostgreSQL custom format dumps (-Fc).
 pub struct CustomHandler {
     processor: DataProcessor,
+    verbose: bool,
 }
 
 impl CustomHandler {
     pub fn new(processor: DataProcessor) -> Self {
-        Self { processor }
+        Self { processor, verbose: false }
+    }
+
+    /// Enable verbose output
+    pub fn verbose(mut self, verbose: bool) -> Self {
+        self.verbose = verbose;
+        self
     }
 
     /// Process a custom format dump from reader to writer.
@@ -35,10 +42,10 @@ impl CustomHandler {
         let mut writer = BufWriter::with_capacity(2 * 1024 * 1024, writer);
 
         // Parse header (bypasses to output)
-        let header = parse_header(&mut reader, &mut writer, initial_bytes)?;
+        let header = parse_header(&mut reader, &mut writer, initial_bytes, self.verbose)?;
 
         // Parse TOC entries (bypasses to output)
-        let entries = parse_toc(&mut reader, &mut writer, &header)?;
+        let entries = parse_toc(&mut reader, &mut writer, &header, self.verbose)?;
 
         // Extract comments from TOC entries to build mutation map
         self.extract_comments(&entries);
@@ -66,10 +73,7 @@ impl CustomHandler {
 
             // Block type 0x01 = DATA
             if block_type[0] == 0x01 {
-                let dump_id = dio.read_int(&mut reader)
-                    .inspect_err(|_| {
-                        eprintln!("Failed to read dump_id after DATA block");
-                    })?;
+                let dump_id = dio.read_int(&mut reader)?;
 
                 // Check if this dump_id is in our data_entries map
                 if let Some(info) = data_entries.get(&dump_id) {

src/format/custom/toc.rs

@@ -64,11 +64,15 @@ pub fn parse_toc<R: Read, W: Write>(
     reader: &mut R,
     writer: &mut W,
     header: &Header,
+    verbose: bool,
 ) -> Result<Vec<TocEntry>> {
     let dio = DumpIO::new(header.int_size, header.offset_size);
 
     // Read TOC count
     let toc_count = dio.read_int_bypass(reader, writer)?;
+    if verbose {
+        eprintln!("[INFO] TOC entries: {}", toc_count);
+    }
     let mut entries = Vec::with_capacity(toc_count.max(0) as usize);
 
     for _ in 0..toc_count {

src/main.rs

@@ -28,6 +28,10 @@ struct Args {
     /// Regex patterns for tables to delete (can be specified multiple times)
     #[arg(long = "delete-table-pattern")]
     delete_table_patterns: Vec<String>,
+
+    /// Enable verbose output (dump version, TOC count, compression info)
+    #[arg(short, long)]
+    verbose: bool,
 }
 
 fn main() {
@@ -77,7 +81,7 @@ fn run() -> Result<()> {
             handler.process(reader, writer, peeked)?;
         }
         DumpFormat::Custom => {
-            let mut handler = CustomHandler::new(processor);
+            let mut handler = CustomHandler::new(processor).verbose(args.verbose);
             handler.process(reader, writer, peeked)?;
         }
     }