Propagate x-databricks-org-id header to DBFS Volume API for SPOG routing

DBFSVolumeClient uses the SDK's ApiClient directly for /api/2.0/fs/*
endpoints. These requests need the x-databricks-org-id header for SPOG
routing, same as telemetry and feature flags.

Add connectionContext.getCustomHeaders() to all 5 HTTP request paths in
DBFSVolumeClient (4 sync via apiClient, 1 async via requestBuilder).

Co-authored-by: Isaac
Signed-off-by: Madhavendra Rathore <[email protected]>

Author: msrathore-db

src/main/java/com/databricks/jdbc/api/impl/volume/DBFSVolumeClient.java

@@ -492,7 +492,7 @@ CreateUploadUrlResponse getCreateUploadUrlResponse(String objectPath)
     CreateUploadUrlRequest request = new CreateUploadUrlRequest(objectPath);
     try {
       Request req = new Request(Request.POST, CREATE_UPLOAD_URL_PATH, apiClient.serialize(request));
-      req.withHeaders(JSON_HTTP_HEADERS);
+      req.withHeaders(JSON_HTTP_HEADERS).withHeaders(connectionContext.getCustomHeaders());
       return apiClient.execute(req, CreateUploadUrlResponse.class);
     } catch (IOException | DatabricksException e) {
       String errorMessage =
@@ -514,7 +514,7 @@ CreateDownloadUrlResponse getCreateDownloadUrlResponse(String objectPath)
     try {
       Request req =
           new Request(Request.POST, CREATE_DOWNLOAD_URL_PATH, apiClient.serialize(request));
-      req.withHeaders(JSON_HTTP_HEADERS);
+      req.withHeaders(JSON_HTTP_HEADERS).withHeaders(connectionContext.getCustomHeaders());
       return apiClient.execute(req, CreateDownloadUrlResponse.class);
     } catch (IOException | DatabricksException e) {
       String errorMessage =
@@ -534,7 +534,7 @@ CreateDeleteUrlResponse getCreateDeleteUrlResponse(String objectPath)
 
     try {
       Request req = new Request(Request.POST, CREATE_DELETE_URL_PATH, apiClient.serialize(request));
-      req.withHeaders(JSON_HTTP_HEADERS);
+      req.withHeaders(JSON_HTTP_HEADERS).withHeaders(connectionContext.getCustomHeaders());
       return apiClient.execute(req, CreateDeleteUrlResponse.class);
     } catch (IOException | DatabricksException e) {
       String errorMessage =
@@ -551,7 +551,7 @@ ListResponse getListResponse(String listPath) throws DatabricksVolumeOperationEx
     ListRequest request = new ListRequest(listPath);
     try {
       Request req = new Request(Request.GET, LIST_PATH);
-      req.withHeaders(JSON_HTTP_HEADERS);
+      req.withHeaders(JSON_HTTP_HEADERS).withHeaders(connectionContext.getCustomHeaders());
       ApiClient.setQuery(req, request);
       return apiClient.execute(req, ListResponse.class);
     } catch (IOException | DatabricksException e) {
@@ -888,6 +888,7 @@ private CompletableFuture<CreateUploadUrlResponse> requestPresignedUrlWithRetry(
         Map<String, String> authHeaders = workspaceClient.config().authenticate();
         authHeaders.forEach(requestBuilder::addHeader);
         JSON_HTTP_HEADERS.forEach(requestBuilder::addHeader);
+        connectionContext.getCustomHeaders().forEach(requestBuilder::addHeader);
 
         requestBuilder.setEntity(
             AsyncEntityProducers.create(requestBody.getBytes(), ContentType.APPLICATION_JSON));