archive_it_client/
wasapi.rs

1use std::path::{Path, PathBuf};
2use std::pin::pin;
3
4use aws_sdk_s3::Client as AwsS3Client;
5use futures_core::Stream;
6use futures_util::{StreamExt, TryStreamExt};
7use http_ferry::Downloader;
8use http_ferry::local::{LocalDir, LocalPath};
9use http_ferry::s3::{S3Dest, S3Location};
10use serde::Serialize;
11use url::Url;
12
13use crate::http::{Transport, header_map};
14use crate::models::wasapi::{Page, WasapiFile};
15use crate::{Config, DownloadOutcome, Error, USER_AGENT};
16
17pub const PRIMARY_LOCATION_SRC: &str = "https://warcs.archive-it.org";
18pub const DEFAULT_WEBDATA_PAGE_SIZE: u32 = 50;
19
20#[derive(Debug, Clone, Serialize)]
21pub struct WebdataQuery {
22    #[serde(skip_serializing_if = "Option::is_none")]
23    pub filename: Option<String>,
24    #[serde(skip_serializing_if = "Option::is_none")]
25    pub filetype: Option<String>,
26    #[serde(skip_serializing_if = "Option::is_none")]
27    pub collection: Option<u64>,
28    #[serde(skip_serializing_if = "Option::is_none")]
29    pub crawl: Option<u64>,
30    #[serde(rename = "crawl-time-after", skip_serializing_if = "Option::is_none")]
31    pub crawl_time_after: Option<String>,
32    #[serde(rename = "crawl-time-before", skip_serializing_if = "Option::is_none")]
33    pub crawl_time_before: Option<String>,
34    #[serde(rename = "crawl-start-after", skip_serializing_if = "Option::is_none")]
35    pub crawl_start_after: Option<String>,
36    #[serde(rename = "crawl-start-before", skip_serializing_if = "Option::is_none")]
37    pub crawl_start_before: Option<String>,
38    #[serde(skip_serializing_if = "Option::is_none")]
39    pub page: Option<u32>,
40    #[serde(skip_serializing_if = "Option::is_none")]
41    pub page_size: Option<u32>,
42}
43
44impl Default for WebdataQuery {
45    fn default() -> Self {
46        Self {
47            filename: None,
48            filetype: None,
49            collection: None,
50            crawl: None,
51            crawl_time_after: None,
52            crawl_time_before: None,
53            crawl_start_after: None,
54            crawl_start_before: None,
55            page: None,
56            page_size: Some(DEFAULT_WEBDATA_PAGE_SIZE),
57        }
58    }
59}
60
61pub struct WasapiClient {
62    transport: Transport,
63    downloader: Downloader,
64    primary_location_src: String,
65}
66
67impl WasapiClient {
68    pub fn new(username: impl Into<String>, password: impl Into<String>) -> Result<Self, Error> {
69        Self::with_config(username, password, Config::wasapi())
70    }
71
72    pub fn with_config(
73        username: impl Into<String>,
74        password: impl Into<String>,
75        cfg: Config,
76    ) -> Result<Self, Error> {
77        let creds = (username.into(), password.into());
78        let transport = Transport::new(cfg.clone(), Some(creds.clone()))?;
79        let download_client = reqwest::Client::builder()
80            .user_agent(USER_AGENT)
81            .default_headers(header_map(&cfg.headers)?)
82            .read_timeout(cfg.download_timeout)
83            .build()?;
84        let downloader = Downloader::builder(download_client)
85            .max_attempts(cfg.max_attempts)
86            .backoff(cfg.backoff)
87            .customize_request(move |req| req.basic_auth(&creds.0, Some(&creds.1)))
88            .build();
89        Ok(Self {
90            transport,
91            downloader,
92            primary_location_src: PRIMARY_LOCATION_SRC.into(),
93        })
94    }
95
96    pub fn with_primary_location_src(mut self, src: impl Into<String>) -> Self {
97        self.primary_location_src = src.into();
98        self
99    }
100
101    pub fn download(
102        &self,
103        file: WasapiFile,
104        path: impl AsRef<Path>,
105    ) -> impl Stream<Item = DownloadOutcome> + Send + '_ {
106        let path = path.as_ref().to_path_buf();
107        http_ferry::drive(
108            &self.downloader,
109            single_file(file).map_err(to_ferry),
110            self.resolver(),
111            LocalPath { path },
112        )
113    }
114
115    pub fn download_collection(
116        &self,
117        query: WebdataQuery,
118        dir: impl Into<PathBuf>,
119    ) -> impl Stream<Item = DownloadOutcome> + Send + '_ {
120        let dir = dir.into();
121        async_stream::stream! {
122            // Preflight the destination once so a bad output path fails the
123            // stream upfront instead of yielding one Failed per file. Also
124            // ensures the directory exists when the collection is empty.
125            if let Err(error) = tokio::fs::create_dir_all(&dir).await {
126                yield DownloadOutcome::StreamFailed {
127                    error: http_ferry::Error::Io(error),
128                };
129                return;
130            }
131            let inner = http_ferry::drive(
132                &self.downloader,
133                self.webdata(query).map_err(to_ferry),
134                self.resolver(),
135                LocalDir { dir },
136            );
137            let mut inner = pin!(inner);
138            while let Some(outcome) = inner.next().await {
139                yield outcome;
140            }
141        }
142    }
143
144    pub fn download_to_s3(
145        &self,
146        file: WasapiFile,
147        s3: AwsS3Client,
148        bucket: impl Into<String>,
149        prefix: Option<String>,
150    ) -> impl Stream<Item = DownloadOutcome<S3Location>> + Send + '_ {
151        http_ferry::drive(
152            &self.downloader,
153            single_file(file).map_err(to_ferry),
154            self.resolver(),
155            S3Dest {
156                client: s3,
157                bucket: bucket.into(),
158                prefix,
159            },
160        )
161    }
162
163    pub fn download_collection_to_s3(
164        &self,
165        query: WebdataQuery,
166        s3: AwsS3Client,
167        bucket: impl Into<String>,
168        prefix: Option<String>,
169    ) -> impl Stream<Item = DownloadOutcome<S3Location>> + Send + '_ {
170        http_ferry::drive(
171            &self.downloader,
172            self.webdata(query).map_err(to_ferry),
173            self.resolver(),
174            S3Dest {
175                client: s3,
176                bucket: bucket.into(),
177                prefix,
178            },
179        )
180    }
181
182    pub async fn list_webdata(&self, query: &WebdataQuery) -> Result<Page<WasapiFile>, Error> {
183        self.transport.get_json("webdata", query).await
184    }
185
186    pub async fn list_webdata_next(
187        &self,
188        prev: &Page<WasapiFile>,
189    ) -> Result<Option<Page<WasapiFile>>, Error> {
190        match prev.next.as_deref() {
191            None => Ok(None),
192            Some(url) => self.transport.get_json(url, &()).await.map(Some),
193        }
194    }
195
196    pub fn primary_location<'a>(&self, file: &'a WasapiFile) -> Option<&'a str> {
197        file.locations
198            .iter()
199            .find(|location| location.starts_with(&self.primary_location_src))
200            .map(String::as_str)
201    }
202
203    pub fn webdata(
204        &self,
205        query: WebdataQuery,
206    ) -> impl Stream<Item = Result<WasapiFile, Error>> + Send + '_ {
207        async_stream::try_stream! {
208            let mut page = self.list_webdata(&query).await?;
209            loop {
210                let files = std::mem::take(&mut page.files);
211                for f in files { yield f; }
212                match self.list_webdata_next(&page).await? {
213                    Some(next) => page = next,
214                    None => break,
215                }
216            }
217        }
218    }
219
220    /// Per-file resolver from a WASAPI item to its primary WARC URL. Runs
221    /// in-loop inside the engine so a file with no primary location yields a
222    /// non-fatal `Failed` rather than tearing down the stream. Host errors are
223    /// mapped into the engine's error type via [`to_ferry`].
224    fn resolver(&self) -> impl FnMut(&WasapiFile) -> Result<Url, http_ferry::Error> + Send + '_ {
225        let primary = self.primary_location_src.as_str();
226        move |file| primary_location_url(primary, file).map_err(to_ferry)
227    }
228}
229
230/// Resolve the WARC URL for `file` from its `locations`, picking the one under
231/// the configured primary source prefix.
232fn primary_location_url(primary_src: &str, file: &WasapiFile) -> Result<Url, Error> {
233    let location = file
234        .locations
235        .iter()
236        .find(|loc| loc.starts_with(primary_src))
237        .ok_or_else(|| Error::PrimaryLocationMissing {
238            filename: file.filename.clone(),
239        })?;
240    Ok(Url::parse(location)?)
241}
242
243fn single_file(file: WasapiFile) -> impl Stream<Item = Result<WasapiFile, Error>> + Send {
244    async_stream::try_stream! { yield file; }
245}
246
247/// Map a host error into the engine's error type at the resolver / item-stream
248/// boundary. Shared infrastructure variants pass through 1:1 so callers can
249/// still match on them (e.g. `http_ferry::Error::Status`); WASAPI-specific
250/// errors like `PrimaryLocationMissing` are boxed into
251/// [`http_ferry::Error::Source`].
252fn to_ferry(err: Error) -> http_ferry::Error {
253    use http_ferry::Error as F;
254    match err {
255        Error::Io(e) => F::Io(e),
256        Error::Request(e) => F::Request(e),
257        Error::Url(e) => F::Url(e),
258        Error::Status(s) => F::Status(s),
259        Error::NotFound(s) => F::NotFound(s),
260        other => F::from_source(other),
261    }
262}
263
264#[cfg(test)]
265mod tests {
266    use super::*;
267    use crate::models::wasapi::Checksums;
268
269    #[test]
270    fn primary_location_uses_configured_source_prefix() {
271        let client = WasapiClient::new("u", "p")
272            .unwrap()
273            .with_primary_location_src("https://example.invalid");
274        let file = WasapiFile {
275            filename: "ARCHIVEIT-1.warc.gz".into(),
276            filetype: "warc".into(),
277            checksums: Checksums {
278                sha1: Some("sha1".into()),
279                md5: Some("md5".into()),
280            },
281            account: 1,
282            size: 1,
283            collection: 1,
284            crawl: Some(1),
285            crawl_time: Some("2020-01-01T00:00:00Z".into()),
286            crawl_start: Some("2020-01-01T00:00:00Z".into()),
287            store_time: "2020-01-01T00:00:00Z".into(),
288            locations: vec![
289                "https://other.example.com/warcs/foo.warc.gz".into(),
290                "https://example.invalid/warcs/foo.warc.gz".into(),
291            ],
292        };
293
294        assert_eq!(
295            client.primary_location(&file),
296            Some("https://example.invalid/warcs/foo.warc.gz")
297        );
298    }
299}
archive_it_client/wasapi.rs

archive_it_client/
wasapi.rs
