[Storage] Decoupled Client Context Manager Methods from Base Client (#41442)

Author: weirongw23-msft

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

@@ -190,6 +190,22 @@ def __init__(
         self._client._config.version = get_api_version(kwargs)  # type: ignore [assignment]
         self._configure_encryption(kwargs)
 
+    def __enter__(self) -> Self:
+        self._client.__enter__()
+        return self
+
+    def __exit__(self, *args) -> None:
+        self._client.__exit__(*args)
+
+    def close(self) -> None:
+        """This method is to close the sockets opened by the client.
+        It need not be used when using with a context manager.
+
+        :return: None
+        :rtype: None
+        """
+        self._client.close()
+
     def _format_url(self, hostname: str) -> str:
         return _format_url(
             container_name=self.container_name,

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

@@ -131,6 +131,22 @@ def __init__(
         self._client._config.version = get_api_version(kwargs)  # type: ignore [assignment]
         self._configure_encryption(kwargs)
 
+    def __enter__(self) -> Self:
+        self._client.__enter__()
+        return self
+
+    def __exit__(self, *args) -> None:
+        self._client.__exit__(*args)
+
+    def close(self) -> None:
+        """This method is to close the sockets opened by the client.
+        It need not be used when using with a context manager.
+
+        :return: None
+        :rtype: None
+        """
+        self._client.close()
+
     def _format_url(self, hostname):
         """Format the endpoint URL according to the current location
         mode hostname.

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

@@ -150,6 +150,22 @@ def __init__(
         self._client = self._build_generated_client()
         self._configure_encryption(kwargs)
 
+    def __enter__(self) -> Self:
+        self._client.__enter__()
+        return self
+
+    def __exit__(self, *args) -> None:
+        self._client.__exit__(*args)
+
+    def close(self) -> None:
+        """This method is to close the sockets opened by the client.
+        It need not be used when using with a context manager.
+
+        :return: None
+        :rtype: None
+        """
+        self._client.close()
+
     def _build_generated_client(self) -> AzureBlobStorage:
         client = AzureBlobStorage(self.url, base_url=self.url, pipeline=self._pipeline)
         client._config.version = self._api_version  # type: ignore [assignment] # pylint: disable=protected-access

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

@@ -127,19 +127,6 @@ def __init__(
         self._sdk_moniker = f"storage-{service}/{VERSION}"
         self._config, self._pipeline = self._create_pipeline(self.credential, sdk_moniker=self._sdk_moniker, **kwargs)
 
-    def __enter__(self):
-        self._client.__enter__()
-        return self
-
-    def __exit__(self, *args):
-        self._client.__exit__(*args)
-
-    def close(self) -> None:
-        """This method is to close the sockets opened by the client.
-        It need not be used when using with a context manager.
-        """
-        self._client.close()
-
     @property
     def url(self) -> str:
         """The full endpoint URL to this entity, including SAS token if used.

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

@@ -51,28 +51,6 @@
 
 class AsyncStorageAccountHostsMixin(object):
 
-    def __enter__(self):
-        raise TypeError("Async client only supports 'async with'.")
-
-    def __exit__(self, *args):
-        pass
-
-    async def __aenter__(self):
-        await self._client.__aenter__()
-        return self
-
-    async def __aexit__(self, *args):
-        await self._client.__aexit__(*args)
-
-    async def close(self) -> None:
-        """This method is to close the sockets opened by the client.
-        It need not be used when using with a context manager.
-
-        :return: None
-        :rtype: None
-        """
-        await self._client.close()
-
     def _format_query_string(
         self,
         sas_token: Optional[str],

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

@@ -196,6 +196,22 @@ def __init__(
         self._client._config.version = get_api_version(kwargs)  # type: ignore [assignment]
         self._configure_encryption(kwargs)
 
+    async def __aenter__(self) -> Self:
+        await self._client.__aenter__()
+        return self
+
+    async def __aexit__(self, *args) -> None:
+        await self._client.__aexit__(*args)
+
+    async def close(self) -> None:
+        """This method is to close the sockets opened by the client.
+        It need not be used when using with a context manager.
+
+        :return: None
+        :rtype: None
+        """
+        await self._client.close()
+
     def _format_url(self, hostname: str) -> str:
         return _format_url(
             container_name=self.container_name,

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

@@ -139,6 +139,22 @@ def __init__(
         self._client._config.version = get_api_version(kwargs)  # type: ignore [assignment]
         self._configure_encryption(kwargs)
 
+    async def __aenter__(self) -> Self:
+        await self._client.__aenter__()
+        return self
+
+    async def __aexit__(self, *args) -> None:
+        await self._client.__aexit__(*args)
+
+    async def close(self) -> None:
+        """This method is to close the sockets opened by the client.
+        It need not be used when using with a context manager.
+
+        :return: None
+        :rtype: None
+        """
+        await self._client.close()
+
     def _format_url(self, hostname):
         """Format the endpoint URL according to the current location
         mode hostname.

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

@@ -143,6 +143,22 @@ def __init__(
         self._client = self._build_generated_client()
         self._configure_encryption(kwargs)
 
+    async def __aenter__(self) -> Self:
+        await self._client.__aenter__()
+        return self
+
+    async def __aexit__(self, *args) -> None:
+        await self._client.__aexit__(*args)
+
+    async def close(self) -> None:
+        """This method is to close the sockets opened by the client.
+        It need not be used when using with a context manager.
+
+        :return: None
+        :rtype: None
+        """
+        await self._client.close()
+
     def _build_generated_client(self) -> AzureBlobStorage:
         client = AzureBlobStorage(self.url, base_url=self.url, pipeline=self._pipeline)
         client._config.version = self._api_version  # type: ignore [assignment] # pylint: disable=protected-access

sdk/storage/azure-storage-file-datalake/azure/storage/filedatalake/_data_lake_service_client.py

@@ -120,18 +120,23 @@ def __init__(
         self._client._config.version = get_api_version(kwargs)  # type: ignore [assignment]
 
     def __enter__(self) -> Self:
+        self._client.__enter__()
         self._blob_service_client.__enter__()
         return self
 
     def __exit__(self, *args: Any) -> None:
-        self._blob_service_client.close()
-        super(DataLakeServiceClient, self).__exit__(*args)
+        self._blob_service_client.__exit__(*args)
+        self._client.__exit__(*args)
 
     def close(self) -> None:
-        """ This method is to close the sockets opened by the client.
+        """This method is to close the sockets opened by the client.
         It need not be used when using with a context manager.
+
+        :return: None
+        :rtype: None
         """
-        self.__exit__()
+        self._blob_service_client.close()
+        self._client.close()
 
     def _format_url(self, hostname: str) -> str:
         """Format the endpoint URL according to hostname.

sdk/storage/azure-storage-file-datalake/azure/storage/filedatalake/_file_system_client.py

@@ -132,6 +132,28 @@ def __init__(
         self._client = self._build_generated_client(self.url)
         self._datalake_client_for_blob_operation = self._build_generated_client(self._container_client.url)
 
+    def __enter__(self) -> Self:
+        self._client.__enter__()
+        self._container_client.__enter__()
+        self._datalake_client_for_blob_operation.__enter__()
+        return self
+
+    def __exit__(self, *args: Any) -> None:
+        self._datalake_client_for_blob_operation.__exit__(*args)
+        self._container_client.__exit__(*args)
+        self._client.__exit__(*args)
+
+    def close(self) -> None:
+        """This method is to close the sockets opened by the client.
+        It need not be used when using with a context manager.
+
+        :return: None
+        :rtype: None
+        """
+        self._datalake_client_for_blob_operation.close()
+        self._container_client.close()
+        self._client.close()
+
     def _build_generated_client(self, url: str) -> AzureDataLakeStorageRESTAPI:
         client = AzureDataLakeStorageRESTAPI(
             url,
@@ -145,17 +167,6 @@ def _build_generated_client(self, url: str) -> AzureDataLakeStorageRESTAPI:
     def _format_url(self, hostname: str) -> str:
         return _format_url(self.scheme, hostname, self.file_system_name, self._query_str)
 
-    def __exit__(self, *args: Any) -> None:
-        self._container_client.close()
-        self._datalake_client_for_blob_operation.close()
-        super(FileSystemClient, self).__exit__(*args)
-
-    def close(self) -> None:
-        """This method is to close the sockets opened by the client.
-        It need not be used when using with a context manager.
-        """
-        self.__exit__()
-
     @classmethod
     def from_connection_string(
         cls, conn_str: str,