use std::{ io::{BufWriter, Write}, path::PathBuf, }; use anyhow::Result; use futures_util::StreamExt; use kordophone::APIInterface; use thiserror::Error; use tokio::pin; mod target { pub static ATTACHMENTS: &str = "attachments"; } #[derive(Debug, Clone)] pub struct Attachment { pub guid: String, pub path: PathBuf, pub downloaded: bool, } #[derive(Debug, Error)] enum AttachmentStoreError { #[error("attachment has already been downloaded")] AttachmentAlreadyDownloaded, #[error("Client error: {0}")] APIClientError(String), } pub struct AttachmentStore { store_path: PathBuf, } impl AttachmentStore { pub fn new(data_dir: &PathBuf) -> AttachmentStore { let store_path = data_dir.join("attachments"); log::info!(target: target::ATTACHMENTS, "Attachment store path: {}", store_path.display()); // Create the attachment store if it doesn't exist std::fs::create_dir_all(&store_path) .expect("Wasn't able to create the attachment store path"); AttachmentStore { store_path: store_path, } } pub fn get_attachment(&self, guid: &String) -> Attachment { let path = self.store_path.join(guid); let path_exists = std::fs::exists(&path).expect( format!( "Wasn't able to check for the existence of an attachment file path at {}", &path.display() ) .as_str(), ); Attachment { guid: guid.to_owned(), path: path, downloaded: path_exists, } } pub async fn download_attachment( &mut self, attachment: &Attachment, mut client_factory: F, ) -> Result<()> where C: APIInterface, F: FnMut() -> Fut, Fut: std::future::Future>, { if attachment.downloaded { log::info!(target: target::ATTACHMENTS, "Attachment already downloaded: {}", attachment.guid); return Err(AttachmentStoreError::AttachmentAlreadyDownloaded.into()); } log::info!(target: target::ATTACHMENTS, "Starting download for attachment: {}", attachment.guid); // Create temporary file first, we'll atomically swap later. assert!(!std::fs::exists(&attachment.path).unwrap()); let file = std::fs::File::create(&attachment.path)?; let mut writer = BufWriter::new(&file); log::trace!(target: target::ATTACHMENTS, "Created attachment file at {}", &attachment.path.display()); let mut client = client_factory().await?; let stream = client .fetch_attachment_data(&attachment.guid) .await .map_err(|e| AttachmentStoreError::APIClientError(format!("{:?}", e)))?; // Since we're async, we need to pin this. pin!(stream); log::trace!(target: target::ATTACHMENTS, "Writing attachment data to disk"); while let Some(Ok(data)) = stream.next().await { writer.write(data.as_ref())?; } log::info!(target: target::ATTACHMENTS, "Completed download for attachment: {}", attachment.guid); Ok(()) } /// Check if an attachment should be downloaded pub fn should_download(&self, attachment_id: &str) -> bool { let attachment = self.get_attachment(&attachment_id.to_string()); !attachment.downloaded } }