Fix FilesExt upload fails when content size is zero

Author: yuanjieding-db

NEXT_CHANGELOG.md

@@ -6,6 +6,8 @@
 
 ### Bug Fixes
 
+- Fix `FilesExt.upload` and `FilesExt.upload_from` would fail when the source content is empty and `use_parallel=True`.
+
 ### Documentation
 
 ### Internal Changes

databricks/sdk/mixins/files.py

@@ -22,8 +22,7 @@
 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
@@ -39,8 +38,7 @@
 from ..service import files
 from ..service._internal import _escape_multi_segment_path_parameter
 from ..service.files import DownloadResponse
-from .files_utils import (CreateDownloadUrlResponse, _ConcatenatedInputStream,
-                          _PresignedUrlDistributor)
+from .files_utils import CreateDownloadUrlResponse, _ConcatenatedInputStream, _PresignedUrlDistributor
 
 if TYPE_CHECKING:
     from _typeshed import Self
@@ -1134,7 +1132,9 @@ def upload(
             f"Upload context: part_size={ctx.part_size}, batch_size={ctx.batch_size}, content_length={ctx.content_length}"
         )
 
-        if ctx.use_parallel:
+        if ctx.use_parallel and (
+            ctx.content_length is None or ctx.content_length >= self._config.files_ext_multipart_upload_min_stream_size
+        ):
             self._parallel_upload_from_stream(ctx, content)
             return UploadStreamResult()
         elif ctx.content_length is not None:
@@ -1206,7 +1206,7 @@ def upload_from(
             use_parallel=use_parallel,
             parallelism=parallelism,
         )
-        if ctx.use_parallel:
+        if ctx.use_parallel and ctx.content_length >= self._config.files_ext_multipart_upload_min_stream_size:
             self._parallel_upload_from_file(ctx)
             return UploadFileResult()
         else:

tests/test_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
@@ -24,9 +24,14 @@
 from databricks.sdk import WorkspaceClient
 from databricks.sdk.core import Config
 from databricks.sdk.environments import Cloud, DatabricksEnvironment
-from databricks.sdk.errors.platform import (AlreadyExists, BadRequest,
-                                            InternalError, NotImplemented,
-                                            PermissionDenied, TooManyRequests)
+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
@@ -1662,6 +1667,7 @@ def processor() -> list:
                 request_json = request.json()
                 etags = {}
 
+                assert len(request_json["parts"]) > 0
                 for part in request_json["parts"]:
                     etags[part["part_number"]] = part["etag"]
 
@@ -1738,10 +1744,17 @@ def to_string(test_case: "MultipartUploadTestCase") -> str:
     [
         # -------------------------- happy cases --------------------------
         MultipartUploadTestCase(
-            "Multipart upload successful: single part",
+            "Multipart upload successful: single part because of small file",
             content_size=1024 * 1024,  # less than part size
-            multipart_upload_part_size=10 * 1024 * 1024,
+            multipart_upload_min_stream_size=10 * 1024 * 1024,
             expected_part_size=1024 * 1024,  # chunk size is used
+            expected_single_shot_upload=True,
+        ),
+        MultipartUploadTestCase(
+            "Multipart upload successful: empty file",
+            content_size=0,  # less than part size
+            multipart_upload_min_stream_size=100 * 1024 * 1024,  # all files smaller than 100M goes to single-shot
+            expected_single_shot_upload=True,
         ),
         MultipartUploadTestCase(
             "Multipart upload successful: multiple parts (aligned)",