[Storage] [STG 104] Unique Put Block IDs

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

@@ -9,10 +9,11 @@
 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 +117,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 [block_id for _, block_id in sorted(range_ids, key=lambda r: r[0])]
     return []
 
 
@@ -257,9 +258,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(uuid4().bytes)
         self.service.stage_block(
             block_id,
             len(chunk_data),
@@ -272,7 +274,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(uuid4().bytes)
             self.service.stage_block(
                 block_id,
                 len(block_stream),
@@ -283,7 +285,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

@@ -11,11 +11,12 @@
 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 +141,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 [block_id for _, block_id in sorted(range_ids, key=lambda r: r[0])]
     return
 
 
@@ -283,9 +284,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(uuid4().bytes)
         await self.service.stage_block(
             block_id,
             len(chunk_data),
@@ -298,7 +300,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(uuid4().bytes)
             await self.service.stage_block(
                 block_id,
                 len(block_stream),
@@ -309,7 +311,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_common_blob.py

@@ -5,8 +5,10 @@
 # --------------------------------------------------------------------------
 # pylint: disable=attribute-defined-outside-init, too-many-public-methods
 
+import base64
 import os
 import tempfile
+import threading
 import uuid
 from datetime import datetime, timedelta
 from enum import Enum
@@ -58,6 +60,11 @@
     upload_blob_to_url,
 )
 from azure.storage.blob._generated.models import RehydratePriority
+from azure.storage.blob._shared.uploads import (
+    BlockBlobChunkUploader,
+    upload_data_chunks,
+    upload_substream_blocks,
+)
 
 # ------------------------------------------------------------------------------
 SMALL_BLOB_SIZE = 1024
@@ -66,6 +73,32 @@
 # ------------------------------------------------------------------------------
 
 
+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
+
+
+def _assert_unique_ordered_block_ids(block_ids, expected_data, staged_blocks):
+    # 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:
+        assert len(base64.b64decode(block_id)) == 16
+    # 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 TestStorageCommonBlob(StorageRecordedTestCase):
     def _setup(self, storage_account_name, key):
         self.bsc = BlobServiceClient(self.account_url(storage_account_name, "blob"), credential=key.secret)
@@ -4020,4 +4053,40 @@ def test_download_blob_returns_access_tier(self, **kwargs):
         assert downloader.properties.blob_tier_change_time is not None
         assert not downloader.properties.blob_tier_inferred
 
+    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,
+        )
+
+        _assert_unique_ordered_block_ids(block_ids, data, service.blocks)
+
+    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,
+        )
+
+        _assert_unique_ordered_block_ids(block_ids, data, service.blocks)
+
     # ------------------------------------------------------------------------------

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

@@ -6,6 +6,7 @@
 # pylint: disable=attribute-defined-outside-init, too-many-public-methods
 
 import asyncio
+import base64
 import os
 import tempfile
 import uuid
@@ -62,6 +63,11 @@
     download_blob_from_url,
     upload_blob_to_url,
 )
+from azure.storage.blob._shared.uploads_async import (
+    BlockBlobChunkUploader,
+    upload_data_chunks,
+    upload_substream_blocks,
+)
 
 # ------------------------------------------------------------------------------
 SMALL_BLOB_SIZE = 1024
@@ -70,6 +76,32 @@
 # ------------------------------------------------------------------------------
 
 
+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
+
+
+def _assert_unique_ordered_block_ids(block_ids, expected_data, staged_blocks):
+    # 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:
+        assert len(base64.b64decode(block_id)) == 16
+    # 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 TestStorageCommonBlobAsync(AsyncStorageRecordedTestCase):
     # --Helpers-----------------------------------------------------------------
     async def _setup(self, storage_account_name, key):
@@ -3928,5 +3960,45 @@ async def test_download_blob_returns_access_tier(self, **kwargs):
         assert downloader.properties.blob_tier_change_time is not None
         assert not downloader.properties.blob_tier_inferred
 
+    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())
+        _assert_unique_ordered_block_ids(block_ids, data, service.blocks)
+
+    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())
+        _assert_unique_ordered_block_ids(block_ids, data, service.blocks)
+
 
 # ------------------------------------------------------------------------------

sdk/storage/azure-storage-file-datalake/azure/storage/filedatalake/_shared/uploads.py

@@ -9,10 +9,11 @@
 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 +117,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 [block_id for _, block_id in sorted(range_ids, key=lambda r: r[0])]
     return []
 
 
@@ -257,9 +258,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(uuid4().bytes)
         self.service.stage_block(
             block_id,
             len(chunk_data),
@@ -272,7 +274,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(uuid4().bytes)
             self.service.stage_block(
                 block_id,
                 len(block_stream),
@@ -283,7 +285,7 @@ def _upload_substream_block(self, index, block_stream):
             )
         finally:
             block_stream.close()
-        return block_id
+        return index, block_id
 
 
 class PageBlobChunkUploader(_ChunkUploader):