[Storage] Adjust client closing logic for Datalake (Azure#30212)

.vscode/cspell.json

@@ -928,7 +928,9 @@
         "thead",
         "tbody",
         "racwdlmeop",
-        "newfs"
+        "newfs",
+        "adlsstorage",
+        "adlskey"
       ]
     },
     {

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

@@ -109,13 +109,14 @@ def __enter__(self):
 
     def __exit__(self, *args):
         self._blob_service_client.close()
+        super(DataLakeServiceClient, self).__exit__(*args)
 
     def close(self):
         # type: () -> None
         """ This method is to close the sockets opened by the client.
         It need not be used when using with a context manager.
         """
-        self._blob_service_client.close()
+        self.__exit__()
 
     def _format_url(self, hostname):
         """Format the endpoint URL according to hostname

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

@@ -130,14 +130,14 @@ def _format_url(self, hostname):
 
     def __exit__(self, *args):
         self._container_client.close()
+        self._datalake_client_for_blob_operation.close()
         super(FileSystemClient, self).__exit__(*args)
 
     def close(self):
         # type: () -> None
         """ This method is to close the sockets opened by the client.
         It need not be used when using with a context manager.
         """
-        self._container_client.close()
         self.__exit__()
 
     @classmethod

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

@@ -117,14 +117,14 @@ def __init__(
 
     def __exit__(self, *args):
         self._blob_client.close()
+        self._datalake_client_for_blob_operation.close()
         super(PathClient, self).__exit__(*args)
 
     def close(self):
         # type: () -> None
         """ This method is to close the sockets opened by the client.
         It need not be used when using with a context manager.
         """
-        self._blob_client.close()
         self.__exit__()
 
     def _format_url(self, hostname):

sdk/storage/azure-storage-file-datalake/azure/storage/filedatalake/aio/_data_lake_service_client_async.py

@@ -98,13 +98,14 @@ async def __aenter__(self):
 
     async def __aexit__(self, *args):
         await self._blob_service_client.close()
+        await super(DataLakeServiceClient, self).__aexit__(*args)
 
     async def close(self):
         # type: () -> None
         """ This method is to close the sockets opened by the client.
         It need not be used when using with a context manager.
         """
-        await self._blob_service_client.close()
+        await self.__aexit__()
 
     async def get_user_delegation_key(self, key_start_time,  # type: datetime
                                       key_expiry_time,  # type: datetime

sdk/storage/azure-storage-file-datalake/azure/storage/filedatalake/aio/_file_system_client_async.py

@@ -113,14 +113,14 @@ def __init__(
 
     async def __aexit__(self, *args):
         await self._container_client.close()
+        await self._datalake_client_for_blob_operation.close()
         await super(FileSystemClient, self).__aexit__(*args)
 
     async def close(self):
         # type: () -> None
         """ This method is to close the sockets opened by the client.
         It need not be used when using with a context manager.
         """
-        await self._container_client.close()
         await self.__aexit__()
 
     @distributed_trace_async

sdk/storage/azure-storage-file-datalake/azure/storage/filedatalake/aio/_path_client_async.py

@@ -88,14 +88,14 @@ def __init__(
 
     async def __aexit__(self, *args):
         await self._blob_client.close()
+        await self._datalake_client_for_blob_operation.close()
         await super(PathClient, self).__aexit__(*args)
 
     async def close(self):
         # type: () -> None
         """ This method is to close the sockets opened by the client.
         It need not be used when using with a context manager.
         """
-        await self._blob_client.close()
         await self.__aexit__()
 
     async def _create(self, resource_type, content_settings=None, metadata=None, **kwargs):

sdk/storage/azure-storage-file-datalake/tests/test_datalake_service_client.py

@@ -5,6 +5,7 @@
 # --------------------------------------------------------------------------
 
 import pytest
+from unittest.mock import MagicMock
 
 from azure.core.credentials import AzureNamedKeyCredential
 from azure.core.exceptions import HttpResponseError
@@ -411,3 +412,43 @@ def test_azure_named_key_credential_access(self, **kwargs):
 
         # Assert
         assert props is not None
+
+    @DataLakePreparer()
+    def test_datalake_clients_properly_close(self, **kwargs):
+        account_name = "adlsstorage"
+        account_key = "adlskey"
+
+        self._setup(account_name, account_key)
+        file_system_client = self.dsc.get_file_system_client(file_system='testfs')
+        dir_client = self.dsc.get_directory_client(file_system='testfs', directory='testdir')
+        file_client = dir_client.get_file_client(file='testfile')
+
+        # Mocks
+        self.dsc._blob_service_client.close = MagicMock()
+        self.dsc._client.__exit__ = MagicMock()
+        file_system_client._client.__exit__ = MagicMock()
+        file_system_client._datalake_client_for_blob_operation.close = MagicMock()
+        dir_client._client.__exit__ = MagicMock()
+        dir_client._datalake_client_for_blob_operation.close = MagicMock()
+        file_client._client.__exit__ = MagicMock()
+        file_client._datalake_client_for_blob_operation.close = MagicMock()
+
+        # Act
+        with self.dsc as dsc:
+            pass
+            with file_system_client as fsc:
+                pass
+                with dir_client as dc:
+                    pass
+                    with file_client as fc:
+                        pass
+
+        # Assert
+        self.dsc._blob_service_client.close.assert_called_once()
+        self.dsc._client.__exit__.assert_called_once()
+        file_system_client._client.__exit__.assert_called_once()
+        file_system_client._datalake_client_for_blob_operation.close.assert_called_once()
+        dir_client._client.__exit__.assert_called_once()
+        dir_client._datalake_client_for_blob_operation.close.assert_called_once()
+        file_client._client.__exit__.assert_called_once()
+        file_client._datalake_client_for_blob_operation.close.assert_called_once()

sdk/storage/azure-storage-file-datalake/tests/test_datalake_service_client_async.py

@@ -5,6 +5,7 @@
 # --------------------------------------------------------------------------
 
 import pytest
+import sys
 
 from azure.core.credentials import AzureNamedKeyCredential
 from azure.core.exceptions import HttpResponseError
@@ -25,6 +26,9 @@
 from devtools_testutils.storage.aio import AsyncStorageRecordedTestCase
 from settings.testcase import DataLakePreparer
 
+if sys.version_info >= (3, 8):
+    from unittest.mock import AsyncMock
+
 # ------------------------------------------------------------------------------
 TEST_FILE_SYSTEM_PREFIX = 'filesystem'
 # ------------------------------------------------------------------------------
@@ -409,3 +413,44 @@ async def test_azure_named_key_credential_access(self, **kwargs):
 
         # Assert
         assert props is not None
+
+    @pytest.mark.skipif(sys.version_info < (3, 8), reason="AsyncMock not introduced until 3.8")
+    @DataLakePreparer()
+    async def test_datalake_clients_properly_close(self, **kwargs):
+        account_name = "adlsstorage"
+        account_key = "adlskey"
+
+        self._setup(account_name, account_key)
+        file_system_client = self.dsc.get_file_system_client(file_system='testfs')
+        dir_client = self.dsc.get_directory_client(file_system='testfs', directory='testdir')
+        file_client = dir_client.get_file_client(file='testfile')
+
+        # Mocks
+        self.dsc._blob_service_client.close = AsyncMock()
+        self.dsc._client.__aexit__ = AsyncMock()
+        file_system_client._client.__aexit__ = AsyncMock()
+        file_system_client._datalake_client_for_blob_operation.close = AsyncMock()
+        dir_client._client.__aexit__ = AsyncMock()
+        dir_client._datalake_client_for_blob_operation.close = AsyncMock()
+        file_client._client.__aexit__ = AsyncMock()
+        file_client._datalake_client_for_blob_operation.close = AsyncMock()
+
+        # Act
+        async with self.dsc as dsc:
+            pass
+            async with file_system_client as fsc:
+                pass
+                async with dir_client as dc:
+                    pass
+                    async with file_client as fc:
+                        pass
+
+        # Assert
+        self.dsc._blob_service_client.close.assert_called_once()
+        self.dsc._client.__aexit__.assert_called_once()
+        file_system_client._client.__aexit__.assert_called_once()
+        file_system_client._datalake_client_for_blob_operation.close.assert_called_once()
+        dir_client._client.__aexit__.assert_called_once()
+        dir_client._datalake_client_for_blob_operation.close.assert_called_once()
+        file_client._client.__aexit__.assert_called_once()
+        file_client._datalake_client_for_blob_operation.close.assert_called_once()