ADD: Add programmatic batch downloads

Author: cjdsellers

CHANGELOG.md

@@ -2,6 +2,8 @@
 
 ## 0.8.0 - TBD
 - Renamed 'dbz' encoding to 'dbn'
+- Added `batch.list_files(...)` method
+- Added `batch.download(...)` method
 - Changed `.to_df(...)` `pretty_ts` default argument to `True`
 - Changed `.to_df(...)` `pretty_px` default argument to `True`
 - Changed `.to_df(...)` `map_symbols` default argument to `True`

databento/common/bento.py

@@ -1,7 +1,8 @@
 import datetime as dt
 import io
 import os.path
-from typing import TYPE_CHECKING, Any, BinaryIO, Callable, Dict, List, Optional
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, BinaryIO, Callable, Dict, List, Optional, Union
 
 import numpy as np
 import pandas as pd
@@ -536,13 +537,13 @@ def replay(self, callback: Callable[[Any], None]) -> None:
             callback(record[0])
 
     @staticmethod
-    def from_file(path: str) -> "FileBento":
+    def from_file(path: Union[Path, str]) -> "FileBento":
         """
         Load the data from a DBN file at the given path.
 
         Parameters
         ----------
-        path : str
+        path : Path or str
             The path to read from.
 
         Returns
@@ -569,7 +570,7 @@ def from_file(path: str) -> "FileBento":
 
         return bento
 
-    def to_file(self, path: str) -> "FileBento":
+    def to_file(self, path: Union[Path, str]) -> "FileBento":
         """
         Write the data to a DBN file at the given path.
 
@@ -591,13 +592,13 @@ def to_file(self, path: str) -> "FileBento":
 
         return bento
 
-    def to_csv(self, path: str) -> None:
+    def to_csv(self, path: Union[Path, str]) -> None:
         """
         Write the data to a file in CSV format.
 
         Parameters
         ----------
-        path : str
+        path : Path or str
             The file path to write to.
 
         Notes
@@ -611,13 +612,13 @@ def to_csv(self, path: str) -> None:
             map_symbols=False,
         ).to_csv(path)
 
-    def to_json(self, path: str) -> None:
+    def to_json(self, path: Union[Path, str]) -> None:
         """
         Write the data to a file in JSON format.
 
         Parameters
         ----------
-        path : str
+        path : Path or str
             The file path to write to.
 
         Notes
@@ -664,7 +665,7 @@ def request_symbology(self, client: "Historical") -> Dict[str, Any]:
     def request_full_definitions(
         self,
         client: "Historical",
-        path: Optional[str] = None,
+        path: Optional[Union[Path, str]] = None,
     ) -> "Bento":
         """
         Request full instrument definitions based on the metadata properties.
@@ -675,7 +676,7 @@ def request_full_definitions(
         ----------
         client : Historical
             The historical client to use for the request (contains the API key).
-        path : str, optional
+        path : Path or str, optional
             The path to stream the data to on disk (will then return a `FileBento`).
 
         Returns
@@ -771,11 +772,11 @@ class FileBento(Bento):
 
     Parameters
     ----------
-    path : str
+    path : Path or str
         The path to the data file.
     """
 
-    def __init__(self, path: str):
+    def __init__(self, path: Union[Path, str]):
         super().__init__()
 
         self._path = path
@@ -790,7 +791,7 @@ def path(self) -> str:
         str
 
         """
-        return self._path
+        return str(self._path)
 
     @property
     def nbytes(self) -> int:

databento/common/parsing.py

@@ -42,7 +42,7 @@ def optional_values_list_to_string(
 
     Returns
     -------
-    str or ``None``
+    str or `None`
 
     """
     if values is None:
@@ -99,7 +99,7 @@ def optional_date_to_string(value: Optional[Union[date, str]]) -> Optional[str]:
 
     Returns
     -------
-    str or ``None``
+    str or `None`
 
     """
     if value is None:
@@ -155,7 +155,7 @@ def optional_datetime_to_string(
 
     Returns
     -------
-    str or ``None``
+    str or `None`
 
     """
     if value is None:

databento/historical/api/batch.py

@@ -1,7 +1,10 @@
+import os
 from datetime import date
+from pathlib import Path
 from typing import Any, Dict, List, Optional, Tuple, Union
 
 import pandas as pd
