use serde::{Deserialize, Serialize};
use std::io::{Read, Write};
use crate::{Result, Error};
pub const DUMP_MAGIC_NUMBER: &[u8; 8] = b"HELIODMP";
pub const DUMP_VERSION: u32 = 1;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum CompressionType {
None,
Zstd,
Gzip,
Brotli,
}
impl CompressionType {
pub fn from_str(s: &str) -> Result<Self> {
match s.to_lowercase().as_str() {
"none" => Ok(Self::None),
"zstd" => Ok(Self::Zstd),
"gzip" | "gz" => Ok(Self::Gzip),
"brotli" | "br" => Ok(Self::Brotli),
_ => Err(Error::config(format!("Unknown compression type: {}", s))),
}
}
pub fn extension(&self) -> &'static str {
match self {
Self::None => "",
Self::Zstd => ".zst",
Self::Gzip => ".gz",
Self::Brotli => ".br",
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DumpMetadata {
pub dump_id: String,
pub version: String,
pub created_at: u64,
pub lsn: u64,
pub compression: CompressionType,
pub table_count: usize,
pub row_count: u64,
pub bytes_uncompressed: u64,
pub bytes_compressed: u64,
pub incremental: bool,
pub previous_lsn: Option<u64>,
}
impl DumpMetadata {
pub fn new(lsn: u64, compression: CompressionType, incremental: bool) -> Self {
Self {
dump_id: uuid::Uuid::new_v4().to_string(),
version: env!("CARGO_PKG_VERSION").to_string(),
created_at: std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|d| d.as_secs())
.unwrap_or(0),
lsn,
compression,
table_count: 0,
row_count: 0,
bytes_uncompressed: 0,
bytes_compressed: 0,
incremental,
previous_lsn: None,
}
}
pub fn write_to<W: Write>(&self, writer: &mut W) -> Result<()> {
writer.write_all(DUMP_MAGIC_NUMBER)
.map_err(|e| Error::io(format!("Failed to write magic number: {}", e)))?;
writer.write_all(&DUMP_VERSION.to_le_bytes())
.map_err(|e| Error::io(format!("Failed to write version: {}", e)))?;
let json = serde_json::to_vec(self)
.map_err(|e| Error::io(format!("Failed to serialize metadata: {}", e)))?;
let len = json.len() as u32;
writer.write_all(&len.to_le_bytes())
.map_err(|e| Error::io(format!("Failed to write metadata length: {}", e)))?;
writer.write_all(&json)
.map_err(|e| Error::io(format!("Failed to write metadata: {}", e)))?;
Ok(())
}
pub fn read_from<R: Read>(reader: &mut R) -> Result<Self> {
let mut magic = [0u8; 8];
reader.read_exact(&mut magic)
.map_err(|e| Error::io(format!("Failed to read magic number: {}", e)))?;
if &magic != DUMP_MAGIC_NUMBER {
return Err(Error::io("Invalid dump file: magic number mismatch"));
}
let mut version_bytes = [0u8; 4];
reader.read_exact(&mut version_bytes)
.map_err(|e| Error::io(format!("Failed to read version: {}", e)))?;
let version = u32::from_le_bytes(version_bytes);
if version != DUMP_VERSION {
return Err(Error::io(format!(
"Incompatible dump version: expected {}, found {}",
DUMP_VERSION, version
)));
}
let mut len_bytes = [0u8; 4];
reader.read_exact(&mut len_bytes)
.map_err(|e| Error::io(format!("Failed to read metadata length: {}", e)))?;
let len = u32::from_le_bytes(len_bytes) as usize;
let mut json = vec![0u8; len];
reader.read_exact(&mut json)
.map_err(|e| Error::io(format!("Failed to read metadata: {}", e)))?;
let metadata: DumpMetadata = serde_json::from_slice(&json)
.map_err(|e| Error::io(format!("Failed to deserialize metadata: {}", e)))?;
Ok(metadata)
}
}
pub struct DumpFormat;
impl DumpFormat {
pub fn compress(data: &[u8], compression: CompressionType) -> Result<Vec<u8>> {
match compression {
CompressionType::None => Ok(data.to_vec()),
CompressionType::Zstd => {
zstd::encode_all(data, 3)
.map_err(|e| Error::io(format!("Zstd compression failed: {}", e)))
}
CompressionType::Gzip => {
use flate2::write::GzEncoder;
use flate2::Compression;
let mut encoder = GzEncoder::new(Vec::new(), Compression::default());
encoder.write_all(data)
.map_err(|e| Error::io(format!("Gzip compression failed: {}", e)))?;
encoder.finish()
.map_err(|e| Error::io(format!("Gzip compression failed: {}", e)))
}
CompressionType::Brotli => {
use brotli::enc::BrotliEncoderParams;
let mut output = Vec::new();
let params = BrotliEncoderParams::default();
brotli::BrotliCompress(
&mut std::io::Cursor::new(data),
&mut output,
¶ms,
).map_err(|e| Error::io(format!("Brotli compression failed: {}", e)))?;
Ok(output)
}
}
}
pub fn decompress(data: &[u8], compression: CompressionType) -> Result<Vec<u8>> {
match compression {
CompressionType::None => Ok(data.to_vec()),
CompressionType::Zstd => {
zstd::decode_all(data)
.map_err(|e| Error::io(format!("Zstd decompression failed: {}", e)))
}
CompressionType::Gzip => {
use flate2::read::GzDecoder;
let mut decoder = GzDecoder::new(data);
let mut output = Vec::new();
decoder.read_to_end(&mut output)
.map_err(|e| Error::io(format!("Gzip decompression failed: {}", e)))?;
Ok(output)
}
CompressionType::Brotli => {
use brotli::BrotliDecompress;
let mut output = Vec::new();
BrotliDecompress(
&mut std::io::Cursor::new(data),
&mut output,
).map_err(|e| Error::io(format!("Brotli decompression failed: {}", e)))?;
Ok(output)
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_metadata_serialization() {
let metadata = DumpMetadata::new(100, CompressionType::Zstd, false);
let mut buffer = Vec::new();
metadata.write_to(&mut buffer).unwrap();
let mut cursor = std::io::Cursor::new(buffer);
let deserialized = DumpMetadata::read_from(&mut cursor).unwrap();
assert_eq!(metadata.lsn, deserialized.lsn);
assert_eq!(metadata.compression, deserialized.compression);
assert_eq!(metadata.incremental, deserialized.incremental);
}
#[test]
fn test_compression_roundtrip() {
let data = b"Hello, World! This is test data for compression.";
for compression in [
CompressionType::None,
CompressionType::Zstd,
CompressionType::Gzip,
CompressionType::Brotli,
] {
let compressed = DumpFormat::compress(data, compression).unwrap();
let decompressed = DumpFormat::decompress(&compressed, compression).unwrap();
assert_eq!(data.as_ref(), decompressed.as_slice());
}
}
}