use std::{collections::HashMap, fs, path::PathBuf};
use crate::record::{Record, RecordSchema};
use crate::Error;
use std::collections::hash_map::Entry::{Occupied, Vacant};
pub trait RecordWriter: Send {
fn write_record(&mut self, record: &Record) -> std::io::Result<()>;
fn finish(&mut self) -> Result<(), Error> {
Ok(())
}
}
pub trait RecordWriterFactory: Send {
type Writer: RecordWriter;
fn make_writer(&mut self, schema: &RecordSchema) -> std::io::Result<Self::Writer>;
}
pub trait FileRecordWriterFactory: Send {
type Writer: RecordWriter;
fn file_name(&self, form_name: String) -> String;
fn make(&mut self, path: &PathBuf, schema: &RecordSchema) -> std::io::Result<Self::Writer>;
fn norm_form_name(&self, name: &str) -> String {
name.replace("/", "-")
}
}
pub struct MultiRecordWriter<F: RecordWriterFactory> {
factory: F,
pub writers: HashMap<RecordSchema, F::Writer>,
}
impl<F: RecordWriterFactory> MultiRecordWriter<F> {
pub fn new(factory: F) -> Self {
Self {
factory,
writers: HashMap::new(),
}
}
pub fn get_writer(&mut self, schema: &RecordSchema) -> std::io::Result<&mut F::Writer> {
Ok(match self.writers.entry(schema.clone()) {
Occupied(e) => e.into_mut(),
Vacant(e) => e.insert(self.factory.make_writer(schema)?),
})
}
}
impl<F: RecordWriterFactory> RecordWriter for MultiRecordWriter<F> {
fn write_record(&mut self, record: &Record) -> std::io::Result<()> {
let writer = self.get_writer(&record.schema)?;
writer.write_record(record)
}
fn finish(&mut self) -> Result<(), Error> {
for (_, writer) in self.writers.iter_mut() {
writer.finish()?;
}
Ok(())
}
}
pub struct MultiFileRecordWriterFactory<F: FileRecordWriterFactory> {
base_path: PathBuf,
factory: F,
}
impl<F: FileRecordWriterFactory> MultiFileRecordWriterFactory<F> {
pub fn new(base_path: PathBuf, factory: F) -> Self {
Self { base_path, factory }
}
}
impl<F: FileRecordWriterFactory> RecordWriterFactory for MultiFileRecordWriterFactory<F> {
type Writer = F::Writer;
fn make_writer(&mut self, schema: &RecordSchema) -> std::io::Result<F::Writer> {
let form_name = self.factory.norm_form_name(&schema.code);
let file_name = self.factory.file_name(form_name);
let path = self.base_path.join(file_name);
fs::create_dir_all(&self.base_path)?;
log::debug!("Creating new FileRecordWriter at: {:?}", path);
let result = self.factory.make(&path, schema)?;
Ok(result)
}
}