Skip to content

Commit db32606

Browse files
tasherif-msfttasherif-msft
authored
[Bug][Datalake] Fixed Datalake Service Client context manager/session closure issue (Azure#16414)
* Fixed Datalake Service Client context manager/session closure issue * test recording * removed redundant code
1 parent 032b09d commit db32606

File tree

6 files changed

+247
-4
lines changed

6 files changed

+247
-4
lines changed

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

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -92,17 +92,19 @@ def __init__(
9292
# ADLS doesn't support secondary endpoint, make sure it's empty
9393
self._hosts[LocationMode.SECONDARY] = ""
9494

95+
def __enter__(self):
96+
self._blob_service_client.__enter__()
97+
return self
98+
9599
def __exit__(self, *args):
96100
self._blob_service_client.close()
97-
super(DataLakeServiceClient, self).__exit__(*args)
98101

99102
def close(self):
100103
# type: () -> None
101104
""" This method is to close the sockets opened by the client.
102105
It need not be used when using with a context manager.
103106
"""
104107
self._blob_service_client.close()
105-
self.__exit__()
106108

107109
def _format_url(self, hostname):
108110
"""Format the endpoint URL according to hostname

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

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -80,17 +80,19 @@ def __init__(
8080
self._client = AzureDataLakeStorageRESTAPI(self.url, pipeline=self._pipeline)
8181
self._loop = kwargs.get('loop', None)
8282

83+
async def __aenter__(self):
84+
await self._blob_service_client.__aenter__()
85+
return self
86+
8387
async def __aexit__(self, *args):
8488
await self._blob_service_client.close()
85-
await super(DataLakeServiceClient, self).__aexit__(*args)
8689

8790
async def close(self):
8891
# type: () -> None
8992
""" This method is to close the sockets opened by the client.
9093
It need not be used when using with a context manager.
9194
"""
9295
await self._blob_service_client.close()
93-
await self.__aexit__()
9496

9597
async def get_user_delegation_key(self, key_start_time, # type: datetime
9698
key_expiry_time, # type: datetime

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

Lines changed: 124 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,124 @@
1+
interactions:
2+
- request:
3+
body: null
4+
headers:
5+
Accept:
6+
- application/xml
7+
Accept-Encoding:
8+
- gzip, deflate
9+
Connection:
10+
- keep-alive
11+
Content-Length:
12+
- '0'
13+
User-Agent:
14+
- azsdk-python-storage-dfs/12.2.3 Python/3.8.5 (Windows-10-10.0.19041-SP0)
15+
x-ms-client-request-id:
16+
- 18d17c7c-61c4-11eb-8696-c8348e5fffbf
17+
x-ms-date:
18+
- Thu, 28 Jan 2021 23:54:02 GMT
19+
x-ms-version:
20+
- '2020-04-08'
21+
method: PUT
22+
uri: https://storagename.blob.core.windows.net/fs1dc581ff4?restype=container
23+
response:
24+
body:
25+
string: ''
26+
headers:
27+
Date:
28+
- Thu, 28 Jan 2021 23:54:02 GMT
29+
ETag:
30+
- '"0x8D8C3E7FD803F14"'
31+
Last-Modified:
32+
- Thu, 28 Jan 2021 23:54:03 GMT
33+
Server:
34+
- Windows-Azure-Blob/1.0 Microsoft-HTTPAPI/2.0
35+
Transfer-Encoding:
36+
- chunked
37+
x-ms-request-id:
38+
- d7930b4a-001e-0017-67d0-f55645000000
39+
x-ms-version:
40+
- '2020-04-08'
41+
status:
42+
code: 201
43+
message: Created
44+
- request:
45+
body: null
46+
headers:
47+
Accept:
48+
- application/xml
49+
Accept-Encoding:
50+
- gzip, deflate
51+
Connection:
52+
- keep-alive
53+
Content-Length:
54+
- '0'
55+
User-Agent:
56+
- azsdk-python-storage-dfs/12.2.3 Python/3.8.5 (Windows-10-10.0.19041-SP0)
57+
x-ms-client-request-id:
58+
- 19180970-61c4-11eb-b060-c8348e5fffbf
59+
x-ms-date:
60+
- Thu, 28 Jan 2021 23:54:02 GMT
61+
x-ms-version:
62+
- '2020-04-08'
63+
method: DELETE
64+
uri: https://storagename.blob.core.windows.net/fs1dc581ff4?restype=container
65+
response:
66+
body:
67+
string: ''
68+
headers:
69+
Date:
70+
- Thu, 28 Jan 2021 23:54:02 GMT
71+
Server:
72+
- Windows-Azure-Blob/1.0 Microsoft-HTTPAPI/2.0
73+
Transfer-Encoding:
74+
- chunked
75+
x-ms-request-id:
76+
- d7930b57-001e-0017-71d0-f55645000000
77+
x-ms-version:
78+
- '2020-04-08'
79+
status:
80+
code: 202
81+
message: Accepted
82+
- request:
83+
body: null
84+
headers:
85+
Accept:
86+
- application/xml
87+
Accept-Encoding:
88+
- gzip, deflate
89+
Connection:
90+
- keep-alive
91+
Content-Length:
92+
- '0'
93+
User-Agent:
94+
- azsdk-python-storage-dfs/12.2.3 Python/3.8.5 (Windows-10-10.0.19041-SP0)
95+
x-ms-client-request-id:
96+
- 192bdf8f-61c4-11eb-a0a8-c8348e5fffbf
97+
x-ms-date:
98+
- Thu, 28 Jan 2021 23:54:02 GMT
99+
x-ms-version:
100+
- '2020-04-08'
101+
method: PUT
102+
uri: https://storagename.blob.core.windows.net/fs2dc581ff4?restype=container
103+
response:
104+
body:
105+
string: ''
106+
headers:
107+
Date:
108+
- Thu, 28 Jan 2021 23:54:03 GMT
109+
ETag:
110+
- '"0x8D8C3E7FDD8A307"'
111+
Last-Modified:
112+
- Thu, 28 Jan 2021 23:54:03 GMT
113+
Server:
114+
- Windows-Azure-Blob/1.0 Microsoft-HTTPAPI/2.0
115+
Transfer-Encoding:
116+
- chunked
117+
x-ms-request-id:
118+
- 2fff5a5b-e01e-002c-04d0-f54f7d000000
119+
x-ms-version:
120+
- '2020-04-08'
121+
status:
122+
code: 201
123+
message: Created
124+
version: 1

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

Lines changed: 90 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,90 @@
1+
interactions:
2+
- request:
3+
body: null
4+
headers:
5+
Accept:
6+
- application/xml
7+
User-Agent:
8+
- azsdk-python-storage-dfs/12.2.3 Python/3.8.5 (Windows-10-10.0.19041-SP0)
9+
x-ms-client-request-id:
10+
- 1407f955-61c4-11eb-a0f0-c8348e5fffbf
11+
x-ms-date:
12+
- Thu, 28 Jan 2021 23:53:54 GMT
13+
x-ms-version:
14+
- '2020-04-08'
15+
method: PUT
16+
uri: https://storagename.blob.core.windows.net/fs1a55d2271?restype=container
17+
response:
18+
body:
19+
string: ''
20+
headers:
21+
Date: Thu, 28 Jan 2021 23:53:54 GMT
22+
ETag: '"0x8D8C3E7F8AA6AA7"'
23+
Last-Modified: Thu, 28 Jan 2021 23:53:54 GMT
24+
Server: Windows-Azure-Blob/1.0 Microsoft-HTTPAPI/2.0
25+
Transfer-Encoding: chunked
26+
x-ms-request-id: 26abe84d-601e-0022-14d0-f50386000000
27+
x-ms-version: '2020-04-08'
28+
status:
29+
code: 201
30+
message: Created
31+
url: https://storagename.blob.core.windows.net/fs1a55d2271?restype=container
32+
- request:
33+
body: null
34+
headers:
35+
Accept:
36+
- application/xml
37+
User-Agent:
38+
- azsdk-python-storage-dfs/12.2.3 Python/3.8.5 (Windows-10-10.0.19041-SP0)
39+
x-ms-client-request-id:
40+
- 14448d10-61c4-11eb-9006-c8348e5fffbf
41+
x-ms-date:
42+
- Thu, 28 Jan 2021 23:53:54 GMT
43+
x-ms-version:
44+
- '2020-04-08'
45+
method: DELETE
46+
uri: https://storagename.blob.core.windows.net/fs1a55d2271?restype=container
47+
response:
48+
body:
49+
string: ''
50+
headers:
51+
Date: Thu, 28 Jan 2021 23:53:54 GMT
52+
Server: Windows-Azure-Blob/1.0 Microsoft-HTTPAPI/2.0
53+
Transfer-Encoding: chunked
54+
x-ms-request-id: 26abe85b-601e-0022-18d0-f50386000000
55+
x-ms-version: '2020-04-08'
56+
status:
57+
code: 202
58+
message: Accepted
59+
url: https://storagename.blob.core.windows.net/fs1a55d2271?restype=container
60+
- request:
61+
body: null
62+
headers:
63+
Accept:
64+
- application/xml
65+
User-Agent:
66+
- azsdk-python-storage-dfs/12.2.3 Python/3.8.5 (Windows-10-10.0.19041-SP0)
67+
x-ms-client-request-id:
68+
- 144f627f-61c4-11eb-b938-c8348e5fffbf
69+
x-ms-date:
70+
- Thu, 28 Jan 2021 23:53:54 GMT
71+
x-ms-version:
72+
- '2020-04-08'
73+
method: PUT
74+
uri: https://storagename.blob.core.windows.net/fs2a55d2271?restype=container
75+
response:
76+
body:
77+
string: ''
78+
headers:
79+
Date: Thu, 28 Jan 2021 23:53:54 GMT
80+
ETag: '"0x8D8C3E7F8E04FCB"'
81+
Last-Modified: Thu, 28 Jan 2021 23:53:55 GMT
82+
Server: Windows-Azure-Blob/1.0 Microsoft-HTTPAPI/2.0
83+
Transfer-Encoding: chunked
84+
x-ms-request-id: 7145ae31-f01e-0011-7ad0-f51774000000
85+
x-ms-version: '2020-04-08'
86+
status:
87+
code: 201
88+
message: Created
89+
url: https://storagename.blob.core.windows.net/fs2a55d2271?restype=container
90+
version: 1

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

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -214,6 +214,16 @@ def test_get_file_system_properties(self):
214214
self.assertIsNotNone(props.has_immutability_policy)
215215
self.assertIsNotNone(props.has_legal_hold)
216216

217+
@record
218+
def test_service_client_session_closes_after_filesystem_creation(self):
219+
# Arrange
220+
dsc2 = DataLakeServiceClient(self.dsc.url, credential=self.settings.STORAGE_DATA_LAKE_ACCOUNT_KEY)
221+
with DataLakeServiceClient(self.dsc.url, credential=self.settings.STORAGE_DATA_LAKE_ACCOUNT_KEY) as ds_client:
222+
fs1 = ds_client.create_file_system(self._get_file_system_reference(prefix="fs1"))
223+
fs1.delete_file_system()
224+
dsc2.create_file_system(self._get_file_system_reference(prefix="fs2"))
225+
dsc2.close()
226+
217227
@record
218228
def test_list_paths(self):
219229
# Arrange

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

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -287,6 +287,21 @@ def test_get_file_system_properties_async(self):
287287
loop = asyncio.get_event_loop()
288288
loop.run_until_complete(self._test_get_file_system_properties_async())
289289

290+
async def _test_service_client_session_closes_after_filesystem_creation(self):
291+
# Arrange
292+
dsc2 = DataLakeServiceClient(self.dsc.url, credential=self.settings.STORAGE_DATA_LAKE_ACCOUNT_KEY)
293+
async with DataLakeServiceClient(
294+
self.dsc.url, credential=self.settings.STORAGE_DATA_LAKE_ACCOUNT_KEY) as ds_client:
295+
fs1 = await ds_client.create_file_system(self._get_file_system_reference(prefix="fs1"))
296+
await fs1.delete_file_system()
297+
await dsc2.create_file_system(self._get_file_system_reference(prefix="fs2"))
298+
await dsc2.close()
299+
300+
@record
301+
def test_service_client_session_closes_after_filesystem_creation(self):
302+
loop = asyncio.get_event_loop()
303+
loop.run_until_complete(self._test_service_client_session_closes_after_filesystem_creation())
304+
290305
async def _test_list_paths_async(self):
291306
# Arrange
292307
file_system = await self._create_file_system()

0 commit comments

Comments
 (0)