Stream to and from Azure (#403)

* Stream to and from Azure (#334)

* Stream uploads for azure

* Download using a stream

* Tests and partials

* changelog

* use pathlib open

* use replace instead of rename

* another try at flaky tests

---------

Co-authored-by: George Boot <884482+georgeboot@users.noreply.github.com>

Author: pjbull

HISTORY.md

@@ -3,6 +3,7 @@
 ## UNRELEASED
 - Implement sliced downloads in GSClient. (Issue [#387](https://github.com/drivendataorg/cloudpathlib/issues/387), PR [#389](https://github.com/drivendataorg/cloudpathlib/pull/389))
 - Implement `as_url` with presigned parameter for all backends. (Issue [#235](https://github.com/drivendataorg/cloudpathlib/issues/235), PR [#236](https://github.com/drivendataorg/cloudpathlib/pull/236))
+- Stream to and from Azure Blob Storage. (PR [#403](https://github.com/drivendataorg/cloudpathlib/pull/403))
 
 ## 0.17.0 (2023-12-21)
 

cloudpathlib/azure/azblobclient.py

@@ -118,6 +118,10 @@ def _get_metadata(self, cloud_path: AzureBlobPath) -> Union["BlobProperties", Di
 
         return properties
 
+    @staticmethod
+    def _partial_filename(local_path) -> Path:
+        return Path(str(local_path) + ".part")
+
     def _download_file(
         self, cloud_path: AzureBlobPath, local_path: Union[str, os.PathLike]
     ) -> Path:
@@ -131,7 +135,17 @@ def _download_file(
 
         local_path.parent.mkdir(exist_ok=True, parents=True)
 
-        local_path.write_bytes(download_stream.content_as_bytes())
+        try:
+            partial_local_path = self._partial_filename(local_path)
+            with partial_local_path.open("wb") as data:
+                download_stream.readinto(data)
+
+            partial_local_path.replace(local_path)
+        except:  # noqa: E722
+            # remove any partial download
+            if partial_local_path.exists():
+                partial_local_path.unlink()
+            raise
 
         return local_path
 
@@ -273,7 +287,8 @@ def _upload_file(
 
         content_settings = ContentSettings(**extra_args)
 
-        blob.upload_blob(Path(local_path).read_bytes(), overwrite=True, content_settings=content_settings)  # type: ignore
+        with Path(local_path).open("rb") as data:
+            blob.upload_blob(data, overwrite=True, content_settings=content_settings)  # type: ignore
 
         return cloud_path
 

cloudpathlib/local/localclient.py

@@ -5,6 +5,7 @@
 from pathlib import Path, PurePosixPath
 import shutil
 from tempfile import TemporaryDirectory
+from time import sleep
 from typing import Callable, Dict, Iterable, List, Optional, Tuple, Union
 
 from ..client import Client
@@ -63,7 +64,17 @@ def _local_to_cloud_path(self, local_path: Union[str, os.PathLike]) -> "LocalPat
     def _download_file(self, cloud_path: "LocalPath", local_path: Union[str, os.PathLike]) -> Path:
         local_path = Path(local_path)
         local_path.parent.mkdir(exist_ok=True, parents=True)
-        shutil.copyfile(self._cloud_path_to_local(cloud_path), local_path)
+
+        try:
+            shutil.copyfile(self._cloud_path_to_local(cloud_path), local_path)
+        except FileNotFoundError:
+            # erroneous FileNotFoundError appears in tests sometimes; patiently insist on the parent directory existing
+            sleep(1.0)
+            local_path.parent.mkdir(exist_ok=True, parents=True)
+            sleep(1.0)
+
+            shutil.copyfile(self._cloud_path_to_local(cloud_path), local_path)
+
         return local_path
 
     def _exists(self, cloud_path: "LocalPath") -> bool:

tests/mock_clients/mock_azureblob.py

@@ -111,7 +111,7 @@ def delete_blob(self):
     def upload_blob(self, data, overwrite, content_settings=None):
         path = self.root / self.key
         path.parent.mkdir(parents=True, exist_ok=True)
-        path.write_bytes(data)
+        path.write_bytes(data.read())
 
         if content_settings is not None:
             self.service_client.metadata_cache[self.root / self.key] = (
@@ -130,6 +130,9 @@ def readall(self):
     def content_as_bytes(self):
         return self.readall()
 
+    def readinto(self, buffer):
+        buffer.write(self.readall())
+
 
 class MockContainerClient:
     def __init__(self, root, container_name):

tests/test_azure_specific.py

@@ -1,12 +1,15 @@
 import os
 
+from azure.storage.blob import StorageStreamDownloader
 import pytest
 
 from urllib.parse import urlparse, parse_qs
 from cloudpathlib import AzureBlobClient, AzureBlobPath
 from cloudpathlib.exceptions import MissingCredentialsError
 from cloudpathlib.local import LocalAzureBlobClient, LocalAzureBlobPath
 
+from .mock_clients.mock_azureblob import MockStorageStreamDownloader
+
 
 @pytest.mark.parametrize("path_class", [AzureBlobPath, LocalAzureBlobPath])
 def test_azureblobpath_properties(path_class, monkeypatch):
@@ -45,3 +48,34 @@ def test_as_url(azure_rig):
     assert "sp" in query_params
     assert "sr" in query_params
     assert "sig" in query_params
+
+
+def test_partial_download(azure_rig, monkeypatch):
+    p: AzureBlobPath = azure_rig.create_cloud_path("dir_0/file0_0.txt")
+
+    # no partial after successful download
+    p.read_text()  # downloads
+    assert p._local.exists()
+    assert not p.client._partial_filename(p._local).exists()
+
+    # remove cache manually
+    p._local.unlink()
+    assert not p._local.exists()
+
+    # no partial after failed download
+    with monkeypatch.context() as m:
+
+        def _patched(self, buffer):
+            buffer.write(b"partial")
+            raise Exception("boom")
+
+        if azure_rig.live_server:
+            m.setattr(StorageStreamDownloader, "readinto", _patched)
+        else:
+            m.setattr(MockStorageStreamDownloader, "readinto", _patched)
+
+        with pytest.raises(Exception):
+            p.read_text()  # downloads; should raise
+
+        assert not p._local.exists()
+        assert not p.client._partial_filename(p._local).exists()