Design Pipeline ownership

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

@@ -28,7 +28,7 @@
 from azure.core.configuration import Configuration
 from azure.core.exceptions import HttpResponseError
 from azure.core.pipeline import Pipeline
-from azure.core.pipeline.transport import RequestsTransport
+from azure.core.pipeline.transport import RequestsTransport, HttpTransport
 from azure.core.pipeline.policies import (
     RedirectPolicy,
     ContentDecodePolicy,
@@ -224,6 +224,29 @@ def _batch_send(
         except StorageErrorException as error:
             process_storage_error(error)
 
+class TransportWrapper(HttpTransport):
+    """Wrapper class that ensures that an inner client created
+    by a `get_client` method does not close the outer transport for the parent
+    when used in a context manager.
+    """
+    def __init__(self, transport):
+        self._transport = transport
+
+    def send(self, request, **kwargs):
+        return self._transport.send(request, **kwargs)
+
+    def open(self):
+        pass
+
+    def close(self):
+        pass
+
+    def __enter__(self):
+        pass
+
+    def __exit__(self, *args):  # pylint: disable=arguments-differ
+        pass
+
 
 def format_shared_key_credential(account, credential):
     if isinstance(credential, six.string_types):

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

@@ -18,6 +18,7 @@
     AsyncRedirectPolicy,
     DistributedTracingPolicy
 )
+from azure.core.pipeline.transport import AsyncHttpTransport
 
 from .constants import STORAGE_OAUTH_SCOPE, DEFAULT_SOCKET_TIMEOUT
 from .authentication import SharedKeyCredentialPolicy
@@ -126,3 +127,27 @@ async def _batch_send(
             return response.parts()  # Return an AsyncIterator
         except StorageErrorException as error:
             process_storage_error(error)
+
+
+class AsyncTransportWrapper(AsyncHttpTransport):
+    """Wrapper class that ensures that an inner client created
+    by a `get_client` method does not close the outer transport for the parent
+    when used in a context manager.
+    """
+    def __init__(self, async_transport):
+        self._transport = async_transport
+
+    async def send(self, request, **kwargs):
+        return await self._transport.send(request, **kwargs)
+
+    async def open(self):
+        pass
+
+    async def close(self):
+        pass
+
+    async def __aenter__(self):
+        pass
+
+    async def __aexit__(self, *args):  # pylint: disable=arguments-differ
+        pass

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

@@ -11,12 +11,13 @@
 )
 
 from azure.core.tracing.decorator import distributed_trace
+from azure.core.pipeline import AsyncPipeline
 from azure.core.tracing.decorator_async import distributed_trace_async
 from azure.core.async_paging import AsyncItemPaged
 
 from .._shared.models import LocationMode
 from .._shared.policies_async import ExponentialRetry
-from .._shared.base_client_async import AsyncStorageAccountHostsMixin
+from .._shared.base_client_async import AsyncStorageAccountHostsMixin, AsyncTransportWrapper
 from .._shared.response_handlers import return_response_headers, process_storage_error
 from .._shared.parser import _to_utc_datetime
 from .._shared.response_handlers import parse_to_internal_user_delegation_key
@@ -487,11 +488,14 @@ def get_container_client(self, container):
             container_name = container.name
         except AttributeError:
             container_name = container
-
+        _pipeline = AsyncPipeline(
+            transport=AsyncTransportWrapper(self._pipeline._transport), # pylint: disable = protected-access
+            policies=self._pipeline._impl_policies # pylint: disable = protected-access
+        )
         return ContainerClient(
             self.url, container_name=container_name,
             credential=self.credential, _configuration=self._config,
-            _pipeline=self._pipeline, _location_mode=self._location_mode, _hosts=self._hosts,
+            _pipeline=_pipeline, _location_mode=self._location_mode, _hosts=self._hosts,
             require_encryption=self.require_encryption, key_encryption_key=self.key_encryption_key,
             key_resolver_function=self.key_resolver_function, loop=self._loop)
 
@@ -537,10 +541,13 @@ def get_blob_client(
             blob_name = blob.name
         except AttributeError:
             blob_name = blob
-
+        _pipeline = AsyncPipeline(
+            transport=AsyncTransportWrapper(self._pipeline._transport), # pylint: disable = protected-access
+            policies=self._pipeline._impl_policies # pylint: disable = protected-access
+        )
         return BlobClient( # type: ignore
             self.url, container_name=container_name, blob_name=blob_name, snapshot=snapshot,
             credential=self.credential, _configuration=self._config,
-            _pipeline=self._pipeline, _location_mode=self._location_mode, _hosts=self._hosts,
+            _pipeline=_pipeline, _location_mode=self._location_mode, _hosts=self._hosts,
             require_encryption=self.require_encryption, key_encryption_key=self.key_encryption_key,
             key_resolver_function=self.key_resolver_function, loop=self._loop)

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

@@ -14,9 +14,10 @@
 from azure.core.tracing.decorator import distributed_trace
 from azure.core.tracing.decorator_async import distributed_trace_async
 from azure.core.async_paging import AsyncItemPaged
+from azure.core.pipeline import AsyncPipeline
 from azure.core.pipeline.transport import HttpRequest, AsyncHttpResponse
 
-from .._shared.base_client_async import AsyncStorageAccountHostsMixin
+from .._shared.base_client_async import AsyncStorageAccountHostsMixin, AsyncTransportWrapper
 from .._shared.policies_async import ExponentialRetry
 from .._shared.request_handlers import add_metadata_headers, serialize_iso
 from .._shared.response_handlers import (
@@ -1000,10 +1001,13 @@ def get_blob_client(
             blob_name = blob.name
         except AttributeError:
             blob_name = blob
-
+        _pipeline = AsyncPipeline(
+            transport=AsyncTransportWrapper(self._pipeline._transport), # pylint: disable = protected-access
+            policies=self._pipeline._impl_policies # pylint: disable = protected-access
+        )
         return BlobClient(
             self.url, container_name=self.container_name, blob_name=blob_name, snapshot=snapshot,
             credential=self.credential, _configuration=self._config,
-            _pipeline=self._pipeline, _location_mode=self._location_mode, _hosts=self._hosts,
+            _pipeline=_pipeline, _location_mode=self._location_mode, _hosts=self._hosts,
             require_encryption=self.require_encryption, key_encryption_key=self.key_encryption_key,
             key_resolver_function=self.key_resolver_function, loop=self._loop)

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

@@ -16,10 +16,11 @@
     from urlparse import urlparse # type: ignore
 
 from azure.core.paging import ItemPaged
+from azure.core.pipeline import Pipeline
 from azure.core.tracing.decorator import distributed_trace
 
 from ._shared.models import LocationMode
-from ._shared.base_client import StorageAccountHostsMixin, parse_connection_str, parse_query
+from ._shared.base_client import StorageAccountHostsMixin, TransportWrapper, parse_connection_str, parse_query
 from ._shared.parser import _to_utc_datetime
 from ._shared.response_handlers import return_response_headers, process_storage_error, \
     parse_to_internal_user_delegation_key
@@ -533,11 +534,14 @@ def get_container_client(self, container):
             container_name = container.name
         except AttributeError:
             container_name = container
-
+        _pipeline = Pipeline(
+            transport=TransportWrapper(self._pipeline._transport), # pylint: disable = protected-access
+            policies=self._pipeline._impl_policies # pylint: disable = protected-access
+        )
         return ContainerClient(
             self.url, container_name=container_name,
             credential=self.credential, _configuration=self._config,
-            _pipeline=self._pipeline, _location_mode=self._location_mode, _hosts=self._hosts,
+            _pipeline=_pipeline, _location_mode=self._location_mode, _hosts=self._hosts,
             require_encryption=self.require_encryption, key_encryption_key=self.key_encryption_key,
             key_resolver_function=self.key_resolver_function)
 
@@ -583,10 +587,13 @@ def get_blob_client(
             blob_name = blob.name
         except AttributeError:
             blob_name = blob
-
+        _pipeline = Pipeline(
+            transport=TransportWrapper(self._pipeline._transport), # pylint: disable = protected-access
+            policies=self._pipeline._impl_policies # pylint: disable = protected-access
+        )
         return BlobClient( # type: ignore
             self.url, container_name=container_name, blob_name=blob_name, snapshot=snapshot,
             credential=self.credential, _configuration=self._config,
-            _pipeline=self._pipeline, _location_mode=self._location_mode, _hosts=self._hosts,
+            _pipeline=_pipeline, _location_mode=self._location_mode, _hosts=self._hosts,
             require_encryption=self.require_encryption, key_encryption_key=self.key_encryption_key,
             key_resolver_function=self.key_resolver_function)

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

@@ -21,9 +21,10 @@
 
 from azure.core.paging import ItemPaged
 from azure.core.tracing.decorator import distributed_trace
+from azure.core.pipeline import Pipeline
 from azure.core.pipeline.transport import HttpRequest
 
-from ._shared.base_client import StorageAccountHostsMixin, parse_connection_str, parse_query
+from ._shared.base_client import StorageAccountHostsMixin, TransportWrapper, parse_connection_str, parse_query
 from ._shared.request_handlers import add_metadata_headers, serialize_iso
 from ._shared.response_handlers import (
     process_storage_error,
@@ -1183,6 +1184,10 @@ def get_blob_client(
             blob_name = blob.name
         except AttributeError:
             blob_name = blob
+        _pipeline = Pipeline(
+            transport=TransportWrapper(self._pipeline._transport), # pylint: disable = protected-access
+            policies=self._pipeline._impl_policies # pylint: disable = protected-access
+        )
         return BlobClient(
             self.url, container_name=self.container_name, blob_name=blob_name, snapshot=snapshot,
             credential=self.credential, _configuration=self._config,

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

@@ -17,6 +17,7 @@
     ResourceNotFoundError,
     ResourceExistsError,
     ClientAuthenticationError)
+from azure.core.pipeline.transport import RequestsTransport
 from azure.storage.blob import (
     upload_blob_to_url,
     download_blob_from_url,
@@ -1902,6 +1903,22 @@ def test_set_blob_permission(self):
         self.assertEqual(permission.delete, True)
         self.assertEqual(permission.write, True)
         self.assertEqual(permission._str, 'wrdx')
+
+    def test_transport_closed_only_once(self):
+        if TestMode.need_recording_file(self.test_mode):
+            return
+        transport = RequestsTransport()
+        url = self._get_account_url()
+        credential = self._get_shared_key_credential()
+        blob_name = self._get_blob_reference()
+        with BlobServiceClient(url, credential=credential, transport=transport) as bsc:
+            bsc.get_service_properties()
+            assert transport.session is not None
+            with bsc.get_blob_client(self.container_name, blob_name) as bc:
+                assert transport.session is not None
+            bsc.get_service_properties()
+            assert transport.session is not None
+
 
 #------------------------------------------------------------------------------
 if __name__ == '__main__':

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

@@ -18,7 +18,7 @@
     ResourceNotFoundError,
     ResourceExistsError,
     ClientAuthenticationError)
-
+from azure.core.pipeline.transport import AsyncioRequestsTransport
 from azure.core.pipeline.transport import AioHttpTransport
 from multidict import CIMultiDict, CIMultiDictProxy
 
@@ -2340,6 +2340,26 @@ def test_upload_to_url_file_with_credential(self):
         loop = asyncio.get_event_loop()
         loop.run_until_complete(self._test_upload_to_url_file_with_credential())
 
+    async def _test_transport_closed_only_once(self):
+        if TestMode.need_recording_file(self.test_mode):
+            return
+        transport = AioHttpTransport()
+        url = self._get_account_url()
+        credential = self._get_shared_key_credential()
+        blob_name = self._get_blob_reference()
+        async with BlobServiceClient(url, credential=credential, transport=transport) as bsc:
+            await bsc.get_service_properties()
+            assert transport.session is not None
+            async with bsc.get_blob_client(self.container_name, blob_name) as bc:
+                assert transport.session is not None
+            await bsc.get_service_properties()
+            assert transport.session is not None
+
+    @record
+    def test_transport_closed_only_once(self):
+        loop = asyncio.get_event_loop()
+        loop.run_until_complete(self._test_transport_closed_only_once())
+
 # ------------------------------------------------------------------------------
 if __name__ == '__main__':
     unittest.main()

sdk/storage/azure-storage-file/azure/storage/file/_shared/base_client.py

@@ -28,7 +28,7 @@
 from azure.core.configuration import Configuration
 from azure.core.exceptions import HttpResponseError
 from azure.core.pipeline import Pipeline
-from azure.core.pipeline.transport import RequestsTransport
+from azure.core.pipeline.transport import RequestsTransport, HttpTransport
 from azure.core.pipeline.policies import (
     RedirectPolicy,
     ContentDecodePolicy,
@@ -225,6 +225,30 @@ def _batch_send(
             process_storage_error(error)
 
 
+class TransportWrapper(HttpTransport):
+    """Wrapper class that ensures that an inner client created
+    by a `get_client` method does not close the outer transport for the parent
+    when used in a context manager.
+    """
+    def __init__(self, transport):
+        self._transport = transport
+
+    def send(self, request, **kwargs):
+        return self._transport.send(request, **kwargs)
+
+    def open(self):
+        pass
+
+    def close(self):
+        pass
+
+    def __enter__(self):
+        pass
+
+    def __exit__(self, *args):  # pylint: disable=arguments-differ
+        pass
+
+
 def format_shared_key_credential(account, credential):
     if isinstance(credential, six.string_types):
         if len(account) < 2:

sdk/storage/azure-storage-file/azure/storage/file/_shared/base_client_async.py

@@ -19,6 +19,7 @@
     DistributedTracingPolicy
 )
 
+from azure.core.pipeline.transport import AsyncHttpTransport
 from .constants import STORAGE_OAUTH_SCOPE, DEFAULT_SOCKET_TIMEOUT
 from .authentication import SharedKeyCredentialPolicy
 from .base_client import create_configuration
@@ -123,3 +124,27 @@ async def _batch_send(
             return response.parts()  # Return an AsyncIterator
         except StorageErrorException as error:
             process_storage_error(error)
+
+
+class AsyncTransportWrapper(AsyncHttpTransport):
+    """Wrapper class that ensures that an inner client created
+    by a `get_client` method does not close the outer transport for the parent
+    when used in a context manager.
+    """
+    def __init__(self, async_transport):
+        self._transport = async_transport
+
+    async def send(self, request, **kwargs):
+        return await self._transport.send(request, **kwargs)
+
+    async def open(self):
+        pass
+
+    async def close(self):
+        pass
+
+    async def __aenter__(self):
+        pass
+
+    async def __aexit__(self, *args):  # pylint: disable=arguments-differ
+        pass