formated

Author: yuanjieding-db

databricks/sdk/config.py

@@ -147,19 +147,19 @@ class Config:
     # For AWS, minimum is 5Mb: https://docs.aws.amazon.com/AmazonS3/latest/userguide/qfacts.html
     # For GCP, minimum is 256 KiB (and also recommended multiple is 256 KiB)
     # boto uses 8Mb: https://boto3.amazonaws.com/v1/documentation/api/latest/reference/customizations/s3.html#boto3.s3.transfer.TransferConfig
-    files_ext_multipart_upload_default_part_size: int = 10 * 1024 * 1024 # 10 MiB
+    files_ext_multipart_upload_default_part_size: int = 10 * 1024 * 1024  # 10 MiB
 
     # List of multipart upload part sizes that can be automatically selected
     files_ext_multipart_upload_part_size_options: list[int] = [
-        10 * 1024 * 1024, # 10 MiB
-        20 * 1024 * 1024, # 20 MiB
-        50 * 1024 * 1024, # 50 MiB
-        100 * 1024 * 1024, # 100 MiB
-        200 * 1024 * 1024, # 200 MiB
-        500 * 1024 * 1024, # 500 MiB
-        1 * 1024 * 1024 * 1024, # 1 GiB
-        2 * 1024 * 1024 * 1024, # 2 GiB
-        4 * 1024 * 1024 * 1024, # 4 GiB
+        10 * 1024 * 1024,  # 10 MiB
+        20 * 1024 * 1024,  # 20 MiB
+        50 * 1024 * 1024,  # 50 MiB
+        100 * 1024 * 1024,  # 100 MiB
+        200 * 1024 * 1024,  # 200 MiB
+        500 * 1024 * 1024,  # 500 MiB
+        1 * 1024 * 1024 * 1024,  # 1 GiB
+        2 * 1024 * 1024 * 1024,  # 2 GiB
+        4 * 1024 * 1024 * 1024,  # 4 GiB
     ]
 
     # Maximum size of a single part in multipart upload.

databricks/sdk/mixins/files.py

@@ -18,11 +18,12 @@
 from dataclasses import dataclass
 from datetime import timedelta
 from io import BytesIO
-from queue import Queue, Empty, Full
-from tempfile import NamedTemporaryFile, mkstemp
-from threading import Thread, Event
+from queue import Empty, Full, Queue
+from tempfile import mkstemp
+from threading import Event, Thread
 from types import TracebackType
-from typing import TYPE_CHECKING, AnyStr, BinaryIO, Callable, Generator, Iterable, Optional, Type, Union
+from typing import (TYPE_CHECKING, AnyStr, BinaryIO, Callable, Generator,
+                    Iterable, Optional, Type, Union)
 from urllib import parse
 
 import requests
@@ -38,8 +39,8 @@
 from ..service import files
 from ..service._internal import _escape_multi_segment_path_parameter
 from ..service.files import DownloadResponse
-
-from .files_utils import _ConcatenatedInputStream, _PresignedUrlDistributor, CreateDownloadUrlResponse
+from .files_utils import (CreateDownloadUrlResponse, _ConcatenatedInputStream,
+                          _PresignedUrlDistributor)
 
 if TYPE_CHECKING:
     from _typeshed import Self
@@ -736,21 +737,18 @@ def __init__(self, message):
 class UploadStreamResult:
     """Result of an upload from stream operation. Currently empty, but can be extended in the future."""
 
-    pass
 
 
 @dataclass
 class UploadFileResult:
     """Result of an upload from file operation. Currently empty, but can be extended in the future."""
 
-    pass
 
 
 @dataclass
 class DownloadFileResult:
     """Result of a download to file operation. Currently empty, but can be extended in the future."""
 
-    pass
 
 
 class FilesExt(files.FilesAPI):

databricks/sdk/mixins/files_utils.py

@@ -1,9 +1,10 @@
 from __future__ import annotations
 
-from dataclasses import dataclass
-from typing import BinaryIO, Optional, Any, Iterable, Callable
 import os
 import threading
+from dataclasses import dataclass
+from typing import Any, BinaryIO, Callable, Iterable, Optional
+
 
 @dataclass
 class CreateDownloadUrlResponse:
@@ -23,6 +24,7 @@ def from_dict(cls, data: dict[str, Any]) -> CreateDownloadUrlResponse:
         parsed_headers = {x["name"]: x["value"] for x in headers}
         return cls(url=data["url"], headers=parsed_headers)
 
+
 class _ConcatenatedInputStream(BinaryIO):
     """This class joins two input streams into one."""
 
@@ -244,6 +246,7 @@ def __exit__(self, __type, __value, __traceback) -> None:
     def __str__(self) -> str:
         return f"Concat: {self._head_stream}, {self._tail_stream}]"
 
+
 class _PresignedUrlDistributor:
     """
     Distributes and manages presigned URLs for downloading files.
@@ -253,7 +256,6 @@ class _PresignedUrlDistributor:
     """
 
     def __init__(self, get_new_url_func: Callable[[], CreateDownloadUrlResponse]):
