MOD: Streaming exceptions to BentoError

Author: nmacholl

databento/historical/api/batch.py

@@ -20,6 +20,7 @@
 from databento.common.enums import Delivery
 from databento.common.enums import Packaging
 from databento.common.enums import SplitDuration
+from databento.common.error import BentoError
 from databento.common.parsing import datetime_to_string
 from databento.common.parsing import optional_datetime_to_string
 from databento.common.parsing import optional_symbols_list_to_list
@@ -371,8 +372,11 @@ def _download_file(
 
             logger.debug("Starting download of file %s", output_path.name)
             with open(output_path, mode=mode) as f:
-                for chunk in response.iter_content(chunk_size=None):
-                    f.write(chunk)
+                try:
+                    for chunk in response.iter_content(chunk_size=None):
+                        f.write(chunk)
+                except Exception as exc:
+                    raise BentoError(f"Error downloading file: {exc}")
             logger.debug("Download of %s completed", output_path.name)
 
     async def download_async(
@@ -512,9 +516,12 @@ async def _download_file_async(
 
                 logger.debug("Starting async download of file %s", output_path.name)
                 with open(output_path, mode=mode) as f:
-                    async for chunk in response.content.iter_chunks():
-                        data: bytes = chunk[0]
-                        f.write(data)
+                    try:
+                        async for chunk in response.content.iter_chunks():
+                            data: bytes = chunk[0]
+                            f.write(data)
+                    except Exception as exc:
+                        raise BentoError(f"Error downloading file: {exc}")
                 logger.debug("Download of %s completed", output_path.name)
 
     def _get_file_download_headers_and_mode(

databento/historical/http.py

@@ -18,6 +18,7 @@
 from databento.common.dbnstore import DBNStore
 from databento.common.error import BentoClientError
 from databento.common.error import BentoDeprecationWarning
+from databento.common.error import BentoError
 from databento.common.error import BentoServerError
 from databento.common.error import BentoWarning
 from databento.common.system import USER_AGENT
@@ -132,8 +133,11 @@ def _stream(
             else:
                 writer = open(path, "x+b")
 
-            for chunk in response.iter_content(chunk_size=None):
-                writer.write(chunk)
+            try:
+                for chunk in response.iter_content(chunk_size=None):
+                    writer.write(chunk)
+            except Exception as exc:
+                raise BentoError(f"Error streaming response: {exc}")
 
             if path is None:
                 writer.seek(0)
@@ -169,8 +173,11 @@ async def _stream_async(
                 else:
                     writer = open(path, "x+b")
 
-                async for chunk in response.content.iter_chunks():
-                    writer.write(chunk[0])
+                try:
+                    async for chunk in response.content.iter_chunks():
+                        writer.write(chunk[0])
+                except Exception as exc:
+                    raise BentoError(f"Error streaming response: {exc}")
 
                 if path is None:
                     writer.seek(0)