+import requests
 from databento.common.enums import (
     Compression,
     Dataset,
@@ -12,6 +15,7 @@
     SplitDuration,
     SType,
 )
+from databento.common.logging import log_error, log_info
 from databento.common.parsing import (
     datetime_to_string,
     optional_datetime_to_string,
@@ -20,7 +24,8 @@
 )
 from databento.common.validation import validate_enum
 from databento.historical.api import API_VERSION
-from databento.historical.http import BentoHttpAPI
+from databento.historical.http import BentoHttpAPI, check_http_error
+from requests.auth import HTTPBasicAuth
 
 
 class BatchHttpAPI(BentoHttpAPI):
@@ -69,7 +74,7 @@ def submit_job(
         symbols : List[Union[str, int]] or str
             The product symbols to filter for. Takes up to 2,000 symbols per request.
             If more than 1 symbol is specified, the data is merged and sorted by time.
-            If 'ALL_SYMBOLS' or ``None`` then will be for **all** symbols.
+            If 'ALL_SYMBOLS' or `None` then will be for **all** symbols.
         schema : Schema or str {'mbo', 'mbp-1', 'mbp-10', 'trades', 'tbbo', 'ohlcv-1s', 'ohlcv-1m', 'ohlcv-1h', 'ohlcv-1d', 'definition', 'statistics', 'status'}, default 'trades'  # noqa
             The data record schema for the request.
         encoding : Encoding or str {'dbn', 'csv', 'json'}, default 'dbn'
@@ -90,7 +95,7 @@ def submit_job(
         stype_out : SType or str, default 'product_id'
             The output symbology type to resolve to.
         limit : int, optional
-            The maximum number of records for the request. If ``None`` then no limit.
+            The maximum number of records for the request. If `None` then no limit.
 
         Returns
         -------
@@ -172,3 +177,138 @@ def list_jobs(
             params=params,
             basic_auth=True,
         ).json()
+
+    def list_files(self, job_id: str) -> List[Dict[str, Any]]:
+        """
+        Request details of all files for a specific batch job.
+
+        Makes a `GET /batch.list_files` HTTP request.
+
+        Parameters
+        ----------
+        job_id : str
+            The batch job identifier.
+
+        Returns
+        -------
+        List[Dict[str, Any]]
+            The file details for the batch job.
+
+        """
+        params: List[Tuple[str, str]] = [
+            ("job_id", job_id),
+        ]
+
+        return self._get(
+            url=self._base_url + ".list_files",
+            params=params,
+            basic_auth=True,
+        ).json()
+
+    def download(
+        self,
+        output_dir: Union[Path, str],
+        job_id: str,
+        filename_to_download: Optional[str] = None,
+    ) -> None:
+        """
+        Download a batch job or a specific file to `{output_dir}/{job_id}/`.
+         - Creates the directories if any do not already exist
+         - Partially downloaded files will be retried using a range request
+
+        Makes one or many `GET /batch.download` HTTP request(s).
+
+        Parameters
+        ----------
+        output_dir: Path or str
+            The directory to download the file(s) to.
+        job_id : str
+            The batch job identifier.
+        filename_to_download : str, optional
+            The specific file to download.
+            If `None` then will download all files for the batch job.
+
+        """
+        params: List[Tuple[str, str]] = [
+            ("job_id", job_id),
+        ]
+
+        job_files: List[Dict[str, Union[str, int]]] = self._get(
+            url=self._base_url + ".list_files",
+            params=params,
+            basic_auth=True,
+        ).json()
+
+        if not job_files:
+            log_error(f"Cannot download batch job {job_id} (no files found).")
+            return
+
+        if filename_to_download:
+            # A specific file is being requested
+            is_file_found = False
+            for details in job_files:
+                if details["filename"] == filename_to_download:
+                    # Reduce files to download only the single file
+                    job_files = [details]
+                    is_file_found = True
+                    break
+            if not is_file_found:
+                log_error(
+                    f"Cannot download batch job {job_id} file "
+                    f"({filename_to_download} not found)",
+                )
+                return
+
+        # Prepare job directory
+        job_dir = os.path.join(output_dir, job_id)
+        os.makedirs(job_dir, exist_ok=True)
+
+        for details in job_files:
+            filename = str(details["filename"])
+            output_path = os.path.join(job_dir, filename)
+            log_info(
+                f"Downloading batch job file {filename} to {output_path} ...",
+            )
+
+            self._download_file(
+                job_id=job_id,
+                filename=filename,
+                filesize=int(details["size"]),
+                output_path=output_path,
+            )
+
+    def _download_file(
+        self,
+        job_id: str,
+        filename: str,
+        filesize: int,
+        output_path: str,
+    ) -> None:
+        params: List[Tuple[str, str]] = [
+            ("job_id", job_id),
+            ("filename", filename),
+        ]
+
+        headers: Dict[str, str] = self._headers.copy()
+
+        # Check if file already exists in partially downloaded state
+        if os.path.isfile(output_path):
+            existing_size = os.path.getsize(output_path)
+            if existing_size < filesize:
+                # Make range request for partial download,
+                # will be from next byte to end of file.
+                headers["Range"] = f"bytes={existing_size}-{filesize - 1}"
+
+        with requests.get(
+            url=self._base_url + ".download",
+            params=params,
+            headers=headers,
+            auth=HTTPBasicAuth(username=self._key, password=None),
+            allow_redirects=True,
+            stream=True,
+        ) as response:
+            check_http_error(response)
+
+            with open(output_path, mode="wb") as f:
+                for chunk in response.iter_content():
+                    f.write(chunk)

databento/historical/api/metadata.py

@@ -302,14 +302,14 @@ def get_record_count(
             If an integer is passed, then this represents nanoseconds since UNIX epoch.
         symbols : List[Union[str, int]] or str, optional
             The product symbols to filter for. Takes up to 2,000 symbols per request.
-            If 'ALL_SYMBOLS' or ``None`` then will be for **all** symbols.
+            If 'ALL_SYMBOLS' or `None` then will be for **all** symbols.
         schema : Schema or str {'mbo', 'mbp-1', 'mbp-10', 'trades', 'tbbo', 'ohlcv-1s', 'ohlcv-1m', 'ohlcv-1h', 'ohlcv-1d', 'definition', 'statistics', 'status'}, default 'trades'  # noqa
             The data record schema for the request.
         stype_in : SType or str, default 'native'
             The input symbology type to resolve from.
         limit : int, optional
             The maximum number of records to include in the query.
-            If ``None`` then no limit.
+            If `None` then no limit.
 
         Returns
         -------
@@ -368,13 +368,13 @@ def get_billable_size(
             If an integer is passed, then this represents nanoseconds since UNIX epoch.
         symbols : List[Union[str, int]] or str, optional
             The product symbols to filter for. Takes up to 2,000 symbols per request.
-            If 'ALL_SYMBOLS' or ``None`` then will be for **all** symbols.
+            If 'ALL_SYMBOLS' or `None` then will be for **all** symbols.
         schema : Schema or str {'mbo', 'mbp-1', 'mbp-10', 'trades', 'tbbo', 'ohlcv-1s', 'ohlcv-1m', 'ohlcv-1h', 'ohlcv-1d', 'definition', 'statistics', 'status'}, default 'trades'  # noqa
             The data record schema for the request.
         stype_in : SType or str, default 'native'
             The input symbology type to resolve from.
         limit : int, optional
-            The maximum number of records to include in the query. If ``None`` then no limit.
+            The maximum number of records to include in the query. If `None` then no limit.
 
         Returns
         -------
@@ -436,13 +436,13 @@ def get_cost(
             The data feed mode for the request.
         symbols : List[Union[str, int]] or str, optional
             The product symbols to filter for. Takes up to 2,000 symbols per request.
-            If 'ALL_SYMBOLS' or ``None`` then will be for **all** symbols.
+            If 'ALL_SYMBOLS' or `None` then will be for **all** symbols.
         schema : Schema or str {'mbo', 'mbp-1', 'mbp-10', 'trades', 'tbbo', 'ohlcv-1s', 'ohlcv-1m', 'ohlcv-1h', 'ohlcv-1d', 'definition', 'statistics', 'status'}, default 'trades'  # noqa
             The data record schema for the request.
         stype_in : SType or str, default 'native'
             The input symbology type to resolve from.
         limit : int, optional
-            The maximum number of records to include in the query. If ``None`` then no limit.
+            The maximum number of records to include in the query. If `None` then no limit.
 
         Returns
         -------