[Datalake] Fixed session closure of filesystem (Azure#14910)

* Fixed session closure of filesystem

* used wrapper

* used wrapper

* added unit test

* fixed all clients

* removed check

* removed check

* recorded test

* recorded async test

Author: tasherif-msft

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

@@ -7,8 +7,9 @@
     from urllib.parse import quote, unquote
 except ImportError:
     from urllib2 import quote, unquote # type: ignore
+from azure.core.pipeline import Pipeline
 from ._deserialize import deserialize_dir_properties
-from ._shared.base_client import parse_connection_str
+from ._shared.base_client import TransportWrapper, parse_connection_str
 from ._data_lake_file_client import DataLakeFileClient
 from ._models import DirectoryProperties
 from ._path_client import PathClient
@@ -505,6 +506,10 @@ def get_file_client(self, file  # type: Union[FileProperties, str]
         except AttributeError:
             file_path = self.path_name + '/' + file
 
+        _pipeline = Pipeline(
+            transport=TransportWrapper(self._pipeline._transport), # pylint: disable = protected-access
+            policies=self._pipeline._impl_policies # pylint: disable = protected-access
+        )
         return DataLakeFileClient(
             self.url, self.file_system_name, file_path=file_path, credential=self._raw_credential,
             _hosts=self._hosts, _configuration=self._config, _pipeline=self._pipeline,
@@ -531,6 +536,10 @@ def get_sub_directory_client(self, sub_directory  # type: Union[DirectoryPropert
         except AttributeError:
             subdir_path = self.path_name + '/' + sub_directory
 
+        _pipeline = Pipeline(
+            transport=TransportWrapper(self._pipeline._transport), # pylint: disable = protected-access
+            policies=self._pipeline._impl_policies # pylint: disable = protected-access
+        )
         return DataLakeDirectoryClient(
             self.url, self.file_system_name, directory_name=subdir_path, credential=self._raw_credential,
             _hosts=self._hosts, _configuration=self._config, _pipeline=self._pipeline,

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

@@ -10,9 +10,10 @@
     from urlparse import urlparse  # type: ignore
 
 from azure.core.paging import ItemPaged
+from azure.core.pipeline import Pipeline
 
 from azure.storage.blob import BlobServiceClient
-from ._shared.base_client import StorageAccountHostsMixin, parse_query, parse_connection_str
+from ._shared.base_client import TransportWrapper, StorageAccountHostsMixin, parse_query, parse_connection_str
 from ._file_system_client import FileSystemClient
 from ._data_lake_directory_client import DataLakeDirectoryClient
 from ._data_lake_file_client import DataLakeFileClient
@@ -325,9 +326,13 @@ def get_file_system_client(self, file_system  # type: Union[FileSystemProperties
         except AttributeError:
             file_system_name = file_system
 
+        _pipeline = Pipeline(
+            transport=TransportWrapper(self._pipeline._transport), # pylint: disable = protected-access
+            policies=self._pipeline._impl_policies # pylint: disable = protected-access
+        )
         return FileSystemClient(self.url, file_system_name, credential=self._raw_credential,
                                 _configuration=self._config,
-                                _pipeline=self._pipeline, _hosts=self._hosts,
+                                _pipeline=_pipeline, _hosts=self._hosts,
                                 require_encryption=self.require_encryption, key_encryption_key=self.key_encryption_key,
                                 key_resolver_function=self.key_resolver_function)
 
@@ -367,9 +372,14 @@ def get_directory_client(self, file_system,  # type: Union[FileSystemProperties,
             directory_name = directory.name
         except AttributeError:
             directory_name = directory
+
+        _pipeline = Pipeline(
+            transport=TransportWrapper(self._pipeline._transport), # pylint: disable = protected-access
+            policies=self._pipeline._impl_policies # pylint: disable = protected-access
+        )
         return DataLakeDirectoryClient(self.url, file_system_name, directory_name=directory_name,
                                        credential=self._raw_credential,
-                                       _configuration=self._config, _pipeline=self._pipeline,
+                                       _configuration=self._config, _pipeline=_pipeline,
                                        _hosts=self._hosts,
                                        require_encryption=self.require_encryption,
                                        key_encryption_key=self.key_encryption_key,
@@ -413,9 +423,13 @@ def get_file_client(self, file_system,  # type: Union[FileSystemProperties, str]
         except AttributeError:
             pass
 
+        _pipeline = Pipeline(
+            transport=TransportWrapper(self._pipeline._transport), # pylint: disable = protected-access
+            policies=self._pipeline._impl_policies # pylint: disable = protected-access
+        )
         return DataLakeFileClient(
             self.url, file_system_name, file_path=file_path, credential=self._raw_credential,
-            _hosts=self._hosts, _configuration=self._config, _pipeline=self._pipeline,
+            _hosts=self._hosts, _configuration=self._config, _pipeline=_pipeline,
             require_encryption=self.require_encryption,
             key_encryption_key=self.key_encryption_key,
             key_resolver_function=self.key_resolver_function)

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

@@ -12,9 +12,10 @@
     from urllib2 import quote  # type: ignore
 
 import six
+from azure.core.pipeline import Pipeline
 from azure.core.paging import ItemPaged
 from azure.storage.blob import ContainerClient
-from ._shared.base_client import StorageAccountHostsMixin, parse_query, parse_connection_str
+from ._shared.base_client import TransportWrapper, StorageAccountHostsMixin, parse_query, parse_connection_str
 from ._serialize import convert_dfs_url_to_blob_url
 from ._models import LocationMode, FileSystemProperties, PublicAccess
 from ._list_paths_helper import PathPropertiesPaged
@@ -737,10 +738,13 @@ def get_directory_client(self, directory  # type: Union[DirectoryProperties, str
             directory_name = directory.name
         except AttributeError:
             directory_name = directory
-
+        _pipeline = Pipeline(
+            transport=TransportWrapper(self._pipeline._transport), # pylint: disable = protected-access
+            policies=self._pipeline._impl_policies # pylint: disable = protected-access
+        )
         return DataLakeDirectoryClient(self.url, self.file_system_name, directory_name=directory_name,
                                        credential=self._raw_credential,
-                                       _configuration=self._config, _pipeline=self._pipeline,
+                                       _configuration=self._config, _pipeline=_pipeline,
                                        _hosts=self._hosts,
                                        require_encryption=self.require_encryption,
                                        key_encryption_key=self.key_encryption_key,
@@ -774,10 +778,13 @@ def get_file_client(self, file_path  # type: Union[FileProperties, str]
             file_path = file_path.name
         except AttributeError:
             pass
-
+        _pipeline = Pipeline(
+            transport=TransportWrapper(self._pipeline._transport), # pylint: disable = protected-access
+            policies=self._pipeline._impl_policies # pylint: disable = protected-access
+        )
         return DataLakeFileClient(
             self.url, self.file_system_name, file_path=file_path, credential=self._raw_credential,
-            _hosts=self._hosts, _configuration=self._config, _pipeline=self._pipeline,
+            _hosts=self._hosts, _configuration=self._config, _pipeline=_pipeline,
             require_encryption=self.require_encryption,
             key_encryption_key=self.key_encryption_key,
             key_resolver_function=self.key_resolver_function)

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

@@ -8,11 +8,13 @@
     from urllib.parse import quote, unquote
 except ImportError:
     from urllib2 import quote, unquote # type: ignore
+from azure.core.pipeline import AsyncPipeline
 from ._data_lake_file_client_async import DataLakeFileClient
 from .._data_lake_directory_client import DataLakeDirectoryClient as DataLakeDirectoryClientBase
 from .._models import DirectoryProperties
 from .._deserialize import deserialize_dir_properties
 from ._path_client_async import PathClient
+from .._shared.base_client_async import AsyncTransportWrapper
 
 
 class DataLakeDirectoryClient(PathClient, DataLakeDirectoryClientBase):
@@ -483,6 +485,10 @@ def get_file_client(self, file  # type: Union[FileProperties, str]
         except AttributeError:
             file_path = self.path_name + '/' + file
 
+        _pipeline = AsyncPipeline(
+            transport=AsyncTransportWrapper(self._pipeline._transport), # pylint: disable = protected-access
+            policies=self._pipeline._impl_policies # pylint: disable = protected-access
+        )
         return DataLakeFileClient(
             self.url, self.file_system_name, file_path=file_path, credential=self._raw_credential,
             _hosts=self._hosts, _configuration=self._config, _pipeline=self._pipeline,
@@ -518,6 +524,10 @@ def get_sub_directory_client(self, sub_directory  # type: Union[DirectoryPropert
         except AttributeError:
             subdir_path = self.path_name + '/' + sub_directory
 
+        _pipeline = AsyncPipeline(
+            transport=AsyncTransportWrapper(self._pipeline._transport), # pylint: disable = protected-access
+            policies=self._pipeline._impl_policies # pylint: disable = protected-access
+        )
         return DataLakeDirectoryClient(
             self.url, self.file_system_name, directory_name=subdir_path, credential=self._raw_credential,
             _hosts=self._hosts, _configuration=self._config, _pipeline=self._pipeline,

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

@@ -6,10 +6,11 @@
 # pylint: disable=invalid-overridden-method
 
 from azure.core.paging import ItemPaged
+from azure.core.pipeline import AsyncPipeline
 
 from azure.storage.blob.aio import BlobServiceClient
 from .._generated.aio import DataLakeStorageClient
-from .._shared.base_client_async import AsyncStorageAccountHostsMixin
+from .._shared.base_client_async import AsyncTransportWrapper, AsyncStorageAccountHostsMixin
 from ._file_system_client_async import FileSystemClient
 from .._data_lake_service_client import DataLakeServiceClient as DataLakeServiceClientBase
 from .._shared.policies_async import ExponentialRetry
@@ -276,6 +277,10 @@ def get_file_system_client(self, file_system  # type: Union[FileSystemProperties
         except AttributeError:
             file_system_name = file_system
 
+        _pipeline = AsyncPipeline(
+            transport=AsyncTransportWrapper(self._pipeline._transport), # pylint: disable = protected-access
+            policies=self._pipeline._impl_policies # pylint: disable = protected-access
+        )
         return FileSystemClient(self.url, file_system_name, credential=self._raw_credential,
                                 _configuration=self._config,
                                 _pipeline=self._pipeline, _hosts=self._hosts,
@@ -318,6 +323,11 @@ def get_directory_client(self, file_system,  # type: Union[FileSystemProperties,
             directory_name = directory.name
         except AttributeError:
             directory_name = directory
+
+        _pipeline = AsyncPipeline(
+            transport=AsyncTransportWrapper(self._pipeline._transport), # pylint: disable = protected-access
+            policies=self._pipeline._impl_policies # pylint: disable = protected-access
+        )
         return DataLakeDirectoryClient(self.url, file_system_name, directory_name=directory_name,
                                        credential=self._raw_credential,
                                        _configuration=self._config, _pipeline=self._pipeline,
@@ -364,6 +374,10 @@ def get_file_client(self, file_system,  # type: Union[FileSystemProperties, str]
         except AttributeError:
             pass
 
+        _pipeline = AsyncPipeline(
+            transport=AsyncTransportWrapper(self._pipeline._transport), # pylint: disable = protected-access
+            policies=self._pipeline._impl_policies # pylint: disable = protected-access
+        )
         return DataLakeFileClient(
             self.url, file_system_name, file_path=file_path, credential=self._raw_credential,
             _hosts=self._hosts, _configuration=self._config, _pipeline=self._pipeline,

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

@@ -13,6 +13,7 @@
 
 from azure.core.tracing.decorator import distributed_trace
 
+from azure.core.pipeline import AsyncPipeline
 from azure.core.async_paging import AsyncItemPaged
 
 from azure.core.tracing.decorator_async import distributed_trace_async
@@ -24,7 +25,7 @@
 from ._data_lake_lease_async import DataLakeLeaseClient
 from .._file_system_client import FileSystemClient as FileSystemClientBase
 from .._generated.aio import DataLakeStorageClient
-from .._shared.base_client_async import AsyncStorageAccountHostsMixin
+from .._shared.base_client_async import AsyncTransportWrapper, AsyncStorageAccountHostsMixin
 from .._shared.policies_async import ExponentialRetry
 from .._models import FileSystemProperties, PublicAccess
 
@@ -698,10 +699,13 @@ def get_directory_client(self, directory  # type: Union[DirectoryProperties, str
             directory_name = directory.name
         except AttributeError:
             directory_name = directory
-
+        _pipeline = AsyncPipeline(
+            transport=AsyncTransportWrapper(self._pipeline._transport), # pylint: disable = protected-access
+            policies=self._pipeline._impl_policies # pylint: disable = protected-access
+        )
         return DataLakeDirectoryClient(self.url, self.file_system_name, directory_name=directory_name,
                                        credential=self._raw_credential,
-                                       _configuration=self._config, _pipeline=self._pipeline,
+                                       _configuration=self._config, _pipeline=_pipeline,
                                        _hosts=self._hosts,
                                        require_encryption=self.require_encryption,
                                        key_encryption_key=self.key_encryption_key,
@@ -736,10 +740,13 @@ def get_file_client(self, file_path  # type: Union[FileProperties, str]
             file_path = file_path.name
         except AttributeError:
             pass
-
+        _pipeline = AsyncPipeline(
+            transport=AsyncTransportWrapper(self._pipeline._transport), # pylint: disable = protected-access
+            policies=self._pipeline._impl_policies # pylint: disable = protected-access
+        )
         return DataLakeFileClient(
             self.url, self.file_system_name, file_path=file_path, credential=self._raw_credential,
-            _hosts=self._hosts, _configuration=self._config, _pipeline=self._pipeline,
+            _hosts=self._hosts, _configuration=self._config, _pipeline=_pipeline,
             require_encryption=self.require_encryption,
             key_encryption_key=self.key_encryption_key,
             key_resolver_function=self.key_resolver_function, loop=self._loop)