[Storage] Optimize single-shot Block Blob upload (#34676)

Author: jalauzon-msft

sdk/storage/azure-storage-blob/CHANGELOG.md

@@ -15,6 +15,7 @@ Python 3.12.
 using async OAuth credentials.
 - Fixed an typing issue which incorrectly typed the `readinto` API. The correct input type is `IO[bytes]`.
 - Fixed a typo in the initialization of `completion_time` for the `CopyProperties` model.
+- Fixed a couple of issues with `upload_blob` when using Iterators/Generators as the data input.
 
 ### Other Changes
 - Passing `prefix` to the following `ContainerClient` APIs now raises a `ValueError`:

sdk/storage/azure-storage-blob/assets.json

@@ -2,5 +2,5 @@
   "AssetsRepo": "Azure/azure-sdk-assets",
   "AssetsRepoPrefixPath": "python",
   "TagPrefix": "python/storage/azure-storage-blob",
-  "Tag": "python/storage/azure-storage-blob_1b66da54e8"
+  "Tag": "python/storage/azure-storage-blob_20d52f0450"
 }

sdk/storage/azure-storage-blob/azure/storage/blob/_blob_client.py

@@ -449,7 +449,6 @@ def _upload_blob_options(  # pylint:disable=too-many-statements
 
         if blob_type == BlobType.BlockBlob:
             kwargs['client'] = self._client.block_blob
-            kwargs['data'] = data
         elif blob_type == BlobType.PageBlob:
             if self.encryption_version == '2.0' and (self.require_encryption or self.key_encryption_key is not None):
                 raise ValueError("Encryption version 2.0 does not currently support page blobs.")
@@ -615,7 +614,7 @@ def upload_blob_from_url(self, source_url, **kwargs):
 
     @distributed_trace
     def upload_blob(
-            self, data: Union[bytes, str, Iterable[AnyStr], IO[AnyStr]],
+            self, data: Union[bytes, str, Iterable[AnyStr], IO[bytes]],
             blob_type: Union[str, BlobType] = BlobType.BlockBlob,
             length: Optional[int] = None,
             metadata: Optional[Dict[str, str]] = None,

sdk/storage/azure-storage-blob/azure/storage/blob/_encryption.py

@@ -15,7 +15,7 @@
     dumps,
     loads,
 )
-from typing import Any, BinaryIO, Callable, Dict, Optional, Tuple, TYPE_CHECKING
+from typing import Any, Callable, Dict, IO, Optional, Tuple, TYPE_CHECKING
 from typing import OrderedDict as TypedOrderedDict
 from typing_extensions import Protocol
 
@@ -220,11 +220,11 @@ class GCMBlobEncryptionStream:
     """
     def __init__(
         self, content_encryption_key: bytes,
-        data_stream: BinaryIO,
+        data_stream: IO[bytes],
     ) -> None:
         """
         :param bytes content_encryption_key: The encryption key to use.
-        :param BinaryIO data_stream: The data stream to read data from.
+        :param IO[bytes] data_stream: The data stream to read data from.
         """
         self.content_encryption_key = content_encryption_key
         self.data_stream = data_stream
@@ -261,28 +261,30 @@ def read(self, size: int = -1) -> bytes:
                     # No more data to read
                     break
 
-                self.current = self._encrypt_region(data)
+                self.current = encrypt_data_v2(data, self.nonce_counter, self.content_encryption_key)
+                # IMPORTANT: Must increment the nonce each time.
+                self.nonce_counter += 1
 
         return result.getvalue()
 
-    def _encrypt_region(self, data: bytes) -> bytes:
-        """
-        Encrypt the given region of data using AES-GCM. The result
-        includes the data in the form: nonce + ciphertext + tag.
 
-        :param bytes data: The data to encrypt.
-        :return: The encrypted bytes.
-        :rtype: bytes
-        """
-        # Each region MUST use a different nonce
-        nonce = self.nonce_counter.to_bytes(_GCM_NONCE_LENGTH, 'big')
-        self.nonce_counter += 1
+def encrypt_data_v2(data: bytes, nonce: int, key: bytes) -> bytes:
+    """
+    Encrypts the given data using the given nonce and key using AES-GCM.
+    The result includes the data in the form: nonce + ciphertext + tag.
 
-        aesgcm = AESGCM(self.content_encryption_key)
+    :param bytes data: The raw data to encrypt.
+    :param int nonce: The nonce to use for encryption.
+    :param bytes key: The encryption key to use for encryption.
+    :return: The encrypted bytes in the form: nonce + ciphertext + tag.
+    :rtype: bytes
+    """
+    nonce_bytes = nonce.to_bytes(_GCM_NONCE_LENGTH, 'big')
+    aesgcm = AESGCM(key)
 
-        # Returns ciphertext + tag
-        ciphertext_with_tag = aesgcm.encrypt(nonce, data, None)
-        return nonce + ciphertext_with_tag
+    # Returns ciphertext + tag
+    ciphertext_with_tag = aesgcm.encrypt(nonce_bytes, data, None)
+    return nonce_bytes + ciphertext_with_tag
 
 
 def is_encryption_v2(encryption_data: Optional[_EncryptionData]) -> bool:

sdk/storage/azure-storage-blob/azure/storage/blob/_upload_helpers.py

@@ -64,7 +64,6 @@ def _any_conditions(modified_access_conditions=None, **kwargs):  # pylint: disab
 
 def upload_block_blob(  # pylint: disable=too-many-locals, too-many-statements
         client=None,
-        data=None,
         stream=None,
         length=None,
         overwrite=None,
@@ -92,12 +91,10 @@ def upload_block_blob(  # pylint: disable=too-many-locals, too-many-statements
 
         # Do single put if the size is smaller than or equal config.max_single_put_size
         if adjusted_count is not None and (adjusted_count <= blob_settings.max_single_put_size):
-            try:
-                data = data.read(length)
-                if not isinstance(data, bytes):
-                    raise TypeError('Blob data should be of type bytes.')
-            except AttributeError:
-                pass
+            data = stream.read(length)
+            if not isinstance(data, bytes):
+                raise TypeError('Blob data should be of type bytes.')
+
             if encryption_options.get('key'):
                 encryption_data, data = encrypt_blob(data, encryption_options['key'], encryption_options['version'])
                 headers['x-ms-meta-encryptiondata'] = encryption_data

sdk/storage/azure-storage-blob/azure/storage/blob/aio/_encryption_async.py

@@ -0,0 +1,72 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for
+# license information.
+# --------------------------------------------------------------------------
+
+import inspect
+import sys
+from io import BytesIO
+from typing import IO
+
+from .._encryption import _GCM_REGION_DATA_LENGTH, encrypt_data_v2
+
+
+class GCMBlobEncryptionStream:
+    """
+    An async stream that performs AES-GCM encryption on the given data as
+    it's streamed. Data is read and encrypted in regions. The stream
+    will use the same encryption key and will generate a guaranteed unique
+    nonce for each encryption region.
+    """
+    def __init__(
+        self, content_encryption_key: bytes,
+        data_stream: IO[bytes],
+    ) -> None:
+        """
+        :param bytes content_encryption_key: The encryption key to use.
+        :param IO[bytes] data_stream: The data stream to read data from.
+        """
+        self.content_encryption_key = content_encryption_key
+        self.data_stream = data_stream
+
+        self.offset = 0
+        self.current = b''
+        self.nonce_counter = 0
+
+    async def read(self, size: int = -1) -> bytes:
+        """
+        Read data from the stream. Specify -1 to read all available data.
+
+        :param int size: The amount of data to read. Defaults to -1 for all data.
+        :return: The bytes read.
+        :rtype: bytes
+        """
+        result = BytesIO()
+        remaining = sys.maxsize if size == -1 else size
+
+        while remaining > 0:
+            # Start by reading from current
+            if len(self.current) > 0:
+                read = min(remaining, len(self.current))
+                result.write(self.current[:read])
+
+                self.current = self.current[read:]
+                self.offset += read
+                remaining -= read
+
+            if remaining > 0:
+                # Read one region of data and encrypt it
+                data = self.data_stream.read(_GCM_REGION_DATA_LENGTH)
+                if inspect.isawaitable(data):
+                    data = await data
+
+                if len(data) == 0:
+                    # No more data to read
+                    break
+
+                self.current = encrypt_data_v2(data, self.nonce_counter, self.content_encryption_key)
+                # IMPORTANT: Must increment the nonce each time.
+                self.nonce_counter += 1
+
+        return result.getvalue()

sdk/storage/azure-storage-blob/azure/storage/blob/aio/_upload_helpers.py

@@ -23,8 +23,8 @@
     AppendPositionAccessConditions,
     ModifiedAccessConditions,
 )
+from ._encryption_async import GCMBlobEncryptionStream
 from .._encryption import (
-    GCMBlobEncryptionStream,
     encrypt_blob,
     get_adjusted_upload_size,
     get_blob_encryptor_and_padder,
@@ -40,7 +40,6 @@
 
 async def upload_block_blob(  # pylint: disable=too-many-locals, too-many-statements
         client=None,
-        data=None,
         stream=None,
         length=None,
         overwrite=None,
@@ -68,17 +67,16 @@ async def upload_block_blob(  # pylint: disable=too-many-locals, too-many-statem
 
         # Do single put if the size is smaller than config.max_single_put_size
         if adjusted_count is not None and (adjusted_count <= blob_settings.max_single_put_size):
-            try:
-                data = data.read(length)
-                if inspect.isawaitable(data):
-                    data = await data
-                if not isinstance(data, bytes):
-                    raise TypeError('Blob data should be of type bytes.')
-            except AttributeError:
-                pass
+            data = stream.read(length)
+            if inspect.isawaitable(data):
+                data = await data
+            if not isinstance(data, bytes):
+                raise TypeError('Blob data should be of type bytes.')
+
             if encryption_options.get('key'):
                 encryption_data, data = encrypt_blob(data, encryption_options['key'], encryption_options['version'])
                 headers['x-ms-meta-encryptiondata'] = encryption_data
+
             response = await client.upload(
                 body=data,
                 content_length=adjusted_count,

sdk/storage/azure-storage-blob/tests/test_blob_encryption_v2.py

@@ -580,6 +580,88 @@ def test_put_blob_multi_region_chunked_size_greater_region(self, **kwargs):
         # Assert
         assert content == data
 
+    @pytest.mark.live_test_only
+    @BlobPreparer()
+    def test_put_blob_other_data_types(self, **kwargs):
+        storage_account_name = kwargs.pop("storage_account_name")
+        storage_account_key = kwargs.pop("storage_account_key")
+
+        self._setup(storage_account_name, storage_account_key)
+        kek = KeyWrapper('key1')
+        bsc = BlobServiceClient(
+            self.account_url(storage_account_name, "blob"),
+            credential=storage_account_key,
+            require_encryption=True,
+            encryption_version='2.0',
+            key_encryption_key=kek)
+
+        blob = bsc.get_blob_client(self.container_name, self._get_blob_reference())
+
+        content = b'Hello World Encrypted!'
+        length = len(content)
+        byte_io = BytesIO(content)
+
+        def generator():
+            yield b'Hello '
+            yield b'World '
+            yield b'Encrypted!'
+
+        def text_generator():
+            yield 'Hello '
+            yield 'World '
+            yield 'Encrypted!'
+
+        data_list = [byte_io, generator(), text_generator()]
+
+        # Act
+        for data in data_list:
+            blob.upload_blob(data, length=length, overwrite=True)
+            result = blob.download_blob().readall()
+
+            # Assert
+            assert content == result
+
+    @pytest.mark.live_test_only
+    @BlobPreparer()
+    def test_put_blob_other_data_types_chunked(self, **kwargs):
+        storage_account_name = kwargs.pop("storage_account_name")
+        storage_account_key = kwargs.pop("storage_account_key")
+
+        self._setup(storage_account_name, storage_account_key)
+        kek = KeyWrapper('key1')
+        bsc = BlobServiceClient(
+            self.account_url(storage_account_name, "blob"),
+            credential=storage_account_key,
+            max_single_put_size=1024,
+            max_block_size=1024,
+            require_encryption=True,
+            encryption_version='2.0',
+            key_encryption_key=kek)
+
+        blob = bsc.get_blob_client(self.container_name, self._get_blob_reference())
+
+        content = b'abcde' * 1030  # 5 KiB + 30
+        byte_io = BytesIO(content)
+
+        def generator():
+            for i in range(0, len(content), 500):
+                yield content[i: i + 500]
+
+        def text_generator():
+            s_content = str(content, encoding='utf-8')
+            for i in range(0, len(s_content), 500):
+                yield s_content[i: i + 500]
+
+        data_list = [byte_io, generator(), text_generator()]
+
+        # Act
+        for data in data_list:
+            blob.upload_blob(data, overwrite=True)
+            result = blob.download_blob().readall()
+
+            # Assert
+            assert content == result
+
     @pytest.mark.live_test_only
     @BlobPreparer()
     def test_get_blob_range_single_region(self, **kwargs):