diff --git a/sdk/storage/azure-storage-file/azure/storage/file/_shared/downloads.py b/sdk/storage/azure-storage-file/azure/storage/file/_shared/downloads.py index 1d46ffc95293..dc713862cb45 100644 --- a/sdk/storage/azure-storage-file/azure/storage/file/_shared/downloads.py +++ b/sdk/storage/azure-storage-file/azure/storage/file/_shared/downloads.py @@ -10,6 +10,7 @@ from azure.core.exceptions import HttpResponseError +from azure.core.tracing.context import tracing_context from .request_handlers import validate_and_format_range_headers from .response_handlers import process_storage_error, parse_length_from_content_range from .encryption import decrypt_blob @@ -454,7 +455,7 @@ def download_to_stream(self, stream, max_connections=1): if max_connections > 1: import concurrent.futures executor = concurrent.futures.ThreadPoolExecutor(max_connections) - list(executor.map(downloader.process_chunk, downloader.get_chunk_offsets())) + list(executor.map(tracing_context.with_current_context(downloader.process_chunk), downloader.get_chunk_offsets())) else: for chunk in downloader.get_chunk_offsets(): downloader.process_chunk(chunk) diff --git a/sdk/storage/azure-storage-file/azure/storage/file/_shared/uploads.py b/sdk/storage/azure-storage-file/azure/storage/file/_shared/uploads.py index 2b269fb1d0ba..68ff3f46e015 100644 --- a/sdk/storage/azure-storage-file/azure/storage/file/_shared/uploads.py +++ b/sdk/storage/azure-storage-file/azure/storage/file/_shared/uploads.py @@ -13,6 +13,7 @@ import six +from azure.core.tracing.context import tracing_context from . import encode_base64, url_quote from .request_handlers import get_length from .response_handlers import return_response_headers @@ -34,7 +35,7 @@ def _parallel_uploads(executor, uploader, pending, running): except StopIteration: break else: - running.add(executor.submit(uploader.process_chunk, next_chunk)) + running.add(executor.submit(tracing_context.with_current_context(uploader.process_chunk), next_chunk)) # Wait for the remaining uploads to finish done, _running = futures.wait(running) @@ -79,7 +80,7 @@ def upload_data_chunks( executor = futures.ThreadPoolExecutor(max_connections) upload_tasks = uploader.get_chunk_streams() running_futures = [ - executor.submit(uploader.process_chunk, u) + executor.submit(tracing_context.with_current_context(uploader.process_chunk), u) for u in islice(upload_tasks, 0, max_connections) ] range_ids = _parallel_uploads(executor, uploader, upload_tasks, running_futures) @@ -115,7 +116,7 @@ def upload_substream_blocks( executor = futures.ThreadPoolExecutor(max_connections) upload_tasks = uploader.get_substream_blocks() running_futures = [ - executor.submit(uploader.process_substream_block, u) + executor.submit(tracing_context.with_current_context(uploader.process_substream_block), u) for u in islice(upload_tasks, 0, max_connections) ] return _parallel_uploads(executor, uploader, upload_tasks, running_futures)