feat: add EtlConfig as well as setting the directory to datadir (#7124)

Co-authored-by: Mikhail Sozin <mikhail.sozin@chainstack.com>
Co-authored-by: Misha <mikawamp@gmail.com>
Co-authored-by: Alexey Shekhirin <a.shekhirin@gmail.com>
This commit is contained in:
joshieDo
2024-03-13 16:06:50 +00:00
committed by GitHub
parent 5d6ac4c815
commit 28f3a2e2d9
18 changed files with 118 additions and 43 deletions

View File

@ -18,7 +18,7 @@ use std::{
cmp::Reverse,
collections::BinaryHeap,
io::{self, BufReader, BufWriter, Read, Seek, SeekFrom, Write},
path::Path,
path::{Path, PathBuf},
};
use rayon::prelude::*;
@ -42,6 +42,8 @@ where
<K as Encode>::Encoded: std::fmt::Debug,
<V as Compress>::Compressed: std::fmt::Debug,
{
/// Parent directory where to create ETL files
parent_dir: Option<PathBuf>,
/// Directory for temporary file storage
dir: Option<TempDir>,
/// Collection of temporary ETL files
@ -66,8 +68,9 @@ where
/// Create a new collector with some capacity.
///
/// Once the capacity (in bytes) is reached, the data is sorted and flushed to disk.
pub fn new(buffer_capacity_bytes: usize) -> Self {
pub fn new(buffer_capacity_bytes: usize, parent_dir: Option<PathBuf>) -> Self {
Self {
parent_dir,
dir: None,
buffer_size_bytes: 0,
files: Vec::new(),
@ -115,7 +118,15 @@ where
/// doesn't exist, it will be created.
fn dir(&mut self) -> io::Result<&TempDir> {
if self.dir.is_none() {
self.dir = Some(TempDir::new()?);
self.dir = match &self.parent_dir {
Some(dir) => {
if !dir.exists() {
std::fs::create_dir_all(dir)?;
}
Some(TempDir::new_in(dir)?)
}
None => Some(TempDir::new()?),
};
}
Ok(self.dir.as_ref().unwrap())
}
@ -273,7 +284,7 @@ mod tests {
let mut entries: Vec<_> =
(0..10_000).map(|id| (TxHash::random(), id as TxNumber)).collect();
let mut collector = Collector::new(1024);
let mut collector = Collector::new(1024, None);
assert!(collector.dir.is_none());
for (k, v) in entries.clone() {