diff --git a/sdk/storage/azure-storage-blob/azure/storage/blob/_shared/base_client.py b/sdk/storage/azure-storage-blob/azure/storage/blob/_shared/base_client.py index 7e416c7e9569..93f83c839ea4 100644 --- a/sdk/storage/azure-storage-blob/azure/storage/blob/_shared/base_client.py +++ b/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): diff --git a/sdk/storage/azure-storage-blob/azure/storage/blob/_shared/base_client_async.py b/sdk/storage/azure-storage-blob/azure/storage/blob/_shared/base_client_async.py index 0fe9eaa90785..215292a64f3e 100644 --- a/sdk/storage/azure-storage-blob/azure/storage/blob/_shared/base_client_async.py +++ b/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 diff --git a/sdk/storage/azure-storage-blob/azure/storage/blob/aio/blob_service_client_async.py b/sdk/storage/azure-storage-blob/azure/storage/blob/aio/blob_service_client_async.py index cca75e3c69c0..2f7064528ac5 100644 --- a/sdk/storage/azure-storage-blob/azure/storage/blob/aio/blob_service_client_async.py +++ b/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) diff --git a/sdk/storage/azure-storage-blob/azure/storage/blob/aio/container_client_async.py b/sdk/storage/azure-storage-blob/azure/storage/blob/aio/container_client_async.py index 684034ed93ff..d8b20b371631 100644 --- a/sdk/storage/azure-storage-blob/azure/storage/blob/aio/container_client_async.py +++ b/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) diff --git a/sdk/storage/azure-storage-blob/azure/storage/blob/blob_service_client.py b/sdk/storage/azure-storage-blob/azure/storage/blob/blob_service_client.py index 57b74da0abc8..c7475d93c9cb 100644 --- a/sdk/storage/azure-storage-blob/azure/storage/blob/blob_service_client.py +++ b/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) diff --git a/sdk/storage/azure-storage-blob/azure/storage/blob/container_client.py b/sdk/storage/azure-storage-blob/azure/storage/blob/container_client.py index 6df9e9b03b11..8f906c8d3b1a 100644 --- a/sdk/storage/azure-storage-blob/azure/storage/blob/container_client.py +++ b/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, diff --git a/sdk/storage/azure-storage-blob/tests/test_common_blob.py b/sdk/storage/azure-storage-blob/tests/test_common_blob.py index 302f2fb55e01..49e33f0c0d8d 100644 --- a/sdk/storage/azure-storage-blob/tests/test_common_blob.py +++ b/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__': diff --git a/sdk/storage/azure-storage-blob/tests/test_common_blob_async.py b/sdk/storage/azure-storage-blob/tests/test_common_blob_async.py index 936fa74af6ae..a5b25ce0f18d 100644 --- a/sdk/storage/azure-storage-blob/tests/test_common_blob_async.py +++ b/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() diff --git a/sdk/storage/azure-storage-file/azure/storage/file/_shared/base_client.py b/sdk/storage/azure-storage-file/azure/storage/file/_shared/base_client.py index 7e416c7e9569..2b9fade672ea 100644 --- a/sdk/storage/azure-storage-file/azure/storage/file/_shared/base_client.py +++ b/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: diff --git a/sdk/storage/azure-storage-file/azure/storage/file/_shared/base_client_async.py b/sdk/storage/azure-storage-file/azure/storage/file/_shared/base_client_async.py index 32a03fcf4f7a..77f6e48ef5ab 100644 --- a/sdk/storage/azure-storage-file/azure/storage/file/_shared/base_client_async.py +++ b/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 diff --git a/sdk/storage/azure-storage-file/azure/storage/file/aio/directory_client_async.py b/sdk/storage/azure-storage-file/azure/storage/file/aio/directory_client_async.py index 63021bca89bc..0f80d74fd1f4 100644 --- a/sdk/storage/azure-storage-file/azure/storage/file/aio/directory_client_async.py +++ b/sdk/storage/azure-storage-file/azure/storage/file/aio/directory_client_async.py @@ -11,7 +11,7 @@ from azure.core.polling import async_poller from azure.core.async_paging import AsyncItemPaged - +from azure.core.pipeline import AsyncPipeline from azure.core.tracing.decorator import distributed_trace from azure.core.tracing.decorator_async import distributed_trace_async from .._parser import _get_file_permission, _datetime_to_str @@ -20,7 +20,7 @@ from .._generated.aio import AzureFileStorage from .._generated.version import VERSION from .._generated.models import StorageErrorException -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 from .._shared.response_handlers import return_response_headers, process_storage_error @@ -112,10 +112,15 @@ def get_file_client(self, file_name, **kwargs): """ if self.directory_path: file_name = self.directory_path.rstrip('/') + "/" + file_name + + _pipeline = AsyncPipeline( + transport=AsyncTransportWrapper(self._pipeline._transport), # pylint: disable = protected-access + policies=self._pipeline._impl_policies # pylint: disable = protected-access + ) return FileClient( self.url, file_path=file_name, share_name=self.share_name, snapshot=self.snapshot, credential=self.credential, _hosts=self._hosts, _configuration=self._config, - _pipeline=self._pipeline, _location_mode=self._location_mode, loop=self._loop, **kwargs) + _pipeline=_pipeline, _location_mode=self._location_mode, loop=self._loop, **kwargs) def get_subdirectory_client(self, directory_name, **kwargs): # type: (str, Any) -> DirectoryClient @@ -138,10 +143,15 @@ def get_subdirectory_client(self, directory_name, **kwargs): :caption: Gets the subdirectory client. """ directory_path = self.directory_path.rstrip('/') + "/" + directory_name + + _pipeline = AsyncPipeline( + transport=AsyncTransportWrapper(self._pipeline._transport), # pylint: disable = protected-access + policies=self._pipeline._impl_policies # pylint: disable = protected-access + ) return DirectoryClient( self.url, share_name=self.share_name, directory_path=directory_path, snapshot=self.snapshot, credential=self.credential, _hosts=self._hosts, _configuration=self._config, - _pipeline=self._pipeline, _location_mode=self._location_mode, loop=self._loop, **kwargs) + _pipeline=_pipeline, _location_mode=self._location_mode, loop=self._loop, **kwargs) @distributed_trace_async async def create_directory(self, **kwargs): # type: ignore diff --git a/sdk/storage/azure-storage-file/azure/storage/file/aio/file_service_client_async.py b/sdk/storage/azure-storage-file/azure/storage/file/aio/file_service_client_async.py index c043a030f4a3..a2050a86e276 100644 --- a/sdk/storage/azure-storage-file/azure/storage/file/aio/file_service_client_async.py +++ b/sdk/storage/azure-storage-file/azure/storage/file/aio/file_service_client_async.py @@ -12,9 +12,10 @@ from azure.core.async_paging import AsyncItemPaged 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 .._shared.base_client_async import AsyncStorageAccountHostsMixin +from .._shared.base_client_async import AsyncStorageAccountHostsMixin, AsyncTransportWrapper from .._shared.response_handlers import process_storage_error from .._shared.policies_async import ExponentialRetry from .._generated.aio import AzureFileStorage @@ -314,6 +315,11 @@ def get_share_client(self, share, snapshot=None): share_name = share.name except AttributeError: share_name = share + + _pipeline = AsyncPipeline( + transport=AsyncTransportWrapper(self._pipeline._transport), # pylint: disable = protected-access + policies=self._pipeline._impl_policies # pylint: disable = protected-access + ) return ShareClient( self.url, share_name=share_name, snapshot=snapshot, credential=self.credential, _hosts=self._hosts, - _configuration=self._config, _pipeline=self._pipeline, _location_mode=self._location_mode, loop=self._loop) + _configuration=self._config, _pipeline=_pipeline, _location_mode=self._location_mode, loop=self._loop) diff --git a/sdk/storage/azure-storage-file/azure/storage/file/aio/share_client_async.py b/sdk/storage/azure-storage-file/azure/storage/file/aio/share_client_async.py index b5e5ce453024..9ccb64e7281e 100644 --- a/sdk/storage/azure-storage-file/azure/storage/file/aio/share_client_async.py +++ b/sdk/storage/azure-storage-file/azure/storage/file/aio/share_client_async.py @@ -10,9 +10,9 @@ from azure.core.tracing.decorator import distributed_trace from azure.core.tracing.decorator_async import distributed_trace_async - +from azure.core.pipeline import AsyncPipeline from .._shared.policies_async import ExponentialRetry -from .._shared.base_client_async import AsyncStorageAccountHostsMixin +from .._shared.base_client_async import AsyncStorageAccountHostsMixin, AsyncTransportWrapper from .._shared.request_handlers import add_metadata_headers, serialize_iso from .._shared.response_handlers import ( return_response_headers, @@ -102,9 +102,14 @@ def get_directory_client(self, directory_path=None): :returns: A Directory Client. :rtype: ~azure.storage.file.aio.directory_client_async.DirectoryClient """ + _pipeline = AsyncPipeline( + transport=AsyncTransportWrapper(self._pipeline._transport), # pylint: disable = protected-access + policies=self._pipeline._impl_policies # pylint: disable = protected-access + ) + return DirectoryClient( self.url, share_name=self.share_name, directory_path=directory_path or "", snapshot=self.snapshot, - credential=self.credential, _hosts=self._hosts, _configuration=self._config, _pipeline=self._pipeline, + credential=self.credential, _hosts=self._hosts, _configuration=self._config, _pipeline=_pipeline, _location_mode=self._location_mode, loop=self._loop) def get_file_client(self, file_path): @@ -117,10 +122,15 @@ def get_file_client(self, file_path): :returns: A File Client. :rtype: ~azure.storage.file.aio.file_client_async.FileClient """ + _pipeline = AsyncPipeline( + transport=AsyncTransportWrapper(self._pipeline._transport), # pylint: disable = protected-access + policies=self._pipeline._impl_policies # pylint: disable = protected-access + ) + return FileClient( self.url, share_name=self.share_name, file_path=file_path, snapshot=self.snapshot, credential=self.credential, _hosts=self._hosts, _configuration=self._config, - _pipeline=self._pipeline, _location_mode=self._location_mode, loop=self._loop) + _pipeline=_pipeline, _location_mode=self._location_mode, loop=self._loop) @distributed_trace_async async def create_share(self, **kwargs): # type: ignore diff --git a/sdk/storage/azure-storage-file/azure/storage/file/directory_client.py b/sdk/storage/azure-storage-file/azure/storage/file/directory_client.py index 9f02423b52cf..039199c48a5e 100644 --- a/sdk/storage/azure-storage-file/azure/storage/file/directory_client.py +++ b/sdk/storage/azure-storage-file/azure/storage/file/directory_client.py @@ -18,12 +18,13 @@ import six from azure.core.polling import LROPoller from azure.core.paging import ItemPaged +from azure.core.pipeline import Pipeline from azure.core.tracing.decorator import distributed_trace from ._generated import AzureFileStorage from ._generated.version import VERSION from ._generated.models import StorageErrorException -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 from ._shared.response_handlers import return_response_headers, process_storage_error from ._shared.parser import _str @@ -217,10 +218,15 @@ def get_file_client(self, file_name, **kwargs): """ if self.directory_path: file_name = self.directory_path.rstrip('/') + "/" + file_name + + _pipeline = Pipeline( + transport=TransportWrapper(self._pipeline._transport), # pylint: disable = protected-access + policies=self._pipeline._impl_policies # pylint: disable = protected-access + ) return FileClient( self.url, file_path=file_name, share_name=self.share_name, napshot=self.snapshot, credential=self.credential, _hosts=self._hosts, _configuration=self._config, - _pipeline=self._pipeline, _location_mode=self._location_mode, **kwargs) + _pipeline=_pipeline, _location_mode=self._location_mode, **kwargs) def get_subdirectory_client(self, directory_name, **kwargs): # type: (str, Any) -> DirectoryClient @@ -243,9 +249,14 @@ def get_subdirectory_client(self, directory_name, **kwargs): :caption: Gets the subdirectory client. """ directory_path = self.directory_path.rstrip('/') + "/" + directory_name + + _pipeline = Pipeline( + transport=TransportWrapper(self._pipeline._transport), # pylint: disable = protected-access + policies=self._pipeline._impl_policies # pylint: disable = protected-access + ) return DirectoryClient( self.url, share_name=self.share_name, directory_path=directory_path, snapshot=self.snapshot, - credential=self.credential, _hosts=self._hosts, _configuration=self._config, _pipeline=self._pipeline, + credential=self.credential, _hosts=self._hosts, _configuration=self._config, _pipeline=_pipeline, _location_mode=self._location_mode, **kwargs) @distributed_trace diff --git a/sdk/storage/azure-storage-file/azure/storage/file/file_service_client.py b/sdk/storage/azure-storage-file/azure/storage/file/file_service_client.py index 591bfefb2fdd..e1a23a172fbd 100644 --- a/sdk/storage/azure-storage-file/azure/storage/file/file_service_client.py +++ b/sdk/storage/azure-storage-file/azure/storage/file/file_service_client.py @@ -16,8 +16,8 @@ from azure.core.paging import ItemPaged from azure.core.tracing.decorator import distributed_trace - -from ._shared.base_client import StorageAccountHostsMixin, parse_connection_str, parse_query +from azure.core.pipeline import Pipeline +from ._shared.base_client import StorageAccountHostsMixin, TransportWrapper, parse_connection_str, parse_query from ._shared.response_handlers import process_storage_error from ._generated import AzureFileStorage from ._generated.models import StorageErrorException, StorageServiceProperties @@ -357,6 +357,11 @@ def get_share_client(self, share, snapshot=None): share_name = share.name except AttributeError: share_name = share + + _pipeline = Pipeline( + transport=TransportWrapper(self._pipeline._transport), # pylint: disable = protected-access + policies=self._pipeline._impl_policies # pylint: disable = protected-access + ) return ShareClient( self.url, share_name=share_name, snapshot=snapshot, credential=self.credential, _hosts=self._hosts, - _configuration=self._config, _pipeline=self._pipeline, _location_mode=self._location_mode) + _configuration=self._config, _pipeline=_pipeline, _location_mode=self._location_mode) diff --git a/sdk/storage/azure-storage-file/azure/storage/file/share_client.py b/sdk/storage/azure-storage-file/azure/storage/file/share_client.py index e41cd1468d4f..4938f769c00d 100644 --- a/sdk/storage/azure-storage-file/azure/storage/file/share_client.py +++ b/sdk/storage/azure-storage-file/azure/storage/file/share_client.py @@ -15,7 +15,8 @@ import six from azure.core.tracing.decorator import distributed_trace -from ._shared.base_client import StorageAccountHostsMixin, parse_connection_str, parse_query +from azure.core.pipeline import Pipeline +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 ( return_response_headers, @@ -205,9 +206,14 @@ def get_directory_client(self, directory_path=None): :returns: A Directory Client. :rtype: ~azure.storage.file.DirectoryClient """ + _pipeline = Pipeline( + transport=TransportWrapper(self._pipeline._transport), # pylint: disable = protected-access + policies=self._pipeline._impl_policies # pylint: disable = protected-access + ) + return DirectoryClient( self.url, share_name=self.share_name, directory_path=directory_path or "", snapshot=self.snapshot, - credential=self.credential, _hosts=self._hosts, _configuration=self._config, _pipeline=self._pipeline, + credential=self.credential, _hosts=self._hosts, _configuration=self._config, _pipeline=_pipeline, _location_mode=self._location_mode) def get_file_client(self, file_path): @@ -220,10 +226,15 @@ def get_file_client(self, file_path): :returns: A File Client. :rtype: ~azure.storage.file.FileClient """ + _pipeline = Pipeline( + transport=TransportWrapper(self._pipeline._transport), # pylint: disable = protected-access + policies=self._pipeline._impl_policies # pylint: disable = protected-access + ) + return FileClient( self.url, share_name=self.share_name, file_path=file_path, snapshot=self.snapshot, credential=self.credential, _hosts=self._hosts, _configuration=self._config, - _pipeline=self._pipeline, _location_mode=self._location_mode) + _pipeline=_pipeline, _location_mode=self._location_mode) @distributed_trace def create_share(self, **kwargs): # type: ignore diff --git a/sdk/storage/azure-storage-file/tests/test_share.py b/sdk/storage/azure-storage-file/tests/test_share.py index 3824ff9f608b..ad441397ab80 100644 --- a/sdk/storage/azure-storage-file/tests/test_share.py +++ b/sdk/storage/azure-storage-file/tests/test_share.py @@ -10,6 +10,7 @@ import pytest import requests +from azure.core.pipeline.transport import RequestsTransport from azure.core.exceptions import ( HttpResponseError, ResourceNotFoundError, @@ -769,6 +770,23 @@ def test_create_permission_for_share(self): # server returned permission self.assertEquals(permission_key, permission_key2) + @record + def test_transport_closed_only_once(self): + if TestMode.need_recording_file(self.test_mode): + return + transport = RequestsTransport() + url = self.get_file_url() + credential = self.get_shared_key_credential() + prefix = TEST_SHARE_PREFIX + share_name = self.get_resource_name(prefix) + with FileServiceClient(url, credential=credential, transport=transport) as fsc: + fsc.get_service_properties() + assert transport.session is not None + with fsc.get_share_client(share_name) as fc: + assert transport.session is not None + fsc.get_service_properties() + assert transport.session is not None + # ------------------------------------------------------------------------------ if __name__ == '__main__': unittest.main() diff --git a/sdk/storage/azure-storage-file/tests/test_share_async.py b/sdk/storage/azure-storage-file/tests/test_share_async.py index 08e2136e0fef..298b573ea9be 100644 --- a/sdk/storage/azure-storage-file/tests/test_share_async.py +++ b/sdk/storage/azure-storage-file/tests/test_share_async.py @@ -11,6 +11,7 @@ import pytest import requests from azure.core.pipeline.transport import AioHttpTransport +from azure.core.pipeline.transport import AsyncioRequestsTransport from multidict import CIMultiDict, CIMultiDictProxy from azure.core.exceptions import ( HttpResponseError, @@ -922,6 +923,27 @@ def test_create_permission_for_share_async(self): loop = asyncio.get_event_loop() loop.run_until_complete(self._test_create_permission_for_share()) + async def _test_transport_closed_only_once_async(self): + if TestMode.need_recording_file(self.test_mode): + return + transport = AioHttpTransport() + url = self.get_file_url() + credential = self.get_shared_key_credential() + prefix = TEST_SHARE_PREFIX + share_name = self.get_resource_name(prefix) + async with FileServiceClient(url, credential=credential, transport=transport) as fsc: + await fsc.get_service_properties() + assert transport.session is not None + async with fsc.get_share_client(share_name) as fc: + assert transport.session is not None + await fsc.get_service_properties() + assert transport.session is not None + + @record + def test_transport_closed_only_once_async(self): + loop = asyncio.get_event_loop() + loop.run_until_complete(self._test_transport_closed_only_once_async()) + # ------------------------------------------------------------------------------ if __name__ == '__main__': unittest.main() diff --git a/sdk/storage/azure-storage-queue/azure/storage/queue/_shared/base_client.py b/sdk/storage/azure-storage-queue/azure/storage/queue/_shared/base_client.py index 7e416c7e9569..2b9fade672ea 100644 --- a/sdk/storage/azure-storage-queue/azure/storage/queue/_shared/base_client.py +++ b/sdk/storage/azure-storage-queue/azure/storage/queue/_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: diff --git a/sdk/storage/azure-storage-queue/azure/storage/queue/_shared/base_client_async.py b/sdk/storage/azure-storage-queue/azure/storage/queue/_shared/base_client_async.py index 32a03fcf4f7a..77f6e48ef5ab 100644 --- a/sdk/storage/azure-storage-queue/azure/storage/queue/_shared/base_client_async.py +++ b/sdk/storage/azure-storage-queue/azure/storage/queue/_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 diff --git a/sdk/storage/azure-storage-queue/azure/storage/queue/aio/queue_service_client_async.py b/sdk/storage/azure-storage-queue/azure/storage/queue/aio/queue_service_client_async.py index 56041f5030e3..2e8e2ef36862 100644 --- a/sdk/storage/azure-storage-queue/azure/storage/queue/aio/queue_service_client_async.py +++ b/sdk/storage/azure-storage-queue/azure/storage/queue/aio/queue_service_client_async.py @@ -15,12 +15,13 @@ from azure.core.async_paging import AsyncItemPaged 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 .._shared.policies_async import ExponentialRetry from ..queue_service_client import QueueServiceClient as QueueServiceClientBase from .._shared.models import LocationMode -from .._shared.base_client_async import AsyncStorageAccountHostsMixin +from .._shared.base_client_async import AsyncStorageAccountHostsMixin, AsyncTransportWrapper from .._shared.response_handlers import process_storage_error from .._generated.aio import AzureQueueStorage from .._generated.models import StorageServiceProperties, StorageErrorException @@ -372,8 +373,14 @@ def get_queue_client(self, queue, **kwargs): queue_name = queue.name except AttributeError: queue_name = queue + + _pipeline = AsyncPipeline( + transport=AsyncTransportWrapper(self._pipeline._transport), # pylint: disable = protected-access + policies=self._pipeline._impl_policies # pylint: disable = protected-access + ) + return QueueClient( self.url, queue_name=queue_name, credential=self.credential, key_resolver_function=self.key_resolver_function, require_encryption=self.require_encryption, - key_encryption_key=self.key_encryption_key, _pipeline=self._pipeline, _configuration=self._config, + key_encryption_key=self.key_encryption_key, _pipeline=_pipeline, _configuration=self._config, _location_mode=self._location_mode, _hosts=self._hosts, loop=self._loop, **kwargs) diff --git a/sdk/storage/azure-storage-queue/azure/storage/queue/queue_service_client.py b/sdk/storage/azure-storage-queue/azure/storage/queue/queue_service_client.py index a7eeec5cd657..259e9f395db2 100644 --- a/sdk/storage/azure-storage-queue/azure/storage/queue/queue_service_client.py +++ b/sdk/storage/azure-storage-queue/azure/storage/queue/queue_service_client.py @@ -14,9 +14,10 @@ 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.response_handlers import process_storage_error from ._generated import AzureQueueStorage from ._generated.models import StorageServiceProperties, StorageErrorException @@ -412,8 +413,14 @@ def get_queue_client(self, queue, **kwargs): queue_name = queue.name except AttributeError: queue_name = queue + + _pipeline = Pipeline( + transport=TransportWrapper(self._pipeline._transport), # pylint: disable = protected-access + policies=self._pipeline._impl_policies # pylint: disable = protected-access + ) + return QueueClient( self.url, queue_name=queue_name, credential=self.credential, key_resolver_function=self.key_resolver_function, require_encryption=self.require_encryption, - key_encryption_key=self.key_encryption_key, _pipeline=self._pipeline, _configuration=self._config, + key_encryption_key=self.key_encryption_key, _pipeline=_pipeline, _configuration=self._config, _location_mode=self._location_mode, _hosts=self._hosts, **kwargs) diff --git a/sdk/storage/azure-storage-queue/tests/test_queue.py b/sdk/storage/azure-storage-queue/tests/test_queue.py index 35f3410b1b55..876cca473254 100644 --- a/sdk/storage/azure-storage-queue/tests/test_queue.py +++ b/sdk/storage/azure-storage-queue/tests/test_queue.py @@ -18,6 +18,7 @@ ) from devtools_testutils import ResourceGroupPreparer, StorageAccountPreparer from azure.mgmt.storage.models import Endpoints +from azure.core.pipeline.transport import RequestsTransport from azure.core.exceptions import ( HttpResponseError, ResourceNotFoundError, @@ -982,6 +983,21 @@ def test_unicode_update_message_unicode_data(self, resource_group, location, sto self.assertIsInstance(message.expires_on, datetime) self.assertIsInstance(message.next_visible_on, datetime) + @ResourceGroupPreparer() + @StorageAccountPreparer(name_prefix='pyacrstorage') + def test_transport_closed_only_once(self, resource_group, location, storage_account, storage_account_key): + if not self.is_live: + return + transport = RequestsTransport() + prefix = TEST_QUEUE_PREFIX + queue_name = self.get_resource_name(prefix) + with QueueServiceClient(self._account_url(storage_account.name), credential=storage_account_key, transport=transport) as qsc: + qsc.get_service_properties() + assert transport.session is not None + with qsc.get_queue_client(queue_name) as qc: + assert transport.session is not None + qsc.get_service_properties() + assert transport.session is not None # ------------------------------------------------------------------------------ if __name__ == '__main__': diff --git a/sdk/storage/azure-storage-queue/tests/test_queue_async.py b/sdk/storage/azure-storage-queue/tests/test_queue_async.py index 874dbd59efd4..fd26c6f22aa5 100644 --- a/sdk/storage/azure-storage-queue/tests/test_queue_async.py +++ b/sdk/storage/azure-storage-queue/tests/test_queue_async.py @@ -16,6 +16,7 @@ ) from devtools_testutils import ResourceGroupPreparer, StorageAccountPreparer from multidict import CIMultiDict, CIMultiDictProxy +from azure.core.pipeline.transport import AsyncioRequestsTransport from azure.core.exceptions import ( HttpResponseError, ResourceNotFoundError, @@ -1044,6 +1045,22 @@ async def test_unicode_update_message_unicode_data(self, resource_group, locatio self.assertIsInstance(message.expires_on, datetime) self.assertIsInstance(message.next_visible_on, datetime) + @ResourceGroupPreparer() + @StorageAccountPreparer(name_prefix='pyacrstorage') + @AsyncQueueTestCase.await_prepared_test + async def test_transport_closed_only_once_async(self, resource_group, location, storage_account, storage_account_key): + if not self.is_live: + return + transport = AioHttpTransport() + prefix = TEST_QUEUE_PREFIX + queue_name = self.get_resource_name(prefix) + async with QueueServiceClient(self._account_url(storage_account.name), credential=storage_account_key, transport=transport) as qsc: + await qsc.get_service_properties() + assert transport.session is not None + async with qsc.get_queue_client(queue_name) as qc: + assert transport.session is not None + await qsc.get_service_properties() + assert transport.session is not None # ------------------------------------------------------------------------------ if __name__ == '__main__':