-
         """
         Initialize the distributor.
 

tests/test_files.py

@@ -10,33 +10,28 @@
 from dataclasses import dataclass
 from datetime import datetime, timedelta, timezone
 from enum import Enum
-from tempfile import mkstemp, NamedTemporaryFile
-from typing import Callable, List, Optional, Type, Union, Any
+from tempfile import NamedTemporaryFile, mkstemp
+from threading import Lock
+from typing import Any, Callable, List, Optional, Type, Union
 from urllib.parse import parse_qs, urlparse
 
 import pytest
 import requests
 import requests_mock
 from requests import RequestException
-from threading import Lock
 
-from .test_files_utils import Utils
 from databricks.sdk import WorkspaceClient
-from databricks.sdk.environments import DatabricksEnvironment, Cloud
 from databricks.sdk.core import Config
-from databricks.sdk.errors.platform import (
-    AlreadyExists,
-    BadRequest,
-    InternalError,
-    PermissionDenied,
-    TooManyRequests,
-    NotImplemented,
-)
-from databricks.sdk.service.files import DownloadResponse
+from databricks.sdk.environments import Cloud, DatabricksEnvironment
+from databricks.sdk.errors.platform import (AlreadyExists, BadRequest,
+                                            InternalError, NotImplemented,
+                                            PermissionDenied, TooManyRequests)
 from databricks.sdk.mixins.files import FallbackToDownloadUsingFilesApi
 from databricks.sdk.mixins.files_utils import CreateDownloadUrlResponse
 from tests.clock import FakeClock
 
+from .test_files_utils import Utils
+
 logger = logging.getLogger(__name__)
 
 
@@ -1275,7 +1270,6 @@ def clear_state(self) -> None:
                 os.remove(file_path)
             except OSError:
                 logger.warning("Failed to remove temp file: %s", file_path)
-                pass
         self.created_temp_files = []
 
     def get_upload_file(self, content: bytes, source_type: "UploadSourceType") -> Union[str, io.BytesIO]:

tests/test_files_utils.py

@@ -1,25 +1,28 @@
 import logging
 import os
+from abc import ABC, abstractmethod
 from io import BytesIO, RawIOBase, UnsupportedOperation
 from typing import BinaryIO, Callable, Optional
+
 import pytest
-from abc import ABC, abstractmethod
 
-from databricks.sdk.mixins.files_utils import _ConcatenatedInputStream, _PresignedUrlDistributor
+from databricks.sdk.mixins.files_utils import (_ConcatenatedInputStream,
+                                               _PresignedUrlDistributor)
 
 logger = logging.getLogger(__name__)
 
+
 class Utils:
     @staticmethod
-    def parse_range_header(range_header: str, content_length: Optional[int]=None) -> tuple[int, int]:
+    def parse_range_header(range_header: str, content_length: Optional[int] = None) -> tuple[int, int]:
         """
         Parses a Range header string and returns the start and end byte positions.
         Example input: "bytes=0-499"
         Example output: (0, 499)
         """
         if not range_header.startswith("bytes="):
             raise ValueError("Invalid Range header format")
-        byte_range = range_header[len("bytes="):]
+        byte_range = range_header[len("bytes=") :]
         start_str, end_str = byte_range.split("-")
         start = int(start_str) if start_str else 0
         end = int(end_str) if end_str else None
@@ -34,6 +37,7 @@ def parse_range_header(range_header: str, content_length: Optional[int]=None) ->
 
         return start, end
 
+
 class NonSeekableBuffer(RawIOBase, BinaryIO):
     """
     A non-seekable buffer that wraps a bytes object. Used for unit tests only.
@@ -268,17 +272,20 @@ def safe_call(buf: BinaryIO, call: Callable[[BinaryIO], any]) -> (any, bool):
         assert buffer.tell() == native_buffer.tell()
         assert read_and_restore(buffer) == read_and_restore(native_buffer)
 
+
 class DummyResponse:
     def __init__(self, value):
         self.value = value
 
+
 def test_get_url_returns_url_and_version():
     distributor = _PresignedUrlDistributor(lambda: DummyResponse("url1"))
     url, version = distributor.get_url()
     assert isinstance(url, DummyResponse)
     assert url.value == "url1"
     assert version == 0
 
+
 def test_get_url_caches_url():
     calls = []
     distributor = _PresignedUrlDistributor(lambda: calls.append(1) or DummyResponse("url2"))
@@ -288,6 +295,7 @@ def test_get_url_caches_url():
     assert version1 == version2
     assert calls.count(1) == 1  # Only called once
 
+
 def test_invalidate_url_changes_url_and_version():
     responses = [DummyResponse("urlA"), DummyResponse("urlB")]
     distributor = _PresignedUrlDistributor(lambda: responses.pop(0))
@@ -298,10 +306,11 @@ def test_invalidate_url_changes_url_and_version():
     assert url2.value == "urlB"
     assert version2 == version1 + 1
 
+
 def test_invalidate_url_wrong_version_does_not_invalidate():
     distributor = _PresignedUrlDistributor(lambda: DummyResponse("urlX"))
     url1, version1 = distributor.get_url()
     distributor.invalidate_url(version1 + 1)  # Wrong version
     url2, version2 = distributor.get_url()
     assert url1 is url2
-    assert version2 == version1
+    assert version2 == version1