MOD: Define a chunk size for requests streaming

Author: nmacholl

CHANGELOG.md

@@ -8,6 +8,7 @@
 
 #### Bug fixes
 - Fixed an issue where `batch.download` and `batch.download_async` would fail if requested files already existed in the output directory
+- Fixed an issue where `batch.download`, `batch.download_async`, and `timeseries.get_range` could use a lot of memory while streaming data
 
 ## 0.33.0 - 2024-04-16
 

databento/common/constants.py

@@ -25,6 +25,8 @@
     x[0]: np.iinfo(x[1]).max for x in InstrumentDefMsg._dtypes if not isinstance(x[1], str)
 }
 
+HTTP_STREAMING_READ_SIZE: Final = 2**12
+
 SCHEMA_STRUCT_MAP: Final[dict[Schema, type[DBNRecord]]] = {
     Schema.DEFINITION: InstrumentDefMsg,
     Schema.IMBALANCE: ImbalanceMsg,

databento/historical/api/batch.py

@@ -24,6 +24,7 @@
 from databento_dbn import SType
 from requests.auth import HTTPBasicAuth
 
+from databento.common.constants import HTTP_STREAMING_READ_SIZE
 from databento.common.enums import Delivery
 from databento.common.enums import Packaging
 from databento.common.enums import SplitDuration
@@ -416,7 +417,7 @@ def _download_batch_file(
                 ) as response:
                     check_http_error(response)
                     with open(output_path, mode=mode) as f:
-                        for chunk in response.iter_content(chunk_size=None):
+                        for chunk in response.iter_content(chunk_size=HTTP_STREAMING_READ_SIZE):
                             f.write(chunk)
             except BentoHttpError as exc:
                 if exc.http_status == 429:

databento/historical/http.py

@@ -8,6 +8,7 @@
 from os import PathLike
 from typing import IO
 from typing import Any
+from typing import Final
 
 import aiohttp
 import requests
@@ -16,6 +17,7 @@
 from requests import Response
 from requests.auth import HTTPBasicAuth
 
+from databento.common.constants import HTTP_STREAMING_READ_SIZE
 from databento.common.dbnstore import DBNStore
 from databento.common.error import BentoClientError
 from databento.common.error import BentoDeprecationWarning
@@ -25,7 +27,7 @@
 from databento.common.system import USER_AGENT
 
 
-WARNING_HEADER_FIELD: str = "X-Warning"
+WARNING_HEADER_FIELD: Final = "X-Warning"
 
 
 class BentoHttpAPI:
@@ -137,7 +139,7 @@ def _stream(
                 writer = open(path, "x+b")
 
             try:
-                for chunk in response.iter_content(chunk_size=None):
+                for chunk in response.iter_content(chunk_size=HTTP_STREAMING_READ_SIZE):
                     writer.write(chunk)
             except Exception as exc:
                 raise BentoError(f"Error streaming response: {exc}") from None