Add SPOG support for ?o= routing in httpPath (#1316)

## Summary

SPOG (Single Panel of Glass) replaces workspace-specific hostnames with
account-level vanity URLs. When httpPath contains `?o=<workspaceId>`,
the driver needs to:
1. Preserve the `?o=` value through URL parsing
2. Extract the clean warehouse ID (without `?o=`) for SEA JSON body
3. Inject `x-databricks-org-id` header on endpoints that don't carry
`?o=` in the URL path (telemetry, feature flags, DBFS volume)

### Changes

**Fix 1 — Property parser** (`DatabricksConnectionContext.java`):
- `split("=")` broke httpPath values containing `=` (e.g.
`httpPath=/sql/1.0/warehouses/xxx?o=yyy` split into 3 parts, dropping
the workspace ID)
- Changed to `indexOf("=")` + `substring` to split on first `=` only

**Fix 2 — Warehouse ID regex** (`DatabricksJdbcConstants.java`):
- `(.+)` captured `abc123?o=999` as warehouse ID, corrupting SEA JSON
body
- Changed to `([^?]+).*` to stop capture at query params
- Trailing `.*` needed because Java `.matches()` requires full-string
match

**Fix 3 — SPOG header extraction** (`DatabricksConnectionContext.java`):
- Uses `URIBuilder` to parse `?o=<workspaceId>` from httpPath
- Injects `x-databricks-org-id` into `customHeaders` map
- Respects explicitly set header (won't override)

**Fix 4 — Header propagation** (telemetry, feature flags, DBFS volume):
- `TelemetryPushClient.java`: Added
`connectionContext.getCustomHeaders()` to telemetry POST
- `DatabricksDriverFeatureFlagsContext.java`: Added
`connectionContext.getCustomHeaders()` to feature flag GET
- `DBFSVolumeClient.java`: Added `connectionContext.getCustomHeaders()`
to all 5 HTTP request paths (4 sync + 1 async), with null guard for test
constructor

### What does NOT get the header

OAuth endpoints (`/oidc/v1/token`, OIDC discovery) — they are
workspace-agnostic and reject `x-databricks-org-id` with HTTP 400. These
use their own HTTP client and never see `customHeaders`.

## Context

Per the [SPOG Peco Clients
doc](https://docs.google.com/document/d/1ZTacOW72Jetr6Qo1iLVogeXKDYoFAE_R8_32KmLawmQ/):
- Thrift: `?o=` in the POST URL handles routing naturally
- SEA, telemetry, feature flags, DBFS volume: Need `x-databricks-org-id`
header since they use separate endpoint paths

Jira: [XTA-15079](https://databricks.atlassian.net/browse/XTA-15079)

## Test plan

**Unit tests added:**
- [x] `testBuildPropertiesMap_preservesQueryParamInHttpPath` — httpPath
with `?o=` preserved
- [x] `testBuildPropertiesMap_handlesValueWithMultipleEquals` — values
with multiple `=` preserved
- [x] `testBuildPropertiesMap_handlesValueWithNoEquals` — key-only
params handled
- [x] `testSpogContext_extractsOrgIdFromHttpPath` —
`x-databricks-org-id` extracted from `?o=`
- [x] `testSpogContext_extractsCleanWarehouseId` — warehouse ID is
`abc123` not `abc123?o=...`
- [x] `testSpogContext_noOrgIdWithoutQueryParam` — no header injected
for legacy URLs
- [x] `testSpogContext_explicitHeaderTakesPrecedence` — explicit
`http.header.x-databricks-org-id` wins over `?o=`
- [x] 3 SPOG URL validation tests in `ValidationUtilTest`

**E2E tested (not committed):**
- [x] DBSQL + PAT: Thrift and SEA on SPOG host
- [x] GP Cluster + PAT: Thrift on SPOG host
- [x] OAuth M2M: Thrift and SEA on SPOG host
- [x] OAuth U2M: Thrift and SEA on SPOG host (with browser login)
- [x] Telemetry: Verified 200 on SPOG with org-id header
- [x] Feature flags: Verified 200 on SPOG with org-id header
- [x] DBFS Volume LIST: Verified on SPOG host
- [x] Legacy host: All auth modes unaffected

NO_CHANGELOG=true

This pull request was AI-assisted by Isaac.

[XTA-15079]:
https://databricks.atlassian.net/browse/XTA-15079?atlOrigin=eyJpIjoiNWRkNTljNzYxNjVmNDY3MDlhMDU5Y2ZhYzA5YTRkZjUiLCJwIjoiZ2l0aHViLWNvbS1KU1cifQ

---------

Signed-off-by: Madhavendra Rathore <madhavendra.rathore@databricks.com>
Signed-off-by: Madhavendra Rathore

Author: msrathore-db

src/main/java/com/databricks/jdbc/api/impl/DatabricksConnectionContext.java

@@ -29,10 +29,12 @@
 import com.google.common.base.Supplier;
 import com.google.common.collect.ImmutableMap;
 import java.net.URI;
+import java.net.URISyntaxException;
 import java.util.*;
 import java.util.Collections;
 import java.util.regex.Matcher;
 import java.util.stream.Collectors;
+import org.apache.http.NameValuePair;
 import org.apache.http.client.utils.URIBuilder;
 
 public class DatabricksConnectionContext implements IDatabricksConnectionContext {
@@ -108,14 +110,14 @@ public static ImmutableMap<String, String> buildPropertiesMap(
     if (!isNullOrEmpty(connectionParamString)) {
       String[] urlParts = connectionParamString.split(DatabricksJdbcConstants.URL_DELIMITER);
       for (String urlPart : urlParts) {
-        String[] pair = urlPart.split(DatabricksJdbcConstants.PAIR_DELIMITER);
-        if (pair.length == 1) {
-          pair = new String[] {pair[0], ""};
-        }
-        if (pair[0].startsWith(DatabricksJdbcUrlParams.HTTP_HEADERS.getParamName())) {
-          parametersBuilder.put(pair[0], pair[1]);
+        // Split on first '=' only — values (like httpPath) may contain '=' (e.g. ?o=123)
+        int delimIdx = urlPart.indexOf(DatabricksJdbcConstants.PAIR_DELIMITER);
+        String key = delimIdx >= 0 ? urlPart.substring(0, delimIdx) : urlPart;
+        String value = delimIdx >= 0 ? urlPart.substring(delimIdx + 1) : "";
+        if (key.startsWith(DatabricksJdbcUrlParams.HTTP_HEADERS.getParamName())) {
+          parametersBuilder.put(key, value);
         } else {
-          parametersBuilder.put(pair[0].toLowerCase(), pair[1]);
+          parametersBuilder.put(key.toLowerCase(), value);
         }
       }
     }
@@ -1166,14 +1168,52 @@ private String getParameter(DatabricksJdbcUrlParams key, String defaultValue) {
     return this.parameters.getOrDefault(key.getParamName().toLowerCase(), defaultValue);
   }
 
+  private static final String ORG_ID_HEADER = "x-databricks-org-id";
+  private static final String ORG_ID_QUERY_PARAM = "o";
+
   private Map<String, String> parseCustomHeaders(ImmutableMap<String, String> parameters) {
     String filterPrefix = DatabricksJdbcUrlParams.HTTP_HEADERS.getParamName();
 
-    return parameters.entrySet().stream()
-        .filter(entry -> entry.getKey().startsWith(filterPrefix))
-        .collect(
-            Collectors.toMap(
-                entry -> entry.getKey().substring(filterPrefix.length()), Map.Entry::getValue));
+    Map<String, String> headers =
+        new HashMap<>(
+            parameters.entrySet().stream()
+                .filter(entry -> entry.getKey().startsWith(filterPrefix))
+                .collect(
+                    Collectors.toMap(
+                        entry -> entry.getKey().substring(filterPrefix.length()),
+                        Map.Entry::getValue)));
+
+    // Extract org ID from ?o= in httpPath for SPOG routing
+    if (!headers.containsKey(ORG_ID_HEADER)) {
+      String httpPath =
+          parameters.getOrDefault(
+              DatabricksJdbcUrlParams.HTTP_PATH.getParamName().toLowerCase(), "");
+      try {
+        for (NameValuePair param :
+            new URIBuilder("http://placeholder" + httpPath).getQueryParams()) {
+          if (ORG_ID_QUERY_PARAM.equals(param.getName())
+              && param.getValue() != null
+              && !param.getValue().isEmpty()) {
+            headers.put(ORG_ID_HEADER, param.getValue());
+            LOGGER.debug(
+                "SPOG header extraction: injecting {}={} (extracted from ?o= in httpPath)",
+                ORG_ID_HEADER,
+                param.getValue());
+            break;
+          }
+        }
+      } catch (URISyntaxException e) {
+        LOGGER.debug(
+            "SPOG header extraction: malformed httpPath, skipping org-id extraction: "
+                + e.getMessage());
+      }
+    } else {
+      LOGGER.debug(
+          "SPOG header extraction: {} already set by caller, not extracting from httpPath",
+          ORG_ID_HEADER);
+    }
+
+    return headers;
   }
 
   @Override

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

@@ -492,7 +492,8 @@ 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 != null ? connectionContext.getCustomHeaders() : Map.of());
       return apiClient.execute(req, CreateUploadUrlResponse.class);
     } catch (IOException | DatabricksException e) {
       String errorMessage =
@@ -514,7 +515,8 @@ 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 != null ? connectionContext.getCustomHeaders() : Map.of());
       return apiClient.execute(req, CreateDownloadUrlResponse.class);
     } catch (IOException | DatabricksException e) {
       String errorMessage =
@@ -534,7 +536,8 @@ 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 != null ? connectionContext.getCustomHeaders() : Map.of());
       return apiClient.execute(req, CreateDeleteUrlResponse.class);
     } catch (IOException | DatabricksException e) {
       String errorMessage =
@@ -551,7 +554,8 @@ 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 != null ? connectionContext.getCustomHeaders() : Map.of());
       ApiClient.setQuery(req, request);
       return apiClient.execute(req, ListResponse.class);
     } catch (IOException | DatabricksException e) {
@@ -888,6 +892,9 @@ private CompletableFuture<CreateUploadUrlResponse> requestPresignedUrlWithRetry(
         Map<String, String> authHeaders = workspaceClient.config().authenticate();
         authHeaders.forEach(requestBuilder::addHeader);
         JSON_HTTP_HEADERS.forEach(requestBuilder::addHeader);
+        if (connectionContext != null) {
+          connectionContext.getCustomHeaders().forEach(requestBuilder::addHeader);
+        }
 
         requestBuilder.setEntity(
             AsyncEntityProducers.create(requestBody.getBytes(), ContentType.APPLICATION_JSON));

src/main/java/com/databricks/jdbc/common/DatabricksJdbcConstants.java

@@ -17,8 +17,10 @@ public final class DatabricksJdbcConstants {
               "(?:/([^;]*))?"
               + // Optional Schema (captured without /)
               "(?:;(.*))?"); // Optional Property=Value pairs (captured without leading ;)
-  public static final Pattern HTTP_WAREHOUSE_PATH_PATTERN = Pattern.compile(".*/warehouses/(.+)");
-  public static final Pattern HTTP_ENDPOINT_PATH_PATTERN = Pattern.compile(".*/endpoints/(.+)");
+  public static final Pattern HTTP_WAREHOUSE_PATH_PATTERN =
+      Pattern.compile(".*/warehouses/([^?&]+).*");
+  public static final Pattern HTTP_ENDPOINT_PATH_PATTERN =
+      Pattern.compile(".*/endpoints/([^?&]+).*");
   public static final Pattern HTTP_CLI_PATTERN = Pattern.compile(".*cliservice(.+)");
   public static final Pattern HTTP_PATH_CLI_PATTERN = Pattern.compile("cliservice");
   public static final Pattern TEST_PATH_PATTERN = Pattern.compile("jdbc:databricks://test");

src/main/java/com/databricks/jdbc/common/safe/DatabricksDriverFeatureFlagsContext.java

@@ -102,6 +102,7 @@ private void refreshAllFeatureFlags() {
           .getDatabricksConfig()
           .authenticate()
           .forEach(request::addHeader);
+      connectionContext.getCustomHeaders().forEach(request::addHeader);
       fetchAndSetFlagsFromServer(httpClient, request);
     } catch (Exception e) {
       LOGGER.trace(

src/main/java/com/databricks/jdbc/telemetry/TelemetryPushClient.java

@@ -59,6 +59,7 @@ public void pushEvent(TelemetryRequest request) throws Exception {
     Map<String, String> authHeaders =
         isAuthenticated ? databricksConfig.authenticate() : Collections.emptyMap();
     authHeaders.forEach(post::addHeader);
+    connectionContext.getCustomHeaders().forEach(post::addHeader);
     try (CloseableHttpResponse response = httpClient.execute(post)) {
       // TODO: check response and add retry for partial failures
       if (!HttpUtil.isSuccessfulHttpResponse(response)) {

src/test/java/com/databricks/jdbc/TestConstants.java

@@ -321,4 +321,17 @@ public class TestConstants {
   public static final List<TSparkArrowBatch> ARROW_BATCH_LIST =
       Collections.singletonList(
           new TSparkArrowBatch().setRowCount(0).setBatch(new byte[] {65, 66, 67}));
+
+  // SPOG URLs with ?o= query parameter in httpPath
+  public static final String VALID_SPOG_URL_WAREHOUSE =
+      "jdbc:databricks://spog.cloud.databricks.com/default;ssl=1;AuthMech=3;"
+          + "httpPath=/sql/1.0/warehouses/abc123?o=6051921418418893;UseThriftClient=1";
+
+  public static final String VALID_SPOG_URL_ENDPOINT =
+      "jdbc:databricks://spog.cloud.databricks.com/default;ssl=1;AuthMech=3;"
+          + "httpPath=/sql/1.0/endpoints/abc123?o=6051921418418893;UseThriftClient=0";
+
+  public static final String VALID_SPOG_URL_WAREHOUSE_NO_EXTRA_PARAMS =
+      "jdbc:databricks://spog.cloud.databricks.com/default;ssl=1;AuthMech=3;"
+          + "httpPath=/sql/1.0/warehouses/abc123?o=6051921418418893";
 }

src/test/java/com/databricks/jdbc/api/impl/DatabricksConnectionContextTest.java

@@ -1358,6 +1358,85 @@ public void testOAuthWebServerTimeoutCustom() throws DatabricksSQLException {
     assertEquals(300, connectionContext.getOAuthWebServerTimeout());
   }
 
+  // ==================== SPOG ?o= Tests ====================
+
+  @Test
+  void testBuildPropertiesMap_preservesQueryParamInHttpPath() {
+    String params = "ssl=1;AuthMech=3;httpPath=/sql/1.0/warehouses/abc123?o=999;UseThriftClient=1";
+    ImmutableMap<String, String> result = buildPropertiesMap(params, new Properties());
+
+    assertEquals("/sql/1.0/warehouses/abc123?o=999", result.get("httppath"));
+    assertEquals("1", result.get("usethriftclient"));
+  }
+
+  @Test
+  void testBuildPropertiesMap_handlesValueWithMultipleEquals() {
+    String params = "httpPath=/sql/1.0/warehouses/abc?o=999&other=foo";
+    ImmutableMap<String, String> result = buildPropertiesMap(params, new Properties());
+
+    assertEquals("/sql/1.0/warehouses/abc?o=999&other=foo", result.get("httppath"));
+  }
+
+  @Test
+  void testBuildPropertiesMap_handlesValueWithNoEquals() {
+    String params = "keyonly";
+    ImmutableMap<String, String> result = buildPropertiesMap(params, new Properties());
+
+    assertEquals("", result.get("keyonly"));
+  }
+
+  @Test
+  void testSpogContext_extractsOrgIdFromHttpPath() throws DatabricksSQLException {
+    Properties props = new Properties();
+    props.put("user", "token");
+    props.put("password", "test-token");
+    IDatabricksConnectionContext ctx =
+        DatabricksConnectionContext.parse(TestConstants.VALID_SPOG_URL_WAREHOUSE, props);
+
+    Map<String, String> headers = ctx.getCustomHeaders();
+    assertEquals("6051921418418893", headers.get("x-databricks-org-id"));
+  }
+
+  @Test
+  void testSpogContext_extractsCleanWarehouseId() throws DatabricksSQLException {
+    Properties props = new Properties();
+    props.put("user", "token");
+    props.put("password", "test-token");
+    IDatabricksConnectionContext ctx =
+        DatabricksConnectionContext.parse(TestConstants.VALID_SPOG_URL_WAREHOUSE, props);
+
+    // Warehouse ID should be "abc123" not "abc123?o=6051921418418893"
+    assertTrue(ctx.getComputeResource() instanceof Warehouse);
+    assertEquals("abc123", ((Warehouse) ctx.getComputeResource()).getWarehouseId());
+  }
+
+  @Test
+  void testSpogContext_noOrgIdWithoutQueryParam() throws DatabricksSQLException {
+    Properties props = new Properties();
+    props.put("user", "token");
+    props.put("password", "test-token");
+    IDatabricksConnectionContext ctx =
+        DatabricksConnectionContext.parse(TestConstants.VALID_URL_1, props);
+
+    Map<String, String> headers = ctx.getCustomHeaders();
+    assertFalse(headers.containsKey("x-databricks-org-id"));
+  }
+
+  @Test
+  void testSpogContext_explicitHeaderTakesPrecedence() throws DatabricksSQLException {
+    String url =
+        "jdbc:databricks://host/default;ssl=1;AuthMech=3;"
+            + "httpPath=/sql/1.0/warehouses/abc123?o=frompath;"
+            + "http.header.x-databricks-org-id=fromheader";
+    Properties props = new Properties();
+    props.put("user", "token");
+    props.put("password", "test-token");
+    IDatabricksConnectionContext ctx = DatabricksConnectionContext.parse(url, props);
+
+    Map<String, String> headers = ctx.getCustomHeaders();
+    assertEquals("fromheader", headers.get("x-databricks-org-id"));
+  }
+
   @Test
   public void testDefaultGetterCoverage() throws DatabricksSQLException {
     IDatabricksConnectionContext ctx =

src/test/java/com/databricks/jdbc/common/util/ValidationUtilTest.java

@@ -128,6 +128,13 @@ private static Stream<Arguments> jdbcUrlValidityTestCases() {
             "Valid URL with invalid compression type",
             true),
         Arguments.of(INVALID_URL_1, "Invalid non-Databricks JDBC URL", false),
-        Arguments.of(INVALID_URL_2, "Invalid malformed JDBC scheme", false));
+        Arguments.of(INVALID_URL_2, "Invalid malformed JDBC scheme", false),
+        Arguments.of(
+            VALID_SPOG_URL_WAREHOUSE, "Valid SPOG URL with ?o= in warehouse httpPath", true),
+        Arguments.of(VALID_SPOG_URL_ENDPOINT, "Valid SPOG URL with ?o= in endpoint httpPath", true),
+        Arguments.of(
+            VALID_SPOG_URL_WAREHOUSE_NO_EXTRA_PARAMS,
+            "Valid SPOG URL with ?o= at end of URL",
+            true));
   }
 }