use super::resolve_archive_path;
use crate::{
processor::{DataSource, SourceContext, StreamingDataSource},
receiver::{RawResponse, Receive, ReceiveMultiple, ReceiveRaw},
};
use eyre::{Result, eyre};
use futures::stream::BoxStream;
use serde::de::DeserializeOwned;
use std::{
fs::File,
io::{BufReader, Read},
path::PathBuf,
sync::Arc,
sync::OnceLock,
time::SystemTime,
};
use tokio::sync::RwLock;
use zip::ZipArchive;
#[derive(Clone)]
pub struct ArchiveFileReceiver {
archive: Arc<RwLock<ZipArchive<File>>>,
filename: String,
subdir: Option<PathBuf>,
modified_date: SystemTime,
source_product: Arc<OnceLock<&'static str>>,
}
impl TryFrom<PathBuf> for ArchiveFileReceiver {
type Error = eyre::Report;
fn try_from(path: PathBuf) -> Result<Self> {
let filename = format!("{}", path.file_name().unwrap_or_default().display());
match path.is_file() {
true => {
tracing::debug!("File is valid: {}", path.display());
let file = File::open(path)?;
let modified_date = file.metadata()?.modified()?;
let archive = ZipArchive::new(file)?;
Ok(Self {
archive: Arc::new(RwLock::new(archive)),
modified_date,
filename,
subdir: None,
source_product: Arc::new(OnceLock::new()),
})
}
false => {
tracing::debug!("File is invalid: {}", path.display());
Err(eyre!("Archive input must be a file: {}", path.display()))
}
}
}
}
impl Receive for ArchiveFileReceiver {
async fn collection_date(&self) -> String {
chrono::DateTime::<chrono::Utc>::from(self.modified_date).to_rfc3339()
}
async fn is_connected(&self) -> bool {
let archive = self.archive.read().await;
let is_empty = archive.is_empty();
if tracing::enabled!(tracing::Level::TRACE) {
let file_names: Vec<String> = archive.file_names().map(|name| name.to_string()).collect();
tracing::trace!("Files in archive: {:?}", file_names);
}
tracing::debug!("Archive {} is valid: {}", &self.filename, !is_empty);
!is_empty
}
fn filename(&self) -> Option<String> {
Some(self.filename.clone())
}
async fn get<T>(&self) -> Result<T>
where
T: DeserializeOwned + DataSource,
{
let mut archive = self.archive.write().await;
let ctx = self.source_context()?;
let source_paths = T::candidate_source_file_paths(&ctx)?;
let mut last_resolve_error = None;
for source_path in source_paths {
match resolve_archive_path(self.subdir.as_ref(), &mut *archive, &source_path) {
Ok(filename) => {
tracing::debug!("Reading {}", filename);
let file = archive.by_name(&filename)?;
let reader = BufReader::new(file);
let data: T = serde_json::from_reader(reader)?;
return Ok(data);
}
Err(e) => {
last_resolve_error = Some(e);
continue;
}
}
}
match last_resolve_error {
Some(e) => Err(e),
None => Err(eyre!("No candidate source files available for {}", T::name())),
}
}
async fn get_stream<T>(&self) -> Result<BoxStream<'static, Result<T::Item>>>
where
T: StreamingDataSource + DeserializeOwned,
T::Item: DeserializeOwned + Send + 'static,
{
let ctx = self.source_context()?;
super::get_stream_from_archive::<File, T>(self.archive.clone(), self.subdir.clone(), ctx).await
}
}
impl ReceiveRaw for ArchiveFileReceiver {
async fn get_raw<T>(&self) -> Result<String>
where
T: DataSource,
{
self.get_raw_response::<T>().await.map(|response| response.body)
}
async fn get_raw_response<T>(&self) -> Result<RawResponse>
where
T: DataSource,
{
let mut archive = self.archive.write().await;
let ctx = self.source_context()?;
let source_paths = T::candidate_source_file_paths(&ctx)?;
let mut last_resolve_error = None;
for source_path in source_paths {
match resolve_archive_path(self.subdir.as_ref(), &mut *archive, &source_path) {
Ok(filename) => {
tracing::debug!("Reading {}", filename);
let file = archive.by_name(&filename)?;
let mut reader = BufReader::new(file);
let mut data = String::new();
reader.read_to_string(&mut data)?;
let response_size_bytes = data.len() as u64;
return Ok(RawResponse {
body: data,
status: None,
response_time_ms: 0,
response_size_bytes,
});
}
Err(e) => {
last_resolve_error = Some(e);
continue;
}
}
}
match last_resolve_error {
Some(e) => Err(e),
None => Err(eyre!("No candidate source files available for {}", T::name())),
}
}
}
impl ReceiveMultiple for ArchiveFileReceiver {
fn set_work_dir(&mut self, work_dir: &str) -> Result<()> {
tracing::trace!("Setting subdir: {}", work_dir);
self.subdir = Some(PathBuf::from(work_dir));
Ok(())
}
}
impl std::fmt::Display for ArchiveFileReceiver {
fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
write!(f, "{}", self.filename)
}
}
impl ArchiveFileReceiver {
pub(crate) fn clone_for_subdir(&self, work_dir: &str) -> Self {
Self {
archive: self.archive.clone(),
filename: self.filename.clone(),
subdir: Some(PathBuf::from(work_dir)),
modified_date: self.modified_date,
source_product: Arc::new(OnceLock::new()),
}
}
pub async fn read_bundle_json<T>(&self, filename: &str) -> Result<T>
where
T: DeserializeOwned,
{
let mut archive = self.archive.write().await;
let filename = resolve_archive_path(self.subdir.as_ref(), &mut *archive, filename)?;
tracing::debug!("Reading bundle file {}", filename);
let file = archive.by_name(&filename)?;
let reader = BufReader::new(file);
serde_json::from_reader(reader).map_err(Into::into)
}
pub fn set_source_product(&self, product: &'static str) -> Result<()> {
match self.source_product.get() {
Some(existing) if *existing != product => Err(eyre!(
"Archive receiver source product already set to {}, cannot change to {}",
existing,
product
)),
Some(_) => Ok(()),
None => self
.source_product
.set(product)
.map_err(|_| eyre!("Failed to initialize archive receiver source product")),
}
}
pub fn source_product(&self) -> Result<&'static str> {
self.source_product
.get()
.copied()
.ok_or_else(|| eyre!("Archive receiver source product is not initialized"))
}
pub fn source_context(&self) -> Result<SourceContext> {
Ok(SourceContext::new(self.source_product()?, None))
}
}