openapi_rs/api/v1/storage/
download.rs

1use crate::common::define::{
2    AsyncResponseFn, BaseRequest, BytesStream, HttpBuilder, HttpFn, HttpStreamBuilder, RequestFn,
3};
4use bytes::Bytes;
5use regex::Regex;
6use reqwest::{Method, Response};
7use serde::{Deserialize, Serialize};
8use std::collections::HashMap;
9
10#[derive(Debug, Default, Clone, Serialize, Deserialize)]
11#[serde(default)]
12pub struct DownloadRequest {
13    #[serde(rename = "Path")]
14    pub path: Option<String>,
15    #[serde(rename = "RangeStart")]
16    pub range_start: Option<isize>,
17    #[serde(rename = "RangeEnd")]
18    pub range_end: Option<isize>,
19}
20
21impl DownloadRequest {
22    pub fn new() -> Self {
23        Default::default()
24    }
25    pub fn with_path(mut self, path: String) -> Self {
26        self.path = Some(path);
27        self
28    }
29    pub fn with_range_start(mut self, range_start: isize) -> Self {
30        self.range_start = Some(range_start);
31        self
32    }
33    pub fn with_range_end(mut self, range_end: isize) -> Self {
34        self.range_end = Some(range_end);
35        self
36    }
37
38    fn request_fn(self) -> RequestFn {
39        let request_fn: RequestFn = Box::new(move || {
40            let mut queries = HashMap::new();
41            if let Some(path) = &self.path {
42                queries.insert("Path".to_string(), path.clone());
43            }
44            if let Some(range_start) = self.range_start {
45                if let Some(range_end) = self.range_end {
46                    queries.insert(
47                        "Range".to_string(),
48                        format!("bytes={}-{}", range_start, range_end),
49                    );
50                }
51            }
52            BaseRequest {
53                method: Method::GET,
54                uri: "/api/storage/download".to_string(),
55                queries: Some(queries),
56                ..Default::default()
57            }
58        });
59        request_fn
60    }
61}
62
63#[derive(Debug, Default, Clone, Serialize, Deserialize)]
64#[serde(default)]
65pub struct DownloadResponse {
66    #[serde(rename = "FileName")]
67    pub file_name: String,
68    #[serde(rename = "FileType")]
69    pub file_type: String,
70    #[serde(rename = "FileSize")]
71    pub file_size: isize,
72    #[serde(rename = "Data", skip)]
73    pub data: Option<Bytes>,
74}
75
76impl HttpBuilder for DownloadRequest {
77    type Response = DownloadResponse;
78
79    fn builder(self) -> HttpFn<Self::Response> {
80        Box::new(move || {
81            let response_fn: AsyncResponseFn<Self::Response> = Box::new(|response: Response| {
82                Box::pin(async move {
83                    let mut download_response = DownloadResponse::default();
84                    let file_name_regex = Regex::new(r#"attachment; filename="(.*?)""#)?;
85                    download_response.file_name = response
86                        .headers()
87                        .get("Content-Disposition")
88                        .and_then(|v| v.to_str().ok())
89                        .and_then(|s| {
90                            file_name_regex
91                                .captures(s)
92                                .and_then(|caps| caps.get(1))
93                                .map(|m| m.as_str().to_owned())
94                        })
95                        .unwrap();
96                    download_response.file_type = response
97                        .headers()
98                        .get("Content-Type")
99                        .unwrap()
100                        .to_str()?
101                        .to_owned();
102                    download_response.file_size = response
103                        .headers()
104                        .get("Content-Length")
105                        .unwrap()
106                        .to_str()?
107                        .parse::<isize>()?;
108                    download_response.data = Some(response.bytes().await?);
109
110                    Ok(download_response)
111                })
112            });
113            (self.request_fn(), response_fn)
114        })
115    }
116}
117
118#[derive(derive_more::Debug, Default)]
119pub struct DownloadStreamResponse {
120    #[debug(skip)]
121    pub stream: Option<BytesStream>,
122}
123
124impl HttpStreamBuilder for DownloadRequest {
125    type Response = DownloadStreamResponse;
126
127    fn stream_builder(self) -> HttpFn<Self::Response> {
128        Box::new(move || {
129            let response_fn: AsyncResponseFn<Self::Response> = Box::new(|response: Response| {
130                Box::pin(async move {
131                    Ok(DownloadStreamResponse {
132                        stream: Some(Box::pin(response.bytes_stream())),
133                    })
134                })
135            });
136            (self.request_fn(), response_fn)
137        })
138    }
139}
140
141#[cfg(test)]
142mod tests {
143    use super::*;
144    use crate::common::client::OpenApiClient;
145    use crate::common::config::{EndpointType, OpenApiConfig};
146    use futures_util::stream::StreamExt;
147    use tracing::info;
148
149    #[tokio::test]
150    async fn test_download() -> anyhow::Result<()> {
151        tracing_subscriber::fmt::init();
152        dotenvy::dotenv()?;
153        let config = OpenApiConfig::new().load_from_env()?;
154        let user_id = config.user_id.clone();
155        let mut client = OpenApiClient::new(config).with_endpoint_type(EndpointType::Cloud);
156
157        let http_fn = DownloadRequest::new()
158            .with_path(format!("/{}/runner.py", user_id))
159            .builder();
160        let response = client.send(http_fn).await?;
161        info!("response: {:#?}", response);
162
163        Ok(())
164    }
165
166    #[tokio::test]
167    async fn test_download_stream() -> anyhow::Result<()> {
168        tracing_subscriber::fmt::init();
169        dotenvy::dotenv()?;
170        let config = OpenApiConfig::new().load_from_env()?;
171        let user_id = config.user_id.clone();
172        let mut client = OpenApiClient::new(config).with_endpoint_type(EndpointType::Cloud);
173
174        let http_fn = DownloadRequest::new()
175            .with_path(format!("/{}/runner.py", user_id))
176            .stream_builder();
177        let mut response = client.send(http_fn).await?;
178        while let Some(data) = response
179            .stream
180            .as_mut()
181            .expect("stream not found")
182            .next()
183            .await
184        {
185            info!("data: {:#?}", data?);
186        }
187
188        Ok(())
189    }
190}
openapi_rs/api/v1/storage/download.rs

openapi_rs/api/v1/storage/
download.rs
