[Bug][Datalake] Fixed Datalake Service Client context manager/session closure issue

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

@@ -92,9 +92,13 @@ def __init__(
         # ADLS doesn't support secondary endpoint, make sure it's empty
         self._hosts[LocationMode.SECONDARY] = ""
 
+    def __enter__(self):
+        self._blob_service_client.__enter__()
+        return self
+
     def __exit__(self, *args):
         self._blob_service_client.close()
-        super(DataLakeServiceClient, self).__exit__(*args)
+        super(BlobServiceClient, self._blob_service_client).__exit__(*args)
 
     def close(self):
         # type: () -> None

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

@@ -80,9 +80,13 @@ def __init__(
         self._client = AzureDataLakeStorageRESTAPI(self.url, pipeline=self._pipeline)
         self._loop = kwargs.get('loop', None)
 
+    async def __aenter__(self):
+        await self._blob_service_client.__aenter__()
+        return self
+
     async def __aexit__(self, *args):
         await self._blob_service_client.close()
-        await super(DataLakeServiceClient, self).__aexit__(*args)
+        await super(BlobServiceClient, self._blob_service_client).__aexit__(*args)
 
     async def close(self):
         # type: () -> None

sdk/storage/azure-storage-file-datalake/tests/recordings/test_file_system.test_service_client_session_closes_after_filesystem_creation.yaml

@@ -0,0 +1,124 @@
+interactions:
+- request:
+    body: null
+    headers:
+      Accept:
+      - application/xml
+      Accept-Encoding:
+      - gzip, deflate
+      Connection:
+      - keep-alive
+      Content-Length:
+      - '0'
+      User-Agent:
+      - azsdk-python-storage-dfs/12.2.3 Python/3.8.5 (Windows-10-10.0.19041-SP0)
+      x-ms-client-request-id:
+      - 18d17c7c-61c4-11eb-8696-c8348e5fffbf
+      x-ms-date:
+      - Thu, 28 Jan 2021 23:54:02 GMT
+      x-ms-version:
+      - '2020-04-08'
+    method: PUT
+    uri: https://sachinssoftdelete3.blob.core.windows.net/fs1dc581ff4?restype=container
+  response:
+    body:
+      string: ''
+    headers:
+      Date:
+      - Thu, 28 Jan 2021 23:54:02 GMT
+      ETag:
+      - '"0x8D8C3E7FD803F14"'
+      Last-Modified:
+      - Thu, 28 Jan 2021 23:54:03 GMT
+      Server:
+      - Windows-Azure-Blob/1.0 Microsoft-HTTPAPI/2.0
+      Transfer-Encoding:
+      - chunked
+      x-ms-request-id:
+      - d7930b4a-001e-0017-67d0-f55645000000
+      x-ms-version:
+      - '2020-04-08'
+    status:
+      code: 201
+      message: Created
+- request:
+    body: null
+    headers:
+      Accept:
+      - application/xml
+      Accept-Encoding:
+      - gzip, deflate
+      Connection:
+      - keep-alive
+      Content-Length:
+      - '0'
+      User-Agent:
+      - azsdk-python-storage-dfs/12.2.3 Python/3.8.5 (Windows-10-10.0.19041-SP0)
+      x-ms-client-request-id:
+      - 19180970-61c4-11eb-b060-c8348e5fffbf
+      x-ms-date:
+      - Thu, 28 Jan 2021 23:54:02 GMT
+      x-ms-version:
+      - '2020-04-08'
+    method: DELETE
+    uri: https://sachinssoftdelete3.blob.core.windows.net/fs1dc581ff4?restype=container
+  response:
+    body:
+      string: ''
+    headers:
+      Date:
+      - Thu, 28 Jan 2021 23:54:02 GMT
+      Server:
+      - Windows-Azure-Blob/1.0 Microsoft-HTTPAPI/2.0
+      Transfer-Encoding:
+      - chunked
+      x-ms-request-id:
+      - d7930b57-001e-0017-71d0-f55645000000
+      x-ms-version:
+      - '2020-04-08'
+    status:
+      code: 202
+      message: Accepted
+- request:
+    body: null
+    headers:
+      Accept:
+      - application/xml
+      Accept-Encoding:
+      - gzip, deflate
+      Connection:
+      - keep-alive
+      Content-Length:
+      - '0'
+      User-Agent:
+      - azsdk-python-storage-dfs/12.2.3 Python/3.8.5 (Windows-10-10.0.19041-SP0)
+      x-ms-client-request-id:
+      - 192bdf8f-61c4-11eb-a0a8-c8348e5fffbf
+      x-ms-date:
+      - Thu, 28 Jan 2021 23:54:02 GMT
+      x-ms-version:
+      - '2020-04-08'
+    method: PUT
+    uri: https://sachinssoftdelete3.blob.core.windows.net/fs2dc581ff4?restype=container
+  response:
+    body:
+      string: ''
+    headers:
+      Date:
+      - Thu, 28 Jan 2021 23:54:03 GMT
+      ETag:
+      - '"0x8D8C3E7FDD8A307"'
+      Last-Modified:
+      - Thu, 28 Jan 2021 23:54:03 GMT
+      Server:
+      - Windows-Azure-Blob/1.0 Microsoft-HTTPAPI/2.0
+      Transfer-Encoding:
+      - chunked
+      x-ms-request-id:
+      - 2fff5a5b-e01e-002c-04d0-f54f7d000000
+      x-ms-version:
+      - '2020-04-08'
+    status:
+      code: 201
+      message: Created
+version: 1

sdk/storage/azure-storage-file-datalake/tests/recordings/test_file_system_async.test_service_client_session_closes_after_filesystem_creation.yaml

@@ -0,0 +1,90 @@
+interactions:
+- request:
+    body: null
+    headers:
+      Accept:
+      - application/xml
+      User-Agent:
+      - azsdk-python-storage-dfs/12.2.3 Python/3.8.5 (Windows-10-10.0.19041-SP0)
+      x-ms-client-request-id:
+      - 1407f955-61c4-11eb-a0f0-c8348e5fffbf
+      x-ms-date:
+      - Thu, 28 Jan 2021 23:53:54 GMT
+      x-ms-version:
+      - '2020-04-08'
+    method: PUT
+    uri: https://sachinssoftdelete3.blob.core.windows.net/fs1a55d2271?restype=container
+  response:
+    body:
+      string: ''
+    headers:
+      Date: Thu, 28 Jan 2021 23:53:54 GMT
+      ETag: '"0x8D8C3E7F8AA6AA7"'
+      Last-Modified: Thu, 28 Jan 2021 23:53:54 GMT
+      Server: Windows-Azure-Blob/1.0 Microsoft-HTTPAPI/2.0
+      Transfer-Encoding: chunked
+      x-ms-request-id: 26abe84d-601e-0022-14d0-f50386000000
+      x-ms-version: '2020-04-08'
+    status:
+      code: 201
+      message: Created
+    url: https://sachinssoftdelete3.blob.core.windows.net/fs1a55d2271?restype=container
+- request:
+    body: null
+    headers:
+      Accept:
+      - application/xml
+      User-Agent:
+      - azsdk-python-storage-dfs/12.2.3 Python/3.8.5 (Windows-10-10.0.19041-SP0)
+      x-ms-client-request-id:
+      - 14448d10-61c4-11eb-9006-c8348e5fffbf
+      x-ms-date:
+      - Thu, 28 Jan 2021 23:53:54 GMT
+      x-ms-version:
+      - '2020-04-08'
+    method: DELETE
+    uri: https://sachinssoftdelete3.blob.core.windows.net/fs1a55d2271?restype=container
+  response:
+    body:
+      string: ''
+    headers:
+      Date: Thu, 28 Jan 2021 23:53:54 GMT
+      Server: Windows-Azure-Blob/1.0 Microsoft-HTTPAPI/2.0
+      Transfer-Encoding: chunked
+      x-ms-request-id: 26abe85b-601e-0022-18d0-f50386000000
+      x-ms-version: '2020-04-08'
+    status:
+      code: 202
+      message: Accepted
+    url: https://sachinssoftdelete3.blob.core.windows.net/fs1a55d2271?restype=container
+- request:
+    body: null
+    headers:
+      Accept:
+      - application/xml
+      User-Agent:
+      - azsdk-python-storage-dfs/12.2.3 Python/3.8.5 (Windows-10-10.0.19041-SP0)
+      x-ms-client-request-id:
+      - 144f627f-61c4-11eb-b938-c8348e5fffbf
+      x-ms-date:
+      - Thu, 28 Jan 2021 23:53:54 GMT
+      x-ms-version:
+      - '2020-04-08'
+    method: PUT
+    uri: https://sachinssoftdelete3.blob.core.windows.net/fs2a55d2271?restype=container
+  response:
+    body:
+      string: ''
+    headers:
+      Date: Thu, 28 Jan 2021 23:53:54 GMT
+      ETag: '"0x8D8C3E7F8E04FCB"'
+      Last-Modified: Thu, 28 Jan 2021 23:53:55 GMT
+      Server: Windows-Azure-Blob/1.0 Microsoft-HTTPAPI/2.0
+      Transfer-Encoding: chunked
+      x-ms-request-id: 7145ae31-f01e-0011-7ad0-f51774000000
+      x-ms-version: '2020-04-08'
+    status:
+      code: 201
+      message: Created
+    url: https://sachinssoftdelete3.blob.core.windows.net/fs2a55d2271?restype=container
+version: 1

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

@@ -214,6 +214,16 @@ def test_get_file_system_properties(self):
         self.assertIsNotNone(props.has_immutability_policy)
         self.assertIsNotNone(props.has_legal_hold)
 
+    @record
+    def test_service_client_session_closes_after_filesystem_creation(self):
+        # Arrange
+        dsc2 = DataLakeServiceClient(self.dsc.url, credential=self.settings.STORAGE_DATA_LAKE_ACCOUNT_KEY)
+        with DataLakeServiceClient(self.dsc.url, credential=self.settings.STORAGE_DATA_LAKE_ACCOUNT_KEY) as ds_client:
+            fs1 = ds_client.create_file_system(self._get_file_system_reference(prefix="fs1"))
+            fs1.delete_file_system()
+        dsc2.create_file_system(self._get_file_system_reference(prefix="fs2"))
+        dsc2.close()
+
     @record
     def test_list_paths(self):
         # Arrange

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

@@ -287,6 +287,21 @@ def test_get_file_system_properties_async(self):
         loop = asyncio.get_event_loop()
         loop.run_until_complete(self._test_get_file_system_properties_async())
 
+    async def _test_service_client_session_closes_after_filesystem_creation(self):
+        # Arrange
+        dsc2 = DataLakeServiceClient(self.dsc.url, credential=self.settings.STORAGE_DATA_LAKE_ACCOUNT_KEY)
+        async with DataLakeServiceClient(
+                self.dsc.url, credential=self.settings.STORAGE_DATA_LAKE_ACCOUNT_KEY) as ds_client:
+            fs1 = await ds_client.create_file_system(self._get_file_system_reference(prefix="fs1"))
+            await fs1.delete_file_system()
+        await dsc2.create_file_system(self._get_file_system_reference(prefix="fs2"))
+        await dsc2.close()
+
+    @record
+    def test_service_client_session_closes_after_filesystem_creation(self):
+        loop = asyncio.get_event_loop()
+        loop.run_until_complete(self._test_service_client_session_closes_after_filesystem_creation())
+
     async def _test_list_paths_async(self):
         # Arrange
         file_system = await self._create_file_system()