Fix typing contract for max_concurrency in Datalake client (#45631)

Author: nateprewitt

sdk/storage/azure-storage-file-datalake/assets.json

@@ -2,5 +2,5 @@
   "AssetsRepo": "Azure/azure-sdk-assets",
   "AssetsRepoPrefixPath": "python",
   "TagPrefix": "python/storage/azure-storage-file-datalake",
-  "Tag": "python/storage/azure-storage-file-datalake_4ab697f017"
+  "Tag": "python/storage/azure-storage-file-datalake_3d29de0db8"
 }

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

@@ -19,6 +19,7 @@
     get_mod_conditions,
     get_path_http_headers
 )
+from ._shared.constants import DEFAULT_MAX_CONCURRENCY
 from ._shared.request_handlers import get_length, read_length
 from ._shared.response_handlers import return_response_headers
 from ._shared.uploads import IterStreamer
@@ -124,7 +125,9 @@ def _upload_options(
     validate_content = kwargs.pop('validate_content', False)
     content_settings = kwargs.pop('content_settings', None)
     metadata = kwargs.pop('metadata', None)
-    max_concurrency = kwargs.pop('max_concurrency', 1)
+    max_concurrency = kwargs.pop('max_concurrency', None)
+    if max_concurrency is None:
+        max_concurrency = DEFAULT_MAX_CONCURRENCY
 
     kwargs['properties'] = add_metadata_headers(metadata)
     kwargs['lease_access_conditions'] = get_access_conditions(kwargs.pop('lease', None))

sdk/storage/azure-storage-file-datalake/azure/storage/filedatalake/_shared/constants.py

@@ -17,4 +17,6 @@
 DEFAULT_OAUTH_SCOPE = "/.default"
 STORAGE_OAUTH_SCOPE = "https://storage.azure.com/.default"
 
+DEFAULT_MAX_CONCURRENCY = 1
+
 SERVICE_HOST_BASE = "core.windows.net"

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

@@ -692,6 +692,27 @@ def test_upload_data_to_existing_file_with_permission_and_umask(self, **kwargs):
         assert data == downloaded_data
         assert prop['permissions'] == 'rwxrwxrwx'
 
+    @DataLakePreparer()
+    @recorded_by_proxy
+    def test_upload_data_with_none_max_concurrency(self, **kwargs):
+        datalake_storage_account_name = kwargs.pop("datalake_storage_account_name")
+        datalake_storage_account_key = kwargs.pop("datalake_storage_account_key")
+
+        self._setUp(datalake_storage_account_name, datalake_storage_account_key)
+        directory_name = self._get_directory_reference()
+
+        # Create a directory to put the file under that
+        directory_client = self.dsc.get_directory_client(self.file_system_name, directory_name)
+        directory_client.create_directory()
+
+        file_client = directory_client.get_file_client('filename')
+        data = self.get_random_bytes(100)
+        # max_concurrency=None should not raise TypeError
+        file_client.upload_data(data, overwrite=True, max_concurrency=None)
+
+        downloaded_data = file_client.download_file().readall()
+        assert data == downloaded_data
+
     @DataLakePreparer()
     @recorded_by_proxy
     def test_read_file(self, **kwargs):
@@ -710,6 +731,24 @@ def test_read_file(self, **kwargs):
         downloaded_data = file_client.download_file().readall()
         assert data == downloaded_data
 
+    @DataLakePreparer()
+    @recorded_by_proxy
+    def test_read_file_with_none_max_concurrency(self, **kwargs):
+        datalake_storage_account_name = kwargs.pop("datalake_storage_account_name")
+        datalake_storage_account_key = kwargs.pop("datalake_storage_account_key")
+
+        self._setUp(datalake_storage_account_name, datalake_storage_account_key)
+        file_client = self._create_file_and_return_client()
+        data = self.get_random_bytes(1024)
+
+        # upload data to file
+        file_client.append_data(data, 0, len(data))
+        file_client.flush_data(len(data))
+
+        # max_concurrency=None should not raise TypeError
+        downloaded_data = file_client.download_file(max_concurrency=None).readall()
+        assert data == downloaded_data
+
     @pytest.mark.live_test_only
     @DataLakePreparer()
     def test_read_file_with_user_delegation_key(self, **kwargs):

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

@@ -727,6 +727,27 @@ async def data_generator():
         result = await (await file_client.download_file()).readall()
         assert result == data*3
 
+    @DataLakePreparer()
+    @recorded_by_proxy_async
+    async def test_upload_data_with_none_max_concurrency(self, **kwargs):
+        datalake_storage_account_name = kwargs.pop("datalake_storage_account_name")
+        datalake_storage_account_key = kwargs.pop("datalake_storage_account_key")
+
+        await self._setUp(datalake_storage_account_name, datalake_storage_account_key)
+
+        # Create a directory to put the file under
+        directory_name = self._get_directory_reference()
+        directory_client = self.dsc.get_directory_client(self.file_system_name, directory_name)
+        await directory_client.create_directory()
+
+        file_client = directory_client.get_file_client('filename')
+        data = self.get_random_bytes(100)
+        # max_concurrency=None should not raise TypeError
+        await file_client.upload_data(data, overwrite=True, max_concurrency=None)
+
+        downloaded_data = await (await file_client.download_file()).readall()
+        assert data == downloaded_data
+
     @DataLakePreparer()
     @recorded_by_proxy_async
     async def test_read_file(self, **kwargs):
@@ -745,6 +766,24 @@ async def test_read_file(self, **kwargs):
         downloaded_data = await (await file_client.download_file()).readall()
         assert data == downloaded_data
 
+    @DataLakePreparer()
+    @recorded_by_proxy_async
+    async def test_read_file_with_none_max_concurrency(self, **kwargs):
+        datalake_storage_account_name = kwargs.pop("datalake_storage_account_name")
+        datalake_storage_account_key = kwargs.pop("datalake_storage_account_key")
+
+        await self._setUp(datalake_storage_account_name, datalake_storage_account_key)
+        file_client = await self._create_file_and_return_client()
+        data = self.get_random_bytes(1024)
+
+        # upload data to file
+        await file_client.append_data(data, 0, len(data))
+        await file_client.flush_data(len(data))
+
+        # max_concurrency=None should not raise TypeError
+        downloaded_data = await (await file_client.download_file(max_concurrency=None)).readall()
+        assert data == downloaded_data
+
     @pytest.mark.live_test_only
     @DataLakePreparer()
     async def test_read_file_with_user_delegation_key(self, **kwargs):