FIX: Python batch download for existing files

Author: nmacholl

CHANGELOG.md

@@ -6,6 +6,9 @@
 - Added `pip-system-certs` dependency for Windows platforms to prevent a connection issue in `requests` when behind a proxy
 - Iteration of the `Live` client will now automatically call `Live.stop` when the iterator is destroyed, such as when a for loop is escaped with an exception or `break` statement.
 
+#### Bug fixes
+- Fixed an issue where `batch.download` and `batch.download_async` would fail if requested files already existed in the output directory
+
 ## 0.33.0 - 2024-04-16
 
 #### Enhancements

databento/historical/api/batch.py

@@ -394,8 +394,16 @@ def _download_batch_file(
             headers: dict[str, str] = self._headers.copy()
             if output_path.exists():
                 existing_size = output_path.stat().st_size
-                headers["Range"] = f"bytes={existing_size}-{batch_download_file.size - 1}"
-                mode = "ab"
+                if existing_size < batch_download_file.size:
+                    headers["Range"] = f"bytes={existing_size}-{batch_download_file.size - 1}"
+                    mode = "ab"
+                elif existing_size == batch_download_file.size:
+                    # File exists and is complete
+                    break
+                else:
+                    raise FileExistsError(
+                        f"Batch file {output_path.name} already exists and has a larger than expected size.",
+                    )
             else:
                 mode = "wb"
             try:
@@ -424,24 +432,26 @@ def _download_batch_file(
                     attempts += 1
                     continue  # try again
                 raise BentoError(f"Error downloading file: {exc}") from None
+            else:
+                break
 
-            logger.debug("Download of %s completed", output_path.name)
-            hash_algo, _, hash_hex = batch_download_file.hash_str.partition(":")
+        logger.debug("Download of %s completed", output_path.name)
+        hash_algo, _, hash_hex = batch_download_file.hash_str.partition(":")
 
-            if hash_algo == "sha256":
-                output_hash = hashlib.sha256(output_path.read_bytes())
-                if output_hash.hexdigest() != hash_hex:
-                    warn_msg = f"Downloaded file failed checksum validation: {output_path.name}"
-                    logger.warning(warn_msg)
-                    warnings.warn(warn_msg, category=BentoWarning)
-            else:
-                logger.warning(
-                    "Skipping %s checksum because %s is not supported",
-                    output_path.name,
-                    hash_algo,
-                )
+        if hash_algo == "sha256":
+            output_hash = hashlib.sha256(output_path.read_bytes())
+            if output_hash.hexdigest() != hash_hex:
+                warn_msg = f"Downloaded file failed checksum validation: {output_path.name}"
+                logger.warning(warn_msg)
+                warnings.warn(warn_msg, category=BentoWarning)
+        else:
+            logger.warning(
+                "Skipping %s checksum because %s is not supported",
+                output_path.name,
+                hash_algo,
+            )
 
-            return output_path
+        return output_path
 
 
 @dataclass

tests/test_historical_batch.py

@@ -276,3 +276,147 @@ def test_batch_download_rate_limit_429(
     assert mocked_batch_list_files.call_args.args == (job_id,)
     assert len(downloaded_files) == 1
     assert downloaded_files[0].read_bytes() == file_content
+
+
+def test_batch_download_file_exists(
+    monkeypatch: pytest.MonkeyPatch,
+    historical_client: Historical,
+    tmp_path: Path,
+) -> None:
+    """
+    Tests batch download by setting up a MagicMock which will return the
+    content "unittest".
+
+    A subsequent call to batch.download should not fail.
+
+    """
+    # Arrange
+    job_id = "GLBX-20220610-5DEFXVTMSM"
+    filename = "glbx-mdp3-20220610.mbo.csv.zst"
+    file_content = b"unittest"
+    file_hash = f"sha256:{hashlib.sha256(file_content).hexdigest()}"
+    file_size = len(file_content)
+
+    # Mock the call to list files so it returns a test manifest
+    monkeypatch.setattr(
+        historical_client.batch,
+        "list_files",
+        mocked_batch_list_files := MagicMock(
+            return_value=[
+                {
+                    "filename": filename,
+                    "hash": file_hash,
+                    "size": file_size,
+                    "urls": {
+                        "https": f"localhost:442/v0/batch/download/TESTUSER/{job_id}/{filename}",
+                        "ftp": "",
+                    },
+                },
+            ],
+        ),
+    )
+
+    # Mock the call for get, so we can simulate a 429 response
+    ok_response = MagicMock()
+    ok_response.__enter__.return_value = MagicMock(
+        status_code=200,
+        iter_content=MagicMock(return_value=iter([file_content])),
+    )
+    monkeypatch.setattr(
+        requests,
+        "get",
+        MagicMock(
+            side_effect=[ok_response],
+        ),
+    )
+
+    # Act
+    historical_client.batch.download(
+        job_id=job_id,
+        output_dir=tmp_path,
+        filename_to_download=filename,
+    )
+
+    downloaded_files = historical_client.batch.download(
+        job_id=job_id,
+        output_dir=tmp_path,
+        filename_to_download=filename,
+    )
+
+    # Assert
+    assert mocked_batch_list_files.call_args.args == (job_id,)
+    assert len(downloaded_files) == 1
+    assert downloaded_files[0].read_bytes() == file_content
+
+
+def test_batch_download_file_larger_than_expected(
+    monkeypatch: pytest.MonkeyPatch,
+    historical_client: Historical,
+    tmp_path: Path,
+) -> None:
+    """
+    Tests batch download by setting up a MagicMock which will return the
+    content "unittest".
+
+    Then, write some extra bytes to that file, and ensure a subsequent
+    call to batch.download will raise an exception.
+
+    """
+    # Arrange
+    job_id = "GLBX-20220610-5DEFXVTMSM"
+    filename = "glbx-mdp3-20220610.mbo.csv.zst"
+    file_content = b"unittest"
+    file_hash = f"sha256:{hashlib.sha256(file_content).hexdigest()}"
+    file_size = len(file_content)
+
+    # Mock the call to list files so it returns a test manifest
+    monkeypatch.setattr(
+        historical_client.batch,
+        "list_files",
+        MagicMock(
+            return_value=[
+                {
+                    "filename": filename,
+                    "hash": file_hash,
+                    "size": file_size,
+                    "urls": {
+                        "https": f"localhost:442/v0/batch/download/TESTUSER/{job_id}/{filename}",
+                        "ftp": "",
+                    },
+                },
+            ],
+        ),
+    )
+
+    # Mock the call for get, so we can simulate a 429 response
+    ok_response = MagicMock()
+    ok_response.__enter__.return_value = MagicMock(
+        status_code=200,
+        iter_content=MagicMock(return_value=iter([file_content])),
+    )
+    monkeypatch.setattr(
+        requests,
+        "get",
+        MagicMock(
+            side_effect=[ok_response],
+        ),
+    )
+
+    # Act
+    downloaded_files = historical_client.batch.download(
+        job_id=job_id,
+        output_dir=tmp_path,
+        filename_to_download=filename,
+    )
+
+    # Increase the existing file size with some junk
+    with downloaded_files[-1].open(mode="ab") as out:
+        out.write(b"junk")
+
+    # Assert
+    with pytest.raises(FileExistsError):
+        historical_client.batch.download(
+            job_id=job_id,
+            output_dir=tmp_path,
+            filename_to_download=filename,
+        )