Add storage proxy support for FilesExt uploads

[parthban-db]

🥞 Stacked PR

Use this link to review incremental changes.

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

---

Summary

Adds optional storage-proxy routing for file uploads, gated behind the experimental_files_ext_enable_storage_proxy config flag. When enabled and the proxy is reachable, uploads bypass the presigned URL coordination APIs and go directly to the storage proxy.

Why

When running inside a Databricks cluster or notebook, a storage proxy is available on the data plane. It can handle file upload operations directly.

The feature is disabled by default and marked experimental, so there is no impact on existing users.

What changed

Interface changes

One new config field:

• experimental_files_ext_enable_storage_proxy: bool = False — enables the storage proxy path.

Behavioral changes

• When the flag is enabled and the probe succeeds, multipart uploads (AWS/Azure) and resumable uploads (GCP) go directly to the proxy instead of fetching presigned URLs from coordination APIs.
• The proxy requires SDK authentication on every request, so an authenticated session (session.auth callback) is used instead of the unauthenticated session used for presigned URL uploads.
• build_abort_url was added to _PresignedUrlRequestBuilder, extracting inline abort URL logic from _abort_multipart_upload.

Internal changes

• _StorageProxyRequestBuilder — new class with the same method signatures as _PresignedUrlRequestBuilder (build_upload_part_urls, build_resumable_upload_url, build_abort_url). Constructs URLs directly to the proxy endpoint instead of calling coordination APIs.
• _create_request_builder() — factory method that returns the proxy builder when the proxy is available, otherwise the presigned URL builder.
• _probe_storage_proxy() — GET to the proxy ping endpoint with SDK auth. Result is cached.
• _create_storage_proxy_session() — cached requests.Session with session.auth callback for SDK credentials. Protected by a lock for thread safety.
• _cloud_provider_session() — updated to return the authenticated proxy session when the proxy is active. Session creation protected by a lock for thread safety.
• _get_hostname() — updated to return the proxy hostname when the flag is enabled and the probe succeeds.
• _abort_multipart_upload — now uses the builder instead of inline create-abort-upload-url API call.
• Test infrastructure: UploadTestCase gains use_storage_proxy parameter. Storage proxy test cases added to existing MultipartUploadTestCase and ResumableUploadTestCase parametrized lists. Presigned URL coordination handlers assert they are never called when storage proxy is active.

How is this tested?

• Unit tests.
• E2E validated on two cloud environments with 200 MB file uploads:
  • AWS (AWS staging): multipart upload via proxy, 20 parallel parts, hash verified.
  • GCP (GCP staging): resumable upload via proxy, 20 sequential chunks, hash verified.
• Both environments confirmed: _StorageProxyRequestBuilder used, no create-upload-part-urls/create-resumable-upload-url calls, session.auth callback active, uploaded and downloaded hashes match.

NO_CHANGELOG=true

[parthban-db]

Range-diff: main (8218fc4 -> c7928c1)

databricks/sdk/mixins/files.py

@@ -45,25 +45,31 @@
 +        results = []
 +        for i in range(count):
 +            part_number = start_part_number + i
-+            query = parse.urlencode({
-+                "session_token": session_token,
-+                "upload_type": "multipart",
-+                "part_number": part_number,
-+            })
-+            results.append(_PresignedUrl(
-+                url=f"{base}?{query}",
-+                headers={"Content-Type": "application/octet-stream"},
-+            ))
++            query = parse.urlencode(
++                {
++                    "session_token": session_token,
++                    "upload_type": "multipart",
++                    "part_number": part_number,
++                }
++            )
++            results.append(
++                _PresignedUrl(
++                    url=f"{base}?{query}",
++                    headers={"Content-Type": "application/octet-stream"},
++                )
++            )
 +        return results
 +
 +    def build_resumable_upload_url(self, path: str, session_token: str) -> _PresignedUrl:
 +        """Builds a URL for resumable upload directly to the storage proxy."""
 +        escaped = _escape_multi_segment_path_parameter(path)
 +        base = f"{self._hostname}/api/2.0/fs/files{escaped}"
-+        query = parse.urlencode({
-+            "session_token": session_token,
-+            "upload_type": "resumable",
-+        })
++        query = parse.urlencode(
++            {
++                "session_token": session_token,
++                "upload_type": "resumable",
++            }
++        )
 +        return _PresignedUrl(
 +            url=f"{base}?{query}",
 +            headers={"Content-Type": "application/octet-stream"},
@@ -73,10 +79,12 @@
 +        """Builds a URL for aborting an upload directly on the storage proxy."""
 +        escaped = _escape_multi_segment_path_parameter(path)
 +        base = f"{self._hostname}/api/2.0/fs/files{escaped}"
-+        query = parse.urlencode({
-+            "action": "abort-upload",
-+            "session_token": session_token,
-+        })
++        query = parse.urlencode(
++            {
++                "action": "abort-upload",
++                "session_token": session_token,
++            }
++        )
 +        return _PresignedUrl(
 +            url=f"{base}?{query}",
 +            headers={"Content-Type": "application/json"},

tests/test_files.py

@@ -346,5 +346,4 @@
      chunk = os.urandom(chunk_size)
 -    # Repeat it until we reach n bytes
 +    # Repeat it until we reach n bytes.
-     return (chunk * (n // chunk_size + 1))[:n]
-+
\ No newline at end of file
+     return (chunk * (n // chunk_size + 1))[:n]
\ No newline at end of file

_{Reproduce locally: git range-diff 9e504b1..8218fc4 ff61559..c7928c1 | 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: 1319
• Commit SHA: 85824ec07a796f348263077a05c5339a190f1131

Checks will be approved automatically on success.