From 2a6b6a525c8daef20bd644186d43974b44caae4a Mon Sep 17 00:00:00 2001 From: axd3v Date: Tue, 21 Apr 2026 23:11:01 +0300 Subject: [PATCH] feat: add pattern-based rules and optimize dump processing - add rules-file support, strict mode, and configurable zstd settings - precompile mutations, conditions, and relations into a shared registry - reduce per-row allocations with scratch buffers and fast hash maps/sets - optimize custom dump I/O and block streaming paths - add verbose processing summary and bump version to 0.2.0 --- Cargo.lock | 56 +++- Cargo.toml | 15 +- src/conditions.rs | 52 ++- src/format/custom/blocks.rs | 366 +++++++++----------- src/format/custom/io.rs | 155 ++++----- src/format/custom/mod.rs | 86 +++-- src/format/plain.rs | 14 +- src/lib.rs | 9 + src/main.rs | 72 +++- src/mutator/identity.rs | 8 +- src/mutator/mod.rs | 113 ++++--- src/processor.rs | 652 ++++++++++++++++++++++++------------ src/relations.rs | 68 ++-- src/types.rs | 142 +++++++- src/unique.rs | 21 +- tests/integration.rs | 8 +- 16 files changed, 1127 insertions(+), 710 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 24ecfd8..aad48ba 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -8,6 +8,20 @@ version = "2.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa" +[[package]] +name = "ahash" +version = "0.8.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a15f179cd60c4584b8a8c596927aadc462e27f2ca70c04e0071964a73ba7a75" +dependencies = [ + "cfg-if", + "getrandom 0.3.4", + "once_cell", + "serde", + "version_check", + "zerocopy", +] + [[package]] name = "aho-corasick" version = "1.1.4" @@ -166,6 +180,15 @@ version = "0.7.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c3e64b0cc0439b12df2fa678eae89a1c56a529fd067a9115f7827f1fffd22b32" +[[package]] +name = "cmake" +version = "0.1.58" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c0f78a02292a74a88ac736019ab962ece0bc380e3f977bf72e376c5d78ff0678" +dependencies = [ + "cc", +] + [[package]] name = "colorchoice" version = "1.0.4" @@ -245,6 +268,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b375d6465b98090a5f25b1c7703f3859783755aa9a80433b36e0379a3ec2f369" dependencies = [ "crc32fast", + "libz-ng-sys", "miniz_oxide", ] @@ -364,6 +388,25 @@ version = "0.2.180" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bcc35a38544a891a5f7c865aca548a982ccb3b8650a5b06d0fd33a10283c56fc" +[[package]] +name = "libmimalloc-sys" +version = "0.1.46" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bc89deee4af0429081d2a518c0431ae068222a5a262a3bc6ff4d8535ec2e02fe" +dependencies = [ + "cc", +] + +[[package]] +name = "libz-ng-sys" +version = "1.1.28" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "be734b33b7bc6a42d92d23e25e69758f866cf564a88d0bf80866fcf5a52c2255" +dependencies = [ + "cmake", + "libc", +] + [[package]] name = "log" version = "0.4.29" @@ -376,6 +419,15 @@ version = "2.7.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f52b00d39961fc5b2736ea853c9cc86238e165017a493d1d5c8eac6bdc4cc273" +[[package]] +name = "mimalloc" +version = "0.1.49" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aca3c01a711f395b4257b81674c0e90e8dd1f1e62c4b7db45f684cc7a4fcb18a" +dependencies = [ + "libmimalloc-sys", +] + [[package]] name = "miniz_oxide" version = "0.8.9" @@ -419,14 +471,16 @@ checksum = "384b8ab6d37215f3c5301a95a4accb5d64aa607f1fcb26a11b5303878451b4fe" [[package]] name = "pg_stage_rs" -version = "0.1.3" +version = "0.2.0" dependencies = [ + "ahash", "chrono", "clap", "crossbeam-channel", "flate2", "hmac", "memchr", + "mimalloc", "num_cpus", "rand", "regex", diff --git a/Cargo.toml b/Cargo.toml index d9cc524..3b264eb 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,9 +1,16 @@ [package] name = "pg_stage_rs" -version = "0.1.3" +version = "0.2.0" edition = "2021" description = "PostgreSQL dump anonymizer - streaming obfuscation for plain and custom formats" +[features] +default = ["mimalloc-allocator", "parallel"] +mimalloc-allocator = ["dep:mimalloc"] +# Opt-in: requires cmake/cc in build env; swaps flate2 backend to zlib-ng for ~1.5-2x zlib throughput. +zlib-ng = ["flate2/zlib-ng"] +parallel = ["dep:crossbeam-channel", "dep:num_cpus"] + [dependencies] clap = { version = "4", features = ["derive"] } serde = { version = "1", features = ["derive"] } @@ -11,8 +18,6 @@ serde_json = "1" flate2 = "1" zstd = { version = "0.13", features = ["zstdmt"] } memchr = "2" -crossbeam-channel = "0.5" -num_cpus = "1.16" uuid = { version = "1", features = ["v4", "v5"] } regex = "1" rand = "0.8" @@ -20,6 +25,10 @@ chrono = { version = "0.4", default-features = false, features = ["clock", "std" hmac = "0.12" sha2 = "0.10" thiserror = "2" +ahash = { version = "0.8", features = ["serde"] } +mimalloc = { version = "0.1", default-features = false, optional = true } +crossbeam-channel = { version = "0.5", optional = true } +num_cpus = { version = "1.16", optional = true } [profile.release] opt-level = 3 diff --git a/src/conditions.rs b/src/conditions.rs index 0eef10e..a2160c5 100644 --- a/src/conditions.rs +++ b/src/conditions.rs @@ -1,46 +1,42 @@ -use regex::Regex; +use std::sync::Arc; -use crate::types::Condition; +use crate::types::{CompiledCondition, CondOp}; +use crate::FastMap; -/// Check if conditions are met for a given row. -/// Returns true if at least one condition matches. -/// Returns true if conditions list is empty. -pub fn check_conditions>( - conditions: &[Condition], - values: &[S], - column_indices: &std::collections::HashMap, +/// Trait giving condition-evaluation access to a row's current (possibly already +/// mutated) values by column index. +pub trait RowRead { + fn len(&self) -> usize; + fn value_at(&self, idx: usize) -> &str; +} + +/// Check if a compiled condition list matches the current row. +/// Returns true if the list is empty, or if at least one condition matches. +pub fn check_conditions( + conditions: &[CompiledCondition], + row: &dyn RowRead, + column_indices: &FastMap, usize>, ) -> bool { if conditions.is_empty() { return true; } - for condition in conditions { - let col_idx = match column_indices.get(&condition.column_name) { - Some(idx) => *idx, + let col_idx = match column_indices.get(condition.column_name.as_ref()) { + Some(&idx) => idx, None => continue, }; - if col_idx >= values.len() { + if col_idx >= row.len() { continue; } - let col_value = values[col_idx].as_ref(); - - let matched = match condition.operation.as_str() { - "equal" => col_value == condition.value, - "not_equal" => col_value != condition.value, - "by_pattern" => { - if let Ok(re) = Regex::new(&condition.value) { - re.is_match(col_value) - } else { - false - } - } - _ => false, + let col_value = row.value_at(col_idx); + let matched = match &condition.op { + CondOp::Equal(v) => col_value == v.as_str(), + CondOp::NotEqual(v) => col_value != v.as_str(), + CondOp::ByPattern(re) => re.is_match(col_value), }; - if matched { return true; } } - false } diff --git a/src/format/custom/blocks.rs b/src/format/custom/blocks.rs index 7f900a7..ea64b0f 100644 --- a/src/format/custom/blocks.rs +++ b/src/format/custom/blocks.rs @@ -12,10 +12,10 @@ use crate::format::custom::header::CompressionMethod; use crate::format::custom::io::DumpIO; use crate::processor::DataProcessor; -const OUTPUT_CHUNK_SIZE: usize = 1024 * 1024; // 1MB for better throughput -const MAX_CHUNK_SIZE: usize = 50 * 1024 * 1024; // 50MB -const READ_BUF_SIZE: usize = 2 * 1024 * 1024; // 2MB read buffer -const COALESCE_TARGET: usize = 256 * 1024; // Accumulate at least 256KB before returning +const OUTPUT_CHUNK_SIZE: usize = 1024 * 1024; +const MAX_CHUNK_SIZE: usize = 50 * 1024 * 1024; +const READ_BUF_SIZE: usize = 2 * 1024 * 1024; +const COALESCE_TARGET: usize = 256 * 1024; /// Streaming reader that coalesces many small chunks into larger reads. /// Critical for -Z0 dumps which can have millions of tiny (~100 byte) chunks. @@ -38,24 +38,25 @@ impl<'a, R: Read> ChunkReader<'a, R> { } } - /// Read multiple chunks until we have at least COALESCE_TARGET bytes or hit end fn fill_buffer(&mut self) -> io::Result<()> { - // Move remaining data to front + // Compact the buffer in-place: move unread bytes to the front once, + // rather than calling drain (which does the same thing + bookkeeping). if self.buf_pos > 0 { - self.buffer.drain(..self.buf_pos); + let len = self.buffer.len(); + self.buffer.copy_within(self.buf_pos..len, 0); + self.buffer.truncate(len - self.buf_pos); self.buf_pos = 0; } - // Read chunks until we have enough data while self.buffer.len() < COALESCE_TARGET { - let chunk_len = self.dio.read_int(self.reader) + let chunk_len = self + .dio + .read_int(self.reader) .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e.to_string()))?; - if chunk_len == 0 { self.done = true; break; } - let len = chunk_len.unsigned_abs() as usize; if len > MAX_CHUNK_SIZE { return Err(io::Error::new( @@ -63,20 +64,16 @@ impl<'a, R: Read> ChunkReader<'a, R> { format!("Chunk size {} exceeds maximum {}", len, MAX_CHUNK_SIZE), )); } - - // Read directly into buffer let start = self.buffer.len(); self.buffer.resize(start + len, 0); self.reader.read_exact(&mut self.buffer[start..])?; } - Ok(()) } } impl Read for ChunkReader<'_, R> { fn read(&mut self, buf: &mut [u8]) -> io::Result { - // Refill if buffer is exhausted or low if self.buf_pos >= self.buffer.len() { if self.done { return Ok(0); @@ -86,8 +83,6 @@ impl Read for ChunkReader<'_, R> { return Ok(0); } } - - // Copy from internal buffer to output let available = self.buffer.len() - self.buf_pos; let to_copy = available.min(buf.len()); buf[..to_copy].copy_from_slice(&self.buffer[self.buf_pos..self.buf_pos + to_copy]); @@ -96,11 +91,12 @@ impl Read for ChunkReader<'_, R> { } } -/// Processes data blocks in a custom format dump. pub struct BlockProcessor<'a> { dio: &'a DumpIO, compression: CompressionMethod, processor: &'a mut DataProcessor, + zstd_level: i32, + zstd_threads: u32, } impl<'a> BlockProcessor<'a> { @@ -108,15 +104,18 @@ impl<'a> BlockProcessor<'a> { dio: &'a DumpIO, compression: CompressionMethod, processor: &'a mut DataProcessor, + zstd_level: i32, + zstd_threads: u32, ) -> Self { Self { dio, compression, processor, + zstd_level, + zstd_threads, } } - /// Process a data block: read chunks, decompress if needed, mutate, compress, write. pub fn process_block( &mut self, reader: &mut R, @@ -131,14 +130,11 @@ 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<()> { - // 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); @@ -148,10 +144,7 @@ impl<'a> BlockProcessor<'a> { if n == 0 { break; } - output_buf.extend_from_slice(&read_buf[..n]); - - // 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)?; @@ -161,150 +154,128 @@ impl<'a> BlockProcessor<'a> { } } - // 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. - /// 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(); + // Reused tail buffer: one allocation per block instead of one per chunk. + let mut line_tail: Vec = Vec::with_capacity(64 * 1024); let mut output_buf: Vec = Vec::with_capacity(OUTPUT_CHUNK_SIZE * 2); loop { - // Read large buffer from ChunkReader (it handles chunk boundaries internally) let n = chunk_reader.read(&mut read_buf)?; if n == 0 { break; } - - // Combine with leftover from previous iteration - let data = if line_tail.is_empty() { + let data_slice: &[u8] = if line_tail.is_empty() { &read_buf[..n] } else { 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_slice) { Some(last_nl) => { - let complete = &data[..=last_nl]; - let tail = &data[last_nl + 1..]; - - self.process_complete_lines(complete, &mut output_buf); - line_tail = tail.to_vec(); + let complete_len = last_nl + 1; + let (complete, tail) = data_slice.split_at(complete_len); + let tail = tail.to_vec(); + process_complete_lines(self.processor, complete, &mut output_buf); + line_tail.clear(); + line_tail.extend_from_slice(&tail); - // Flush output when large enough if output_buf.len() >= OUTPUT_CHUNK_SIZE { - self.flush_uncompressed(writer, &mut output_buf)?; + flush_uncompressed(self.dio, writer, &mut output_buf)?; } } None => { if line_tail.is_empty() { - line_tail = read_buf[..n].to_vec(); + line_tail.extend_from_slice(&read_buf[..n]); } - // else: data already IS line_tail (extended above), keep as-is + // else: line_tail already extended above, leave as-is. } } } - // Process remaining tail (last line without newline) if !line_tail.is_empty() { if let Some(mutated) = self.processor.process_line(&line_tail) { - output_buf.extend_from_slice(&mutated); + output_buf.extend_from_slice(mutated); } } - // Write remaining output if !output_buf.is_empty() { - self.flush_uncompressed(writer, &mut output_buf)?; + flush_uncompressed(self.dio, writer, &mut output_buf)?; } - // Terminator self.dio.write_int(writer, 0)?; Ok(()) } - /// Streaming processing for zlib-compressed blocks. - /// Uses ChunkReader for on-demand chunk reading to minimize memory usage. fn process_block_zlib( &mut self, reader: &mut R, writer: &mut W, ) -> Result<()> { - // Use streaming chunk reader instead of loading entire block into memory let chunk_reader = ChunkReader::new(reader, self.dio); - - // Stream: decompress → process lines → compress → write chunks let mut decoder = ZlibDecoder::new(chunk_reader); - let mut encoder = ZlibEncoder::new(Vec::with_capacity(OUTPUT_CHUNK_SIZE), Compression::new(6)); + let mut encoder = + ZlibEncoder::new(Vec::with_capacity(OUTPUT_CHUNK_SIZE), Compression::new(6)); let mut read_buf = vec![0u8; READ_BUF_SIZE]; - let mut line_tail: Vec = Vec::new(); + let mut line_tail: Vec = Vec::with_capacity(64 * 1024); loop { - let n = decoder.read(&mut read_buf) + let n = decoder + .read(&mut read_buf) .map_err(|e| PgStageError::CompressionError(format!("Zlib decompression failed: {}", e)))?; if n == 0 { break; } - - let data = if line_tail.is_empty() { + let data_slice: &[u8] = if line_tail.is_empty() { &read_buf[..n] } else { line_tail.extend_from_slice(&read_buf[..n]); line_tail.as_slice() }; - - match memrchr(b'\n', data) { + match memrchr(b'\n', data_slice) { Some(last_nl) => { - let complete = &data[..=last_nl]; - let tail = &data[last_nl + 1..]; - - // Process complete lines directly into encoder - self.process_complete_lines_to_writer(complete, &mut encoder)?; - - line_tail = tail.to_vec(); - - // Flush compressed output when large enough - self.flush_encoder_chunks(writer, &mut encoder)?; + let complete_len = last_nl + 1; + let (complete, tail) = data_slice.split_at(complete_len); + let tail = tail.to_vec(); + process_complete_lines_to_writer(self.processor, complete, &mut encoder)?; + line_tail.clear(); + line_tail.extend_from_slice(&tail); + flush_encoder_chunks(self.dio, writer, encoder.get_mut())?; } None => { if line_tail.is_empty() { - line_tail = read_buf[..n].to_vec(); + line_tail.extend_from_slice(&read_buf[..n]); } - // else: data already IS line_tail (extended above), keep as-is } } } - // Process remaining tail if !line_tail.is_empty() { if let Some(mutated) = self.processor.process_line(&line_tail) { - encoder.write_all(&mutated) + encoder + .write_all(mutated) .map_err(|e| PgStageError::CompressionError(format!("Zlib compression failed: {}", e)))?; } } - // Finalize encoder and write remaining compressed data - let remaining = encoder.finish() + let remaining = encoder + .finish() .map_err(|e| PgStageError::CompressionError(format!("Zlib compression finish failed: {}", e)))?; if !remaining.is_empty() { for chunk in remaining.chunks(OUTPUT_CHUNK_SIZE) { @@ -312,83 +283,68 @@ impl<'a> BlockProcessor<'a> { writer.write_all(chunk)?; } } - - // Terminator self.dio.write_int(writer, 0)?; Ok(()) } - /// Streaming processing for zstd-compressed blocks. - /// Uses ChunkReader for on-demand chunk reading to minimize memory usage. fn process_block_zstd( &mut self, reader: &mut R, writer: &mut W, ) -> Result<()> { - // Use streaming chunk reader instead of loading entire block into memory let chunk_reader = ChunkReader::new(reader, self.dio); - - // Stream: decompress → process lines → compress → write chunks - // Use compression level 1 for speed (was 3) let mut decoder = ZstdDecoder::new(chunk_reader) .map_err(|e| PgStageError::CompressionError(format!("Zstd decoder init failed: {}", e)))?; - - // Use multithread zstd compression for better performance on large data - let mut encoder = ZstdEncoder::new(Vec::with_capacity(OUTPUT_CHUNK_SIZE), 1) + let mut encoder = ZstdEncoder::new(Vec::with_capacity(OUTPUT_CHUNK_SIZE), self.zstd_level) .map_err(|e| PgStageError::CompressionError(format!("Zstd encoder init failed: {}", e)))?; - // Enable multithreaded compression (0 = auto-detect CPU count) - encoder.multithread(0) + encoder + .multithread(self.zstd_threads) .map_err(|e| PgStageError::CompressionError(format!("Zstd multithread init failed: {}", e)))?; let mut read_buf = vec![0u8; READ_BUF_SIZE]; - let mut line_tail: Vec = Vec::new(); + let mut line_tail: Vec = Vec::with_capacity(64 * 1024); loop { - let n = decoder.read(&mut read_buf) + let n = decoder + .read(&mut read_buf) .map_err(|e| PgStageError::CompressionError(format!("Zstd decompression failed: {}", e)))?; if n == 0 { break; } - - let data = if line_tail.is_empty() { + let data_slice: &[u8] = if line_tail.is_empty() { &read_buf[..n] } else { line_tail.extend_from_slice(&read_buf[..n]); line_tail.as_slice() }; - - match memrchr(b'\n', data) { + match memrchr(b'\n', data_slice) { Some(last_nl) => { - let complete = &data[..=last_nl]; - let tail = &data[last_nl + 1..]; - - // Process complete lines directly into encoder - self.process_complete_lines_to_writer(complete, &mut encoder)?; - - line_tail = tail.to_vec(); - - // Flush compressed output when large enough - self.flush_zstd_encoder_chunks(writer, &mut encoder)?; + let complete_len = last_nl + 1; + let (complete, tail) = data_slice.split_at(complete_len); + let tail = tail.to_vec(); + process_complete_lines_to_writer(self.processor, complete, &mut encoder)?; + line_tail.clear(); + line_tail.extend_from_slice(&tail); + flush_encoder_chunks_zstd(self.dio, writer, encoder.get_mut())?; } None => { if line_tail.is_empty() { - line_tail = read_buf[..n].to_vec(); + line_tail.extend_from_slice(&read_buf[..n]); } - // else: data already IS line_tail (extended above), keep as-is } } } - // Process remaining tail if !line_tail.is_empty() { if let Some(mutated) = self.processor.process_line(&line_tail) { - encoder.write_all(&mutated) + encoder + .write_all(mutated) .map_err(|e| PgStageError::CompressionError(format!("Zstd compression failed: {}", e)))?; } } - // Finalize encoder and write remaining compressed data - let remaining = encoder.finish() + let remaining = encoder + .finish() .map_err(|e| PgStageError::CompressionError(format!("Zstd compression finish failed: {}", e)))?; if !remaining.is_empty() { for chunk in remaining.chunks(OUTPUT_CHUNK_SIZE) { @@ -396,115 +352,111 @@ impl<'a> BlockProcessor<'a> { writer.write_all(chunk)?; } } - - // Terminator self.dio.write_int(writer, 0)?; Ok(()) } +} - /// Process complete lines (each ending with \n) and append results to output buffer. - fn process_complete_lines(&mut self, data: &[u8], output: &mut Vec) { - // Fast path: if table is being deleted, skip all data - if self.processor.is_delete() { - return; - } - - // Fast path: if no mutations, copy data directly - if !self.processor.has_mutations() { - output.extend_from_slice(data); - return; - } - - // Slow path: process line by line - use memchr::memchr; - let mut start = 0; - while start < data.len() { - let end = memchr(b'\n', &data[start..]) - .map(|p| start + p) - .unwrap_or(data.len()); - - let line = &data[start..end]; - if let Some(mutated) = self.processor.process_line(line) { - output.extend_from_slice(&mutated); - if end < data.len() { - output.push(b'\n'); - } +fn process_complete_lines(processor: &mut DataProcessor, data: &[u8], output: &mut Vec) { + if processor.is_delete() { + return; + } + if !processor.has_mutations() { + output.extend_from_slice(data); + return; + } + use memchr::memchr; + let mut start = 0; + while start < data.len() { + let end = memchr(b'\n', &data[start..]) + .map(|p| start + p) + .unwrap_or(data.len()); + let line = &data[start..end]; + if let Some(mutated) = processor.process_line(line) { + output.extend_from_slice(mutated); + if end < data.len() { + output.push(b'\n'); } - - start = end + 1; } + start = end + 1; } +} - /// Process complete lines and write results to a Write impl (encoder). - fn process_complete_lines_to_writer(&mut self, data: &[u8], writer: &mut W) -> Result<()> { - // Fast path: if table is being deleted, skip all data - if self.processor.is_delete() { - return Ok(()); - } - - // Fast path: if no mutations, write data directly - if !self.processor.has_mutations() { - writer.write_all(data) +fn process_complete_lines_to_writer( + processor: &mut DataProcessor, + data: &[u8], + writer: &mut W, +) -> Result<()> { + if processor.is_delete() { + return Ok(()); + } + if !processor.has_mutations() { + writer + .write_all(data) + .map_err(|e| PgStageError::CompressionError(format!("Write failed: {}", e)))?; + return Ok(()); + } + use memchr::memchr; + let mut start = 0; + while start < data.len() { + let end = memchr(b'\n', &data[start..]) + .map(|p| start + p) + .unwrap_or(data.len()); + let line = &data[start..end]; + if let Some(mutated) = processor.process_line(line) { + writer + .write_all(mutated) .map_err(|e| PgStageError::CompressionError(format!("Write failed: {}", e)))?; - return Ok(()); - } - - // Slow path: process line by line - use memchr::memchr; - let mut start = 0; - while start < data.len() { - let end = memchr(b'\n', &data[start..]) - .map(|p| start + p) - .unwrap_or(data.len()); - - let line = &data[start..end]; - if let Some(mutated) = self.processor.process_line(line) { - writer.write_all(&mutated) + if end < data.len() { + writer + .write_all(b"\n") .map_err(|e| PgStageError::CompressionError(format!("Write failed: {}", e)))?; - if end < data.len() { - writer.write_all(b"\n") - .map_err(|e| PgStageError::CompressionError(format!("Write failed: {}", e)))?; - } } - - start = end + 1; } - Ok(()) + start = end + 1; } + Ok(()) +} - /// Write all data in output_buf as uncompressed chunks and clear the buffer. - fn flush_uncompressed(&self, writer: &mut W, output_buf: &mut Vec) -> Result<()> { - 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(); - Ok(()) +fn flush_uncompressed( + dio: &DumpIO, + writer: &mut W, + output_buf: &mut Vec, +) -> Result<()> { + for chunk in output_buf.chunks(OUTPUT_CHUNK_SIZE) { + dio.write_int(writer, chunk.len() as i32)?; + writer.write_all(chunk)?; } + output_buf.clear(); + Ok(()) +} - /// Flush accumulated compressed bytes from zlib encoder's inner buffer as chunks. - fn flush_encoder_chunks(&self, writer: &mut W, encoder: &mut ZlibEncoder>) -> Result<()> { - let inner = encoder.get_mut(); - if inner.len() >= OUTPUT_CHUNK_SIZE { - for chunk in inner.chunks(OUTPUT_CHUNK_SIZE) { - self.dio.write_int(writer, chunk.len() as i32)?; - writer.write_all(chunk)?; - } - inner.clear(); +fn flush_encoder_chunks( + dio: &DumpIO, + writer: &mut W, + inner: &mut Vec, +) -> Result<()> { + if inner.len() >= OUTPUT_CHUNK_SIZE { + for chunk in inner.chunks(OUTPUT_CHUNK_SIZE) { + dio.write_int(writer, chunk.len() as i32)?; + writer.write_all(chunk)?; } - Ok(()) + inner.clear(); } + Ok(()) +} - /// Flush accumulated compressed bytes from zstd encoder's inner buffer as chunks. - fn flush_zstd_encoder_chunks(&self, writer: &mut W, encoder: &mut ZstdEncoder<'_, Vec>) -> Result<()> { - let inner = encoder.get_mut(); - if inner.len() >= OUTPUT_CHUNK_SIZE { - for chunk in inner.chunks(OUTPUT_CHUNK_SIZE) { - self.dio.write_int(writer, chunk.len() as i32)?; - writer.write_all(chunk)?; - } - inner.clear(); +fn flush_encoder_chunks_zstd( + dio: &DumpIO, + writer: &mut W, + inner: &mut Vec, +) -> Result<()> { + if inner.len() >= OUTPUT_CHUNK_SIZE { + for chunk in inner.chunks(OUTPUT_CHUNK_SIZE) { + dio.write_int(writer, chunk.len() as i32)?; + writer.write_all(chunk)?; } - Ok(()) + inner.clear(); } -} \ No newline at end of file + Ok(()) +} diff --git a/src/format/custom/io.rs b/src/format/custom/io.rs index a23fc39..ef7ee1a 100644 --- a/src/format/custom/io.rs +++ b/src/format/custom/io.rs @@ -18,36 +18,15 @@ impl DumpIO { Self { int_size, offset_size } } - /// Read a single byte from the reader. pub fn read_byte(reader: &mut R) -> Result { let mut buf = [0u8; 1]; reader.read_exact(&mut buf)?; Ok(buf[0]) } - /// Read a signed integer encoded as `1 byte sign + int_size bytes`. + /// Read a signed integer: 1 sign byte + int_size little-endian magnitude bytes. pub fn read_int(&self, reader: &mut R) -> Result { - // Sign byte - let mut sign_buf = [0u8; 1]; - reader.read_exact(&mut sign_buf)?; - let sign = sign_buf[0]; - - // Magnitude bytes (little-endian) — stack buffer, max 8 bytes - let mut buf = [0u8; 8]; - reader.read_exact(&mut buf[..self.int_size])?; - - let mut value: i32 = 0; - let mut shift = 0; - for &b in &buf[..self.int_size] { - value |= (b as i32) << shift; - shift += 8; - } - - if sign != 0 { - value = -value; - } - - Ok(value) + self.read_int_inner(reader, None::<&mut std::io::Sink>) } /// Read an int and also write its raw bytes to the bypass output. @@ -56,78 +35,53 @@ impl DumpIO { reader: &mut R, writer: &mut W, ) -> Result { - // Sign byte - let mut sign_buf = [0u8; 1]; - reader.read_exact(&mut sign_buf)?; - writer.write_all(&sign_buf)?; - let sign = sign_buf[0]; - - // Magnitude bytes — stack buffer - let mut buf = [0u8; 8]; - reader.read_exact(&mut buf[..self.int_size])?; - writer.write_all(&buf[..self.int_size])?; - - let mut value: i32 = 0; - let mut shift = 0; - for &b in &buf[..self.int_size] { - value |= (b as i32) << shift; - shift += 8; - } - - if sign != 0 { - value = -value; - } - - Ok(value) + self.read_int_inner(reader, Some(writer)) } - /// Debug version: read int and log raw bytes + /// Debug version — reads, bypasses, and logs raw bytes to stderr. pub fn read_int_bypass_debug( &self, reader: &mut R, writer: &mut W, label: &str, ) -> Result { - // Sign byte - let mut sign_buf = [0u8; 1]; - reader.read_exact(&mut sign_buf)?; - writer.write_all(&sign_buf)?; - let sign = sign_buf[0]; - - // Magnitude bytes — stack buffer - let mut buf = [0u8; 8]; - reader.read_exact(&mut buf[..self.int_size])?; - writer.write_all(&buf[..self.int_size])?; - - let mut value: i32 = 0; - let mut shift = 0; - for &b in &buf[..self.int_size] { - value |= (b as i32) << shift; - shift += 8; - } - - if sign != 0 { - value = -value; - } - - // Debug output + // Read sign+magnitude onto a stack buffer, bypass to writer, then decode. + let mut stack = [0u8; 9]; + let total = 1 + self.int_size; + reader.read_exact(&mut stack[..total])?; + writer.write_all(&stack[..total])?; + let value = decode_int(stack[0], &stack[1..1 + self.int_size]); eprintln!( "[DEBUG] {} raw bytes: sign={:02X} magnitude={:02X?} -> value={}", - label, sign, &buf[..self.int_size], value + label, + stack[0], + &stack[1..1 + self.int_size], + value ); - Ok(value) } + fn read_int_inner( + &self, + reader: &mut R, + writer: Option<&mut W>, + ) -> Result { + let mut stack = [0u8; 9]; + let total = 1 + self.int_size; + reader.read_exact(&mut stack[..total])?; + if let Some(w) = writer { + w.write_all(&stack[..total])?; + } + Ok(decode_int(stack[0], &stack[1..1 + self.int_size])) + } + /// Write a signed integer as `1 byte sign + int_size bytes`. pub fn write_int(&self, writer: &mut W, val: i32) -> Result<()> { let (sign, v_abs) = if val < 0 { - (1u8, val.wrapping_neg()) // Use wrapping_neg to handle i32::MIN edge cases safely + (1u8, val.wrapping_neg()) } else { (0u8, val) }; - - // Write sign + magnitude in one syscall (max 9 bytes: 1 sign + 8 magnitude) let mut buf = [0u8; 9]; buf[0] = sign; let mut current = v_abs; @@ -136,7 +90,6 @@ impl DumpIO { current >>= 8; } writer.write_all(&buf[..1 + self.int_size])?; - Ok(()) } @@ -148,8 +101,7 @@ impl DumpIO { } let mut buf = vec![0u8; len as usize]; reader.read_exact(&mut buf)?; - let s = String::from_utf8_lossy(&buf).to_string(); - Ok(Some(s)) + Ok(Some(String::from_utf8_lossy(&buf).into_owned())) } /// Read a string and bypass raw bytes to output. @@ -165,18 +117,14 @@ impl DumpIO { let mut buf = vec![0u8; len as usize]; reader.read_exact(&mut buf)?; writer.write_all(&buf)?; - let s = String::from_utf8_lossy(&buf).to_string(); - Ok(Some(s)) + Ok(Some(String::from_utf8_lossy(&buf).into_owned())) } - /// Read an offset value as raw bytes (no sign prefix), little-endian. + /// Read an offset as `offset_size` little-endian bytes, in a single read. pub fn read_offset(&self, reader: &mut R) -> Result { - let mut offset: i64 = 0; - for i in 0..self.offset_size { - let byte = Self::read_byte(reader)? as i64; - offset |= byte << (i * 8); - } - Ok(offset) + let mut buf = [0u8; 8]; + reader.read_exact(&mut buf[..self.offset_size])?; + Ok(decode_offset(&buf[..self.offset_size])) } /// Read an offset and bypass raw bytes to output. @@ -185,15 +133,10 @@ impl DumpIO { reader: &mut R, writer: &mut W, ) -> Result { - let mut offset: i64 = 0; - for i in 0..self.offset_size { - let mut buf = [0u8; 1]; - reader.read_exact(&mut buf)?; - writer.write_all(&buf)?; - let byte = buf[0] as i64; - offset |= byte << (i * 8); - } - Ok(offset) + let mut buf = [0u8; 8]; + reader.read_exact(&mut buf[..self.offset_size])?; + writer.write_all(&buf[..self.offset_size])?; + Ok(decode_offset(&buf[..self.offset_size])) } /// Read exactly n bytes. @@ -214,4 +157,22 @@ impl DumpIO { writer.write_all(&buf)?; Ok(buf) } -} \ No newline at end of file +} + +#[inline] +fn decode_int(sign: u8, magnitude: &[u8]) -> i32 { + let mut value: i32 = 0; + for (i, &b) in magnitude.iter().enumerate() { + value |= (b as i32) << (i * 8); + } + if sign != 0 { -value } else { value } +} + +#[inline] +fn decode_offset(bytes: &[u8]) -> i64 { + let mut value: i64 = 0; + for (i, &b) in bytes.iter().enumerate() { + value |= (b as i64) << (i * 8); + } + value +} diff --git a/src/format/custom/mod.rs b/src/format/custom/mod.rs index dab82b2..c4dbd64 100644 --- a/src/format/custom/mod.rs +++ b/src/format/custom/mod.rs @@ -11,51 +11,57 @@ use crate::format::custom::header::parse_header; use crate::format::custom::io::DumpIO; use crate::format::custom::toc::{parse_toc, Section, TocEntry}; use crate::processor::DataProcessor; +use crate::FastMap; /// Handler for PostgreSQL custom format dumps (-Fc). pub struct CustomHandler { processor: DataProcessor, verbose: bool, + zstd_level: i32, + zstd_threads: u32, } impl CustomHandler { pub fn new(processor: DataProcessor) -> Self { - Self { processor, verbose: false } + Self { + processor, + verbose: false, + zstd_level: 1, + zstd_threads: 0, + } } - /// Enable verbose output pub fn verbose(mut self, verbose: bool) -> Self { self.verbose = verbose; self } - /// Process a custom format dump from reader to writer. - /// `initial_bytes` contains the bytes already read for format detection. + pub fn zstd_level(mut self, level: i32) -> Self { + self.zstd_level = level; + self + } + + pub fn zstd_threads(mut self, threads: u32) -> Self { + self.zstd_threads = threads; + self + } + pub fn process( &mut self, reader: R, writer: W, initial_bytes: &[u8], ) -> Result<()> { - // Large buffers for better throughput on big dumps (2MB each) let mut reader = BufReader::with_capacity(2 * 1024 * 1024, reader); 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, self.verbose)?; - - // Parse TOC entries (bypasses to output) let entries = parse_toc(&mut reader, &mut writer, &header, self.verbose)?; - // Extract comments from TOC entries to build mutation map self.extract_comments(&entries); - - // Build a map of dump_id -> table info for data blocks let data_entries = self.build_data_map(&entries); - let dio = DumpIO::new(header.int_size, header.offset_size); - // Process data blocks loop { let mut block_type = [0u8; 1]; match reader.read_exact(&mut block_type) { @@ -64,74 +70,84 @@ impl CustomHandler { Err(e) => return Err(e.into()), } - // Block type 0x04 = END if block_type[0] == 0x04 { writer.write_all(&block_type)?; writer.flush()?; break; } - // Block type 0x01 = DATA if block_type[0] == 0x01 { 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) { - // Set up processor for this table if !info.copy_stmt.is_empty() { self.processor.setup_table(&info.copy_stmt); } - if self.processor.has_mutations() || self.processor.is_delete() { - // Process with mutations writer.write_all(&block_type)?; dio.write_int(&mut writer, dump_id)?; - let mut bp = BlockProcessor::new(&dio, header.compression, &mut self.processor); + let mut bp = BlockProcessor::new( + &dio, + header.compression, + &mut self.processor, + self.zstd_level, + self.zstd_threads, + ); bp.process_block(&mut reader, &mut writer)?; } else { - // No mutations: pass through writer.write_all(&block_type)?; dio.write_int(&mut writer, dump_id)?; - let bp = BlockProcessor::new(&dio, header.compression, &mut self.processor); + let bp = BlockProcessor::new( + &dio, + header.compression, + &mut self.processor, + self.zstd_level, + self.zstd_threads, + ); bp.pass_through_block(&mut reader, &mut writer)?; } - self.processor.reset_table(); } else { - // Entry not in data_entries map - pass through writer.write_all(&block_type)?; dio.write_int(&mut writer, dump_id)?; - let bp = BlockProcessor::new(&dio, header.compression, &mut self.processor); + let bp = BlockProcessor::new( + &dio, + header.compression, + &mut self.processor, + self.zstd_level, + self.zstd_threads, + ); bp.pass_through_block(&mut reader, &mut writer)?; } } else { - // Other block types (BLOBS, etc.) - pass through writer.write_all(&block_type)?; - // Read and write dump_id for other block types too let dump_id = dio.read_int(&mut reader)?; dio.write_int(&mut writer, dump_id)?; - let bp = BlockProcessor::new(&dio, header.compression, &mut self.processor); + let bp = BlockProcessor::new( + &dio, + header.compression, + &mut self.processor, + self.zstd_level, + self.zstd_threads, + ); bp.pass_through_block(&mut reader, &mut writer)?; } } writer.flush()?; + self.processor.emit_summary(); Ok(()) } - /// Extract COMMENT ON statements from TOC definitions to build mutation map. fn extract_comments(&mut self, entries: &[TocEntry]) { for entry in entries { if entry.desc == "COMMENT" { - // Parse the definition for mutation comments self.processor.parse_comment(&entry.defn); } } } - /// Build a map from dump_id to data entry info. - fn build_data_map(&self, entries: &[TocEntry]) -> std::collections::HashMap { - let mut map = std::collections::HashMap::new(); + fn build_data_map(&self, entries: &[TocEntry]) -> FastMap { + let mut map = FastMap::new(); for entry in entries { if entry.section == Section::Data || entry.desc == "TABLE DATA" { map.insert( @@ -148,4 +164,4 @@ impl CustomHandler { struct DataEntryInfo { copy_stmt: String, -} \ No newline at end of file +} diff --git a/src/format/plain.rs b/src/format/plain.rs index aa42652..86ff5f8 100644 --- a/src/format/plain.rs +++ b/src/format/plain.rs @@ -21,12 +21,10 @@ impl PlainHandler { writer: W, initial_bytes: &[u8], ) -> Result<()> { - // Large buffers for better throughput (2MB each) let mut writer = BufWriter::with_capacity(2 * 1024 * 1024, writer); let mut is_data = false; let mut comment_buf: Option = None; - // If we have initial bytes, chain them with the reader let combined = std::io::Cursor::new(initial_bytes.to_vec()).chain(reader); let buf_reader = BufReader::with_capacity(2 * 1024 * 1024, combined); @@ -35,7 +33,6 @@ impl PlainHandler { if is_data { if line == "\\." { - // End of COPY data if !self.processor.is_delete() { writer.write_all(b"\\.\n")?; } @@ -44,20 +41,18 @@ impl PlainHandler { continue; } - // Process data line if let Some(mutated) = self.processor.process_line(line.as_bytes()) { - writer.write_all(&mutated)?; + writer.write_all(mutated)?; writer.write_all(b"\n")?; } continue; } - // Handle multiline COMMENT accumulation if let Some(ref mut buf) = comment_buf { buf.push('\n'); buf.push_str(&line); if line.ends_with("';") { - let full_comment = buf.clone(); + let full_comment = std::mem::take(buf); comment_buf = None; self.processor.parse_comment(&full_comment); writer.write_all(full_comment.as_bytes())?; @@ -66,7 +61,6 @@ impl PlainHandler { continue; } - // Detect start of a multiline COMMENT ON COLUMN/TABLE with 'anon: if (line.starts_with("COMMENT ON COLUMN ") || line.starts_with("COMMENT ON TABLE ")) && line.contains("'anon: ") && !line.ends_with("';") @@ -75,10 +69,8 @@ impl PlainHandler { continue; } - // Try to parse as single-line comment self.processor.parse_comment(&line); - // Try to parse as COPY statement if self.processor.setup_table(&line) { if !self.processor.is_delete() { writer.write_all(line.as_bytes())?; @@ -88,12 +80,12 @@ impl PlainHandler { continue; } - // Pass through other lines unchanged writer.write_all(line.as_bytes())?; writer.write_all(b"\n")?; } writer.flush()?; + self.processor.emit_summary(); Ok(()) } } diff --git a/src/lib.rs b/src/lib.rs index 4716919..a8ba38e 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -6,3 +6,12 @@ pub mod processor; pub mod relations; pub mod types; pub mod unique; + +/// Project-wide hashmap/set type aliases. +/// +/// pg_stage_rs hashes keys it fully controls (table/column names, column values +/// already seen during this very dump), so HashDoS protection is unnecessary. +/// aHash is 3-5x faster than the std SipHash on short string keys that dominate +/// this workload. +pub type FastMap = ahash::AHashMap; +pub type FastSet = ahash::AHashSet; diff --git a/src/main.rs b/src/main.rs index d03cc0d..bb33ffb 100644 --- a/src/main.rs +++ b/src/main.rs @@ -3,13 +3,17 @@ use std::io::{self, Read}; use clap::Parser; use regex::Regex; -use pg_stage_rs::error::Result; +use pg_stage_rs::error::{PgStageError, Result}; use pg_stage_rs::format::custom::CustomHandler; use pg_stage_rs::format::plain::PlainHandler; use pg_stage_rs::format::{detect_format, DumpFormat}; use pg_stage_rs::processor::DataProcessor; use pg_stage_rs::types::Locale; +#[cfg(feature = "mimalloc-allocator")] +#[global_allocator] +static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc; + #[derive(Parser, Debug)] #[command(name = "pg_stage_rs", version, about = "PostgreSQL dump anonymizer")] struct Args { @@ -29,7 +33,24 @@ struct Args { #[arg(long = "delete-table-pattern")] delete_table_patterns: Vec, - /// Enable verbose output (dump version, TOC count, compression info) + /// Optional rules file (JSON) with pattern-based mutations for many schemas. + /// See README §"Pattern rules". + #[arg(long = "rules-file")] + rules_file: Option, + + /// Zstd compression level for the output dump (1-22). Lower is faster. + #[arg(long = "zstd-level", default_value_t = 1)] + zstd_level: i32, + + /// Zstd compression threads (0 = auto-detect CPU count). + #[arg(long = "zstd-threads", default_value_t = 0)] + zstd_threads: u32, + + /// Fail fast on invalid JSON in COMMENT mutations instead of logging a warning. + #[arg(long)] + strict: bool, + + /// Enable verbose output (dump version, TOC count, compression info, progress) #[arg(short, long)] verbose: bool, } @@ -44,14 +65,30 @@ fn main() { fn run() -> Result<()> { let args = Args::parse(); - let locale: Locale = args.locale.parse().unwrap(); - let delimiter = args.delimiter.bytes().next().unwrap_or(b'\t'); + if !matches!(args.locale.to_lowercase().as_str(), "en" | "ru" | "russian" | "english") { + eprintln!( + "pg_stage_rs warning: unknown locale '{}', falling back to 'en'", + args.locale + ); + } + let locale: Locale = args.locale.parse().unwrap_or(Locale::En); + + let delimiter = args.delimiter.bytes().next().ok_or_else(|| { + PgStageError::InvalidParameter("--delimiter must be a non-empty string".to_string()) + })?; let delete_patterns: Vec = args .delete_table_patterns .iter() - .filter_map(|p| Regex::new(p).ok()) - .collect(); + .map(|p| { + Regex::new(p).map_err(|e| { + PgStageError::InvalidParameter(format!( + "invalid --delete-table-pattern regex '{}': {}", + p, e + )) + }) + }) + .collect::>>()?; let stdin = io::stdin(); let stdout = io::stdout(); @@ -67,13 +104,27 @@ fn run() -> Result<()> { match fmt.as_str() { "plain" | "p" => DumpFormat::Plain, "custom" | "c" => DumpFormat::Custom, - _ => detect_format(peeked)?, + other => { + return Err(PgStageError::InvalidParameter(format!( + "unknown --format '{}', expected plain|p|custom|c", + other + ))) + } } } else { detect_format(peeked)? }; - let processor = DataProcessor::new(locale, delimiter, delete_patterns); + let mut processor = DataProcessor::new(locale, delimiter, delete_patterns); + processor.set_strict(args.strict); + processor.set_verbose(args.verbose); + + if let Some(rules_path) = &args.rules_file { + let text = std::fs::read_to_string(rules_path).map_err(|e| { + PgStageError::InvalidParameter(format!("cannot read --rules-file '{}': {}", rules_path, e)) + })?; + processor.load_rules(&text)?; + } match format { DumpFormat::Plain => { @@ -81,7 +132,10 @@ fn run() -> Result<()> { handler.process(reader, writer, peeked)?; } DumpFormat::Custom => { - let mut handler = CustomHandler::new(processor).verbose(args.verbose); + let mut handler = CustomHandler::new(processor) + .verbose(args.verbose) + .zstd_level(args.zstd_level) + .zstd_threads(args.zstd_threads); handler.process(reader, writer, peeked)?; } } diff --git a/src/mutator/identity.rs b/src/mutator/identity.rs index 3d9efb4..45146e8 100644 --- a/src/mutator/identity.rs +++ b/src/mutator/identity.rs @@ -23,13 +23,7 @@ pub fn uuid5_by_source_value(ctx: &mut MutationContext) -> Result { PgStageError::InvalidParameter(format!("Invalid UUID namespace '{}': {}", namespace_str, e)) })?; - // Get source value from obfuscated_values - let source_value = ctx - .obfuscated_values - .get(source_column) - .cloned() - .unwrap_or_default(); - + let source_value = ctx.obfuscated_values.get(source_column).unwrap_or(""); let today = Utc::now().format("%Y-%m-%d").to_string(); let name = format!("{}-{}", source_value, today); let uuid5 = Uuid::new_v5(&namespace, name.as_bytes()); diff --git a/src/mutator/mod.rs b/src/mutator/mod.rs index b9c657a..9f88330 100644 --- a/src/mutator/mod.rs +++ b/src/mutator/mod.rs @@ -8,22 +8,31 @@ pub mod network; pub mod numeric; pub mod simple; -use std::collections::HashMap; - use rand::rngs::ThreadRng; -use crate::error::{PgStageError, Result}; +use crate::error::Result; use crate::types::Locale; use crate::unique::UniqueTracker; +use crate::FastMap; + +/// Monomorphic function pointer type used by the dispatch table. +pub type MutationFn = fn(&mut MutationContext) -> Result; + +/// Read-only accessor for already-obfuscated values in the current row. +/// Used by mutations like `uuid5_by_source_value` that derive their output +/// from another column's (already obfuscated) value. +pub trait ObfuscatedLookup { + fn get(&self, column: &str) -> Option<&str>; +} pub struct MutationContext<'a> { - pub kwargs: &'a HashMap, - pub current_value: &'a str, // Changed from String to &str to avoid allocation + pub kwargs: &'a FastMap, + pub current_value: &'a str, pub rng: &'a mut ThreadRng, pub unique_tracker: &'a mut UniqueTracker, pub locale: Locale, - pub secrets: &'a HashMap, - pub obfuscated_values: &'a HashMap, + pub secrets: &'a FastMap, + pub obfuscated_values: &'a dyn ObfuscatedLookup, } impl<'a> MutationContext<'a> { @@ -39,52 +48,46 @@ impl<'a> MutationContext<'a> { } } -pub fn dispatch_mutation(name: &str, ctx: &mut MutationContext) -> Result { - match name { - // Names - "first_name" => names::first_name(ctx), - "last_name" => names::last_name(ctx), - "full_name" => names::full_name(ctx), - "middle_name" => names::middle_name(ctx), - - // Contact - "email" => contact::email(ctx), - "phone_number" => contact::phone_number(ctx), - "address" => contact::address(ctx), - "deterministic_phone_number" => contact::deterministic_phone(ctx), - - // Numeric - "numeric_smallint" => numeric::smallint(ctx), - "numeric_integer" => numeric::integer(ctx), - "numeric_bigint" => numeric::bigint(ctx), - "numeric_decimal" => numeric::decimal(ctx), - "numeric_real" => numeric::real(ctx), - "numeric_double_precision" => numeric::double_precision(ctx), - "numeric_smallserial" => numeric::smallserial(ctx), - "numeric_serial" => numeric::serial(ctx), - "numeric_bigserial" => numeric::bigserial(ctx), - - // DateTime - "date" => datetime::date(ctx), - - // Network - "uri" => network::uri(ctx), - "ipv4" => network::ipv4(ctx), - "ipv6" => network::ipv6(ctx), - - // Identity - "uuid4" => identity::uuid4(ctx), - "uuid5_by_source_value" => identity::uuid5_by_source_value(ctx), - - // Simple - "null" => simple::null(ctx), - "empty_string" => simple::empty_string(ctx), - "fixed_value" => simple::fixed_value(ctx), - "random_choice" => simple::random_choice(ctx), - - // Mask - "string_by_mask" => mask::string_by_mask(ctx), - - _ => Err(PgStageError::UnknownMutation(name.to_string())), - } +/// Resolve a mutation name to its function pointer at parse time (once). +/// Returns `None` for unknown names — callers turn that into an error. +pub fn resolve_mutation(name: &str) -> Option { + Some(match name { + "first_name" => names::first_name, + "last_name" => names::last_name, + "full_name" => names::full_name, + "middle_name" => names::middle_name, + + "email" => contact::email, + "phone_number" => contact::phone_number, + "address" => contact::address, + "deterministic_phone_number" => contact::deterministic_phone, + + "numeric_smallint" => numeric::smallint, + "numeric_integer" => numeric::integer, + "numeric_bigint" => numeric::bigint, + "numeric_decimal" => numeric::decimal, + "numeric_real" => numeric::real, + "numeric_double_precision" => numeric::double_precision, + "numeric_smallserial" => numeric::smallserial, + "numeric_serial" => numeric::serial, + "numeric_bigserial" => numeric::bigserial, + + "date" => datetime::date, + + "uri" => network::uri, + "ipv4" => network::ipv4, + "ipv6" => network::ipv6, + + "uuid4" => identity::uuid4, + "uuid5_by_source_value" => identity::uuid5_by_source_value, + + "null" => simple::null, + "empty_string" => simple::empty_string, + "fixed_value" => simple::fixed_value, + "random_choice" => simple::random_choice, + + "string_by_mask" => mask::string_by_mask, + + _ => return None, + }) } diff --git a/src/processor.rs b/src/processor.rs index 3cc79ca..7d9d68e 100644 --- a/src/processor.rs +++ b/src/processor.rs @@ -1,68 +1,120 @@ -use std::borrow::Cow; -use std::collections::HashMap; +use std::sync::Arc; use rand::rngs::ThreadRng; use rand::thread_rng; use regex::Regex; -use crate::conditions::check_conditions; -use crate::mutator::{dispatch_mutation, MutationContext}; +use crate::conditions::{check_conditions, RowRead}; +use crate::error::{PgStageError, Result}; +use crate::mutator::{MutationContext, ObfuscatedLookup}; use crate::relations::RelationTracker; -use crate::types::{Locale, MutationMap, MutationSpec, TableMutationMap, TableMutationSpec}; +use crate::types::{ + ColumnPatternRule, CompiledMutationSpec, Locale, MutationMap, MutationSpec, RulesFile, + TableMutationMap, TableMutationSpec, TablePatternRule, +}; use crate::unique::UniqueTracker; +use crate::FastMap; -pub struct DataProcessor { +/// Compiled mutation registry, filled during parse-time. +/// Separated from per-row runtime state so it can, in the future, be shared +/// read-only between worker threads. +#[derive(Default)] +pub struct MutationRegistry { pub mutation_map: MutationMap, pub table_mutations: TableMutationMap, + pub table_pattern_rules: Vec<(Regex, TableMutationSpec)>, + pub column_pattern_rules: Vec<(Regex, Regex, Vec)>, +} + +impl MutationRegistry { + pub fn new() -> Self { + Self::default() + } + + pub fn table_delete(&self, table: &str, extra_patterns: &[Regex]) -> bool { + if let Some(spec) = self.table_mutations.get(table) { + if spec.mutation_name == "delete" { + return true; + } + } + for (re, spec) in &self.table_pattern_rules { + if spec.mutation_name == "delete" && re.is_match(table) { + return true; + } + } + for re in extra_patterns { + if re.is_match(table) { + return true; + } + } + false + } +} + +pub struct DataProcessor { + pub registry: MutationRegistry, pub locale: Locale, pub delimiter: u8, pub delete_patterns: Vec, - // Current table state - current_table: String, - current_columns: Vec, - column_indices: HashMap, - current_mutations: HashMap>, + strict: bool, + verbose: bool, + + pub rows_processed: u64, + pub mutations_applied: u64, + + current_table: Arc, + current_columns: Vec>, + column_indices: FastMap, usize>, + current_mutations: FastMap, Vec>, + sorted_col_indices: Vec, is_delete_table: bool, - sorted_columns: Vec, - // Shared state + // Per-row scratch — cleared, not reallocated, each row. + scratch_spans: Vec<(u32, u32)>, + scratch_replacements: Vec>>, + scratch_output: Vec, + rng: ThreadRng, unique_tracker: UniqueTracker, relation_tracker: RelationTracker, - secrets: HashMap, + secrets: FastMap, - // Regex patterns comment_column_re: Regex, comment_table_re: Regex, copy_re: Regex, + + json_errors: u64, + unknown_mutation_errors: u64, } impl DataProcessor { pub fn new(locale: Locale, delimiter: u8, delete_patterns: Vec) -> Self { - let secrets = { - let mut m = HashMap::new(); - if let Ok(v) = std::env::var("SECRET_KEY") { - m.insert("SECRET_KEY".to_string(), v); - } - if let Ok(v) = std::env::var("SECRET_KEY_NONCE") { - m.insert("SECRET_KEY_NONCE".to_string(), v); - } - m - }; - + let mut secrets = FastMap::new(); + if let Ok(v) = std::env::var("SECRET_KEY") { + secrets.insert("SECRET_KEY".to_string(), v); + } + if let Ok(v) = std::env::var("SECRET_KEY_NONCE") { + secrets.insert("SECRET_KEY_NONCE".to_string(), v); + } Self { - mutation_map: HashMap::new(), - table_mutations: HashMap::new(), + registry: MutationRegistry::new(), locale, delimiter, delete_patterns, - current_table: String::new(), + strict: false, + verbose: false, + rows_processed: 0, + mutations_applied: 0, + current_table: Arc::from(""), current_columns: Vec::new(), - column_indices: HashMap::new(), - current_mutations: HashMap::new(), + column_indices: FastMap::new(), + current_mutations: FastMap::new(), + sorted_col_indices: Vec::new(), is_delete_table: false, - sorted_columns: Vec::new(), + scratch_spans: Vec::new(), + scratch_replacements: Vec::new(), + scratch_output: Vec::new(), rng: thread_rng(), unique_tracker: UniqueTracker::new(), relation_tracker: RelationTracker::new(), @@ -76,39 +128,133 @@ impl DataProcessor { ) .unwrap(), copy_re: Regex::new(r"COPY ([\d\w_\.]+) \(([#\w\W]+)\) FROM stdin;").unwrap(), + json_errors: 0, + unknown_mutation_errors: 0, } } - /// Parse a COMMENT ON COLUMN or COMMENT ON TABLE line. - /// Returns true if a comment was parsed. + pub fn set_strict(&mut self, strict: bool) { + self.strict = strict; + } + + pub fn set_verbose(&mut self, verbose: bool) { + self.verbose = verbose; + } + + pub fn load_rules(&mut self, text: &str) -> Result<()> { + let file: RulesFile = serde_json::from_str(text) + .map_err(|e| PgStageError::InvalidParameter(format!("invalid rules file: {}", e)))?; + for TablePatternRule { table, mutation } in file.table_patterns { + let re = Regex::new(&table).map_err(|e| { + PgStageError::InvalidParameter(format!("invalid table pattern '{}': {}", table, e)) + })?; + self.registry.table_pattern_rules.push((re, mutation)); + } + for ColumnPatternRule { table, column, mutations } in file.column_patterns { + let table_re = Regex::new(&table).map_err(|e| { + PgStageError::InvalidParameter(format!( + "invalid rule table pattern '{}': {}", + table, e + )) + })?; + let col_re = Regex::new(&column).map_err(|e| { + PgStageError::InvalidParameter(format!( + "invalid rule column pattern '{}': {}", + column, e + )) + })?; + let compiled: Vec = mutations + .into_iter() + .map(CompiledMutationSpec::compile) + .collect::>>()?; + self.registry + .column_pattern_rules + .push((table_re, col_re, compiled)); + } + Ok(()) + } + + pub fn parse_warnings(&self) -> (u64, u64) { + (self.json_errors, self.unknown_mutation_errors) + } + + /// Parse a COMMENT ON COLUMN / COMMENT ON TABLE line. Returns true if a + /// comment was recognized (even if its JSON failed to parse — that error + /// is reported via stderr based on `strict`/`verbose`). pub fn parse_comment(&mut self, line: &str) -> bool { if let Some(caps) = self.comment_column_re.captures(line) { let full_name = caps.get(1).unwrap().as_str(); let json_str = caps.get(2).unwrap().as_str(); - // Parse table.column from full_name (e.g., "public.users.email") let parts: Vec<&str> = full_name.rsplitn(2, '.').collect(); if parts.len() < 2 { return false; } - let column_name = parts[0].to_string(); - let table_name = parts[1].to_string(); - - if let Ok(specs) = serde_json::from_str::>(json_str) { - self.mutation_map - .entry(table_name) - .or_default() - .insert(column_name, specs); + let column_name: Arc = Arc::from(parts[0]); + let table_name: Arc = Arc::from(parts[1]); + + match serde_json::from_str::>(json_str) { + Ok(specs) => { + let mut compiled = Vec::with_capacity(specs.len()); + for spec in specs { + match CompiledMutationSpec::compile(spec) { + Ok(c) => compiled.push(c), + Err(e) => { + self.unknown_mutation_errors += 1; + if self.strict { + eprintln!( + "pg_stage_rs error: compile failed for {}: {}", + full_name, e + ); + } else if self.verbose { + eprintln!( + "pg_stage_rs warning: compile failed for {}: {}", + full_name, e + ); + } + } + } + } + self.registry + .mutation_map + .entry(table_name) + .or_default() + .insert(column_name, compiled); + } + Err(e) => { + self.json_errors += 1; + if self.strict { + eprintln!( + "pg_stage_rs error: invalid JSON in COMMENT ON COLUMN {}: {}", + full_name, e + ); + } else if self.verbose { + eprintln!( + "pg_stage_rs warning: invalid JSON in COMMENT ON COLUMN {}: {}", + full_name, e + ); + } + } } return true; } if let Some(caps) = self.comment_table_re.captures(line) { - let table_name = caps.get(1).unwrap().as_str().to_string(); + let table_name: Arc = Arc::from(caps.get(1).unwrap().as_str()); let json_str = caps.get(2).unwrap().as_str(); - - if let Ok(spec) = serde_json::from_str::(json_str) { - self.table_mutations.insert(table_name, spec); + match serde_json::from_str::(json_str) { + Ok(spec) => { + self.registry.table_mutations.insert(table_name, spec); + } + Err(e) => { + self.json_errors += 1; + if self.strict || self.verbose { + eprintln!( + "pg_stage_rs warning: invalid JSON in COMMENT ON TABLE {}: {}", + table_name, e + ); + } + } } return true; } @@ -116,158 +262,238 @@ impl DataProcessor { false } - /// Set up the processor for a new table based on COPY statement. - /// Returns true if line was a COPY statement. pub fn setup_table(&mut self, line: &str) -> bool { - if let Some(caps) = self.copy_re.captures(line) { - let table_name = caps.get(1).unwrap().as_str().to_string(); - let columns_str = caps.get(2).unwrap().as_str(); - - self.current_columns = columns_str - .split(", ") - .map(|s| s.trim().to_string()) - .collect(); - - self.column_indices.clear(); - for (i, col) in self.current_columns.iter().enumerate() { - self.column_indices.insert(col.clone(), i); - } + let Some(caps) = self.copy_re.captures(line) else { + return false; + }; + let table_name_str = caps.get(1).unwrap().as_str(); + let columns_str = caps.get(2).unwrap().as_str(); + + self.current_columns.clear(); + self.column_indices.clear(); + self.current_mutations.clear(); - // Check if table should be deleted - self.is_delete_table = self.should_delete_table(&table_name); + for (i, raw) in columns_str.split(", ").enumerate() { + let col: Arc = Arc::from(raw.trim()); + self.column_indices.insert(Arc::clone(&col), i); + self.current_columns.push(col); + } - // Get mutations for this table - self.current_mutations = self - .mutation_map - .get(&table_name) - .cloned() - .unwrap_or_default(); + let table_name: Arc = Arc::from(table_name_str); + self.current_table = Arc::clone(&table_name); - // Sort columns: non-source-dependent first - self.sorted_columns = self.sort_columns_by_dependency(); + self.is_delete_table = self + .registry + .table_delete(&table_name, &self.delete_patterns); - self.current_table = table_name; - self.unique_tracker.clear(); - return true; + if let Some(cols) = self.registry.mutation_map.get(&table_name) { + for (col, specs) in cols.iter() { + self.current_mutations + .entry(Arc::clone(col)) + .or_default() + .extend(specs.iter().cloned()); + } } - false + + for (table_re, col_re, specs) in &self.registry.column_pattern_rules { + if !table_re.is_match(&table_name) { + continue; + } + for col in self.current_columns.iter() { + if col_re.is_match(col) { + self.current_mutations + .entry(Arc::clone(col)) + .or_default() + .extend(specs.iter().cloned()); + } + } + } + + self.build_sorted_indices(); + self.unique_tracker.clear(); + true + } + + fn build_sorted_indices(&mut self) { + self.sorted_col_indices.clear(); + let mut dependent = Vec::new(); + for (i, col) in self.current_columns.iter().enumerate() { + if let Some(specs) = self.current_mutations.get(col) { + let has_source = specs.iter().any(|s| s.has_source_column()); + if has_source { + dependent.push(i); + } else { + self.sorted_col_indices.push(i); + } + } + } + self.sorted_col_indices.extend(dependent); + } + + pub fn reset_table(&mut self) { + self.current_table = Arc::from(""); + self.current_columns.clear(); + self.column_indices.clear(); + self.current_mutations.clear(); + self.sorted_col_indices.clear(); + self.is_delete_table = false; + } + + pub fn has_mutations(&self) -> bool { + !self.current_mutations.is_empty() } - /// Process a single data line (tab-separated values). - /// Returns None if the line should be deleted. - /// Returns Some(mutated_line) otherwise. - /// - /// Note: Callers should check has_mutations() and is_delete() first - /// for bulk operations to avoid per-line overhead. - #[inline] - pub fn process_line(&mut self, line: &[u8]) -> Option> { + pub fn is_delete(&self) -> bool { + self.is_delete_table + } + + pub fn relation_tracker_size(&self) -> usize { + self.relation_tracker.len() + } + + pub fn unique_tracker_size(&self) -> usize { + self.unique_tracker.len() + } + + /// Process a single data line. Returns `None` if the table is being + /// deleted, else `Some(bytes)` where `bytes` is valid until the next call + /// to any `&mut self` method. + pub fn process_line<'a>(&'a mut self, line: &'a [u8]) -> Option<&'a [u8]> { if self.is_delete_table { return None; } + self.rows_processed = self.rows_processed.wrapping_add(1); if self.current_mutations.is_empty() { - return Some(line.to_vec()); + return Some(line); } - // Split line by delimiter - let line_str = match std::str::from_utf8(line) { - Ok(s) => s, - Err(_) => return Some(line.to_vec()), - }; - - let delimiter_char = self.delimiter as char; - let values: Vec<&str> = line_str.split(delimiter_char).collect(); - if values.len() != self.current_columns.len() { - return Some(line.to_vec()); + if std::str::from_utf8(line).is_err() { + return Some(line); } - // Use Cow to avoid allocating Strings for unmodified columns - let mut result_values: Vec> = values.iter().map(|&s| Cow::Borrowed(s)).collect(); - - // Pre-allocate with expected capacity to reduce reallocations - let mut obfuscated_values: HashMap = HashMap::with_capacity(self.current_mutations.len()); + self.scratch_spans.clear(); + self.scratch_replacements.clear(); + self.scratch_replacements + .resize_with(self.current_columns.len(), || None); + + let delim = self.delimiter; + let mut start: u32 = 0; + for (i, &b) in line.iter().enumerate() { + if b == delim { + self.scratch_spans.push((start, i as u32)); + start = i as u32 + 1; + } + } + self.scratch_spans.push((start, line.len() as u32)); - // Iterate by index to avoid cloning sorted_columns - for col_sort_idx in 0..self.sorted_columns.len() { - let col_name = &self.sorted_columns[col_sort_idx]; + if self.scratch_spans.len() != self.current_columns.len() { + return Some(line); + } - let col_idx = match self.column_indices.get(col_name.as_str()) { - Some(&idx) => idx, - None => continue, - }; + self.run_mutations(line); + self.build_output(line); + Some(&self.scratch_output) + } - let specs = match self.current_mutations.get(col_name.as_str()) { - Some(s) => s, - None => continue, + fn run_mutations(&mut self, line: &[u8]) { + let Self { + current_columns, + column_indices, + current_mutations, + sorted_col_indices, + scratch_spans, + scratch_replacements, + rng, + unique_tracker, + relation_tracker, + secrets, + locale, + mutations_applied, + .. + } = self; + + for &col_idx in sorted_col_indices.iter() { + let col_name: &Arc = ¤t_columns[col_idx]; + let Some(specs) = current_mutations.get(col_name) else { + continue; }; - // Try each mutation spec in order for spec in specs.iter() { - // Check conditions — borrow of result_values ends after this call (NLL) - if !check_conditions(&spec.conditions, result_values.as_slice(), &self.column_indices) { + let row = ScratchRow { + line, + spans: scratch_spans, + replacements: scratch_replacements, + }; + if !check_conditions(&spec.conditions, &row, column_indices) { continue; } - // Check relations (inlined to avoid &self method call) if !spec.relations.is_empty() { - let mut relation_found = false; - for relation in &spec.relations { - if let Some(&from_idx) = self.column_indices.get(&relation.from_column_name) { - let fk_value = result_values[from_idx].as_ref(); - if let Some(existing) = self.relation_tracker.lookup( - &relation.table_name, - &relation.to_column_name, - fk_value, - ) { - let val = existing.clone(); - obfuscated_values.insert(col_name.clone(), val.clone()); - result_values[col_idx] = Cow::Owned(val); - relation_found = true; - break; - } + let mut found: Option = None; + for rel in &spec.relations { + let from_idx = match column_indices.get(rel.from_column_name.as_ref()) { + Some(&i) => i, + None => continue, + }; + let fk_view = current_value(line, scratch_spans, scratch_replacements, from_idx); + if let Some(existing) = + relation_tracker.lookup(&rel.table_name, &rel.to_column_name, fk_view) + { + found = Some(existing.to_string()); + break; } } - if relation_found { + if let Some(val) = found { + scratch_replacements[col_idx] = Some(Box::from(val.as_str())); break; } } - // Get current value as &str to avoid allocation - let current_value = result_values[col_idx].as_ref(); - - // Dispatch mutation + let cur = current_value(line, scratch_spans, scratch_replacements, col_idx); + let lookup = ScratchLookup { + column_indices, + replacements: scratch_replacements, + }; let mut ctx = MutationContext { - kwargs: &spec.mutation_kwargs, - current_value, - rng: &mut self.rng, - unique_tracker: &mut self.unique_tracker, - locale: self.locale, - secrets: &self.secrets, - obfuscated_values: &obfuscated_values, + kwargs: spec.mutation_kwargs.as_ref(), + current_value: cur, + rng, + unique_tracker, + locale: *locale, + secrets, + obfuscated_values: &lookup, }; - match dispatch_mutation(&spec.mutation_name, &mut ctx) { + match spec.call(&mut ctx) { Ok(new_val) => { - // Store relation (inlined to avoid &mut self method call) if !spec.relations.is_empty() { - for relation in &spec.relations { - if let Some(&from_idx) = self.column_indices.get(&relation.from_column_name) { - self.relation_tracker.store( - &relation.table_name, - &relation.to_column_name, - result_values[from_idx].as_ref(), + for rel in &spec.relations { + if let Some(&from_idx) = + column_indices.get(rel.from_column_name.as_ref()) + { + let fk_view = current_value( + line, + scratch_spans, + scratch_replacements, + from_idx, + ); + relation_tracker.store( + &rel.table_name, + &rel.to_column_name, + fk_view, &new_val, ); } } } - obfuscated_values.insert(col_name.clone(), new_val.clone()); - result_values[col_idx] = Cow::Owned(new_val); + scratch_replacements[col_idx] = Some(Box::from(new_val.as_str())); + *mutations_applied = mutations_applied.wrapping_add(1); break; } Err(e) => { eprintln!( - "Warning: mutation '{}' failed for column '{}': {}", + "pg_stage_rs warning: mutation '{}' failed for column '{}': {}", spec.mutation_name, col_name, e ); continue; @@ -275,75 +501,85 @@ impl DataProcessor { } } } + } - // Build output directly into bytes without intermediate String - let mut output = Vec::with_capacity(line.len()); - for (i, val) in result_values.iter().enumerate() { + fn build_output(&mut self, line: &[u8]) { + self.scratch_output.clear(); + self.scratch_output.reserve(line.len() + 16); + for (i, span) in self.scratch_spans.iter().enumerate() { if i > 0 { - output.push(self.delimiter); + self.scratch_output.push(self.delimiter); + } + match &self.scratch_replacements[i] { + Some(s) => self.scratch_output.extend_from_slice(s.as_bytes()), + None => self + .scratch_output + .extend_from_slice(&line[span.0 as usize..span.1 as usize]), } - output.extend_from_slice(val.as_ref().as_bytes()); } - Some(output) - } - - /// Reset table state (called when COPY data ends) - pub fn reset_table(&mut self) { - self.current_table.clear(); - self.current_columns.clear(); - self.current_mutations.clear(); - self.is_delete_table = false; - } - - /// Check if current table has any mutations configured - pub fn has_mutations(&self) -> bool { - !self.current_mutations.is_empty() } - /// Check if current table should be deleted - pub fn is_delete(&self) -> bool { - self.is_delete_table + pub fn emit_summary(&self) { + if !self.verbose { + return; + } + eprintln!( + "[INFO] processed rows: {}, mutations applied: {}, unique values tracked: {}, relations tracked: {}", + self.rows_processed, + self.mutations_applied, + self.unique_tracker.len(), + self.relation_tracker.len(), + ); + if self.json_errors > 0 || self.unknown_mutation_errors > 0 { + eprintln!( + "[WARN] parse warnings: {} invalid JSON comments, {} unknown mutations", + self.json_errors, self.unknown_mutation_errors + ); + } } +} - fn should_delete_table(&self, table_name: &str) -> bool { - // Check table mutations - if let Some(spec) = self.table_mutations.get(table_name) { - if spec.mutation_name == "delete" { - return true; - } +#[inline] +fn current_value<'a>( + line: &'a [u8], + spans: &[(u32, u32)], + replacements: &'a [Option>], + idx: usize, +) -> &'a str { + match &replacements[idx] { + Some(b) => b.as_ref(), + None => { + let (s, e) = spans[idx]; + // SAFETY: caller verified the whole line is valid UTF-8 and the + // delimiter byte is ASCII, so each span is a valid UTF-8 slice. + unsafe { std::str::from_utf8_unchecked(&line[s as usize..e as usize]) } } + } +} - // Check delete patterns - for pattern in &self.delete_patterns { - if pattern.is_match(table_name) { - return true; - } - } +struct ScratchRow<'a> { + line: &'a [u8], + spans: &'a [(u32, u32)], + replacements: &'a [Option>], +} - false +impl<'a> RowRead for ScratchRow<'a> { + fn len(&self) -> usize { + self.spans.len() } + fn value_at(&self, idx: usize) -> &str { + current_value(self.line, self.spans, self.replacements, idx) + } +} - fn sort_columns_by_dependency(&self) -> Vec { - let mut independent = Vec::new(); - let mut dependent = Vec::new(); - - for col_name in self.current_columns.iter() { - if let Some(specs) = self.current_mutations.get(col_name) { - let has_source = specs.iter().any(|s| { - s.mutation_kwargs.contains_key("source_column") - }); - if has_source { - dependent.push(col_name.clone()); - } else { - independent.push(col_name.clone()); - } - } else { - independent.push(col_name.clone()); - } - } +struct ScratchLookup<'a> { + column_indices: &'a FastMap, usize>, + replacements: &'a [Option>], +} - independent.extend(dependent); - independent +impl<'a> ObfuscatedLookup for ScratchLookup<'a> { + fn get(&self, column: &str) -> Option<&str> { + let idx = *self.column_indices.get(column)?; + self.replacements.get(idx)?.as_deref() } - } diff --git a/src/relations.rs b/src/relations.rs index e986561..6456bb0 100644 --- a/src/relations.rs +++ b/src/relations.rs @@ -1,17 +1,21 @@ -use std::collections::HashMap; +use std::sync::Arc; -use uuid::Uuid; +use crate::FastMap; /// Tracks FK relationships to ensure consistent obfuscation across tables. /// -/// Maps: table_name -> to_column_name -> fk_value -> relation_key (UUID) -/// And: relation_key -> obfuscated_value +/// Layout: `by_table[table][column][fk_value] = obfuscated_value`. +/// +/// `table` and `column` keys are `Arc`, shared with `CompiledRelation` and +/// `MutationRegistry`, so a name that appears in thousands of schemas is stored +/// in memory exactly once per unique string. +/// +/// `fk_value` and the stored obfuscated value are `Box` (no String +/// capacity overhead). #[derive(Debug, Default)] pub struct RelationTracker { - /// table_name -> column_name -> fk_value -> relation_uuid - fk_map: HashMap>>, - /// relation_uuid -> obfuscated_value - values: HashMap, + by_table: FastMap, FastMap, FastMap, Box>>>, + count: usize, } impl RelationTracker { @@ -19,36 +23,34 @@ impl RelationTracker { Self::default() } - /// Look up if a relation already has an obfuscated value. - /// Returns the existing obfuscated value if found. - pub fn lookup( - &self, - table_name: &str, - to_column_name: &str, - fk_value: &str, - ) -> Option<&String> { - self.fk_map - .get(table_name) - .and_then(|cols| cols.get(to_column_name)) - .and_then(|fks| fks.get(fk_value)) - .and_then(|key| self.values.get(key)) + pub fn len(&self) -> usize { + self.count + } + + pub fn is_empty(&self) -> bool { + self.count == 0 + } + + pub fn lookup(&self, table: &Arc, column: &Arc, fk_value: &str) -> Option<&str> { + self.by_table + .get(table.as_ref())? + .get(column.as_ref())? + .get(fk_value) + .map(|v| v.as_ref()) } - /// Store a new relation mapping. pub fn store( &mut self, - table_name: &str, - to_column_name: &str, + table: &Arc, + column: &Arc, fk_value: &str, - obfuscated_value: &str, + obfuscated: &str, ) { - let key = Uuid::new_v4().to_string(); - self.fk_map - .entry(table_name.to_string()) - .or_default() - .entry(to_column_name.to_string()) - .or_default() - .insert(fk_value.to_string(), key.clone()); - self.values.insert(key, obfuscated_value.to_string()); + let outer = self.by_table.entry(Arc::clone(table)).or_default(); + let mid = outer.entry(Arc::clone(column)).or_default(); + let inserted = mid.insert(Box::from(fk_value), Box::from(obfuscated)); + if inserted.is_none() { + self.count += 1; + } } } diff --git a/src/types.rs b/src/types.rs index 3bd7e02..3b9fea5 100644 --- a/src/types.rs +++ b/src/types.rs @@ -1,6 +1,11 @@ +use regex::Regex; use serde::Deserialize; -use std::collections::HashMap; use std::str::FromStr; +use std::sync::Arc; + +use crate::error::{PgStageError, Result}; +use crate::mutator::{resolve_mutation, MutationFn, MutationContext}; +use crate::FastMap; #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum Locale { @@ -38,7 +43,7 @@ pub struct Relation { pub struct MutationSpec { pub mutation_name: String, #[serde(default)] - pub mutation_kwargs: HashMap, + pub mutation_kwargs: FastMap, #[serde(default)] pub conditions: Vec, #[serde(default)] @@ -50,8 +55,135 @@ pub struct TableMutationSpec { pub mutation_name: String, } -/// Maps table_name -> column_name -> Vec -pub type MutationMap = HashMap>>; +/// Condition operation resolved at parse time. +#[derive(Debug, Clone)] +pub enum CondOp { + Equal(String), + NotEqual(String), + ByPattern(Regex), +} + +#[derive(Debug, Clone)] +pub struct CompiledCondition { + pub column_name: Arc, + pub op: CondOp, +} + +#[derive(Debug, Clone)] +pub struct CompiledRelation { + pub table_name: Arc, + pub to_column_name: Arc, + pub from_column_name: Arc, +} + +/// Mutation spec with the function resolved and regex/op compiled. +#[derive(Clone)] +pub struct CompiledMutationSpec { + pub mutation_name: Arc, + pub mutation_fn: MutationFn, + pub mutation_kwargs: Arc>, + pub conditions: Vec, + pub relations: Vec, +} + +impl std::fmt::Debug for CompiledMutationSpec { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("CompiledMutationSpec") + .field("mutation_name", &self.mutation_name) + .field("conditions", &self.conditions) + .field("relations", &self.relations) + .finish() + } +} + +impl CompiledMutationSpec { + pub fn compile(spec: MutationSpec) -> Result { + let mutation_fn: MutationFn = resolve_mutation(&spec.mutation_name) + .ok_or_else(|| PgStageError::UnknownMutation(spec.mutation_name.clone()))?; + let conditions = spec + .conditions + .into_iter() + .map(|c| { + let op = match c.operation.as_str() { + "equal" => CondOp::Equal(c.value), + "not_equal" => CondOp::NotEqual(c.value), + "by_pattern" => { + let re = Regex::new(&c.value).map_err(|e| { + PgStageError::InvalidParameter(format!( + "invalid regex in condition.by_pattern '{}': {}", + c.value, e + )) + })?; + CondOp::ByPattern(re) + } + other => { + return Err(PgStageError::InvalidParameter(format!( + "unknown condition operation '{}'", + other + ))) + } + }; + Ok(CompiledCondition { + column_name: Arc::from(c.column_name.as_str()), + op, + }) + }) + .collect::>>()?; + let relations = spec + .relations + .into_iter() + .map(|r| CompiledRelation { + table_name: Arc::from(r.table_name.as_str()), + to_column_name: Arc::from(r.to_column_name.as_str()), + from_column_name: Arc::from(r.from_column_name.as_str()), + }) + .collect(); + Ok(Self { + mutation_name: Arc::from(spec.mutation_name.as_str()), + mutation_fn, + mutation_kwargs: Arc::new(spec.mutation_kwargs), + conditions, + relations, + }) + } + + pub fn has_source_column(&self) -> bool { + self.mutation_kwargs.contains_key("source_column") + } + + /// Run the compiled mutation. + pub fn call(&self, ctx: &mut MutationContext) -> Result { + (self.mutation_fn)(ctx) + } +} + +/// Maps table_name -> column_name -> Vec +pub type MutationMap = FastMap, FastMap, Vec>>; /// Maps table_name -> TableMutationSpec (e.g., delete) -pub type TableMutationMap = HashMap; +pub type TableMutationMap = FastMap, TableMutationSpec>; + +/// File format for --rules-file: pattern-based mutations that apply to many schemas. +#[derive(Debug, Clone, Deserialize, Default)] +pub struct RulesFile { + #[serde(default)] + pub table_patterns: Vec, + #[serde(default)] + pub column_patterns: Vec, +} + +#[derive(Debug, Clone, Deserialize)] +pub struct TablePatternRule { + /// Regex matched against fully-qualified "schema.table". + pub table: String, + pub mutation: TableMutationSpec, +} + +#[derive(Debug, Clone, Deserialize)] +pub struct ColumnPatternRule { + /// Regex matched against fully-qualified "schema.table". + pub table: String, + /// Regex matched against the column name. + pub column: String, + pub mutations: Vec, +} diff --git a/src/unique.rs b/src/unique.rs index 4ac67a9..d88cde5 100644 --- a/src/unique.rs +++ b/src/unique.rs @@ -1,25 +1,32 @@ -use std::collections::HashSet; - use crate::error::{PgStageError, Result}; +use crate::FastSet; const MAX_RETRIES: u32 = 1000; #[derive(Debug, Default)] pub struct UniqueTracker { - values: HashSet, + values: FastSet>, } impl UniqueTracker { pub fn new() -> Self { Self { - values: HashSet::new(), + values: FastSet::new(), } } - /// Try to insert a value. Returns Ok(true) if inserted (unique), - /// Ok(false) if already exists. + pub fn len(&self) -> usize { + self.values.len() + } + + /// Try to insert. Returns true if the value was new. + /// Avoids allocating a `String` when the value is already present. pub fn try_insert(&mut self, value: &str) -> bool { - self.values.insert(value.to_string()) + if self.values.contains(value) { + return false; + } + self.values.insert(Box::from(value)); + true } /// Generate a unique value using the provided generator function. diff --git a/tests/integration.rs b/tests/integration.rs index ea49899..03420f0 100644 --- a/tests/integration.rs +++ b/tests/integration.rs @@ -412,8 +412,8 @@ fn test_processor_parse_comment() { let mut proc = make_processor(); let comment = "COMMENT ON COLUMN public.users.email IS 'anon: [{\"mutation_name\": \"email\", \"mutation_kwargs\": {\"unique\": true}}]';"; assert!(proc.parse_comment(comment)); - assert!(proc.mutation_map.contains_key("public.users")); - assert!(proc.mutation_map["public.users"].contains_key("email")); + assert!(proc.registry.mutation_map.contains_key("public.users")); + assert!(proc.registry.mutation_map["public.users"].contains_key("email")); } #[test] @@ -421,8 +421,8 @@ fn test_processor_parse_table_comment() { let mut proc = make_processor(); let comment = "COMMENT ON TABLE public.logs IS 'anon: {\"mutation_name\": \"delete\"}';"; assert!(proc.parse_comment(comment)); - assert!(proc.table_mutations.contains_key("public.logs")); - assert_eq!(proc.table_mutations["public.logs"].mutation_name, "delete"); + assert!(proc.registry.table_mutations.contains_key("public.logs")); + assert_eq!(proc.registry.table_mutations["public.logs"].mutation_name, "delete"); } #[test]