Add changes for Kill switch integration for SEA rollout (#1067)

## Description
Added changes to integrate with Kill Switch to prepare for SEA rollout.
Changes include lazy evaluation of clientType to cover different
scenarios for Thrift vs SEA rollout

This PR also adds comprehensive test coverage for the lazy clientType
inference feature. The implementation makes
clientType evaluation lazy using a Supplier pattern, avoiding
unnecessary feature flag service calls during connection parsing.

Test Coverage Added

  Total: 32 new tests across 4 test files (101 tests run, all passing ✅)

  1. DatabricksConnectionContextTest (+15 tests)

Added comprehensive tests for lazy clientType inference and decision
logic:

  Lazy Initialization Tests:
- Caching behavior verification - ensures clientType is computed once
and cached
- Thread safety validation - concurrent access returns consistent
results

  Client Type Selection Logic:
  - Explicit useThriftClient parameter handling (0 and 1 values)
  - Arrow disabled scenarios (forces THRIFT)
  - CloudFetch disabled scenarios (forces THRIFT)
  - Combined Arrow + CloudFetch disabled scenarios

  Feature Flag Integration:
  - Feature flag enabled → SEA client
  - Feature flag disabled → THRIFT client
  - Invalid feature flag values → safe default (THRIFT)
  - Empty feature flag values → safe default (THRIFT)
  - Missing feature flag → safe default (THRIFT)

  Override Behavior:
  - setClientType() before first access
  - setClientType() after lazy evaluation

  Parameterized Decision Matrix:
  - Comprehensive test covering all combinations of:
    - AllPurposeCluster vs Warehouse
    - useThriftClient (null/0/1)
    - enableArrow (0/1)
    - enableCloudFetch (0/1)
    - Feature flag state (true/false)
  - 10 different scenarios validated

  2. DatabricksSessionTest (+2 tests)

  Session-level integration tests:
  - Session opens correctly with lazy SEA client type
  - Session opens correctly with lazy THRIFT client type for clusters

  3. DatabricksDriverFeatureFlagsContextFactoryTest (NEW FILE, 8 tests)

  Factory pattern and workspace isolation tests:
  - Singleton pattern per workspace validation
  - Different workspaces get isolated contexts
  - Reference counting increments/decrements correctly
  - Context persists until last reference removed
  - Null safety handling
  - Flag sharing across connections to same workspace
  - Flag isolation across different workspaces

  4. DatabricksDriverFeatureFlagsContextTest (+7 tests)

  Additional feature flag integration tests:
  - SQL exec flag specific handling
  - Multiple feature flags in response parsing
  - HTTP error handling (404, 403, 503)
  - Case-sensitive flag name matching

  Code Fixes

  1. DatabricksConnectionContext.java

Fixed getClientTypeFromContext() to remove premature default value
check:
- Before: The method was checking getParameter(USE_THRIFT_CLIENT) which
returns default value "1", causing it to always return THRIFT before
checking feature flags
- After: Removed the redundant default check, allowing feature flag
evaluation when user hasn't explicitly set useThriftClient
- Impact: Feature flags can now properly control SEA client enablement
when prerequisites (Arrow, CloudFetch) are met

  2. DatabricksDriverFeatureFlagsContextFactory.java

  Made setFeatureFlagsContext() method public:
  - Already marked @VisibleForTesting
  - Needed for test mocking and validation
  - Allows tests to inject feature flag values for deterministic testing

  Test Coverage Highlights

  What's Tested ✅

1. Lazy Evaluation: ClientType is not computed during connection parsing
2. Thread Safety: Concurrent access to getClientType() is safe and
consistent
  3. Caching: Once computed, clientType is cached and reused
  4. Decision Logic: All branches of clientType selection are covered
5. Feature Flag Integration: All feature flag states
(enabled/disabled/invalid/missing) are tested
6. Error Handling: HTTP errors, malformed responses, and edge cases
default to safe behavior
7. Workspace Isolation: Feature flags are properly isolated per
workspace
  8. Reference Counting: Factory correctly manages context lifecycle

  Decision Matrix Validated

| Cluster Type | useThrift | Arrow | CloudFetch | FeatureFlag | Expected
|

|--------------|-----------|-------|------------|-------------|----------|
| APC | * | * | * | * | THRIFT |
| Warehouse | 1 | * | * | * | THRIFT |
| Warehouse | 0 | * | * | * | SEA |
| Warehouse | null | 0 | * | * | THRIFT |
| Warehouse | null | * | 0 | * | THRIFT |
| Warehouse | null | 1 | 1 | true | SEA |
| Warehouse | null | 1 | 1 | false | THRIFT |

## Testing
<!-- Describe how the changes have been tested-->
mvn test
-Dtest=DatabricksConnectionContextTest,DatabricksSessionTest,DatabricksDriverFeatureFlagsContextTest,DatabricksDriverFeatureFlagsContextFactoryTest

  Result: All 101 tests passing ✅

## Additional Notes to the Reviewer
<!-- Share any additional context or insights that may help the reviewer
understand the changes better. This could include challenges faced,
limitations, or compromises made during the development process.
Also, mention any areas of the code that you would like the reviewer to
focus on specifically. -->

---------

Co-authored-by: Claude <noreply@anthropic.com>

Author: gopalldb

NEXT_CHANGELOG.md

@@ -3,6 +3,7 @@
 ## [Unreleased]
 
 ### Added
+- Added Feature-flag integration for SQL Exec API rollout
 
 ### Updated
 - Minimized OAuth requests by reducing calls in feature flags and telemetry.

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

@@ -10,6 +10,7 @@
 
 import com.databricks.jdbc.api.internal.IDatabricksConnectionContext;
 import com.databricks.jdbc.common.*;
+import com.databricks.jdbc.common.safe.DatabricksDriverFeatureFlagsContextFactory;
 import com.databricks.jdbc.common.util.ValidationUtil;
 import com.databricks.jdbc.exception.DatabricksDriverException;
 import com.databricks.jdbc.exception.DatabricksParsingException;
@@ -23,6 +24,7 @@
 import com.databricks.sdk.core.utils.Cloud;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Strings;
+import com.google.common.base.Supplier;
 import com.google.common.collect.ImmutableMap;
 import java.net.URI;
 import java.util.*;
@@ -34,13 +36,17 @@ public class DatabricksConnectionContext implements IDatabricksConnectionContext
 
   private static final JdbcLogger LOGGER =
       JdbcLoggerFactory.getLogger(DatabricksConnectionContext.class);
+
+  private static final String SQL_EXEC_FLAG_NAME =
+      "databricks.partnerplatform.clientConfigsFeatureFlags.enableSqlExecForJdbc";
+
   private final String host;
   @VisibleForTesting final int port;
   private final String schema;
   private final String connectionURL;
   private final IDatabricksComputeResource computeResource;
   private final Map<String, String> customHeaders;
-  private DatabricksClientType clientType;
+  private Supplier<DatabricksClientType> clientTypeSupplier;
   @VisibleForTesting final ImmutableMap<String, String> parameters;
   @VisibleForTesting final String connectionUuid;
 
@@ -59,7 +65,18 @@ private DatabricksConnectionContext(
     this.customHeaders = parseCustomHeaders(parameters);
     this.computeResource = buildCompute();
     this.connectionUuid = UUID.randomUUID().toString();
-    this.clientType = getClientTypeFromContext();
+    this.clientTypeSupplier =
+        new Supplier<>() {
+          private DatabricksClientType cType;
+
+          @Override
+          public synchronized DatabricksClientType get() {
+            if (cType == null) {
+              cType = getClientTypeFromContext();
+            }
+            return cType;
+          }
+        };
   }
 
   private DatabricksConnectionContext(
@@ -422,20 +439,40 @@ public DatabricksClientType getClientTypeFromContext() {
     if (computeResource instanceof AllPurposeCluster) {
       return DatabricksClientType.THRIFT;
     }
-    String useThriftClient = getParameter(DatabricksJdbcUrlParams.USE_THRIFT_CLIENT);
-    if (useThriftClient != null && useThriftClient.equals("1")) {
+    // First we will check if user as explicitly provided a client type, and we will honour that
+    String useThriftClient = getParameterIgnoreDefault(DatabricksJdbcUrlParams.USE_THRIFT_CLIENT);
+    if (useThriftClient != null) {
+      if (useThriftClient.equals("1")) {
+        return DatabricksClientType.THRIFT;
+      } else if (useThriftClient.equals("0")) {
+        return DatabricksClientType.SEA;
+      }
+    }
+    // Now, user has not provided a value, we will decide based on our checks
+    // Check if Arrow is disabled - Thrift is required for inline mode
+    if (!Objects.equals(getParameter(DatabricksJdbcUrlParams.ENABLE_ARROW), "1")) {
+      return DatabricksClientType.THRIFT;
+    }
+    // Check if CloudFetch is disabled - Thrift is required for inline mode
+    if (!Objects.equals(getParameter(DatabricksJdbcUrlParams.ENABLE_CLOUD_FETCH), "1")) {
       return DatabricksClientType.THRIFT;
     }
-    return DatabricksClientType.SEA;
+    // Check feature flag to determine if SEA client should be enabled
+    if (DatabricksDriverFeatureFlagsContextFactory.getInstance(this)
+        .isFeatureEnabled(SQL_EXEC_FLAG_NAME)) {
+      return DatabricksClientType.SEA;
+    }
+    // Default to THRIFT if feature flag is not enabled or cannot be determined
+    return DatabricksClientType.THRIFT;
   }
 
   @Override
   public DatabricksClientType getClientType() {
-    return clientType;
+    return clientTypeSupplier.get();
   }
 
   public void setClientType(DatabricksClientType clientType) {
-    this.clientType = clientType;
+    this.clientTypeSupplier = () -> clientType;
   }
 
   @Override
@@ -1036,6 +1073,10 @@ private String getParameter(DatabricksJdbcUrlParams key) {
     return this.parameters.getOrDefault(key.getParamName().toLowerCase(), key.getDefaultValue());
   }
 
+  private String getParameterIgnoreDefault(DatabricksJdbcUrlParams key) {
+    return this.parameters.getOrDefault(key.getParamName().toLowerCase(), null);
+  }
+
   private String getParameter(DatabricksJdbcUrlParams key, String defaultValue) {
     return this.parameters.getOrDefault(key.getParamName().toLowerCase(), defaultValue);
   }

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

@@ -62,19 +62,10 @@ public class DatabricksSession implements IDatabricksSession {
    */
   public DatabricksSession(IDatabricksConnectionContext connectionContext)
       throws DatabricksSQLException {
-    if (connectionContext.getClientType() == DatabricksClientType.THRIFT) {
-      this.databricksClient =
-          DatabricksMetricsTimedProcessor.createProxy(
-              new DatabricksThriftServiceClient(connectionContext));
-    } else {
-      this.databricksClient =
-          DatabricksMetricsTimedProcessor.createProxy(new DatabricksSdkClient(connectionContext));
-      this.databricksMetadataClient =
-          DatabricksMetricsTimedProcessor.createProxy(
-              new DatabricksMetadataSdkClient(databricksClient));
-    }
     this.isSessionOpen = false;
     this.sessionInfo = null;
+    this.databricksClient = null;
+    this.databricksMetadataClient = null;
     this.computeResource = connectionContext.getComputeResource();
     this.catalog = connectionContext.getCatalog();
     this.schema = connectionContext.getSchema();
@@ -139,6 +130,22 @@ public boolean isOpen() {
   @Override
   public void open() throws DatabricksSQLException {
     LOGGER.debug("public void open()");
+
+    // Skip for tests, it would be already set
+    if (databricksClient == null) {
+      if (connectionContext.getClientType() == DatabricksClientType.THRIFT) {
+        this.databricksClient =
+            DatabricksMetricsTimedProcessor.createProxy(
+                new DatabricksThriftServiceClient(connectionContext));
+      } else {
+        this.databricksClient =
+            DatabricksMetricsTimedProcessor.createProxy(new DatabricksSdkClient(connectionContext));
+        this.databricksMetadataClient =
+            DatabricksMetricsTimedProcessor.createProxy(
+                new DatabricksMetadataSdkClient(databricksClient));
+      }
+    }
+
     synchronized (this) {
       if (!isSessionOpen) {
         try {

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

@@ -172,6 +172,7 @@ public enum DatabricksJdbcUrlParams {
   ENABLE_METRIC_VIEW_METADATA("EnableMetricViewMetadata", "Enable metric view metadata", "0"),
   ENABLE_MULTIPLE_CATALOG_SUPPORT(
       "enableMultipleCatalogSupport", "Enable multiple catalog support", "1"),
+  ENABLE_CLOUD_FETCH("EnableQueryResultDownload", "Enable Cloud Fetch", "1"),
   ENABLE_SEA_SYNC_METADATA(
       "EnableSeaSyncMetadata",
       "Enable x-databricks-sea-can-run-fully-sync header for synchronous metadata requests in SEA mode",

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

@@ -63,7 +63,7 @@ public static void removeInstance(IDatabricksConnectionContext connectionContext
   }
 
   @VisibleForTesting
-  static void setFeatureFlagsContext(
+  public static void setFeatureFlagsContext(
       IDatabricksConnectionContext connectionContext, Map<String, String> featureFlags) {
     String key = keyOf(connectionContext);
     contextMap.put(
@@ -73,6 +73,6 @@ static void setFeatureFlagsContext(
   }
 
   private static String keyOf(IDatabricksConnectionContext context) {
-    return context.getHostForOAuth();
+    return context.getHost();
   }
 }