copy from universe

Author: yuanjieding-db

databricks/sdk/__init__.py

@@ -6,7 +6,7 @@
 import pathlib
 import sys
 import urllib.parse
-from typing import Dict, Iterable, List, Optional
+from typing import Dict, Iterable, Optional, List
 
 import requests
 
@@ -110,7 +110,7 @@ class Config:
 
     disable_async_token_refresh: bool = ConfigAttribute(env="DATABRICKS_DISABLE_ASYNC_TOKEN_REFRESH")
 
-    enable_experimental_files_api_client: bool = ConfigAttribute(env="DATABRICKS_ENABLE_EXPERIMENTAL_FILES_API_CLIENT")
+    disable_experimental_files_api_client: bool = ConfigAttribute(env="DATABRICKS_DISABLE_EXPERIMENTAL_FILES_API_CLIENT")
 
     files_ext_client_download_streaming_chunk_size: int = 2 * 1024 * 1024  # 2 MiB
 

databricks/sdk/config.py

@@ -796,6 +796,9 @@ def download(
 
         :returns: :class:`DownloadResponse`
         """
+        if self._config.disable_experimental_files_api_client:
+            _LOG.info("Disable experimental files API client, will use the original download method.")
+            return super().download(file_path)
 
         initial_response: DownloadResponse = self._open_download_stream(
             file_path=file_path, start_byte_offset=0, if_unmodified_since_timestamp=None
@@ -829,6 +832,11 @@ def download_to(
 
         :returns: :class:`DownloadFileResult`
         """
+        if self._config.disable_experimental_files_api_client:
+            raise NotImplementedError(
+                "Experimental files API features are disabled, download_to is not supported. Please use download instead."
+            )
+
         # The existence of the target file is checked before starting the download. This is a best-effort check
         # to avoid overwriting an existing file. However, there is nothing preventing a file from being created
         # at the destination path after this check and before the file is written, and no way to prevent other
@@ -1086,6 +1094,11 @@ def upload(
         :returns: :class:`UploadStreamResult`
         """
 
+        if self._config.disable_experimental_files_api_client:
+            _LOG.info("Disable experimental files API client, will use the original upload method.")
+            super().upload(file_path=file_path, contents=content, overwrite=overwrite)
+            return UploadStreamResult()
+
         _LOG.debug(f"Uploading file from BinaryIO stream")
         if parallelism is not None and not use_parallel:
             raise ValueError("parallelism can only be set if use_parallel is True")
@@ -1163,6 +1176,10 @@ def upload_from(
 
         :returns: :class:`UploadFileResult`
         """
+        if self._config.disable_experimental_files_api_client:
+            raise NotImplementedError(
+                "Experimental files API features are disabled, upload_from is not supported. Please use upload instead."
+            )
 
         _LOG.debug(f"Uploading file from local path: {source_path}")
 

databricks/sdk/mixins/files.py

@@ -13,7 +13,7 @@
 from enum import Enum
 from tempfile import NamedTemporaryFile
 from threading import Lock
-from typing import Any, Callable, Dict, List, Optional, Type, Union
+from typing import Any, Callable, List, Optional, Type, Union, Dict
 from urllib.parse import parse_qs, urlparse
 
 import pytest
@@ -203,7 +203,7 @@ def run(self, config: Config, monkeypatch) -> None:
         if self.use_parallel and platform.system() == "Windows":
             pytest.skip("Skipping parallel download tests on Windows")
         config = config.copy()
-        config.enable_experimental_files_api_client = self.enable_new_client
+        config.disable_experimental_files_api_client = not self.enable_new_client
         config.files_ext_client_download_max_total_recovers = self.max_recovers_total
         config.files_ext_client_download_max_total_recovers_without_progressing = self.max_recovers_without_progressing
         config.enable_presigned_download_api = False
@@ -950,7 +950,6 @@ def run_one_case(self, config: Config, monkeypatch, download_mode: DownloadMode,
             logger.debug("Parallel download is not supported on Windows. Falling back to sequential download.")
             return
         config = config.copy()
-        config.enable_experimental_files_api_client = True
         config.enable_presigned_download_api = True
         config._clock = FakeClock()
         if self.parallel_download_min_file_size is not None:
@@ -1323,7 +1322,6 @@ def run_one_case(self, config: Config, use_parallel: bool, source_type: "UploadS
         logger.debug(f"Running test case: {self.name}, source_type={source_type}, use_parallel={use_parallel}")
         config = config.copy()
         config._clock = FakeClock()
-        config.enable_experimental_files_api_client = True
 
         if self.cloud:
             config.databricks_environment = DatabricksEnvironment(self.cloud, "")

tests/test_files.py

@@ -2,7 +2,7 @@
 import os
 from abc import ABC, abstractmethod
 from io import BytesIO, RawIOBase, UnsupportedOperation
-from typing import BinaryIO, Callable, List, Optional, Tuple
+from typing import BinaryIO, Callable, Optional, Tuple, List
 
 import pytest