feat: Add support for disabling CloudFetch via EnableQueryResultDownload parameter (databricks#1183)

## Summary

Add support for disabling CloudFetch via `EnableQueryResultDownload=0`
connection parameter to use inline Arrow results instead.

## Changes

- Add `isCloudFetchEnabled()` method to `IDatabricksConnectionContext`
interface
- Implement the method in `DatabricksConnectionContext` using existing
`EnableQueryResultDownload` parameter
- Update `DatabricksThriftServiceClient` to respect this setting when
making execute requests
- Add unit tests for the new functionality

## Usage

To disable CloudFetch and use inline Arrow results:

```
jdbc:databricks://host:port/default;EnableQueryResultDownload=0;...
```

---------

Signed-off-by: Vikrant Puppala <vikrant.puppala@databricks.com>

Author: vikrantpuppala

NEXT_CHANGELOG.md

@@ -3,6 +3,7 @@
 ## [Unreleased]
 
 ### Added
+- Added support for disabling CloudFetch via `EnableQueryResultDownload=0` to use inline Arrow results instead.
 
 ### Updated
 - Implemented lazy loading for inline Arrow results, fetching arrow batches on demand instead of all at once. This improves memory usage and initial response time for large result sets when using the Thrift protocol with Arrow format.

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

@@ -469,7 +469,7 @@ public DatabricksClientType getClientTypeFromContext() {
       return DatabricksClientType.THRIFT;
     }
     // Check if CloudFetch is disabled - Thrift is required for inline mode
-    if (!Objects.equals(getParameter(DatabricksJdbcUrlParams.ENABLE_CLOUD_FETCH), "1")) {
+    if (!isCloudFetchEnabled()) {
       return DatabricksClientType.THRIFT;
     }
     // Check feature flag to determine if SEA client should be enabled
@@ -1198,6 +1198,11 @@ public boolean isStreamingChunkProviderEnabled() {
     return getParameter(DatabricksJdbcUrlParams.ENABLE_STREAMING_CHUNK_PROVIDER).equals("1");
   }
 
+  @Override
+  public boolean isCloudFetchEnabled() {
+    return getParameter(DatabricksJdbcUrlParams.ENABLE_CLOUD_FETCH).equals("1");
+  }
+
   @Override
   public int getLinkPrefetchWindow() {
     return Integer.parseInt(getParameter(DatabricksJdbcUrlParams.LINK_PREFETCH_WINDOW));

src/main/java/com/databricks/jdbc/api/internal/IDatabricksConnectionContext.java

@@ -426,6 +426,16 @@ public interface IDatabricksConnectionContext {
   /** Returns whether streaming chunk provider is enabled for result fetching. */
   boolean isStreamingChunkProviderEnabled();
 
+  /**
+   * Returns whether CloudFetch (URL-based result download) is enabled.
+   *
+   * <p>When enabled (default), the server may return URL_BASED_SET results that are downloaded from
+   * cloud storage. When disabled, the server returns ARROW_BASED_SET with inline Arrow data.
+   *
+   * @return true if CloudFetch is enabled, false otherwise
+   */
+  boolean isCloudFetchEnabled();
+
   /**
    * Returns the number of chunk links to prefetch ahead of consumption.
    *

src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftServiceClient.java

@@ -232,7 +232,8 @@ private TExecuteStatementReq getRequest(
       request.setCanDecompressLZ4Result(true);
     }
     if (ProtocolFeatureUtil.supportsCloudFetch(serverProtocolVersion)) {
-      request.setCanDownloadResult(true);
+      // Use EnableQueryResultDownload param to control CloudFetch vs inline Arrow
+      request.setCanDownloadResult(this.connectionContext.isCloudFetchEnabled());
     }
     if (ProtocolFeatureUtil.supportsAdvancedArrowTypes(serverProtocolVersion)) {
       arrowNativeTypes

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

@@ -1210,4 +1210,47 @@ public void testEnableTokenFederation() throws DatabricksSQLException {
             DatabricksConnectionContext.parse(TestConstants.VALID_URL_1, props);
     assertFalse(ctx.isTokenFederationEnabled());
   }
+
+  @Test
+  public void testIsCloudFetchEnabled() throws DatabricksSQLException {
+    // Test default value (should be enabled by default)
+    DatabricksConnectionContext ctx =
+        (DatabricksConnectionContext)
+            DatabricksConnectionContext.parse(TestConstants.VALID_URL_1, properties);
+    assertTrue(ctx.isCloudFetchEnabled()); // Default should be true
+
+    // Test via URL parameter - enabled
+    String urlWithCloudFetchEnabled =
+        "jdbc:databricks://sample-host.18.azuredatabricks.net:9999/default;httpPath=/sql/1.0/warehouses/999999999;EnableQueryResultDownload=1";
+    ctx =
+        (DatabricksConnectionContext)
+            DatabricksConnectionContext.parse(urlWithCloudFetchEnabled, properties);
+    assertTrue(ctx.isCloudFetchEnabled());
+
+    // Test via URL parameter - disabled
+    String urlWithCloudFetchDisabled =
+        "jdbc:databricks://sample-host.18.azuredatabricks.net:9999/default;httpPath=/sql/1.0/warehouses/999999999;EnableQueryResultDownload=0";
+    ctx =
+        (DatabricksConnectionContext)
+            DatabricksConnectionContext.parse(urlWithCloudFetchDisabled, properties);
+    assertFalse(ctx.isCloudFetchEnabled());
+
+    // Test via Properties - enabled
+    Properties props = new Properties();
+    props.setProperty("password", "passwd");
+    props.setProperty("EnableQueryResultDownload", "1");
+    ctx =
+        (DatabricksConnectionContext)
+            DatabricksConnectionContext.parse(TestConstants.VALID_URL_1, props);
+    assertTrue(ctx.isCloudFetchEnabled());
+
+    // Test via Properties - disabled
+    props = new Properties();
+    props.setProperty("password", "passwd");
+    props.setProperty("EnableQueryResultDownload", "0");
+    ctx =
+        (DatabricksConnectionContext)
+            DatabricksConnectionContext.parse(TestConstants.VALID_URL_1, props);
+    assertFalse(ctx.isCloudFetchEnabled());
+  }
 }

src/test/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftServiceClientTest.java

@@ -12,6 +12,7 @@
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyLong;
 import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.lenient;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
@@ -164,6 +165,9 @@ private static Stream<Arguments> protocolVersionProvider() {
   void testGetRequestWithDifferentProtocolVersions(TProtocolVersion protocolVersion)
       throws SQLException {
     when(connectionContext.shouldEnableArrow()).thenReturn(true);
+    // Use lenient() because isCloudFetchEnabled() is only called for protocols that support
+    // CloudFetch
+    lenient().when(connectionContext.isCloudFetchEnabled()).thenReturn(true);
     DatabricksThriftServiceClient client =
         new DatabricksThriftServiceClient(thriftAccessor, connectionContext);
     when(session.getSessionInfo()).thenReturn(SESSION_INFO);
@@ -243,6 +247,7 @@ void testGetRequestWithDifferentProtocolVersions(TProtocolVersion protocolVersio
   @Test
   void testExecute() throws SQLException {
     when(connectionContext.shouldEnableArrow()).thenReturn(true);
+    when(connectionContext.isCloudFetchEnabled()).thenReturn(true);
     DatabricksThriftServiceClient client =
         new DatabricksThriftServiceClient(thriftAccessor, connectionContext);
     when(session.getSessionInfo()).thenReturn(SESSION_INFO);
@@ -281,9 +286,52 @@ void testExecute() throws SQLException {
     assertEquals(resultSet, actualResultSet);
   }
 
+  @Test
+  void testExecuteWithCloudFetchDisabled() throws SQLException {
+    when(connectionContext.shouldEnableArrow()).thenReturn(true);
+    when(connectionContext.isCloudFetchEnabled()).thenReturn(false);
+    DatabricksThriftServiceClient client =
+        new DatabricksThriftServiceClient(thriftAccessor, connectionContext);
+    when(session.getSessionInfo()).thenReturn(SESSION_INFO);
+    when(parentStatement.getStatement()).thenReturn(statement);
+    when(parentStatement.getMaxRows()).thenReturn(10);
+    when(statement.getQueryTimeout()).thenReturn(10);
+    TSparkArrowTypes arrowNativeTypes =
+        new TSparkArrowTypes()
+            .setComplexTypesAsArrow(true)
+            .setIntervalTypesAsArrow(true)
+            .setNullTypeAsArrow(true)
+            .setDecimalAsArrow(true)
+            .setTimestampAsArrow(true);
+    TExecuteStatementReq executeStatementReq =
+        new TExecuteStatementReq()
+            .setStatement(TEST_STRING)
+            .setSessionHandle(SESSION_HANDLE)
+            .setCanReadArrowResult(true)
+            .setQueryTimeout(10)
+            .setResultRowLimit(10)
+            .setCanDecompressLZ4Result(true)
+            .setCanDownloadResult(false)
+            .setParameters(Collections.emptyList())
+            .setRunAsync(true)
+            .setUseArrowNativeTypes(arrowNativeTypes);
+    when(thriftAccessor.execute(executeStatementReq, parentStatement, session, StatementType.SQL))
+        .thenReturn(resultSet);
+    DatabricksResultSet actualResultSet =
+        client.executeStatement(
+            TEST_STRING,
+            CLUSTER_COMPUTE,
+            Collections.emptyMap(),
+            StatementType.SQL,
+            session,
+            parentStatement);
+    assertEquals(resultSet, actualResultSet);
+  }
+
   @Test
   void testExecuteAsync() throws SQLException {
     when(connectionContext.shouldEnableArrow()).thenReturn(true);
+    when(connectionContext.isCloudFetchEnabled()).thenReturn(true);
     DatabricksThriftServiceClient client =
         new DatabricksThriftServiceClient(thriftAccessor, connectionContext);
     when(session.getSessionInfo()).thenReturn(SESSION_INFO);
@@ -630,6 +678,7 @@ void testListFunctionsWithSQLEnabled() throws SQLException {
         new DatabricksThriftServiceClient(thriftAccessor, connectionContext);
     when(connectionContext.enableShowCommandsForGetFunctions()).thenReturn(true);
     when(connectionContext.shouldEnableArrow()).thenReturn(true);
+    when(connectionContext.isCloudFetchEnabled()).thenReturn(true);
     when(session.getSessionInfo()).thenReturn(SESSION_INFO);
     TSparkArrowTypes arrowNativeTypes =
         new TSparkArrowTypes()