[PECOBLR-450] Successfull identification of request type in the execute function #922

nishkarsh-db · 2025-08-04T08:58:58Z

Description

Implemented request identification by explicitly passing the request type (e.g. THRIFT_METADATA, UC_VOLUME_PUT_INTO) (an enum) in the DatabricksHttpClient execute function, for all clients (Thrift, UCVolume, DBFS, Telemetry, AWS)

Testing

The old tests are successfully executing. New tests are yet to be added.

… for retry handling

…n the DatabricksHttpClient

shivam2680 · 2025-08-08T10:29:38Z

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

+   * Helper method to set HTTP request configuration for volume operations based on operation type.
+   * Only works with Thrift client connections; SDK clients use their own retry logic.
+   */
+  private void setVolumeOperationConfig(HTTPRequestType config) {


rename to setVolumeOperationRequestType

shivam2680 · 2025-08-08T10:30:48Z

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

+      }
+    } catch (Exception e) {
+      // Catch all exceptions (SQLException, NullPointerException, etc.)
+      // This ensures volume operations continue to work even if HTTP config setup fails


How will the retry logic work in this case?

It wont work in this case, the request type will be NOT_SET which will be treated as non idempotent.

shivam2680 · 2025-08-08T10:31:48Z

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


    String listFilesSQLQuery = createListQuery(catalog, schema, volume, folder);

+    setVolumeOperationConfig(HTTPRequestType.UC_VOLUME_LIST);


Better move this inside try block of createStatement. This indicates you are update request type just before sending request.

Applies to the entire class

shivam2680 · 2025-08-08T10:32:32Z

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


    boolean volumeOperationStatus = false;

+    setVolumeOperationConfig(HTTPRequestType.UC_VOLUME_GET_TO);


what is UC_VOLUME_GET_TO? should this be UC_VOLUME_GET?

shivam2680 · 2025-08-08T10:33:33Z

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

      if (isAllowedInputStreamForVolumeOperation) {
-        responseStream = databricksHttpClient.execute(httpGet);
+        responseStream =
+            databricksHttpClient.execute(httpGet, HTTPRequestType.DBFS_VOLUME_EXECUTE_GET);


I think we can use same request types across UcVolumeClient and DBFSVolumeClient

shivam2680 · 2025-08-08T10:34:54Z

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

+  UC_VOLUME_LIST,
+  UC_VOLUME_SHOW_VOLUMES,
+  UC_VOLUME_GET_TO,
+  UC_VOLUME_PUT_INTO,
+  UC_VOLUME_REMOVE,
+  DBFS_VOLUME_EXECUTE_GET,
+  DBFS_VOLUME_EXECUTE_PUT,
+  DBFS_VOLUME_EXECUTE_DELETE,


seems like we are having redundant enums for dbfs and uc volume clients? Please unify

shivam2680 · 2025-08-08T10:35:20Z

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

+
+// enumerates the types of HTTP requests
+
+public enum RequestIdempotency {


Do we need this now?

No we dont need this now, altough a similar enum will be needed in Retry Handler. WIll delete this file.

shivam2680 · 2025-08-08T10:36:13Z

src/main/java/com/databricks/jdbc/dbclient/impl/http/DatabricksHttpClient.java

  @Override
-  public CloseableHttpResponse execute(HttpUriRequest request, boolean supportGzipEncoding)
+  public void setCurrentRequestType(HTTPRequestType currentHTTPRequestType) {
+    if (this.currentHTTPRequestType == HTTPRequestType.NOT_SET) {


Please mention in comments, in what case NOT_SET becomes important. Document the UC volumes example we discussed

shivam2680 · 2025-08-08T10:37:57Z

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


  TCancelOperationResp cancelOperation(TCancelOperationReq req) throws DatabricksHttpException {
    try {
+      setHttpRequestType(HTTPRequestType.THRIFT_CANCEL_OPERATION);


redundant setting of request type. you've already set the request type in getThriftResponse

shivam2680 · 2025-08-08T10:49:18Z

src/main/java/com/databricks/jdbc/dbclient/impl/http/DatabricksHttpClient.java

-            .setDefaultRequestConfig(makeRequestConfig(connectionContext.getSocketTimeout()))
-            .setRetryHandler(retryHandler)
-            .addInterceptorFirst(retryHandler);


You can't remove this until the new retry mechanism is ready. Please undo this chage

shivam2680 · 2025-08-08T10:49:38Z

src/main/java/com/databricks/jdbc/dbclient/impl/http/DatabricksHttpClient.java

-    DatabricksHttpRetryHandler retryHandler =
-        type.equals(HttpClientType.COMMON)
-            ? new DatabricksHttpRetryHandler(connectionContext)
-            : new UCVolumeHttpRetryHandler(connectionContext);


I feel we can rid of HttpClientType now that we have request type.

shivam2680 · 2025-08-08T10:53:02Z

src/main/java/com/databricks/jdbc/dbclient/impl/http/DatabricksHttpClient.java

+
+  @Override
+  public CloseableHttpResponse execute(
+      HttpUriRequest request, HTTPRequestType config, boolean supportGzipEncoding)


maintain consistent naming for params. requestType instead of config

shivam2680 · 2025-08-08T10:53:43Z

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

+  public void setHttpRequestType(HTTPRequestType config) {
+    this.httpClient.setCurrentRequestType(config);
+  }


maintain consistent naming for params.

shivam2680 · 2025-08-08T10:54:17Z

src/main/java/com/databricks/jdbc/dbclient/impl/http/DatabricksHttpClient.java

      request.setHeader("Content-Encoding", "gzip");
    }
+
+    currentHTTPRequestType = HTTPRequestType.NOT_SET;


The state management feels a bit weird. Why are you updating to NOT_SET here?

Because in setRequestType we are setting only when its previously NOT_SET

shivam2680 · 2025-08-08T10:55:10Z

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

    // Execute the request and handle the response
    long httpRequestStartTime = System.currentTimeMillis();
-    try (CloseableHttpResponse response = httpClient.execute(request)) {
+    HTTPRequestType currentRequestConfig = httpClient.getCurrentRequestType();


rename to currentRequestType

samikshya-db · 2025-08-17T16:33:08Z

src/test/java/com/databricks/jdbc/common/safe/DatabricksDriverFeatureFlagsContextTest.java

      context.fetchAndSetFlagsFromServer(httpClientMock, request);
      assertTrue(context.isFeatureEnabled(FEATURE_FLAG_NAME));
-      verify(httpClientMock).execute(request);
+      verify(httpClientMock).execute(eq(request), isA(HTTPRequestType.class));


any reason we are not specifying the actual HTTPRequestType across tests?

New tests for request Type verification are going to be added anyways, so didn't want to change the existing tests behaviour.

vikrantpuppala · 2025-09-02T04:36:49Z

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

+package com.databricks.jdbc.common;
+
+public enum HTTPRequestType {
+  NOT_SET,


HTTPRequestType should be UNKNOWN instead of NOT_SET (I guess the weird nomenclature is because you're actually using RequestType as a state for the http client).

vikrantpuppala · 2025-09-02T04:37:00Z

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

+  VOLUME_DELETE,
+  AUTH,
+  TELEMETRY_PUSH,
+  OTHER


what's this used for?

vikrantpuppala · 2025-09-02T04:39:12Z

src/main/java/com/databricks/jdbc/dbclient/impl/http/DatabricksHttpClient.java

+  private final IDatabricksConnectionContext connectionContext;
  private IdleConnectionEvictor idleConnectionEvictor;
  private CloseableHttpAsyncClient asyncClient;
+  private HTTPRequestType currentHTTPRequestType;


treating this as a class level property makes this class thread unsafe, this is a request level property and the implementation must be in accordance

vikrantpuppala · 2025-09-02T04:41:08Z

src/main/java/com/databricks/jdbc/dbclient/impl/http/DatabricksHttpClient.java

+  public void setCurrentRequestType(HTTPRequestType currentHTTPRequestType) {
+    /* The NOT_SET check is needed while executing UC volume requests, as they go through Thrift ExecuteStatement,
+    So, this function will be called 2 times (from UC Volume Client as well as Thrift Client). In order to avoid that conflict,
+    this check is required.


this seems like a hacky way of doing things to bypass the request type being set at the class level

vikrantpuppala · 2025-09-02T04:47:22Z

src/main/java/com/databricks/jdbc/dbclient/impl/http/HttpRequestTypeBasedRetryHandler.java

+    REQUEST_IDEMPOTENCY_MAP.put(HTTPRequestType.VOLUME_PUT, RequestIdempotency.NON_IDEMPOTENT);
+    REQUEST_IDEMPOTENCY_MAP.put(HTTPRequestType.NOT_SET, RequestIdempotency.NON_IDEMPOTENT);
+    REQUEST_IDEMPOTENCY_MAP.put(HTTPRequestType.OTHER, RequestIdempotency.NON_IDEMPOTENT);
+  }


can we define this mapping as a part of HTTPRequestType enum

vikrantpuppala · 2025-09-02T05:04:11Z

src/main/java/com/databricks/jdbc/dbclient/impl/http/HttpRequestTypeBasedRetryHandler.java

+      throws DatabricksHttpException {
+
+    HttpContext httpContext = new BasicHttpContext();
+    initializeRetryState(httpContext);


why do we maintain the retry counts/duration in the http context? what's the benefit of that?

vikrantpuppala · 2025-09-02T05:07:34Z

src/main/java/com/databricks/jdbc/dbclient/impl/http/HttpRequestTypeBasedRetryHandler.java

+        Thread.currentThread().interrupt();
+        throw new RuntimeException("Thread interrupted during retry", e);
+      } catch (RuntimeException e) {
+        throwHttpException(e, request);


throw RuntimeException for InterruptedException and HttpException for RuntimeException? what's the logic here?

vikrantpuppala · 2025-09-02T05:10:27Z

src/main/java/com/databricks/jdbc/dbclient/impl/http/IRetryStrategy.java

+
+public interface IRetryStrategy {
+  /* Tells whether a given HTTP status code is retriable or not */
+  boolean isStatusCodeRetriable(int statusCode, IDatabricksConnectionContext connectionContext);


we should instead just do isRetriable, it should be oblivious to clients on why we chose to retry (status code, etc.), additionally, we should only retry if we have 503 + retry-after header, all those checks should be done by this class, take a look at: https://github.com/databricks/databricks-sdk-java/blob/d1de7b8a9a397fcf34e93120e43d8bc2f459d004/databricks-sdk-java/src/main/java/com/databricks/sdk/core/retry/RetryStrategy.java

vikrantpuppala · 2025-09-02T05:14:26Z

src/main/java/com/databricks/jdbc/dbclient/impl/http/HttpRequestTypeBasedRetryHandler.java

+
+        // Last attempt check
+        if (attempt > MAX_RETRIES) {
+          return null;


why do we return null? should we just not return the last response?

vikrantpuppala · 2025-09-02T08:06:21Z

src/main/java/com/databricks/jdbc/dbclient/impl/http/IdempotentRetryStrategy.java

+          HttpStatus.SC_REQUEST_TOO_LONG,
+          HttpStatus.SC_REQUEST_URI_TOO_LONG,
+          HttpStatus.SC_UNSUPPORTED_MEDIA_TYPE,
+          HttpStatus.SC_REQUESTED_RANGE_NOT_SATISFIABLE);


do we also have non retryable exceptions?

nishkarsh-db added 2 commits August 4, 2025 12:56

successfully passing config in the execute function for Thrift Client…

cbb6a7e

… for retry handling

Changed the execute function to have an extra parameter RequestType i…

a619c95

…n the DatabricksHttpClient

nishkarsh-db changed the title ~~successfully passing config in the execute function for Thrift Client…~~ Successfull identification of request type in the execute function Aug 8, 2025

nishkarsh-db marked this pull request as ready for review August 8, 2025 10:23

shivam2680 requested review from vikrantpuppala and shivam2680 August 8, 2025 10:27

nishkarsh-db changed the title ~~Successfull identification of request type in the execute function~~ [PECOBLR-450] Successfull identification of request type in the execute function Aug 8, 2025

shivam2680 requested changes Aug 8, 2025

View reviewed changes

shivam2680 reviewed Aug 8, 2025

View reviewed changes

nishkarsh-db added 2 commits August 12, 2025 12:38

added Retry Handlerand removed ClientType enum

74ef808

handled exceptions thrown while HTTP request execution

66748a6

samikshya-db reviewed Aug 17, 2025

View reviewed changes

vikrantpuppala reviewed Sep 2, 2025

View reviewed changes


		String listFilesSQLQuery = createListQuery(catalog, schema, volume, folder);

		setVolumeOperationConfig(HTTPRequestType.UC_VOLUME_LIST);


		boolean volumeOperationStatus = false;

		setVolumeOperationConfig(HTTPRequestType.UC_VOLUME_GET_TO);


		// enumerates the types of HTTP requests

		public enum RequestIdempotency {

[PECOBLR-450] Successfull identification of request type in the execute function #922

Are you sure you want to change the base?

[PECOBLR-450] Successfull identification of request type in the execute function #922

Uh oh!

Conversation

nishkarsh-db commented Aug 4, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Testing

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

shivam2680 Aug 8, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

shivam2680 Aug 8, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

nishkarsh-db commented Aug 4, 2025 •

edited

Loading

shivam2680 Aug 8, 2025 •

edited

Loading

shivam2680 Aug 8, 2025 •

edited

Loading