db-ingest: Add zstd compression

OpenSource-Deadlock-Tools · Oct 9, 2024 · fecbfbd · fecbfbd
1 parent feb1db6
commit fecbfbd
Show file tree

Hide file tree

Showing 5 changed files with 66 additions and 5 deletions.
diff --git a/db-ingest/Cargo.lock b/db-ingest/Cargo.lock
diff --git a/db-ingest/Cargo.toml b/db-ingest/Cargo.toml
@@ -10,7 +10,7 @@ rust-s3 = "0.35.1"
 tokio = { version = "1.40.0", features = ["rt-multi-thread", "macros"] }
 log = "0.4.22"
 futures-lite = "2.3.0"
-async-compression = { version = "0.4.13", features = ["bzip2", "tokio"] }
+async-compression = { version = "0.4.13", features = ["bzip2", "tokio", "zstd"] }
 bytes = "1.7.2"
 valveprotos = { git = "https://github.com/OpenSource-Deadlock-Tools/valveprotos-rs", rev = "71ab6d7de2cd43f567397f65821efd80f5aa0b71", features = ["deadlock"] }
 prost = "0.13.3"

diff --git a/db-ingest/src/models/compression.rs b/db-ingest/src/models/compression.rs
@@ -1,6 +1,6 @@
 use crate::models::error::ParseError;
-use async_compression::tokio::bufread::BzDecoder;
-use async_compression::tokio::write::BzEncoder;
+use async_compression::tokio::bufread::{BzDecoder, ZstdDecoder};
+use async_compression::tokio::write::{BzEncoder, ZstdEncoder};
 use std::fmt::Display;
 use std::str::FromStr;
 use std::{vec, write};
@@ -11,6 +11,7 @@ pub enum Compression {
     #[default]
     Uncompressed,
     Bzip2,
+    Zstd,
 }
 
 impl Compression {
@@ -25,6 +26,14 @@ impl Compression {
                     .map_err(ParseError::Decompress)?;
                 Ok(decompressed)
             }
+            Self::Zstd => {
+                let mut decompressed = vec![];
+                ZstdDecoder::new(data.as_ref())
+                    .read_to_end(&mut decompressed)
+                    .await
+                    .map_err(ParseError::Decompress)?;
+                Ok(decompressed)
+            }
         }
     }
     pub async fn compress(&self, data: &[u8]) -> Result<Vec<u8>, ParseError> {
@@ -38,6 +47,14 @@ impl Compression {
                     .map_err(ParseError::Decompress)?;
                 Ok(encoder.into_inner())
             }
+            Self::Zstd => {
+                let mut encoder = ZstdEncoder::new(Vec::new());
+                encoder
+                    .write_all(data)
+                    .await
+                    .map_err(ParseError::Decompress)?;
+                Ok(encoder.into_inner())
+            }
         }
     }
 }
@@ -47,6 +64,7 @@ impl Display for Compression {
         match self {
             Self::Uncompressed => write!(f, ""),
             Self::Bzip2 => write!(f, "bz2"),
+            Self::Zstd => write!(f, "zstd"),
         }
     }
 }

diff --git a/db-ingest/src/parsers/metadata_content_parser.rs b/db-ingest/src/parsers/metadata_content_parser.rs
@@ -1,3 +1,4 @@
+use crate::models::compression::Compression;
 use crate::models::error::ParseError;
 use crate::models::file_data::FileData;
 use crate::models::file_type::FileType;
@@ -23,7 +24,7 @@ impl Parser<MatchInfo> for MetaDataContentParser {
         Ok(ParseResult {
             file_type: FileType::MetadataContent,
             data: data.to_vec(),
-            compression: file_data.compression,
+            compression: Compression::Zstd,
             parsed_data,
         })
     }

diff --git a/db-ingest/src/parsers/metadata_parser.rs b/db-ingest/src/parsers/metadata_parser.rs
@@ -1,6 +1,7 @@
 use crate::parsers::parser::Parser;
 use prost::Message;
 
+use crate::models::compression::Compression;
 use crate::models::error::ParseError;
 use crate::models::file_data::FileData;
 use crate::models::file_type::FileType;
@@ -28,7 +29,7 @@ impl Parser<MatchInfo> for MetaDataParser {
         Ok(ParseResult {
             file_type: FileType::MetadataContent,
             data: match_details.to_vec(),
-            compression: file_data.compression,
+            compression: Compression::Zstd,
             parsed_data,
         })
     }