Add workspace ID validation to storage proxy probe#1324
Merged
parthban-db merged 2 commits intomainfrom Mar 13, 2026
Merged
Conversation
This was referenced Mar 11, 2026
parthban-db
temporarily deployed
to
test-trigger-is
GitHub Actions
Inactive
parthban-db
temporarily deployed
to
test-trigger-is
GitHub Actions
Inactive
github-merge-queue bot
pushed a commit
that referenced
this pull request
Mar 11, 2026
…1295) ## 🥞 Stacked PR Use this [link](https://github.com/databricks/databricks-sdk-py/pull/1295/files) to review incremental changes. - [**stack/refactor-files-client-3**](#1295) [[Files changed](https://github.com/databricks/databricks-sdk-py/pull/1295/files)] - [stack/refactor-files-client-4](#1319) [[Files changed](https://github.com/databricks/databricks-sdk-py/pull/1319/files/9e504b191ad065bf8adbeb003fa493ba51bcc165..8218fc4ae84ff4b9ec2f838b26faa0c3a957f2e5)] - [stack/refactor-files-client-5](#1324) [[Files changed](https://github.com/databricks/databricks-sdk-py/pull/1324/files/8218fc4ae84ff4b9ec2f838b26faa0c3a957f2e5..2c4f3e63287e86b0984a4dfb323547cdaacfa058)] --------- ## 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
parthban-db
force-pushed
the
parthban-db/stack/refactor-files-client-5
branch
from
2c4f3e6 to
93372a6
Compare
parthban-db
temporarily deployed
to
test-trigger-is
GitHub Actions
Inactive
parthban-db
temporarily deployed
to
test-trigger-is
GitHub Actions
Inactive
parthban-db
force-pushed
the
parthban-db/stack/refactor-files-client-5
branch
from
93372a6 to
14e1367
Compare
parthban-db
temporarily deployed
to
test-trigger-is
GitHub Actions
Inactive
parthban-db
temporarily deployed
to
test-trigger-is
GitHub Actions
Inactive
github-merge-queue bot
pushed a commit
that referenced
this pull request
Mar 13, 2026
## 🥞 Stacked PR Use this [link](https://github.com/databricks/databricks-sdk-py/pull/1319/files) to review incremental changes. - [**stack/refactor-files-client-4**](#1319) [[Files changed](https://github.com/databricks/databricks-sdk-py/pull/1319/files)] - [stack/refactor-files-client-5](#1324) [[Files changed](https://github.com/databricks/databricks-sdk-py/pull/1324/files/85824ec07a796f348263077a05c5339a190f1131..14e13676c5e12adff8f68d5f30aca295a423acdd)] - [stack/refactor-files-client-6](#1329) [[Files changed](https://github.com/databricks/databricks-sdk-py/pull/1329/files/14e13676c5e12adff8f68d5f30aca295a423acdd..80c211ffa59afea71f5d01697d95381e0d126865)] --------- ## 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
force-pushed
the
parthban-db/stack/refactor-files-client-5
branch
from
14e1367 to
214368b
Compare
parthban-db
temporarily deployed
to
test-trigger-is
GitHub Actions
Inactive
parthban-db
temporarily deployed
to
test-trigger-is
GitHub Actions
Inactive
|
If integration tests don't run automatically, an authorized user can run them manually by following the instructions below: Trigger: Inputs:
Checks will be approved automatically on success. |
rauchy
approved these changes
Mar 13, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
🥞 Stacked PR
Use this link to review incremental changes.
Summary
Adds workspace ID validation to the storage proxy probe so the proxy is only used when it can serve the target workspace. Prevents false positives when a
WorkspaceClientis configured to point at a different workspace than the one hosting the data plane.Why
The previous probe (
GET /DatabricksInternal/Probes/ping) only checked if the proxy was reachable. It didn't verify whether the proxy could serve the specific workspace the client is targeting. This meant aWorkspaceClientconfigured with a different workspace's host and token could incorrectly route uploads through the local proxy, which cannot serve that workspace.What changed
Interface changes
None.
Behavioral changes
GET /api/2.0/preview/scim/v2/Me) and passes it as?ew=<workspace_id>to the probe URL. The proxy returns 200 only if it can serve that workspace, 403 otherwise.Internal changes
_resolve_workspace_id()— new method that calls SCIM Me to get the workspace ID from theX-Databricks-Org-Idresponse header._probe_storage_proxy()— updated to resolve workspace ID first, skip if unavailable, use HEAD method, and passewparameter.ewparameter is present in the probe URL.How is this tested?
ewreturns 200,_StorageProxyRequestBuilderused, 200 MB hash verified._PresignedUrlRequestBuilderused, 200 MB hash verified.ew=<gcp_id>to local AWS proxy, proxy returns 403, correctly falls back to presigned URLs.NO_CHANGELOG=true