//! Shared support for splitting a large generated SQL export into multiple //! parts packaged inside a single `.zip` archive, instead of one unbounded //! `.sql` file. Used by both the whole-database export (`database_export.rs`) //! and the single-table SQL export (`table_export.rs`). //! //! # Cut boundary //! The writer never splits mid-write: it only checks the size threshold //! *before* starting a new `write()` call and, if exceeded, rotates to the //! next zip entry first. Every caller in this codebase writes one complete //! SQL statement (or comment line) per `write`/`writeln!` call, so a part //! boundary can land between two INSERT batches of the same table but never //! inside a single statement -- every part is independently valid, replayable //! SQL. use std::io::Write; use serde::Serialize; use zip::write::SimpleFileOptions; use zip::{CompressionMethod, ZipWriter}; /// Default part size when a caller enables splitting without specifying a /// size (should not normally happen -- the frontend always sends an explicit /// value -- but keeps the backend safe against malformed requests). pub const DEFAULT_SPLIT_PART_MAX_MB: u32 = 100; pub const MIN_SPLIT_PART_MAX_MB: u32 = 1; pub const MAX_SPLIT_PART_MAX_MB: u32 = 4096; /// Clamp a user-supplied "max MB per part" value to a sane range so a /// malformed or hostile request cannot create an unbounded number of zip /// entries (too-small) or a single-part export that defeats the point /// (too-large is still allowed up to a generous ceiling). pub fn clamp_split_part_max_mb(value: u32) -> u32 { value.clamp(MIN_SPLIT_PART_MAX_MB, MAX_SPLIT_PART_MAX_MB) } fn split_part_max_bytes(max_mb: u32) -> u64 { u64::from(clamp_split_part_max_mb(max_mb)) * 1024 * 1024 } #[derive(Debug, Serialize)] #[serde(rename_all = "camelCase")] struct SplitExportManifest<'a> { format: &'static str, generated_at: String, source_file_name: &'a str, part_max_bytes: u64, total_parts: u32, parts: &'a [String], } /// Writes SQL export content across multiple `part-N.sql` entries inside a /// single `.zip` archive, rotating to a new entry once the current one /// reaches `part_max_bytes`. Call [`SplitZipExportWriter::finish`] to write /// the trailing `manifest.json` entry and flush the archive to disk. pub struct SplitZipExportWriter { zip: ZipWriter, options: SimpleFileOptions, part_max_bytes: u64, current_part_bytes: u64, part_index: u32, part_stem: String, part_extension: String, part_names: Vec, } impl SplitZipExportWriter { /// `part_stem`/`part_extension` name each entry as `{stem}-part-{n}.{extension}` /// (e.g. `mydb-part-1.sql`). pub fn create( zip_path: &std::path::Path, max_mb: u32, part_stem: &str, part_extension: &str, ) -> Result { Self::create_with_overwrite(zip_path, max_mb, part_stem, part_extension, false) } pub fn create_with_overwrite( zip_path: &std::path::Path, max_mb: u32, part_stem: &str, part_extension: &str, prevent_overwrite: bool, ) -> Result { let file = std::fs::OpenOptions::new() .write(true) .create(true) .truncate(!prevent_overwrite) .create_new(prevent_overwrite) .open(zip_path) .map_err(|error| { if prevent_overwrite && error.kind() == std::io::ErrorKind::AlreadyExists { format!("Backup file already exists: {}", zip_path.display()) } else { format!("Failed to create zip file: {error}") } })?; let zip = ZipWriter::new(file); let options = SimpleFileOptions::default().compression_method(CompressionMethod::Deflated); let mut writer = Self { zip, options, part_max_bytes: split_part_max_bytes(max_mb), current_part_bytes: 0, part_index: 0, part_stem: sanitize_zip_entry_component(part_stem), part_extension: sanitize_zip_entry_component(part_extension), part_names: Vec::new(), }; writer.start_next_part()?; Ok(writer) } fn start_next_part(&mut self) -> Result<(), String> { self.part_index += 1; let name = format!("{}-part-{}.{}", self.part_stem, self.part_index, self.part_extension); self.zip .start_file(&name, self.options) .map_err(|error| format!("Failed to start zip entry {name}: {error}"))?; self.part_names.push(name); self.current_part_bytes = 0; Ok(()) } /// Writes the trailing `manifest.json` entry (listing every part in /// order) and finalizes the zip archive. Consumes the writer because the /// underlying `zip::ZipWriter::finish` does. pub fn finish(mut self, source_file_name: &str) -> Result<(), String> { let manifest = SplitExportManifest { format: "dbx-sql-export-parts-v1", generated_at: chrono::Local::now().to_rfc3339(), source_file_name, part_max_bytes: self.part_max_bytes, total_parts: self.part_index, parts: &self.part_names, }; let manifest_json = serde_json::to_vec_pretty(&manifest) .map_err(|error| format!("Failed to encode export manifest: {error}"))?; let stored_options = SimpleFileOptions::default().compression_method(CompressionMethod::Stored); self.zip .start_file("manifest.json", stored_options) .map_err(|error| format!("Failed to start manifest entry: {error}"))?; self.zip.write_all(&manifest_json).map_err(|error| format!("Failed to write export manifest: {error}"))?; let mut file = self.zip.finish().map_err(|error| format!("Failed to finalize zip archive: {error}"))?; file.flush().map_err(|error| format!("Failed to flush zip archive: {error}")) } } impl Write for SplitZipExportWriter { fn write(&mut self, buffer: &[u8]) -> std::io::Result { if self.current_part_bytes >= self.part_max_bytes { self.start_next_part().map_err(std::io::Error::other)?; } let written = self.zip.write(buffer)?; self.current_part_bytes += written as u64; Ok(written) } fn flush(&mut self) -> std::io::Result<()> { self.zip.flush() } } /// Zip entry names must not contain path separators or other characters that /// could be misread as a directory traversal; export file stems come from /// user-controlled database/table names, so sanitize defensively. fn sanitize_zip_entry_component(value: &str) -> String { let sanitized: String = value .chars() .map(|c| if c.is_ascii_alphanumeric() || c == '-' || c == '_' || c == '.' { c } else { '_' }) .collect(); let trimmed = sanitized.trim_matches('_').to_string(); if trimmed.is_empty() { "export".to_string() } else { trimmed } } #[cfg(test)] mod tests { use super::*; use std::io::Read; fn read_zip_entries(path: &std::path::Path) -> Vec<(String, Vec)> { let file = std::fs::File::open(path).unwrap(); let mut archive = zip::ZipArchive::new(file).unwrap(); let mut entries = Vec::new(); for index in 0..archive.len() { let mut entry = archive.by_index(index).unwrap(); let mut contents = Vec::new(); entry.read_to_end(&mut contents).unwrap(); entries.push((entry.name().to_string(), contents)); } entries } #[test] fn clamps_split_part_max_mb_to_a_sane_range() { assert_eq!(clamp_split_part_max_mb(0), MIN_SPLIT_PART_MAX_MB); assert_eq!(clamp_split_part_max_mb(50), 50); assert_eq!(clamp_split_part_max_mb(u32::MAX), MAX_SPLIT_PART_MAX_MB); } #[test] fn rotates_to_a_new_part_once_the_threshold_is_reached() { let dir = std::env::temp_dir().join(format!("dbx-split-zip-test-{}", uuid::Uuid::new_v4())); std::fs::create_dir_all(&dir).unwrap(); let zip_path = dir.join("export.zip"); { // 1 MB minimum threshold; each write is far smaller so every // write should land in the very first part with no rotation. let mut writer = SplitZipExportWriter::create(&zip_path, 1, "mydb", "sql").unwrap(); for index in 0..5 { writeln!(writer, "INSERT INTO t VALUES ({index});").unwrap(); } writer.finish("mydb.sql").unwrap(); } let entries = read_zip_entries(&zip_path); let sql_entries: Vec<_> = entries.iter().filter(|(name, _)| name.ends_with(".sql")).collect(); assert_eq!(sql_entries.len(), 1); assert_eq!(sql_entries[0].0, "mydb-part-1.sql"); let manifest_entry = entries.iter().find(|(name, _)| name == "manifest.json").unwrap(); let manifest: serde_json::Value = serde_json::from_slice(&manifest_entry.1).unwrap(); assert_eq!(manifest["totalParts"], 1); assert_eq!(manifest["parts"].as_array().unwrap().len(), 1); let _ = std::fs::remove_dir_all(&dir); } #[test] fn splits_into_multiple_parts_and_every_part_is_valid_statements() { let dir = std::env::temp_dir().join(format!("dbx-split-zip-test-{}", uuid::Uuid::new_v4())); std::fs::create_dir_all(&dir).unwrap(); let zip_path = dir.join("export.zip"); // Force a rotation by using a threshold smaller than a single write. let long_statement = format!("INSERT INTO t VALUES ('{}');", "x".repeat(200)); { let mut writer = SplitZipExportWriter::create(&zip_path, MIN_SPLIT_PART_MAX_MB, "mydb", "sql").unwrap(); // Each write() call is one full statement; the threshold check // happens before the call, so a small per-write payload combined // with a below-1MB threshold still needs enough total bytes to // exceed 1MB and trigger a second part. Use the internal minimum // by writing repeatedly. for index in 0..6000 { writeln!(writer, "{long_statement} -- row {index}").unwrap(); } writer.finish("mydb.sql").unwrap(); } let entries = read_zip_entries(&zip_path); let mut sql_entries: Vec<_> = entries.iter().filter(|(name, _)| name.ends_with(".sql")).collect(); sql_entries.sort_by(|a, b| a.0.cmp(&b.0)); assert!(sql_entries.len() >= 2, "expected at least 2 parts, got {}", sql_entries.len()); for (name, contents) in &sql_entries { let text = String::from_utf8(contents.clone()).unwrap(); assert!(!text.is_empty(), "{name} must not be empty"); for line in text.lines().filter(|line| !line.trim().is_empty()) { assert!(line.contains(';'), "{name} contains a line that is not a complete statement: {line}"); } } let manifest_entry = entries.iter().find(|(name, _)| name == "manifest.json").unwrap(); let manifest: serde_json::Value = serde_json::from_slice(&manifest_entry.1).unwrap(); assert_eq!(manifest["totalParts"].as_u64().unwrap() as usize, sql_entries.len()); assert_eq!(manifest["parts"].as_array().unwrap().len(), sql_entries.len()); let _ = std::fs::remove_dir_all(&dir); } #[test] fn sanitizes_zip_entry_stem_from_untrusted_names() { assert_eq!(sanitize_zip_entry_component("../../etc/passwd"), ".._.._etc_passwd"); assert_eq!(sanitize_zip_entry_component(""), "export"); assert_eq!(sanitize_zip_entry_component("my db (prod)"), "my_db__prod"); } }