From 2e0054a03697207b3299b441f91d0199d5f7d848 Mon Sep 17 00:00:00 2001 From: axd3v Date: Thu, 29 Jan 2026 19:46:02 +0300 Subject: [PATCH 1/2] Optimized to coalesce many small chunks into larger output chunks. --- Cargo.toml | 2 +- src/format/custom/blocks.rs | 83 +++++++++++++++++++++---------------- src/format/custom/header.rs | 6 +-- src/format/custom/mod.rs | 5 +-- 4 files changed, 52 insertions(+), 44 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 27948e6..d9cc524 100644 --- a/Cargo.toml +++ b/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" diff --git a/src/format/custom/blocks.rs b/src/format/custom/blocks.rs index c31df93..d809c90 100644 --- a/src/format/custom/blocks.rs +++ b/src/format/custom/blocks.rs @@ -114,76 +114,84 @@ 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( &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 = 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( &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 = Vec::new(); let mut output_buf: Vec = 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 { @@ -191,7 +199,10 @@ impl<'a> BlockProcessor<'a> { } } 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 } } } diff --git a/src/format/custom/header.rs b/src/format/custom/header.rs index fafe6a5..0e1b947 100644 --- a/src/format/custom/header.rs +++ b/src/format/custom/header.rs @@ -73,7 +73,7 @@ pub fn parse_header( writer.write_all(&[vrev])?; #[cfg(debug_assertions)] - eprintln!("pg_dump version: {}.{}.{}", vmaj, vmin, vrev); + eprintln!("[DEBUG] pg_dump format version: {}.{}.{}", vmaj, vmin, vrev); // custom.py validation: < 1.12 or > 1.16 is unsupported if vmaj < 1 || (vmaj == 1 && vmin < 12) { @@ -98,7 +98,7 @@ pub fn parse_header( writer.write_all(&[offset_size as u8])?; #[cfg(debug_assertions)] - eprintln!("int_size={}, offset_size={}", int_size, offset_size); + eprintln!("[DEBUG] int_size={}, offset_size={}", int_size, offset_size); // Validate sizes if int_size == 0 || int_size > 8 || offset_size == 0 || offset_size > 8 { @@ -163,7 +163,7 @@ pub fn parse_header( }; #[cfg(debug_assertions)] - eprintln!("Compression: {:?}", compression); + eprintln!("[DEBUG] 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. diff --git a/src/format/custom/mod.rs b/src/format/custom/mod.rs index 959b723..5b365c8 100644 --- a/src/format/custom/mod.rs +++ b/src/format/custom/mod.rs @@ -66,10 +66,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) { From 06f3ce963f16a3b8c0c2f6ea378f0583f2b01340 Mon Sep 17 00:00:00 2001 From: axd3v Date: Thu, 29 Jan 2026 19:56:05 +0300 Subject: [PATCH 2/2] Added verbose flag --- Cargo.lock | 2 +- README.md | 9 +++++++++ src/format/custom/header.rs | 37 ++++++++++++++++++++++--------------- src/format/custom/mod.rs | 13 ++++++++++--- src/format/custom/toc.rs | 4 ++++ src/main.rs | 6 +++++- 6 files changed, 51 insertions(+), 20 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index c6f38fa..24ecfd8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -419,7 +419,7 @@ checksum = "384b8ab6d37215f3c5301a95a4accb5d64aa607f1fcb26a11b5303878451b4fe" [[package]] name = "pg_stage_rs" -version = "0.1.2" +version = "0.1.3" dependencies = [ "chrono", "clap", diff --git a/README.md b/README.md index 5787c1a..4025c84 100644 --- a/README.md +++ b/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 diff --git a/src/format/custom/header.rs b/src/format/custom/header.rs index 0e1b947..77f68a2 100644 --- a/src/format/custom/header.rs +++ b/src/format/custom/header.rs @@ -38,6 +38,7 @@ pub fn parse_header( reader: &mut R, writer: &mut W, initial_bytes: &[u8], + verbose: bool, ) -> Result
{ // Write initial bytes (the magic we already consumed for detection) writer.write_all(initial_bytes)?; @@ -72,8 +73,9 @@ pub fn parse_header( let vrev = DumpIO::read_byte(reader)?; writer.write_all(&[vrev])?; - #[cfg(debug_assertions)] - eprintln!("[DEBUG] pg_dump format 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( let offset_size = DumpIO::read_byte(reader)? as usize; writer.write_all(&[offset_size as u8])?; - #[cfg(debug_assertions)] - eprintln!("[DEBUG] 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( } }; - #[cfg(debug_assertions)] - eprintln!("[DEBUG] 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( } // 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, diff --git a/src/format/custom/mod.rs b/src/format/custom/mod.rs index 5b365c8..dab82b2 100644 --- a/src/format/custom/mod.rs +++ b/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); diff --git a/src/format/custom/toc.rs b/src/format/custom/toc.rs index 809cf15..a5f0809 100644 --- a/src/format/custom/toc.rs +++ b/src/format/custom/toc.rs @@ -64,11 +64,15 @@ pub fn parse_toc( reader: &mut R, writer: &mut W, header: &Header, + verbose: bool, ) -> Result> { 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 { diff --git a/src/main.rs b/src/main.rs index 29b4078..d03cc0d 100644 --- a/src/main.rs +++ b/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, + + /// 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)?; } }