databento/historical/
batch.rs

1//! The historical batch download API.
2
3use std::{
4    cmp::Ordering,
5    collections::HashMap,
6    fmt::{self, Write},
7    num::NonZeroU64,
8    path::{Path, PathBuf},
9    str::FromStr,
10};
11
12use dbn::{Compression, Encoding, SType, Schema};
13use futures::StreamExt;
14use hex::ToHex;
15use reqwest::RequestBuilder;
16use serde::{de, Deserialize, Deserializer};
17use sha2::{Digest, Sha256};
18use time::OffsetDateTime;
19use tokio::{
20    fs::File,
21    io::{AsyncReadExt, BufWriter},
22};
23use tracing::{debug, error, info, info_span, instrument, warn, Instrument};
24
25use crate::{
26    deserialize::{deserialize_date_time, deserialize_opt_date_time},
27    historical::{check_http_error, AddToForm, Limit, ReqwestForm},
28    Error, Symbols,
29};
30
31use super::{handle_response, DateTimeRange};
32
33/// A client for the batch group of Historical API endpoints.
34#[derive(Debug)]
35pub struct BatchClient<'a> {
36    pub(crate) inner: &'a mut super::Client,
37}
38
39struct SplitSize(Option<NonZeroU64>);
40impl AddToForm<SplitSize> for ReqwestForm {
41    fn add_to_form(mut self, SplitSize(split_size): &SplitSize) -> Self {
42        if let Some(split_size) = split_size {
43            self.push(("split_size", split_size.to_string()));
44        }
45        self
46    }
47}
48
49impl AddToForm<SplitDuration> for ReqwestForm {
50    fn add_to_form(mut self, split_duration: &SplitDuration) -> Self {
51        self.push(("split_duration", split_duration.to_string()));
52        self
53    }
54}
55
56impl BatchClient<'_> {
57    /// Submits a new batch job and returns a description and identifiers for the job.
58    ///
59    /// <div class="warning">
60    /// Calling this method will incur a cost.
61    /// </div>
62    ///
63    /// # Errors
64    /// This function returns an error when it fails to communicate with the Databento API
65    /// or the API indicates there's an issue with the request.
66    #[instrument(name = "batch.submit_job")]
67    pub async fn submit_job(&mut self, params: &SubmitJobParams) -> crate::Result<BatchJob> {
68        let form = vec![
69            ("dataset", params.dataset.to_string()),
70            ("schema", params.schema.to_string()),
71            ("encoding", params.encoding.to_string()),
72            ("compression", params.compression.to_string()),
73            ("pretty_px", params.pretty_px.to_string()),
74            ("pretty_ts", params.pretty_ts.to_string()),
75            (
76                "map_symbols",
77                params
78                    .map_symbols
79                    .unwrap_or(params.encoding != Encoding::Dbn)
80                    .to_string(),
81            ),
82            ("split_symbols", params.split_symbols.to_string()),
83            ("delivery", params.delivery.to_string()),
84            ("stype_in", params.stype_in.to_string()),
85            ("stype_out", params.stype_out.to_string()),
86            ("symbols", params.symbols.to_api_string()),
87        ]
88        .add_to_form(&params.date_time_range)
89        .add_to_form(&Limit(params.limit))
90        .add_to_form(&SplitSize(params.split_size))
91        .add_to_form(&params.split_duration);
92        let builder = self.post("submit_job")?.form(&form);
93        let resp = builder.send().await?;
94        handle_response(resp).await
95    }
96
97    /// Lists previous batch jobs with filtering by `params`.
98    ///
99    /// # Errors
100    /// This function returns an error when it fails to communicate with the Databento API
101    /// or the API indicates there's an issue with the request.
102    #[instrument(name = "batch.list_jobs")]
103    pub async fn list_jobs(&mut self, params: &ListJobsParams) -> crate::Result<Vec<BatchJob>> {
104        let mut builder = self.get("list_jobs")?;
105        if let Some(ref states) = params.states {
106            let states_str = states.iter().fold(String::new(), |mut acc, s| {
107                if acc.is_empty() {
108                    s.as_str().to_owned()
109                } else {
110                    write!(acc, ",{}", s.as_str()).unwrap();
111                    acc
112                }
113            });
114            builder = builder.query(&[("states", states_str)]);
115        }
116        if let Some(ref since) = params.since {
117            builder = builder.query(&[("since", &since.unix_timestamp_nanos().to_string())]);
118        }
119        let resp = builder.send().await?;
120        handle_response(resp).await
121    }
122
123    /// Gets the details of a batch job with ID `job_id`.
124    ///
125    /// # Errors
126    /// This function returns an error when it fails to communicate with the Databento API
127    /// or the API indicates there's an issue with the request.
128    #[instrument(name = "batch.get_job_details")]
129    pub async fn get_job_details(&mut self, job_id: &str) -> crate::Result<BatchJob> {
130        let resp = self
131            .get("get_job_details")?
132            .query(&[("job_id", job_id)])
133            .send()
134            .await?;
135        handle_response(resp).await
136    }
137
138    /// Lists all files associated with the batch job with ID `job_id`.
139    ///
140    /// # Errors
141    /// This function returns an error when it fails to communicate with the Databento API
142    /// or the API indicates there's an issue with the request.
143    #[instrument(name = "batch.list_files")]
144    pub async fn list_files(&mut self, job_id: &str) -> crate::Result<Vec<BatchFileDesc>> {
145        let resp = self
146            .get("list_files")?
147            .query(&[("job_id", job_id)])
148            .send()
149            .await?;
150        handle_response(resp).await
151    }
152
153    /// Downloads the file specified in `params` or all files associated with the job ID.
154    ///
155    /// # Errors
156    /// This function returns an error when it fails to communicate with the Databento API
157    /// or the API indicates there's an issue with the request. It will also return an
158    /// error if it encounters an issue downloading a file.
159    #[instrument(name = "batch.download")]
160    pub async fn download(&mut self, params: &DownloadParams) -> crate::Result<Vec<PathBuf>> {
161        let job_dir = params.output_dir.join(&params.job_id);
162        if job_dir.exists() {
163            if !job_dir.is_dir() {
164                return Err(Error::bad_arg(
165                    "output_dir",
166                    "exists but is not a directory",
167                ));
168            }
169        } else {
170            tokio::fs::create_dir_all(&job_dir).await?;
171        }
172        let job_files = self.list_files(&params.job_id).await?;
173        if let Some(filename_to_download) = params.filename_to_download.as_ref() {
174            let Some(file_desc) = job_files
175                .iter()
176                .find(|file| file.filename == *filename_to_download)
177            else {
178                return Err(Error::bad_arg(
179                    "filename_to_download",
180                    "not found for batch job",
181                ));
182            };
183            let output_path = job_dir.join(filename_to_download);
184            let https_url = file_desc
185                .urls
186                .get("https")
187                .ok_or_else(|| Error::internal("Missing https URL for batch file"))?;
188            self.download_file(https_url, &output_path, &file_desc.hash, file_desc.size)
189                .await?;
190            Ok(vec![output_path])
191        } else {
192            let mut paths = Vec::with_capacity(job_files.len());
193            for file_desc in job_files.iter() {
194                let output_path = params
195                    .output_dir
196                    .join(&params.job_id)
197                    .join(&file_desc.filename);
198                let https_url = file_desc
199                    .urls
200                    .get("https")
201                    .ok_or_else(|| Error::internal("Missing https URL for batch file"))?;
202                self.download_file(https_url, &output_path, &file_desc.hash, file_desc.size)
203                    .await?;
204                paths.push(output_path);
205            }
206            Ok(paths)
207        }
208    }
209
210    #[instrument(name = "batch.download_file")]
211    async fn download_file(
212        &mut self,
213        url: &str,
214        path: &Path,
215        hash: &str,
216        exp_size: u64,
217    ) -> crate::Result<()> {
218        const MAX_RETRIES: usize = 5;
219        let url = reqwest::Url::parse(url)
220            .map_err(|e| Error::internal(format!("Unable to parse URL: {e:?}")))?;
221
222        let Some((hash_algo, exp_hash_hex)) = hash.split_once(':') else {
223            return Err(Error::internal("Unexpected hash string format {hash:?}"));
224        };
225        let mut hasher = if hash_algo == "sha256" {
226            Some(Sha256::new())
227        } else {
228            warn!(
229                hash_algo,
230                "Skipping checksum with unsupported hash algorithm"
231            );
232            None
233        };
234
235        let span = info_span!("BatchDownload", %url, path=%path.display());
236        async move {
237            let mut retries = 0;
238            'retry: loop {
239                let mut req = self.inner.get_with_path(url.path())?;
240                match Self::check_if_exists(path, exp_size, &mut hasher).await? {
241                    Header::Skip => {
242                        return Ok(());
243                    }
244                    Header::Range(Some((key, val))) => {
245                        req = req.header(key, val);
246                    }
247                    Header::Range(None) => {}
248                }
249                let resp = req.send().await?;
250                let mut stream = check_http_error(resp).await?.bytes_stream();
251                info!("Downloading file");
252                let mut output = BufWriter::new(
253                    tokio::fs::OpenOptions::new()
254                        .create(true)
255                        .append(true)
256                        .write(true)
257                        .open(path)
258                        .await?,
259                );
260                while let Some(chunk) = stream.next().await {
261                    let chunk = match chunk {
262                        Ok(chunk) => chunk,
263                        Err(err) if retries < MAX_RETRIES => {
264                            retries += 1;
265                            error!(?err, retries, "Retrying download");
266                            continue 'retry;
267                        }
268                        Err(err) => {
269                            return Err(crate::Error::from(err));
270                        }
271                    };
272                    if retries > 0 {
273                        retries = 0;
274                        info!("Resumed download");
275                    }
276                    if let Some(hasher) = hasher.as_mut() {
277                        hasher.update(&chunk)
278                    }
279                    tokio::io::copy(&mut chunk.as_ref(), &mut output).await?;
280                }
281                debug!("Completed download");
282                Self::verify_hash(hasher, exp_hash_hex).await;
283                return Ok(());
284            }
285        }
286        .instrument(span)
287        .await
288    }
289
290    async fn check_if_exists(
291        path: &Path,
292        exp_size: u64,
293        hasher: &mut Option<Sha256>,
294    ) -> crate::Result<Header> {
295        let Ok(metadata) = tokio::fs::metadata(path).await else {
296            return Ok(Header::Range(None));
297        };
298        let actual_size = metadata.len();
299        match actual_size.cmp(&exp_size) {
300            Ordering::Less => {
301                debug!(
302                    prev_downloaded_bytes = actual_size,
303                    total_bytes = exp_size,
304                    "Found existing file, resuming download"
305                );
306                if let Some(hasher) = hasher {
307                    let mut buf = vec![0; 1 << 23];
308                    let mut file = File::open(path).await?;
309                    loop {
310                        let read_size = file.read(&mut buf).await?;
311                        if read_size == 0 {
312                            break;
313                        }
314                        hasher.update(&buf[..read_size]);
315                    }
316                }
317            }
318            Ordering::Equal => {
319                debug!("Skipping download as file already exists and matches expected size");
320                return Ok(Header::Skip);
321            }
322            Ordering::Greater => {
323                return Err(crate::Error::Io(std::io::Error::other(format!(
324                                    "Batch file {} already exists with size {actual_size} which is larger than expected size {exp_size}",
325                                    path.file_name().unwrap().display(),
326                                ))));
327            }
328        }
329        Ok(Header::Range(Some((
330            "Range",
331            format!("bytes={}-", metadata.len()),
332        ))))
333    }
334
335    async fn verify_hash(hasher: Option<Sha256>, exp_hash_hex: &str) {
336        let Some(hasher) = hasher else {
337            return;
338        };
339        let hash_hex = hasher.finalize().encode_hex::<String>();
340        if hash_hex != exp_hash_hex {
341            warn!(
342                hash_hex,
343                exp_hash_hex, "Downloaded file failed checksum verification"
344            );
345        } else {
346            debug!("Successfully verified checksum");
347        }
348    }
349
350    const PATH_PREFIX: &'static str = "batch";
351
352    fn get(&mut self, slug: &str) -> crate::Result<RequestBuilder> {
353        self.inner.get(&format!("{}.{slug}", Self::PATH_PREFIX))
354    }
355
356    fn post(&mut self, slug: &str) -> crate::Result<RequestBuilder> {
357        self.inner.post(&format!("{}.{slug}", Self::PATH_PREFIX))
358    }
359}
360
361/// The duration of time at which batch files will be split.
362#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
363pub enum SplitDuration {
364    /// One file per day.
365    #[default]
366    Day,
367    /// One file per week. A week starts on Sunday UTC.
368    Week,
369    /// One file per month.
370    Month,
371    /// One file per year.
372    Year,
373    /// No time-based splitting.
374    None,
375}
376
377/// How the batch job will be delivered.
378#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
379pub enum Delivery {
380    /// Via download from the Databento portal.
381    #[default]
382    Download,
383}
384
385/// The state of a batch job.
386#[derive(Clone, Copy, Debug, PartialEq, Eq)]
387pub enum JobState {
388    /// The job has been queued for processing.
389    Queued,
390    /// The job has begun processing.
391    Processing,
392    /// The job has finished processing and is ready for delivery.
393    Done,
394    /// The job is no longer available.
395    Expired,
396}
397
398/// The parameters for [`BatchClient::submit_job()`]. Use [`SubmitJobParams::builder()`] to
399/// get a builder type with all the preset defaults.
400#[derive(Debug, Clone, bon::Builder, PartialEq, Eq)]
401pub struct SubmitJobParams {
402    /// The dataset code.
403    #[builder(with = |d: impl ToString| d.to_string())]
404    pub dataset: String,
405    /// The symbols to filter for.
406    #[builder(into)]
407    pub symbols: Symbols,
408    /// The data record schema.
409    pub schema: Schema,
410    /// The request range with an inclusive start and an exclusive end.
411    ///
412    /// Filters on `ts_recv` if it exists in the schema, otherwise `ts_event`.
413    #[builder(into)]
414    pub date_time_range: DateTimeRange,
415    /// The data encoding. Defaults to [`Dbn`](Encoding::Dbn).
416    #[builder(default = Encoding::Dbn)]
417    pub encoding: Encoding,
418    /// The data compression mode. Defaults to [`Zstd`](Compression::Zstd).
419    #[builder(default = Compression::Zstd)]
420    pub compression: Compression,
421    /// If `true`, prices will be formatted to the correct scale (using the fixed-
422    /// precision scalar 1e-9). Only valid for [`Encoding::Csv`] and [`Encoding::Json`].
423    #[builder(default)]
424    pub pretty_px: bool,
425    /// If `true`, timestamps will be formatted as ISO 8601 strings. Only valid for
426    /// [`Encoding::Csv`] and [`Encoding::Json`].
427    #[builder(default)]
428    pub pretty_ts: bool,
429    /// If `true`, a symbol field will be included with each text-encoded
430    /// record. Defaults to `true` for [`Encoding::Csv`] and [`Encoding::Json`] encodings
431    /// when `None`, and `false` for [`Encoding::Dbn`].
432    pub map_symbols: Option<bool>,
433    /// If `true`, files will be split by raw symbol. Cannot be requested with [`Symbols::All`].
434    #[builder(default)]
435    pub split_symbols: bool,
436    /// The maximum time duration before batched data is split into multiple
437    /// files.
438    ///
439    /// [`None`](SplitDuration::None) means the data will not be split by time. Defaults
440    /// to [`Day`](SplitDuration::Day).
441    #[builder(default)]
442    pub split_duration: SplitDuration,
443    /// The optional maximum size (in bytes) of each batched data file before being split.
444    /// Must be an integer between 1e9 and 10e9 inclusive (1GB - 10GB). Defaults to `None`.
445    pub split_size: Option<NonZeroU64>,
446    /// The delivery mechanism for the batched data files once processed.
447    /// Only [`Download`](Delivery::Download) is supported at this time.
448    #[builder(default)]
449    pub delivery: Delivery,
450    /// The symbology type of the input `symbols`. Defaults to
451    /// [`RawSymbol`](dbn::enums::SType::RawSymbol).
452    #[builder(default = SType::RawSymbol)]
453    pub stype_in: SType,
454    /// The symbology type of the output `symbols`. Defaults to
455    /// [`InstrumentId`](dbn::enums::SType::InstrumentId).
456    ///
457    /// Must be a valid symbology combination with [`stype_in`](Self::stype_in).
458    /// See [symbology combinations](https://databento.com/docs/standards-and-conventions/symbology#supported-symbology-combinations).
459    #[builder(default = SType::InstrumentId)]
460    pub stype_out: SType,
461    /// The optional maximum number of records to return. Defaults to no limit.
462    pub limit: Option<NonZeroU64>,
463}
464
465/// The description of a submitted batch job.
466#[derive(Debug, Clone, Deserialize)]
467pub struct BatchJob {
468    /// The unique job ID.
469    pub id: String,
470    /// The user ID of the user who submitted the job.
471    pub user_id: Option<String>,
472    /// The cost of the job in US dollars. Will be `None` until the job is processed.
473    pub cost_usd: Option<f64>,
474    /// The dataset code.
475    pub dataset: String,
476    /// The list of symbols specified in the request.
477    pub symbols: Symbols,
478    /// The symbology type of the input `symbols`.
479    pub stype_in: SType,
480    /// The symbology type of the output `symbols`.
481    pub stype_out: SType,
482    /// The data record schema.
483    pub schema: Schema,
484    /// The inclusive start of the request range.
485    #[serde(deserialize_with = "deserialize_date_time")]
486    pub start: OffsetDateTime,
487    /// The exclusive end of the request range.
488    #[serde(deserialize_with = "deserialize_date_time")]
489    pub end: OffsetDateTime,
490    /// The maximum number of records to return.
491    pub limit: Option<NonZeroU64>,
492    /// The data encoding.
493    pub encoding: Encoding,
494    /// The data compression mode.
495    #[serde(deserialize_with = "deserialize_compression")]
496    pub compression: Compression,
497    /// If prices are formatted to the correct scale (using the fixed-precision scalar 1e-9).
498    pub pretty_px: bool,
499    /// If timestamps are formatted as ISO 8601 strings.
500    pub pretty_ts: bool,
501    /// If a symbol field is included with each text-encoded record.
502    pub map_symbols: bool,
503    /// If files are split by raw symbol.
504    pub split_symbols: bool,
505    /// The maximum time interval for an individual file before splitting into multiple
506    /// files.
507    pub split_duration: SplitDuration,
508    /// The maximum size for an individual file before splitting into multiple files.
509    pub split_size: Option<NonZeroU64>,
510    /// The delivery mechanism of the batch data.
511    pub delivery: Delivery,
512    /// The number of data records (`None` until the job is processed).
513    pub record_count: Option<u64>,
514    /// The size of the raw binary data used to process the batch job (used for billing purposes).
515    pub billed_size: Option<u64>,
516    /// The total size of the result of the batch job after splitting and compression.
517    pub actual_size: Option<u64>,
518    /// The total size of the result of the batch job after any packaging (including metadata).
519    pub package_size: Option<u64>,
520    /// The current status of the batch job.
521    pub state: JobState,
522    /// The timestamp of when Databento received the batch job.
523    #[serde(deserialize_with = "deserialize_date_time")]
524    pub ts_received: OffsetDateTime,
525    /// The timestamp of when the batch job was queued.
526    #[serde(deserialize_with = "deserialize_opt_date_time")]
527    pub ts_queued: Option<OffsetDateTime>,
528    /// The timestamp of when the batch job began processing.
529    #[serde(deserialize_with = "deserialize_opt_date_time")]
530    pub ts_process_start: Option<OffsetDateTime>,
531    /// The timestamp of when the batch job finished processing.
532    #[serde(deserialize_with = "deserialize_opt_date_time")]
533    pub ts_process_done: Option<OffsetDateTime>,
534    /// The timestamp of when the batch job will expire from the Download center.
535    #[serde(deserialize_with = "deserialize_opt_date_time")]
536    pub ts_expiration: Option<OffsetDateTime>,
537    /// The progress percentage of the batch job (0-100). `None` for jobs that
538    /// were just submitted.
539    #[serde(default)]
540    pub progress: Option<u8>,
541}
542
543/// The parameters for [`BatchClient::list_jobs()`]. Use [`ListJobsParams::builder()`] to
544/// get a builder type with all the preset defaults.
545#[derive(Debug, Clone, Default, bon::Builder, PartialEq, Eq)]
546pub struct ListJobsParams {
547    /// The optional filter for job states. If `None`, defaults to all except `Expired`.
548    pub states: Option<Vec<JobState>>,
549    /// The optional filter for timestamp submitted (will not include jobs prior to
550    /// this time).
551    pub since: Option<OffsetDateTime>,
552}
553
554/// The file details for a batch job.
555#[derive(Debug, Clone, Deserialize)]
556pub struct BatchFileDesc {
557    /// The file name.
558    pub filename: String,
559    /// The size of the file in bytes.
560    pub size: u64,
561    /// The SHA256 hash of the file.
562    pub hash: String,
563    /// A map of download protocol to URL.
564    pub urls: HashMap<String, String>,
565}
566
567/// The parameters for [`BatchClient::download()`]. Use [`DownloadParams::builder()`] to
568/// get a builder type with all the preset defaults.
569#[derive(Debug, Clone, bon::Builder, PartialEq, Eq)]
570pub struct DownloadParams {
571    /// The directory to download the file(s) to.
572    #[builder(into)]
573    pub output_dir: PathBuf,
574    /// The batch job identifier.
575    #[builder(with = |id: impl ToString| id.to_string())]
576    pub job_id: String,
577    /// `None` means all files associated with the job will be downloaded.
578    #[builder(with = |f: impl ToString| f.to_string())]
579    pub filename_to_download: Option<String>,
580}
581
582impl SplitDuration {
583    /// Converts the enum to its `str` representation.
584    pub const fn as_str(&self) -> &'static str {
585        match self {
586            SplitDuration::Day => "day",
587            SplitDuration::Week => "week",
588            SplitDuration::Month => "month",
589            SplitDuration::Year => "year",
590            SplitDuration::None => "none",
591        }
592    }
593}
594
595impl fmt::Display for SplitDuration {
596    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
597        f.write_str(self.as_str())
598    }
599}
600
601impl FromStr for SplitDuration {
602    type Err = crate::Error;
603
604    fn from_str(s: &str) -> Result<Self, Self::Err> {
605        match s {
606            "day" => Ok(SplitDuration::Day),
607            "week" => Ok(SplitDuration::Week),
608            "month" => Ok(SplitDuration::Month),
609            "year" => Ok(SplitDuration::Year),
610            "none" => Ok(SplitDuration::None),
611            _ => Err(crate::Error::bad_arg(
612                "s",
613                format!(
614                    "{s} does not correspond with any {} variant",
615                    std::any::type_name::<Self>()
616                ),
617            )),
618        }
619    }
620}
621
622impl<'de> Deserialize<'de> for SplitDuration {
623    fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
624        let opt = Option::<String>::deserialize(deserializer)?;
625        match opt {
626            Some(str) => FromStr::from_str(&str).map_err(de::Error::custom),
627            // The API returns `null` instead of `"none"` for no time-based splitting
628            None => Ok(SplitDuration::None),
629        }
630    }
631}
632
633impl Delivery {
634    /// Converts the enum to its `str` representation.
635    pub const fn as_str(&self) -> &'static str {
636        match self {
637            Delivery::Download => "download",
638        }
639    }
640}
641
642impl fmt::Display for Delivery {
643    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
644        f.write_str(self.as_str())
645    }
646}
647
648impl FromStr for Delivery {
649    type Err = crate::Error;
650
651    fn from_str(s: &str) -> Result<Self, Self::Err> {
652        match s {
653            "download" => Ok(Delivery::Download),
654            _ => Err(crate::Error::bad_arg(
655                "s",
656                format!(
657                    "{s} does not correspond with any {} variant",
658                    std::any::type_name::<Self>()
659                ),
660            )),
661        }
662    }
663}
664
665impl<'de> Deserialize<'de> for Delivery {
666    fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
667        let str = String::deserialize(deserializer)?;
668        FromStr::from_str(&str).map_err(de::Error::custom)
669    }
670}
671
672impl JobState {
673    /// Converts the enum to its `str` representation.
674    pub const fn as_str(&self) -> &'static str {
675        match self {
676            JobState::Queued => "queued",
677            JobState::Processing => "processing",
678            JobState::Done => "done",
679            JobState::Expired => "expired",
680        }
681    }
682}
683
684impl fmt::Display for JobState {
685    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
686        f.write_str(self.as_str())
687    }
688}
689
690impl FromStr for JobState {
691    type Err = crate::Error;
692
693    fn from_str(s: &str) -> Result<Self, Self::Err> {
694        match s {
695            "queued" => Ok(JobState::Queued),
696            "processing" => Ok(JobState::Processing),
697            "done" => Ok(JobState::Done),
698            "expired" => Ok(JobState::Expired),
699            _ => Err(crate::Error::bad_arg(
700                "s",
701                format!(
702                    "{s} does not correspond with any {} variant",
703                    std::any::type_name::<Self>()
704                ),
705            )),
706        }
707    }
708}
709
710impl<'de> Deserialize<'de> for JobState {
711    fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
712        let str = String::deserialize(deserializer)?;
713        FromStr::from_str(&str).map_err(de::Error::custom)
714    }
715}
716
717// Handles Compression::None being serialized as null in JSON
718fn deserialize_compression<'de, D: serde::Deserializer<'de>>(
719    deserializer: D,
720) -> Result<Compression, D::Error> {
721    let opt = Option::<Compression>::deserialize(deserializer)?;
722    Ok(opt.unwrap_or(Compression::None))
723}
724
725enum Header {
726    Skip,
727    Range(Option<(&'static str, String)>),
728}
729
730#[cfg(test)]
731mod tests {
732    use dbn::Dataset;
733    use reqwest::StatusCode;
734    use serde_json::json;
735    use time::macros::datetime;
736    use wiremock::{
737        matchers::{basic_auth, method, path, query_param, query_param_is_missing},
738        Mock, MockServer, ResponseTemplate,
739    };
740
741    use super::*;
742    use crate::{
743        body_contains,
744        historical::test_infra::{client, API_KEY},
745        historical::API_VERSION,
746    };
747
748    #[tokio::test]
749    async fn test_submit_job() -> crate::Result<()> {
750        const START: time::OffsetDateTime = datetime!(2023 - 06 - 14 00:00 UTC);
751        const END: time::OffsetDateTime = datetime!(2023 - 06 - 17 00:00 UTC);
752        const SCHEMA: Schema = Schema::Trades;
753
754        let mock_server = MockServer::start().await;
755        Mock::given(method("POST"))
756            .and(basic_auth(API_KEY, ""))
757            .and(path(format!("/v{API_VERSION}/batch.submit_job")))
758            .and(body_contains("dataset", "XNAS.ITCH"))
759            .and(body_contains("schema", "trades"))
760            .and(body_contains("symbols", "TSLA"))
761            .and(body_contains(
762                "start",
763                START.unix_timestamp_nanos().to_string(),
764            ))
765            .and(body_contains("encoding", "dbn"))
766            .and(body_contains("compression", "zstd"))
767            .and(body_contains("map_symbols", "false"))
768            .and(body_contains("end", END.unix_timestamp_nanos().to_string()))
769            // // default
770            .and(body_contains("stype_in", "raw_symbol"))
771            .and(body_contains("stype_out", "instrument_id"))
772            .respond_with(
773                ResponseTemplate::new(StatusCode::OK.as_u16()).set_body_json(json!({
774                    "id": "123",
775                    "user_id": "test_user",
776                    "cost_usd": 10.50,
777                    "dataset": "XNAS.ITCH",
778                    "symbols": ["TSLA"],
779                    "stype_in": "raw_symbol",
780                    "stype_out": "instrument_id",
781                    "schema": SCHEMA.as_str(),
782                    "start": "2023-06-14T00:00:00.000000000Z",
783                    "end": "2023-06-17 00:00:00.000000+00:00",
784                    "limit": null,
785                    "encoding": "dbn",
786                    "compression": "zstd",
787                    "pretty_px": false,
788                    "pretty_ts": false,
789                    "map_symbols": false,
790                    "split_symbols": false,
791                    "split_duration": "day",
792                    "split_size": null,
793                    "delivery": "download",
794                    "state": "queued",
795                     "ts_received": "2023-07-19T23:00:04.095538123Z",
796                     "ts_queued": null,
797                     "ts_process_start": null,
798                     "ts_process_done": null,
799                     "ts_expiration": null
800                })),
801            )
802            .mount(&mock_server)
803            .await;
804        let mut target = client(&mock_server);
805        let job_desc = target
806            .batch()
807            .submit_job(
808                &SubmitJobParams::builder()
809                    .dataset(dbn::Dataset::XnasItch)
810                    .schema(SCHEMA)
811                    .symbols("TSLA")
812                    .date_time_range(START..END)
813                    .build(),
814            )
815            .await?;
816        assert_eq!(job_desc.dataset, dbn::Dataset::XnasItch.as_str());
817        Ok(())
818    }
819
820    #[tokio::test]
821    async fn test_submit_job_param_map_symbols() -> crate::Result<()> {
822        const START: time::OffsetDateTime = datetime!(2023 - 06 - 14 00:00 UTC);
823        const END: time::OffsetDateTime = datetime!(2023 - 06 - 17 00:00 UTC);
824
825        // When not explicitly set, map_symbols is None (resolved at request time
826        // based on encoding)
827        let params = SubmitJobParams::builder()
828            .dataset(Dataset::GlbxMdp3)
829            .encoding(Encoding::Dbn)
830            .symbols("ESM5")
831            .schema(Schema::Mbo)
832            .date_time_range(START..END)
833            .build();
834        assert_eq!(params.encoding, Encoding::Dbn);
835        assert!(params.map_symbols.is_none());
836
837        let params = SubmitJobParams::builder()
838            .dataset(Dataset::GlbxMdp3)
839            .encoding(Encoding::Csv)
840            .symbols("ESM5")
841            .schema(Schema::Mbo)
842            .date_time_range(START..END)
843            .build();
844        assert_eq!(params.encoding, Encoding::Csv);
845        assert!(params.map_symbols.is_none());
846
847        // When explicitly set, map_symbols preserves the value
848        let params = SubmitJobParams::builder()
849            .dataset(Dataset::GlbxMdp3)
850            .encoding(Encoding::Json)
851            .symbols("ESM5")
852            .schema(Schema::Mbo)
853            .date_time_range(START..END)
854            .map_symbols(false)
855            .build();
856        assert_eq!(params.encoding, Encoding::Json);
857        assert_eq!(params.map_symbols, Some(false));
858
859        Ok(())
860    }
861
862    #[tokio::test]
863    async fn test_list_jobs() -> crate::Result<()> {
864        const SCHEMA: Schema = Schema::Trades;
865
866        let mock_server = MockServer::start().await;
867        Mock::given(method("GET"))
868            .and(basic_auth(API_KEY, ""))
869            .and(path(format!("/v{API_VERSION}/batch.list_jobs")))
870            .and(query_param_is_missing("states"))
871            .and(query_param_is_missing("since"))
872            .respond_with(
873                ResponseTemplate::new(StatusCode::OK.as_u16()).set_body_json(json!([{
874                    "id": "123",
875                    "user_id": "test_user",
876                    "cost_usd": 10.50,
877                    "dataset": "XNAS.ITCH",
878                    "symbols": "TSLA",
879                    "stype_in": "raw_symbol",
880                    "stype_out": "instrument_id",
881                    "schema": SCHEMA.as_str(),
882                    // test both time formats
883                    "start": "2023-06-14 00:00:00+00:00",
884                    "end": "2023-06-17T00:00:00.012345678Z",
885                    "limit": null,
886                    "encoding": "json",
887                    "compression": "zstd",
888                    "pretty_px": true,
889                    "pretty_ts": false,
890                    "map_symbols": true,
891                    "split_symbols": false,
892                    "split_duration": "day",
893                    "split_size": null,
894                    "delivery": "download",
895                    "state": "processing",
896                     "ts_received": "2023-07-19 23:00:04.095538+00:00",
897                     "ts_queued": "2023-07-19T23:00:08.095538123Z",
898                     "ts_process_start": "2023-07-19 23:01:04.000000+00:00",
899                     "ts_process_done": null,
900                     "ts_expiration": null
901                },
902                {
903                    "id": "XNAS-20250602-5KM3HL5BUW",
904                    "user_id": "AA89XSlBV",
905                    "cost_usd": 0.0,
906                    "dataset": "XNAS.ITCH",
907                    "symbols": "MSFT",
908                    "stype_in": "raw_symbol",
909                    "stype_out": "instrument_id",
910                    "schema": "trades",
911                    "start": "2022-06-10T12:30:00.000000000Z",
912                    "end": "2022-06-10T14:00:00.000000000Z",
913                    "limit": 1000,
914                    "encoding": "csv",
915                    "compression": null,
916                    "pretty_px": false,
917                    "pretty_ts": false,
918                    "map_symbols": true,
919                    "split_symbols": false,
920                    "split_duration": null,
921                    "split_size": null,
922                    "packaging": null,
923                    "delivery": "download",
924                    "record_count": 1000,
925                    "billed_size": 48000,
926                    "actual_size": 94000,
927                    "package_size": 97690,
928                    "state": "done",
929                    "ts_received": "2025-06-02T15:51:19.251582000Z",
930                    "ts_queued": "2025-06-02T15:51:20.997673000Z",
931                    "ts_process_start": "2025-06-02T15:51:45.312317000Z",
932                    "ts_process_done": "2025-06-02T15:51:46.324860000Z",
933                    "ts_expiration": "2025-07-02T16:00:00.000000000Z",
934                    "progress": 100
935                }])),
936            )
937            .mount(&mock_server)
938            .await;
939        let mut target = client(&mock_server);
940        let job_descs = target.batch().list_jobs(&ListJobsParams::default()).await?;
941        assert_eq!(job_descs.len(), 2);
942        let mut job_desc = &job_descs[0];
943        assert_eq!(
944            job_desc.ts_queued.unwrap(),
945            datetime!(2023-07-19 23:00:08.095538123 UTC)
946        );
947        assert_eq!(
948            job_desc.ts_process_start.unwrap(),
949            datetime!(2023-07-19 23:01:04 UTC)
950        );
951        assert_eq!(job_desc.encoding, Encoding::Json);
952        assert!(job_desc.pretty_px);
953        assert!(!job_desc.pretty_ts);
954        assert!(job_desc.map_symbols);
955        assert_eq!(job_desc.split_duration, SplitDuration::Day);
956        assert!(job_desc.progress.is_none());
957
958        job_desc = &job_descs[1];
959        assert_eq!(
960            job_desc.ts_queued.unwrap(),
961            datetime!(2025-06-02 15:51:20.997673000 UTC)
962        );
963        assert_eq!(
964            job_desc.ts_process_start.unwrap(),
965            datetime!(2025-06-02 15:51:45.312317000 UTC)
966        );
967        assert_eq!(job_desc.start, datetime!(2022-06-10 12:30:00.000000000 UTC));
968        assert_eq!(job_desc.end, datetime!(2022-06-10 14:00:00.000000000 UTC));
969        assert_eq!(job_desc.encoding, Encoding::Csv);
970        assert!(!job_desc.pretty_px);
971        assert!(!job_desc.pretty_ts);
972        assert!(job_desc.map_symbols);
973        assert!(!job_desc.split_symbols);
974        assert_eq!(job_desc.split_duration, SplitDuration::None);
975        assert_eq!(job_desc.progress, Some(100));
976
977        Ok(())
978    }
979
980    #[tokio::test]
981    async fn test_get_job_details() -> crate::Result<()> {
982        const SCHEMA: Schema = Schema::Trades;
983        const JOB_ID: &str = "XNAS-20250602-5KM3HL5BUW";
984
985        let mock_server = MockServer::start().await;
986        Mock::given(method("GET"))
987            .and(basic_auth(API_KEY, ""))
988            .and(path(format!("/v{API_VERSION}/batch.get_job_details")))
989            .and(query_param("job_id", JOB_ID))
990            .respond_with(
991                ResponseTemplate::new(StatusCode::OK.as_u16()).set_body_json(json!({
992                    "id": JOB_ID,
993                    "user_id": "test_user",
994                    "cost_usd": 10.50,
995                    "dataset": "XNAS.ITCH",
996                    "symbols": "TSLA",
997                    "stype_in": "raw_symbol",
998                    "stype_out": "instrument_id",
999                    "schema": SCHEMA.as_str(),
1000                    "start": "2023-06-14T00:00:00.000000000Z",
1001                    "end": "2023-06-17T00:00:00.000000000Z",
1002                    "limit": null,
1003                    "encoding": "dbn",
1004                    "compression": "zstd",
1005                    "pretty_px": false,
1006                    "pretty_ts": false,
1007                    "map_symbols": false,
1008                    "split_symbols": false,
1009                    "split_duration": "day",
1010                    "split_size": null,
1011                    "delivery": "download",
1012                    "state": "done",
1013                    "ts_received": "2023-07-19T23:00:04.095538123Z",
1014                    "ts_queued": "2023-07-19T23:00:08.095538123Z",
1015                    "ts_process_start": "2023-07-19T23:01:04.000000000Z",
1016                    "ts_process_done": "2023-07-19T23:02:04.000000000Z",
1017                    "ts_expiration": "2023-08-19T23:00:00.000000000Z",
1018                    "progress": 100
1019                })),
1020            )
1021            .mount(&mock_server)
1022            .await;
1023        let mut target = client(&mock_server);
1024        let job_desc = target.batch().get_job_details(JOB_ID).await?;
1025        assert_eq!(job_desc.id, JOB_ID);
1026        assert_eq!(job_desc.dataset, "XNAS.ITCH");
1027        assert_eq!(job_desc.state, JobState::Done);
1028        assert_eq!(job_desc.progress, Some(100));
1029        Ok(())
1030    }
1031
1032    #[test]
1033    fn test_deserialize_compression() {
1034        #[derive(serde::Deserialize)]
1035        struct Test {
1036            #[serde(deserialize_with = "deserialize_compression")]
1037            compression: Compression,
1038        }
1039
1040        const JSON: &str =
1041            r#"[{"compression":null}, {"compression":"none"}, {"compression":"zstd"}]"#;
1042        let res: Vec<Test> = serde_json::from_str(JSON).unwrap();
1043        assert_eq!(
1044            res.into_iter().map(|t| t.compression).collect::<Vec<_>>(),
1045            vec![Compression::None, Compression::None, Compression::Zstd]
1046        );
1047    }
1048}
databento/historical/batch.rs

databento/historical/
batch.rs
