Last active
July 19, 2023 10:31
-
-
Save sibukixxx/31021b48b9c64230dcd48aac06edcd03 to your computer and use it in GitHub Desktop.
download.rs
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| use async_trait::async_trait; | |
| use reqwest::Client; | |
| use rusoto_core::Region; | |
| use rusoto_s3::{GetObjectRequest, S3Client, S3}; | |
| use std::fs::File; | |
| use std::io::{Read, Write}; | |
| use std::path::Path; | |
| use tokio::io::AsyncReadExt; | |
| use url::Url; | |
| use data_encoding::BASE64; | |
| // ... previous code ... | |
| pub struct S3DownloadClient { | |
| s3_client: S3Client, | |
| } | |
| #[async_trait] | |
| impl DownloadClient for S3DownloadClient { | |
| async fn match_uri(&self, uri: &str) -> bool { | |
| uri.starts_with("s3://") | |
| } | |
| async fn get(&self, uri: &str) -> Vec<u8> { | |
| let url = Url::parse(uri).unwrap(); | |
| let bucket = url.host_str().unwrap(); | |
| let key = &url[url::Position::BeforePath..]; | |
| let req = GetObjectRequest { | |
| bucket: bucket.to_string(), | |
| key: key.to_string(), | |
| ..Default::default() | |
| }; | |
| let result = self.s3_client.get_object(req).await.unwrap(); | |
| let mut buf = Vec::new(); | |
| result.body.unwrap().into_async_read().read_to_end(&mut buf).await.unwrap(); | |
| buf | |
| } | |
| async fn download(&self, uri: &str, target: &str) { | |
| let content = self.get(uri).await; | |
| let mut file = File::create(target).unwrap(); | |
| file.write_all(&content).unwrap(); | |
| } | |
| } | |
| pub struct DataUriDownloadClient; | |
| #[async_trait] | |
| impl DownloadClient for DataUriDownloadClient { | |
| async fn match_uri(&self, uri: &str) -> bool { | |
| uri.starts_with("data:") | |
| } | |
| async fn get(&self, uri: &str) -> Vec<u8> { | |
| let comma_index = uri.find(',').unwrap(); | |
| let data = &uri[(comma_index + 1)..]; | |
| BASE64.decode(data.as_bytes()).unwrap() | |
| } | |
| async fn download(&self, uri: &str, target: &str) { | |
| let content = self.get(uri).await; | |
| let mut file = File::create(target).unwrap(); | |
| file.write_all(&content).unwrap(); | |
| } | |
| } | |
| pub struct LocalFileDownloadClient; | |
| #[async_trait] | |
| impl DownloadClient for LocalFileDownloadClient { | |
| async fn match_uri(&self, uri: &str) -> bool { | |
| Path::new(uri).exists() | |
| } | |
| async fn get(&self, uri: &str) -> Vec<u8> { | |
| let mut file = File::open(uri).unwrap(); | |
| let mut content = Vec::new(); | |
| file.read_to_end(&mut content).unwrap(); | |
| content | |
| } | |
| async fn download(&self, uri: &str, target: &str) { | |
| let content = self.get(uri).await; | |
| let mut file = File::create(target).unwrap(); | |
| file.write_all(&content).unwrap(); | |
| } | |
| } | |
| pub struct UniversalDownloadClient { | |
| clients: Vec<Box<dyn DownloadClient>>, | |
| } | |
| impl UniversalDownloadClient { | |
| pub fn new() -> UniversalDownloadClient { | |
| let clients: Vec<Box<dyn DownloadClient>> = vec![ | |
| Box::new(HttpDownloadClient { http_client: Client::new() }), | |
| Box::new(S3DownloadClient { s3_client: S3Client::new(Region::default()) }), | |
| Box::new(DataUriDownloadClient), | |
| Box::new(LocalFileDownloadClient), | |
| ]; | |
| UniversalDownloadClient { clients } | |
| } | |
| async fn get_client(&self, uri: &str) -> Result<&dyn DownloadClient, Box<dyn Error>> { | |
| for client in &self.clients { | |
| if client.match_uri(uri).await { | |
| return Ok(client.as_ref()); | |
| } | |
| } | |
| Err(Box::new(std::io::Error::new(std::io::ErrorKind::Other, format!("Unsupported URI: {}", uri)))) | |
| } | |
| pub async fn get(&self, uri: &str) -> Result<Vec<u8>, Box<dyn Error>> { | |
| let client = self.get_client(uri).await?; | |
| client.get(uri).await | |
| } | |
| pub async fn download(&self, uri: &str, target: &str) -> Result<(), Box<dyn Error>> { | |
| let client = self.get_client(uri).await?; | |
| client.download(uri, target).await | |
| } | |
| } |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment