[Storage] [STG 104] Unique Put Block IDs

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_42c2d90038"
+  "Tag": "python/storage/azure-storage-blob_05e6598c4a"
 }

sdk/storage/azure-storage-blob/azure/storage/blob/_shared/uploads.py

@@ -4,15 +4,17 @@
 # license information.
 # --------------------------------------------------------------------------
 
+import os
 from concurrent import futures
 from io import BytesIO, IOBase, SEEK_CUR, SEEK_END, SEEK_SET, UnsupportedOperation
 from itertools import islice
 from math import ceil
 from threading import Lock
+from uuid import uuid4
 
 from azure.core.tracing.common import with_current_context
 
-from . import encode_base64, url_quote
+from . import encode_base64
 from .request_handlers import get_length
 from .response_handlers import return_response_headers
 
@@ -116,7 +118,7 @@ def upload_substream_blocks(
     else:
         range_ids = [uploader.process_substream_block(b) for b in uploader.get_substream_blocks()]
     if any(range_ids):
-        return sorted(range_ids)
+        return [r[1] for r in sorted(range_ids, key=lambda r: r[0])]
     return []
 
 
@@ -257,9 +259,10 @@ def __init__(self, *args, **kwargs):
         self.current_length = None
 
     def _upload_chunk(self, chunk_offset, chunk_data):
-        # TODO: This is incorrect, but works with recording.
+        # Generate a unique block ID for each staged block. The chunk offset is
+        # still returned so the block list can be committed in the correct order.
         index = f"{chunk_offset:032d}"
-        block_id = encode_base64(url_quote(encode_base64(index)))
+        block_id = encode_base64(f"{uuid4().int:048d}")
         self.service.stage_block(
             block_id,
             len(chunk_data),
@@ -272,7 +275,7 @@ def _upload_chunk(self, chunk_offset, chunk_data):
 
     def _upload_substream_block(self, index, block_stream):
         try:
-            block_id = f"BlockId{(index//self.chunk_size):05}"
+            block_id = encode_base64(os.urandom(9))
             self.service.stage_block(
                 block_id,
                 len(block_stream),
@@ -283,7 +286,7 @@ def _upload_substream_block(self, index, block_stream):
             )
         finally:
             block_stream.close()
-        return block_id
+        return index, block_id
 
 
 class PageBlobChunkUploader(_ChunkUploader):

sdk/storage/azure-storage-blob/azure/storage/blob/_shared/uploads_async.py

@@ -6,16 +6,18 @@
 
 import asyncio  # pylint: disable=do-not-import-asyncio
 import inspect
+import os
 import threading
 from io import UnsupportedOperation
 from itertools import islice
 from math import ceil
 from typing import AsyncGenerator, Union
+from uuid import uuid4
 
-from . import encode_base64, url_quote
+from . import encode_base64
 from .request_handlers import get_length
 from .response_handlers import return_response_headers
-from .uploads import SubStream, IterStreamer  # pylint: disable=unused-import
+from .uploads import IterStreamer, SubStream  # pylint: disable=unused-import
 
 
 async def _async_parallel_uploads(uploader, pending, running):
@@ -140,7 +142,7 @@ async def upload_substream_blocks(
         for block in uploader.get_substream_blocks():
             range_ids.append(await uploader.process_substream_block(block))
     if any(range_ids):
-        return sorted(range_ids)
+        return [r[1] for r in sorted(range_ids, key=lambda r: r[0])]
     return
 
 
@@ -283,9 +285,10 @@ def __init__(self, *args, **kwargs):
         self.current_length = None
 
     async def _upload_chunk(self, chunk_offset, chunk_data):
-        # TODO: This is incorrect, but works with recording.
+        # Generate a unique block ID for each staged block. The chunk offset is
+        # still returned so the block list can be committed in the correct order.
         index = f"{chunk_offset:032d}"
-        block_id = encode_base64(url_quote(encode_base64(index)))
+        block_id = encode_base64(f"{uuid4().int:048d}")
         await self.service.stage_block(
             block_id,
             len(chunk_data),
@@ -298,7 +301,7 @@ async def _upload_chunk(self, chunk_offset, chunk_data):
 
     async def _upload_substream_block(self, index, block_stream):
         try:
-            block_id = f"BlockId{(index//self.chunk_size):05}"
+            block_id = encode_base64(os.urandom(9))
             await self.service.stage_block(
                 block_id,
                 len(block_stream),
@@ -309,7 +312,7 @@ async def _upload_substream_block(self, index, block_stream):
             )
         finally:
             block_stream.close()
-        return block_id
+        return index, block_id
 
 
 class PageBlobChunkUploader(_ChunkUploader):

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

@@ -9,11 +9,11 @@
 import pytest
 
 from devtools_testutils import (
+    add_body_regex_sanitizer,
     add_general_regex_sanitizer,
     add_header_regex_sanitizer,
     add_oauth_response_sanitizer,
     add_uri_regex_sanitizer,
-    test_proxy,
 )
 
 
@@ -36,3 +36,11 @@ def add_sanitizers(test_proxy):
         regex=r"(?<=[?&]sktid=)[^&#]+",
         value="00000000-0000-0000-0000-000000000000",
     )
+    add_uri_regex_sanitizer(
+        regex=r"(?<=[?&]blockid=)[^&#]+",
+        value="00000000-0000-0000-0000-000000000000",
+    )
+    add_body_regex_sanitizer(
+        regex=r"<Latest>[^<]+</Latest>",
+        value="<Latest>00000000-0000-0000-0000-000000000000</Latest>",
+    )

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

@@ -6,8 +6,11 @@
 # pylint: disable=attribute-defined-outside-init, too-many-public-methods
 
 import tempfile
+import threading
+from base64 import b64decode
 from datetime import datetime, timedelta
 from io import BytesIO
+from unittest import mock
 
 import pytest
 import requests
@@ -16,6 +19,7 @@
 from devtools_testutils.storage import StorageRecordedTestCase
 from fake_credentials import CPK_KEY_HASH, CPK_KEY_VALUE
 from settings.testcase import BlobPreparer
+from test_content_validation import _deterministic_urandom
 from test_helpers import _build_base_file_share_headers, _create_file_share_oauth, NonSeekableStream, ProgressTracker
 
 from azure.core.exceptions import HttpResponseError, ResourceExistsError, ResourceModifiedError, ResourceNotFoundError
@@ -33,6 +37,12 @@
     ImmutabilityPolicy,
     StandardBlobTier,
 )
+from azure.storage.blob._shared.uploads import (
+    BlockBlobChunkUploader,
+    upload_data_chunks,
+    upload_substream_blocks,
+)
+
 from azure.storage.blob._shared.validation import calculate_content_md5
 
 # ------------------------------------------------------------------------------
@@ -43,6 +53,34 @@
 # ------------------------------------------------------------------------------
 
 
+def _assert_unique_ordered_block_ids(block_ids, expected_data, staged_blocks, expected_length):
+    # The data was actually split into multiple blocks.
+    assert len(block_ids) > 1
+    # Every block ID is unique (the point of the feature).
+    assert len(set(block_ids)) == len(block_ids)
+    # All block IDs are equal-length, valid base64 (Azure requires equal length per blob).
+    assert len({len(block_id) for block_id in block_ids}) == 1
+    for block_id in block_ids:
+        # Length is fixed per upload path to match the previous scheme (chunk=64, substream=12).
+        assert len(block_id) == expected_length
+        b64decode(block_id)  # must be valid base64
+    # Committing the blocks in the returned order must reproduce the original content.
+    assert b"".join(staged_blocks[block_id] for block_id in block_ids) == expected_data
+
+
+class _RecordingBlockBlobService:
+    """Minimal fake service that records the block IDs and data passed to stage_block."""
+
+    def __init__(self):
+        self._lock = threading.Lock()
+        self.blocks = {}  # block_id -> staged bytes
+
+    def stage_block(self, block_id, length, data, **kwargs):  # pylint: disable=unused-argument
+        content = data.read() if hasattr(data, "read") else data
+        with self._lock:
+            self.blocks[block_id] = content
+
+
 class TestStorageBlockBlob(StorageRecordedTestCase):
     # --Helpers-----------------------------------------------------------------
     def _setup(self, storage_account_name, key, container_name="utcontainer"):
@@ -1780,7 +1818,8 @@ def test_create_blob_with_md5_chunked(self, **kwargs):
         data = self.get_random_bytes(LARGE_BLOB_SIZE)
 
         # Act
-        blob.upload_blob(data, validate_content=True)
+        with mock.patch("os.urandom", _deterministic_urandom()):
+            blob.upload_blob(data, validate_content=True)
 
         # Assert
 
@@ -2123,5 +2162,43 @@ def test_smart_access_tier(self, **kwargs):
             assert blob.blob_tier == StandardBlobTier.SMART
             assert blob.smart_access_tier is not None
 
+    def test_block_blob_upload_generates_unique_ordered_block_ids(self):
+        # Each staged block must get a unique block ID, and the block list must be
+        # committed in byte-offset order even when chunks finish out of order under concurrency.
+        data = b"".join(bytes([i]) * 8 for i in range(20))  # 20 distinct 8-byte chunks
+        service = _RecordingBlockBlobService()
+
+        block_ids = upload_data_chunks(
+            service=service,
+            uploader_class=BlockBlobChunkUploader,
+            total_size=len(data),
+            chunk_size=8,
+            max_concurrency=4,
+            stream=BytesIO(data),
+            validate_content=False,
+            progress_hook=None,
+        )
+
+        # Chunk path uses a 48-digit zero-padded UUID -> 64-char block IDs.
+        _assert_unique_ordered_block_ids(block_ids, data, service.blocks, 64)
+
+    def test_block_blob_substream_upload_generates_unique_ordered_block_ids(self):
+        data = b"".join(bytes([i]) * 8 for i in range(20))
+        service = _RecordingBlockBlobService()
+
+        block_ids = upload_substream_blocks(
+            service=service,
+            uploader_class=BlockBlobChunkUploader,
+            total_size=len(data),
+            chunk_size=8,
+            max_concurrency=4,
+            stream=BytesIO(data),
+            validate_content=False,
+            progress_hook=None,
+        )
+
+        # Substream path uses os.urandom(9) -> 12-char block IDs (matches old length).
+        _assert_unique_ordered_block_ids(block_ids, data, service.blocks, 12)
+
 
 # ------------------------------------------------------------------------------

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

@@ -5,9 +5,11 @@
 # --------------------------------------------------------------------------
 # pylint: disable=attribute-defined-outside-init, too-many-public-methods
 
+import asyncio
 import tempfile
 from datetime import datetime, timedelta
 from io import BytesIO
+from unittest import mock
 
 import aiohttp
 import pytest
@@ -16,6 +18,8 @@
 from devtools_testutils.storage.aio import AsyncStorageRecordedTestCase
 from fake_credentials import CPK_KEY_HASH, CPK_KEY_VALUE
 from settings.testcase import BlobPreparer
+from test_block_blob import _assert_unique_ordered_block_ids
+from test_content_validation import _deterministic_urandom
 from test_helpers_async import (
     NonSeekableStream,
     ProgressTracker,
@@ -36,6 +40,11 @@
     ImmutabilityPolicy,
     StandardBlobTier,
 )
+from azure.storage.blob._shared.uploads_async import (
+    BlockBlobChunkUploader,
+    upload_data_chunks,
+    upload_substream_blocks,
+)
 from azure.storage.blob._shared.validation import calculate_content_md5
 from azure.storage.blob.aio import BlobClient, BlobServiceClient
 
@@ -47,6 +56,19 @@
 # ------------------------------------------------------------------------------
 
 
+class _RecordingBlockBlobServiceAsync:
+    """Minimal fake async service that records the block IDs and data passed to stage_block."""
+
+    def __init__(self):
+        self.blocks = {}  # block_id -> staged bytes
+
+    async def stage_block(self, block_id, length, data=None, body=None, **kwargs):  # pylint: disable=unused-argument
+        content = data if data is not None else body
+        if hasattr(content, "read"):
+            content = content.read()
+        self.blocks[block_id] = content
+
+
 class TestStorageBlockBlobAsync(AsyncStorageRecordedTestCase):
     # --Helpers-----------------------------------------------------------------
     async def _setup(self, storage_account_name, key, container_name="utcontainer"):
@@ -1923,7 +1945,8 @@ async def test_create_blob_with_md5_chunked(self, **kwargs):
         data = self.get_random_bytes(LARGE_BLOB_SIZE)
 
         # Act
-        await blob.upload_blob(data, validate_content=True)
+        with mock.patch("os.urandom", _deterministic_urandom()):
+            await blob.upload_blob(data, validate_content=True)
 
         # Assert
 
@@ -2267,5 +2290,47 @@ async def test_smart_access_tier(self, **kwargs):
             assert blob.blob_tier == StandardBlobTier.SMART
             assert blob.smart_access_tier is not None
 
+    def test_block_blob_upload_generates_unique_ordered_block_ids(self):
+        # Each staged block must get a unique block ID, and the block list must be
+        # committed in byte-offset order even when chunks finish out of order under concurrency.
+        data = b"".join(bytes([i]) * 8 for i in range(20))  # 20 distinct 8-byte chunks
+        service = _RecordingBlockBlobServiceAsync()
+
+        async def run():
+            return await upload_data_chunks(
+                service=service,
+                uploader_class=BlockBlobChunkUploader,
+                total_size=len(data),
+                chunk_size=8,
+                max_concurrency=4,
+                stream=BytesIO(data),
+                validate_content=False,
+                progress_hook=None,
+            )
+
+        block_ids = asyncio.run(run())
+        # Chunk path uses a 48-digit zero-padded UUID -> 64-char block IDs.
+        _assert_unique_ordered_block_ids(block_ids, data, service.blocks, 64)
+
+    def test_block_blob_substream_upload_generates_unique_ordered_block_ids(self):
+        data = b"".join(bytes([i]) * 8 for i in range(20))
+        service = _RecordingBlockBlobServiceAsync()
+
+        async def run():
+            return await upload_substream_blocks(
+                service=service,
+                uploader_class=BlockBlobChunkUploader,
+                total_size=len(data),
+                chunk_size=8,
+                max_concurrency=4,
+                stream=BytesIO(data),
+                validate_content=False,
+                progress_hook=None,
+            )
+
+        block_ids = asyncio.run(run())
+        # Substream path uses os.urandom(9) -> 12-char block IDs (matches old length).
+        _assert_unique_ordered_block_ids(block_ids, data, service.blocks, 12)
+
 
 # ------------------------------------------------------------------------------