Refactor FilesExt: extract presigned URL coordination into builder

[parthban-db]

🥞 Stacked PR

Use this link to review incremental changes.

• stack/refactor-files-client-3 [Files changed]
  • stack/refactor-files-client-4 [Files changed]
    • stack/refactor-files-client-5 [Files changed]

---

Summary

Extracts the presigned URL coordination logic (create-upload-part-urls and create-resumable-upload-url API calls and response parsing) into a dedicated _PresignedUrlRequestBuilder class, replacing inline code in three upload methods.

Why

The create-upload-part-urls API call, response validation, and header parsing was duplicated across _do_upload_one_part, _perform_multipart_upload, and _perform_resumable_upload (with a similar pattern for create-resumable-upload-url). Each copy built the request body, called _api.do(), validated the response structure, and converted the headers list to a dict — all inline. This made the upload methods longer than necessary and meant any change to the coordination logic required updating multiple places.

_PresignedUrlRequestBuilder consolidates this into a single class. It also prepares for storage-proxy routing (#1278), where a different builder implementation can construct URLs directly instead of calling the presigned URL APIs.

What changed

Interface changes

None. All changes are to private methods.

Behavioral changes

• Malformed presigned URL responses (ValueError/KeyError from response parsing) now trigger fallback to single-shot upload on the first part, instead of propagating as hard errors. Previously the _api.do() call was wrapped in try/except but the response parsing was outside it. Now both are encapsulated in the builder, so parsing errors are also caught. This is more resilient — a broken coordination response should not abort the upload when a simpler path exists.

Internal changes

• _PresignedUrl — new dataclass holding a resolved presigned URL and its associated headers.
• _PresignedUrlRequestBuilder — new class with build_upload_part_urls() and build_resumable_upload_url(). Encapsulates the coordination API calls and response parsing.
• _do_upload_one_part — replaced inline create-upload-part-urls call and response parsing with builder.build_upload_part_urls(..., count=1)[0].
• _perform_multipart_upload — replaced inline create-upload-part-urls batch call and response parsing with builder.build_upload_part_urls(..., count=batch_size).
• _perform_resumable_upload — replaced inline create-resumable-upload-url call and response parsing with builder.build_resumable_upload_url().

How is this tested?

Unit tests.

NO_CHANGELOG=true

[Divyansh-db]

LGTM, left some comments

[Divyansh-db]

can we keep current_part_number as part of _PresignedUrl to validate this assertion?

[Divyansh-db]

Suggested change

	ctx.target_path, session_token, part_index, 1, self._get_upload_url_expire_time()
	ctx.target_path, session_token, part_index, 1, self._get_upload_url_expire_time(),

nit, The parameters of a function should be linted in a uniform format (as https://github.com/databricks/databricks-sdk-py/pull/1295/changes#diff-469b96ba8e49ef7cbe8d3b4577d2a0ddf3a3552638b76d87bfccb936fda65d18R1849-R1854)

[Divyansh-db]

This was not set in the existing code

[parthban-db]

Range-diff: main (27164f7 -> 18f57c6)

databricks/sdk/mixins/files.py

@@ -10,6 +10,7 @@
 +
 +    url: str
 +    headers: dict[str, str]
++    part_number: int | None = None
 +
 +
 +class _PresignedUrlRequestBuilder:
@@ -46,11 +47,11 @@
 +        results = []
 +        for part_info in upload_part_urls:
 +            url = part_info["url"]
-+            _ = part_info["part_number"]
++            part_number = part_info["part_number"]
 +            headers: dict[str, str] = {"Content-Type": "application/octet-stream"}
 +            for h in part_info.get("headers", []):
 +                headers[h["name"]] = h["value"]
-+            results.append(_PresignedUrl(url=url, headers=headers))
++            results.append(_PresignedUrl(url=url, headers=headers, part_number=part_number))
 +        return results
 +
 +    def build_resumable_upload_url(self, path: str, session_token: str) -> _PresignedUrl:
@@ -69,7 +70,7 @@
 +        url = url_node.get("url")
 +        if not url:
 +            raise ValueError(f"Unexpected server response: {response}")
-+        headers: dict[str, str] = {"Content-Type": "application/octet-stream"}
++        headers: dict[str, str] = {}
 +        for h in url_node.get("headers", []):
 +            headers[h["name"]] = h["value"]
 +        return _PresignedUrl(url=url, headers=headers)
@@ -104,7 +105,11 @@
 -                    "POST", url=f"{hostname}/api/2.0/fs/create-upload-part-urls", headers=headers, body=body
 -                )
 +                presigned = builder.build_upload_part_urls(
-+                    ctx.target_path, session_token, part_index, 1, self._get_upload_url_expire_time()
++                    ctx.target_path,
++                    session_token,
++                    part_index,
++                    1,
++                    self._get_upload_url_expire_time(),
 +                )[0]
              except Exception as e:
                  if is_first_part:
@@ -195,6 +200,7 @@
 -                for h in required_headers:
 -                    headers[h["name"]] = h["value"]
 -
++                assert current_part_number == presigned.part_number
                  actual_chunk_length = min(actual_buffer_length, ctx.part_size)
                  _LOG.debug(
                      f"Uploading part {current_part_number}: [{chunk_offset}, {chunk_offset + actual_chunk_length - 1}]"

_{Reproduce locally: git range-diff 127b960..27164f7 127b960..18f57c6 | Disable: git config gitstack.push-range-diff false}

[github-actions]

If integration tests don't run automatically, an authorized user can run them manually by following the instructions below:

Trigger:
go/deco-tests-run/sdk-py

Inputs:

• PR number: 1295
• Commit SHA: 18f57c6c1f3589180a131f70aaead1eea39950eb

Checks will be approved automatically on success.