use super::{
cdx::{self, IndexClient},
digest::compute_digest,
downloader::Downloader,
Item,
};
use bytes::Buf;
use chrono::Utc;
use csv::{ReaderBuilder, WriterBuilder};
use flate2::{Compression, GzBuilder};
use futures::{StreamExt, TryStreamExt};
use std::collections::HashSet;
use std::fs::{create_dir_all, File};
use std::io::{BufRead, BufReader, Read, Write};
use std::path::{Path, PathBuf};
#[derive(thiserror::Error, Debug)]
pub enum Error {
#[error("I/O error: {0:?}")]
Io(#[from] std::io::Error),
#[error("CDX error: {0:?}")]
IndexClient(#[from] cdx::Error),
#[error("HTTP client error: {0:?}")]
Client(#[from] reqwest::Error),
#[error("CSV writing error: {0:?}")]
Csv(#[from] csv::Error),
#[error("Item parsing error: {0:?}")]
Item(#[from] super::item::Error),
}
pub struct Session {
base: PathBuf,
known_digests: Option<PathBuf>,
parallelism: usize,
index_client: IndexClient,
client: Downloader,
}
impl Session {
const TIMESTAMP_FMT: &'static str = "%Y%m%d%H%M%S";
pub fn new<P1: AsRef<Path>, P2: AsRef<Path>>(
base: P1,
known_digests: Option<P2>,
parallelism: usize,
) -> Result<Session, Error> {
Ok(Session {
base: base.as_ref().to_path_buf(),
known_digests: known_digests.map(|path| path.as_ref().to_path_buf()),
parallelism,
index_client: IndexClient::default(),
client: Downloader::default(),
})
}
pub fn new_timestamped<P: AsRef<Path>>(
known_digests: Option<P>,
parallelism: usize,
) -> Result<Session, Error> {
Self::new(
Utc::now().format(Self::TIMESTAMP_FMT).to_string(),
known_digests,
parallelism,
)
}
pub async fn save_cdx_results(&self, queries: &[String]) -> Result<(), Error> {
create_dir_all(&self.base)?;
let mut query_log = File::create(self.base.join("queries.txt"))?;
query_log.write_all(format!("{}\n", queries.join("\n")).as_bytes())?;
let results: Vec<Result<Vec<Item>, String>> = futures::stream::iter(queries.iter())
.map(|query| Ok(self.index_client.search(query, None, None)))
.try_buffer_unordered(self.parallelism)
.map(|result| match result {
Err(cdx::Error::BlockedQuery(query)) => Ok(Err(query)),
Err(other) => Err(other),
Ok(items) => Ok(Ok(items)),
})
.err_into::<Error>()
.try_collect()
.await?;
let mut blocked: Vec<String> = vec![];
let mut items: Vec<Item> = Vec::with_capacity(results.len());
for result in results {
match result {
Ok(batch) => items.extend(batch),
Err(query) => blocked.push(query),
}
}
if !blocked.is_empty() {
let mut blocked_log = File::create(self.base.join("blocked.txt"))?;
blocked.sort();
blocked_log.write_all(format!("{}\n", blocked.join("\n")).as_bytes())?;
}
items.sort();
items.dedup();
let originals_item_log = File::create(self.base.join("originals.csv"))?;
let redirects_item_log = File::create(self.base.join("redirects.csv"))?;
let mut originals_csv = WriterBuilder::new().from_writer(originals_item_log);
let mut redirects_csv = WriterBuilder::new().from_writer(redirects_item_log);
for item in &items {
if item.status == Some(302) {
redirects_csv.write_record(item.to_record())?;
} else {
originals_csv.write_record(item.to_record())?;
}
}
Ok(())
}
pub async fn resolve_redirects(&self) -> Result<(), Error> {
let redirects_item_log = File::open(self.base.join("redirects.csv"))?;
let mut items = Self::read_csv(redirects_item_log)?;
items.sort();
create_dir_all(self.base.join("data"))?;
create_dir_all(self.base.join("invalid"))?;
let mut digests = HashSet::new();
items.retain(|item| digests.insert(item.digest.clone()));
if let Some(path) = &self.known_digests {
let file = File::open(path)?;
for line in BufReader::new(file).lines() {
digests.remove(line?.trim());
}
}
items.retain(|item| digests.remove(&item.digest));
log::info!("Resolving {} items", items.len());
let results = futures::stream::iter(items.iter())
.map(|item| async move {
log::info!("Resolving: {}", item.url);
(
item,
self.client
.resolve_redirect(&item.url, &item.timestamp(), &item.digest)
.await,
)
})
.buffer_unordered(self.parallelism)
.map(|(item, result)| async move {
let resolution = result.map_err(|_| item)?;
if resolution.valid_digest {
let mut items = self
.index_client
.search(&resolution.url, Some(&resolution.timestamp), None)
.await
.map_err(|_| item)?;
let actual_item = items.pop().ok_or(item)?;
let output =
File::create(self.base.join("data").join(format!("{}.gz", item.digest)))
.map_err(|_| item)?;
let mut gz = GzBuilder::new()
.filename(item.make_filename())
.write(output, Compression::default());
gz.write_all(&resolution.content).map_err(|_| item)?;
gz.finish().map_err(|_| item)?;
Ok(actual_item)
} else {
Err(item)
}
})
.buffer_unordered(self.parallelism)
.collect::<Vec<_>>()
.await;
create_dir_all(self.base.join("errors"))?;
let redirects_error_log = File::create(self.base.join("errors").join("redirects.csv"))?;
let mut redirects_error_csv = WriterBuilder::new().from_writer(redirects_error_log);
let extras_item_log = File::create(self.base.join("extras.csv"))?;
let mut extras_item_csv = WriterBuilder::new().from_writer(extras_item_log);
for result in results {
match result {
Ok(item) => {
extras_item_csv.write_record(item.to_record())?;
}
Err(item) => {
redirects_error_csv.write_record(item.to_record())?;
}
}
}
Ok(())
}
pub async fn download_items(&self) -> Result<(usize, usize, usize, usize), Error> {
let originals_file = File::open(self.base.join("originals.csv"))?;
let mut items = Self::read_csv(originals_file)?;
let extras_file = File::open(self.base.join("extras.csv"))?;
items.extend(Self::read_csv(extras_file)?);
items.sort();
let total_count = items.len();
let mut digests = HashSet::new();
items.retain(|item| digests.insert(item.digest.clone()));
if let Some(path) = &self.known_digests {
let file = File::open(path)?;
for line in BufReader::new(file).lines() {
digests.remove(line?.trim());
}
}
items.retain(|item| digests.remove(&item.digest));
log::info!("Downloading {} items", items.len());
let results = futures::stream::iter(items)
.map(|item| async {
let content = self
.client
.download_item(&item)
.await
.map_err(|_| item.clone())?;
let expected = item.digest.clone();
let computed = compute_digest(&mut content.clone().reader()).unwrap();
if computed == expected {
let output =
File::create(self.base.join("data").join(format!("{}.gz", expected)))
.map_err(|_| item.clone())?;
let mut gz = GzBuilder::new()
.filename(item.make_filename())
.write(output, Compression::default());
gz.write_all(&content).map_err(|_| item.clone())?;
gz.finish().map_err(|_| item)?;
Ok(None)
} else {
let output =
File::create(self.base.join("invalid").join(format!("{}.gz", computed)))
.map_err(|_| item.clone())?;
let mut gz = GzBuilder::new()
.filename(item.make_filename())
.write(output, Compression::default());
gz.write_all(&content).map_err(|_| item.clone())?;
gz.finish().map_err(|_| item)?;
Ok(Some((expected, computed)))
}
})
.buffer_unordered(self.parallelism)
.collect::<Vec<Result<Option<(String, String)>, Item>>>()
.await;
let error_log = File::create(self.base.join("errors").join("items.csv"))?;
let mut error_csv = WriterBuilder::new().from_writer(error_log);
let invalid_log = File::create(self.base.join("errors").join("invalid.csv"))?;
let mut invalid_csv = WriterBuilder::new().from_writer(invalid_log);
let mut success_count = 0;
let mut invalid_count = 0;
let mut error_count = 0;
for result in results {
match result {
Ok(None) => {
success_count += 1;
}
Ok(Some((expected, computed))) => {
invalid_count += 1;
invalid_csv.write_record(vec![expected, computed])?;
}
Err(item) => {
error_count += 1;
error_csv.write_record(item.to_record())?;
}
}
}
Ok((
success_count,
invalid_count,
total_count - success_count - error_count - invalid_count,
error_count,
))
}
fn read_csv<R: Read>(reader: R) -> Result<Vec<Item>, Error> {
let mut csv_reader = ReaderBuilder::new().has_headers(false).from_reader(reader);
csv_reader
.records()
.map(|record| {
let row = record?;
Ok(Item::parse_optional_record(
row.get(0),
row.get(1),
row.get(2),
row.get(3),
row.get(4),
row.get(5),
)?)
})
.collect()
}
}