use crate::exchange::http::{MAX_ERROR_BODY_DOWNLOAD_BYTES, read_body_capped};
use crate::exchange::lse::error::{LseError, extract_detail};
use crate::exchange::lse::market::{LseDataset, candle_interval_str};
use crate::exchange::lse::vault::LseVaultClient;
use crate::subscription::candle::CandleInterval;
use chrono::{DateTime, NaiveDate, Utc};
use futures_util::StreamExt;
use rustrade_instrument::exchange::ExchangeId;
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use std::path::{Path, PathBuf};
use std::time::Duration;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tracing::{debug, info, warn};
const HASH_READ_BUFFER_BYTES: usize = 64 * 1024;
const DOWNLOAD_BUFFER_BYTES: usize = 256 * 1024;
const DOWNLOAD_TIMEOUT: Duration = Duration::from_secs(6 * 60 * 60);
const ALL_SYMBOLS_LITERAL: &str = "all";
const EXPORT_FORMAT: &str = "parquet";
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum LseExportTimeframe {
Tick,
Candle(CandleInterval),
}
impl LseExportTimeframe {
pub fn as_lse_str(&self) -> Result<&'static str, LseError> {
match self {
Self::Tick => Ok("tick"),
Self::Candle(interval) => {
candle_interval_str(*interval).ok_or(LseError::UnsupportedInterval {
interval: *interval,
})
}
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct LseExportRange {
start: NaiveDate,
end: NaiveDate,
}
impl LseExportRange {
pub fn new(start: NaiveDate, end: NaiveDate) -> Result<Self, LseError> {
if start >= end {
return Err(LseError::InvalidInput {
message: format!(
"export range start {start} must be before end {end} (end is exclusive)"
),
});
}
Ok(Self { start, end })
}
pub fn start(&self) -> NaiveDate {
self.start
}
pub fn end(&self) -> NaiveDate {
self.end
}
#[must_use]
pub fn is_wholly_after(&self, today: NaiveDate) -> bool {
self.start
.checked_sub_days(chrono::Days::new(1))
.is_some_and(|slack| slack > today)
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct LseExportRequest {
dataset: LseDataset,
symbol: String,
timeframe: LseExportTimeframe,
range: LseExportRange,
}
impl LseExportRequest {
pub fn new(
dataset: LseDataset,
symbol: impl Into<String>,
timeframe: LseExportTimeframe,
range: LseExportRange,
) -> Result<Self, LseError> {
let symbol: String = symbol.into();
let symbol = symbol.trim().to_owned();
if symbol.is_empty() {
return Err(LseError::InvalidInput {
message: "export symbol must not be blank; the provider rejects a missing symbol \
with a 400, which still consumes an export"
.to_owned(),
});
}
let _ = timeframe.as_lse_str()?;
let equity_symbology = matches!(dataset, LseDataset::Stocks | LseDataset::Etf);
if !(equity_symbology && symbol == "ALL")
&& symbol.eq_ignore_ascii_case(ALL_SYMBOLS_LITERAL)
{
let hint = if equity_symbology {
format!(
" (If you meant Allstate, its symbol is exactly {:?}, uppercase.)",
ALL_SYMBOLS_LITERAL.to_uppercase()
)
} else {
format!(
" No casing of it is a symbol on dataset {:?}, whose symbols are not bare \
equity tickers.",
dataset.as_catalog_str()
)
};
return Err(LseError::InvalidInput {
message: format!(
"symbol {symbol:?} is a literal the provider matches against the symbol \
column, not a request for every symbol: an export naming it returns a valid \
but EMPTY artifact, with no error, and still consumes one of five hourly \
exports - name a real symbol instead. Measured on both the candle and the \
tick path; every artifact this provider will produce is single-symbol.{hint}"
),
});
}
if let LseExportTimeframe::Candle(interval) = timeframe
&& !dataset.is_candle_class()
{
return Err(LseError::InvalidInput {
message: format!(
"dataset {:?} is tick-only on the export path, so resolution {interval} \
cannot be exported for it - it serves candles over REST, but export \
timeframes follow the provider's candle-class split; use \
LseExportTimeframe::Tick",
dataset.as_catalog_str()
),
});
}
Ok(Self {
dataset,
symbol,
timeframe,
range,
})
}
pub fn dataset(&self) -> LseDataset {
self.dataset
}
pub fn symbol(&self) -> &str {
&self.symbol
}
pub fn timeframe(&self) -> LseExportTimeframe {
self.timeframe
}
pub fn range(&self) -> LseExportRange {
self.range
}
fn to_body(&self) -> Result<ExportRequestBody<'_>, LseError> {
Ok(ExportRequestBody {
dataset: self.dataset.as_catalog_str(),
symbol: &self.symbol,
timeframe: self.timeframe.as_lse_str()?,
start: self.range.start().to_string(),
end: self.range.end().to_string(),
format: EXPORT_FORMAT,
})
}
}
#[derive(Debug, Serialize)]
struct ExportRequestBody<'a> {
dataset: &'a str,
symbol: &'a str,
timeframe: &'a str,
start: String,
end: String,
format: &'a str,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum LseExportStatus {
Queued,
Running,
Ready,
Failed,
Expired,
Other(String),
}
impl LseExportStatus {
pub fn is_terminal(&self) -> bool {
matches!(self, Self::Ready | Self::Failed | Self::Expired)
}
pub fn as_str(&self) -> &str {
match self {
Self::Queued => "queued",
Self::Running => "running",
Self::Ready => "ready",
Self::Failed => "failed",
Self::Expired => "expired",
Self::Other(other) => other,
}
}
}
impl From<&str> for LseExportStatus {
fn from(value: &str) -> Self {
match value {
"queued" => Self::Queued,
"running" => Self::Running,
"ready" => Self::Ready,
"failed" => Self::Failed,
"expired" => Self::Expired,
other => Self::Other(other.to_owned()),
}
}
}
impl std::fmt::Display for LseExportStatus {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(self.as_str())
}
}
impl<'de> Deserialize<'de> for LseExportStatus {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: serde::Deserializer<'de>,
{
Ok(Self::from(String::deserialize(deserializer)?.as_str()))
}
}
#[derive(Debug, Clone, Deserialize)]
#[non_exhaustive]
pub struct LseExportJob {
pub job_id: String,
pub status: LseExportStatus,
#[serde(default)]
pub est_bytes: Option<u64>,
}
#[derive(Debug, Clone, Deserialize)]
#[non_exhaustive]
pub struct LseExportJobStatus {
pub id: String,
pub status: LseExportStatus,
#[serde(default)]
pub rows: Option<u64>,
#[serde(default)]
pub bytes: Option<u64>,
#[serde(default)]
pub sha256: Option<String>,
#[serde(default)]
pub error: Option<String>,
#[serde(default)]
pub expires_at: Option<DateTime<Utc>>,
#[serde(default)]
pub table_name: Option<String>,
#[serde(default)]
pub dataset: Option<String>,
#[serde(default)]
pub symbol: Option<String>,
#[serde(default)]
pub timeframe: Option<String>,
#[serde(default)]
pub start: Option<String>,
#[serde(default)]
pub end: Option<String>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct LseExport {
path: PathBuf,
dataset: LseDataset,
symbol: String,
timeframe: LseExportTimeframe,
range: LseExportRange,
}
impl LseExport {
pub fn new(
path: impl Into<PathBuf>,
dataset: LseDataset,
symbol: impl Into<String>,
timeframe: LseExportTimeframe,
range: LseExportRange,
) -> Self {
Self {
path: path.into(),
dataset,
symbol: symbol.into(),
timeframe,
range,
}
}
pub fn path(&self) -> &Path {
&self.path
}
pub fn dataset(&self) -> LseDataset {
self.dataset
}
pub fn symbol(&self) -> &str {
&self.symbol
}
pub fn timeframe(&self) -> LseExportTimeframe {
self.timeframe
}
pub fn range(&self) -> LseExportRange {
self.range
}
pub fn exchange_id(&self) -> ExchangeId {
self.dataset.exchange_id()
}
}
impl LseVaultClient {
pub async fn submit_export(
&self,
request: &LseExportRequest,
) -> Result<LseExportJob, LseError> {
let today = chrono::Utc::now().date_naive();
if request.range().is_wholly_after(today) {
return Err(LseError::InvalidInput {
message: format!(
"export range {} to {} (end exclusive) lies wholly in the future, and today is \
{today}: the provider would answer with a valid but EMPTY artifact, with no \
error, and still consume one of five hourly exports - check the year",
request.range().start(),
request.range().end()
),
});
}
let body = request.to_body()?;
info!(
dataset = body.dataset,
symbol = body.symbol,
timeframe = body.timeframe,
start = body.start,
end = body.end,
"submitting vault export job; this consumes one of the hourly export allowance"
);
self.post_json("export", &body).await
}
pub async fn export_status(&self, job_id: &str) -> Result<LseExportJobStatus, LseError> {
self.get_json(&format!("export/{job_id}"), &[]).await
}
pub async fn await_export(
&self,
job_id: &str,
poll_interval: Duration,
timeout: Duration,
) -> Result<LseExportJobStatus, LseError> {
let now = tokio::time::Instant::now();
let deadline = now
.checked_add(timeout)
.unwrap_or_else(|| now + Duration::from_secs(86_400 * 365));
loop {
let status = self.export_status(job_id).await?;
match &status.status {
LseExportStatus::Ready => return Ok(status),
LseExportStatus::Failed | LseExportStatus::Expired => {
return Err(LseError::ExportFailed {
job_id: job_id.to_owned(),
status: status.status.to_string(),
message: status.error.unwrap_or_default(),
});
}
pending => {
debug!(job_id, status = %pending, "export job not ready");
}
}
if tokio::time::Instant::now() >= deadline {
return Err(LseError::ExportTimeout {
job_id: job_id.to_owned(),
status: status.status.to_string(),
});
}
let next_poll = tokio::time::Instant::now()
.checked_add(poll_interval)
.map_or(deadline, |instant| deadline.min(instant));
tokio::time::sleep_until(next_poll).await;
}
}
pub async fn download_export(
&self,
job: &LseExportJobStatus,
destination: impl AsRef<Path>,
request: &LseExportRequest,
) -> Result<LseExport, LseError> {
if job.status != LseExportStatus::Ready {
return Err(LseError::InvalidInput {
message: format!(
"export job {} is {}, not ready; nothing to download",
job.id, job.status
),
});
}
verify_job_covers_request(job, request)?;
let destination = destination.as_ref();
let part = part_path(destination, &job.id);
if let Some(parent) = destination
.parent()
.filter(|parent| !parent.as_os_str().is_empty())
{
tokio::fs::create_dir_all(parent)
.await
.map_err(|source| LseError::Io {
message: format!("creating {}", parent.display()),
source,
})?;
}
let verifiable = job.bytes.is_some() || job.sha256.is_some();
let (mut hasher, mut downloaded) = if verifiable {
hash_existing_part(&part).await?
} else {
(Sha256::new(), 0)
};
let mut complete = downloaded > 0 && job.bytes.is_some_and(|total| downloaded >= total);
if !complete {
let transferred = self.transfer_export(job, &part, hasher, downloaded).await?;
hasher = transferred.hasher;
downloaded = transferred.downloaded;
complete = transferred.exhausted
|| job
.bytes
.is_some_and(|total| transferred.downloaded >= total);
}
let digest = hex::encode(hasher.finalize());
if let Some(expected) = job.bytes.filter(|expected| *expected != downloaded) {
let corrupt = downloaded > expected;
return Err(self
.integrity_mismatch(
&part,
format!("{expected} bytes"),
format!("{downloaded} bytes"),
corrupt,
)
.await);
}
if let Some(expected) = job
.sha256
.as_ref()
.filter(|expected| !expected.eq_ignore_ascii_case(&digest))
{
return Err(self
.integrity_mismatch(&part, expected.clone(), digest, complete)
.await);
}
if job.bytes.is_none() || job.sha256.is_none() {
warn!(
job_id = job.id,
path = %destination.display(),
checked_length = job.bytes.is_some(),
checked_sha256 = job.sha256.is_some(),
"export job reports no integrity metadata for one or both checks; accepting the \
artifact with that verification skipped"
);
}
tokio::fs::rename(&part, destination)
.await
.map_err(|source| LseError::Io {
message: format!("renaming {} to {}", part.display(), destination.display()),
source,
})?;
info!(
job_id = job.id,
path = %destination.display(),
bytes = downloaded,
rows = job.rows.unwrap_or_default(),
"export downloaded and verified"
);
Ok(LseExport::new(
destination,
request.dataset(),
request.symbol(),
request.timeframe(),
request.range(),
))
}
async fn transfer_export(
&self,
job: &LseExportJobStatus,
part: &Path,
mut hasher: Sha256,
mut downloaded: u64,
) -> Result<Transferred, LseError> {
let permit = self.enter_gate().await;
let url = format!("{}/export/{}/download", self.base_url(), job.id);
let mut builder = self
.http()
.get(&url)
.timeout(DOWNLOAD_TIMEOUT)
.header(reqwest::header::ACCEPT_ENCODING, "identity");
if downloaded > 0 {
builder = builder.header(reqwest::header::RANGE, format!("bytes={downloaded}-"));
}
let response = builder.send().await?;
let status = response.status();
if status == reqwest::StatusCode::TOO_MANY_REQUESTS {
drop(permit);
return Err(self.quota_exceeded().await);
}
if downloaded > 0 && status == reqwest::StatusCode::RANGE_NOT_SATISFIABLE {
warn!(
job_id = job.id,
downloaded,
"range request rejected as unsatisfiable; the existing part file already holds the \
whole artifact"
);
return Ok(Transferred {
hasher,
downloaded,
exhausted: true,
});
}
if !status.is_success() {
let body = read_body_capped(response, MAX_ERROR_BODY_DOWNLOAD_BYTES).await?;
return Err(LseError::Api {
status: status.as_u16(),
message: extract_detail(&body),
});
}
let resuming = downloaded > 0 && status == reqwest::StatusCode::PARTIAL_CONTENT;
if resuming {
verify_resume_offset(&response, downloaded)?;
}
if downloaded > 0 && !resuming {
warn!(
job_id = job.id,
downloaded, "range request answered in full; restarting the download"
);
hasher = Sha256::new();
downloaded = 0;
}
let file = tokio::fs::OpenOptions::new()
.create(true)
.write(true)
.truncate(!resuming)
.append(resuming)
.open(part)
.await
.map_err(|source| LseError::Io {
message: format!("opening {}", part.display()),
source,
})?;
let mut file = tokio::io::BufWriter::with_capacity(DOWNLOAD_BUFFER_BYTES, file);
let mut stream = response.bytes_stream();
while let Some(chunk) = stream.next().await {
let chunk = chunk?;
hasher.update(&chunk);
downloaded += chunk.len() as u64;
file.write_all(&chunk)
.await
.map_err(|source| LseError::Io {
message: format!("writing {}", part.display()),
source,
})?;
if job.bytes.is_some_and(|total| downloaded > total) {
warn!(
job_id = job.id,
downloaded,
expected = job.bytes,
"export download exceeded the length the job reports; abandoning the transfer \
rather than writing an unbounded body to disk"
);
break;
}
}
file.flush().await.map_err(|source| LseError::Io {
message: format!("flushing {}", part.display()),
source,
})?;
Ok(Transferred {
hasher,
downloaded,
exhausted: false,
})
}
pub(crate) async fn post_json<B, T>(&self, path: &str, body: &B) -> Result<T, LseError>
where
B: Serialize + ?Sized,
T: serde::de::DeserializeOwned,
{
let permit = self.enter_gate().await;
let url = format!("{}/{path}", self.base_url());
let response = self.http().post(&url).json(body).send().await?;
let status = response.status();
if status == reqwest::StatusCode::TOO_MANY_REQUESTS {
drop(permit);
return Err(self.quota_exceeded().await);
}
if !status.is_success() {
let body = read_body_capped(response, MAX_ERROR_BODY_DOWNLOAD_BYTES).await?;
return Err(LseError::Api {
status: status.as_u16(),
message: extract_detail(&body),
});
}
let body = response.text().await?;
debug!(len = body.len(), path, "vault response received");
serde_json::from_str(&body).map_err(|error| LseError::Deserialize {
message: format!("{path}: {error}"),
})
}
async fn integrity_mismatch(
&self,
part: &Path,
expected: String,
actual: String,
discard: bool,
) -> LseError {
if discard {
warn!(
path = %part.display(),
%expected,
%actual,
"the partial file is corrupt rather than incomplete -- it is longer than the \
artifact, or complete at the wrong digest -- so it cannot be resumed; discarding \
it so a re-call restarts"
);
if let Err(error) = tokio::fs::remove_file(part).await {
warn!(
path = %part.display(),
%error,
"could not remove the unusable partial file; remove it by hand, or this job \
will keep failing verification"
);
}
}
LseError::IntegrityMismatch {
path: part.to_path_buf(),
expected,
actual,
discarded: discard,
}
}
async fn quota_exceeded(&self) -> LseError {
let status = match self.usage().await {
Ok(status) => Some(status),
Err(error) => {
warn!(
%error,
"export allowance exhausted, and the follow-up usage call failed; reporting \
the rejection without an allowance position"
);
None
}
};
LseError::QuotaExceeded { status }
}
}
struct Transferred {
hasher: Sha256,
downloaded: u64,
exhausted: bool,
}
fn verify_job_covers_request(
job: &LseExportJobStatus,
request: &LseExportRequest,
) -> Result<(), LseError> {
let mismatch = |field: &str, requested: String, reported: &str| LseError::ExportJobMismatch {
job_id: job.id.clone(),
field: field.to_owned(),
requested,
reported: reported.to_owned(),
};
if let Some(reported) = job.dataset.as_deref() {
let requested = request.dataset().as_catalog_str();
if reported != requested {
return Err(mismatch("dataset", requested.to_owned(), reported));
}
}
if let Some(reported) = job.symbol.as_deref()
&& !reported.eq_ignore_ascii_case(request.symbol())
{
return Err(mismatch("symbol", request.symbol().to_owned(), reported));
}
if let Some(reported) = job.timeframe.as_deref()
&& let Ok(requested) = request.timeframe().as_lse_str()
&& reported != requested
{
return Err(mismatch("timeframe", requested.to_owned(), reported));
}
if let Some(reported) = job.start.as_deref() {
let requested = request.range().start().to_string();
if reported != requested {
return Err(mismatch("start", requested, reported));
}
}
if let Some(reported) = job.end.as_deref() {
let requested = request.range().end().to_string();
if reported != requested {
return Err(mismatch("end", requested, reported));
}
}
Ok(())
}
fn part_path(destination: &Path, job_id: &str) -> PathBuf {
let mut name = destination.as_os_str().to_owned();
name.push(".");
name.push(
job_id
.chars()
.map(|character| {
if character.is_ascii_alphanumeric() || matches!(character, '.' | '_' | '-') {
character
} else {
'_'
}
})
.collect::<String>(),
);
name.push(".part");
PathBuf::from(name)
}
fn verify_resume_offset(response: &reqwest::Response, expected_start: u64) -> Result<(), LseError> {
let status = reqwest::StatusCode::PARTIAL_CONTENT.as_u16();
let header = response
.headers()
.get(reqwest::header::CONTENT_RANGE)
.and_then(|value| value.to_str().ok())
.ok_or_else(|| LseError::Api {
status,
message: "partial-content response carries no usable Content-Range header".to_owned(),
})?;
let first = header
.trim()
.strip_prefix("bytes ")
.and_then(|range| range.split('-').next())
.and_then(|first| first.trim().parse::<u64>().ok())
.ok_or_else(|| LseError::Api {
status,
message: format!("unparseable Content-Range on a partial-content response: {header}"),
})?;
if first != expected_start {
return Err(LseError::Api {
status,
message: format!(
"partial-content response resumes at byte {first}, not the requested \
{expected_start}"
),
});
}
Ok(())
}
async fn hash_existing_part(part: &Path) -> Result<(Sha256, u64), LseError> {
let mut hasher = Sha256::new();
let mut file = match tokio::fs::File::open(part).await {
Ok(file) => file,
Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok((hasher, 0)),
Err(source) => {
return Err(LseError::Io {
message: format!("opening {}", part.display()),
source,
});
}
};
let mut total = 0_u64;
let mut buffer = vec![0_u8; HASH_READ_BUFFER_BYTES];
loop {
let read = file
.read(&mut buffer)
.await
.map_err(|source| LseError::Io {
message: format!("reading {}", part.display()),
source,
})?;
if read == 0 {
break;
}
hasher.update(&buffer[..read]);
total += read as u64;
}
debug!(path = %part.display(), bytes = total, "resuming a partial export download");
Ok((hasher, total))
}