[PECO-1565] [PECO-1563] Thrift - add other auth types + refresh token (#173)

* add other auth types

* Refresh token

* Add more tests

* Modify naming

Author: samikshya-db

src/main/java/com/databricks/jdbc/client/impl/sdk/DatabricksSdkClient.java

@@ -69,7 +69,6 @@ public DatabricksSdkClient(
       ApiClient apiClient)
       throws DatabricksParsingException {
     this.connectionContext = connectionContext;
-    // Handle more auth types
     this.databricksConfig =
         new DatabricksConfig()
             .setHost(connectionContext.getHostUrl())

src/main/java/com/databricks/jdbc/client/impl/thrift/commons/DatabricksThriftAccessor.java

@@ -34,37 +34,31 @@ public DatabricksThriftAccessor(IDatabricksConnectionContext connectionContext)
         new DatabricksHttpTTransport(
             DatabricksHttpClient.getInstance(connectionContext),
             connectionContext.getEndpointURL());
-    // TODO : add other auth in followup PRs
-    this.databricksConfig =
-        new DatabricksConfig()
-            .setHost(connectionContext.getHostUrl())
-            .setToken(connectionContext.getToken());
+    this.databricksConfig = new OAuthAuthenticator(connectionContext).getDatabricksConfig();
     Map<String, String> authHeaders = databricksConfig.authenticate();
     transport.setCustomHeaders(authHeaders);
     TBinaryProtocol protocol = new TBinaryProtocol(transport);
     this.thriftClient = new TCLIService.Client(protocol);
   }
 
   @VisibleForTesting
-  public DatabricksThriftAccessor(TCLIService.Client client) {
-    this.databricksConfig = null;
+  DatabricksThriftAccessor(TCLIService.Client client, DatabricksConfig config) {
+    this.databricksConfig = config;
     this.thriftClient = client;
   }
 
   public TBase getThriftResponse(
       TBase request, CommandName commandName, IDatabricksStatement parentStatement)
       throws DatabricksSQLException {
     /*Todo list :
-     *  1. Poll until we get a success status
-     *  2. Test out metadata operations.
-     *  3. Add token refresh
-     *  4. Handle cloud-fetch
-     *  5. Handle compression
+     *  1. Test out metadata operations.
+     *  2. Handle compression
      * */
     LOGGER.debug(
         "Fetching thrift response for request {}, CommandName {}",
         request.toString(),
         commandName.name());
+    refreshHeadersIfRequired();
     try {
       switch (commandName) {
         case OPEN_SESSION:
@@ -107,11 +101,12 @@ public TBase getThriftResponse(
 
   public TFetchResultsResp getResultSetResp(TOperationHandle operationHandle, String context)
       throws DatabricksHttpException {
+    refreshHeadersIfRequired();
     return getResultSetResp(
         TStatusCode.SUCCESS_STATUS, operationHandle, context, DEFAULT_ROW_LIMIT, false);
   }
 
-  public TFetchResultsResp getResultSetResp(
+  private TFetchResultsResp getResultSetResp(
       TStatusCode responseCode,
       TOperationHandle operationHandle,
       String context,
@@ -169,6 +164,7 @@ public DatabricksResultSet execute(
       IDatabricksSession session,
       StatementType statementType)
       throws SQLException {
+    refreshHeadersIfRequired();
     int maxRows = (parentStatement == null) ? DEFAULT_ROW_LIMIT : parentStatement.getMaxRows();
     TSparkGetDirectResults directResults =
         new TSparkGetDirectResults().setMaxBytes(DEFAULT_BYTE_LIMIT).setMaxRows(maxRows);
@@ -338,10 +334,15 @@ private TFetchResultsResp listColumns(TGetColumnsReq request)
         false);
   }
 
-  public TGetResultSetMetadataResp getResultSetMetadata(TOperationHandle operationHandle)
+  private TGetResultSetMetadataResp getResultSetMetadata(TOperationHandle operationHandle)
       throws TException {
     TGetResultSetMetadataReq resultSetMetadataReq =
         new TGetResultSetMetadataReq().setOperationHandle(operationHandle);
     return thriftClient.GetResultSetMetadata(resultSetMetadataReq);
   }
+
+  private void refreshHeadersIfRequired() {
+    ((DatabricksHttpTTransport) thriftClient.getInputProtocol().getTransport())
+        .setCustomHeaders(databricksConfig.authenticate());
+  }
 }

src/main/java/com/databricks/jdbc/core/OAuthAuthenticator.java

@@ -14,56 +14,55 @@ public OAuthAuthenticator(IDatabricksConnectionContext connectionContext) {
   }
 
   public WorkspaceClient getWorkspaceClient() throws DatabricksParsingException {
+    return new WorkspaceClient(getDatabricksConfig());
+  }
+
+  public DatabricksConfig getDatabricksConfig() throws DatabricksParsingException {
     if (this.connectionContext.getAuthMech().equals(IDatabricksConnectionContext.AuthMech.PAT)) {
-      return authenticateAccessToken();
+      return createAccessTokenConfig();
     }
     // TODO(Madhav): Revisit these to set JDBC values
     else if (this.connectionContext
         .getAuthMech()
         .equals(IDatabricksConnectionContext.AuthMech.OAUTH)) {
       switch (this.connectionContext.getAuthFlow()) {
         case TOKEN_PASSTHROUGH:
-          return authenticateAccessToken();
+          return createAccessTokenConfig();
         case CLIENT_CREDENTIALS:
-          return authenticateM2M();
+          return createM2MConfig();
         case BROWSER_BASED_AUTHENTICATION:
-          return authenticateU2M();
+          return createU2MConfig();
       }
     }
-    return authenticateAccessToken();
+    return createAccessTokenConfig();
   }
 
-  public WorkspaceClient authenticateU2M() throws DatabricksParsingException {
+  public DatabricksConfig createU2MConfig() throws DatabricksParsingException {
     DatabricksConfig config =
         new DatabricksConfig()
             .setAuthType(DatabricksJdbcConstants.U2M_AUTH_TYPE)
-            .setHost(this.connectionContext.getHostForOAuth())
-            .setClientId(this.connectionContext.getClientId())
-            .setClientSecret(this.connectionContext.getClientSecret())
+            .setHost(connectionContext.getHostForOAuth())
+            .setClientId(connectionContext.getClientId())
+            .setClientSecret(connectionContext.getClientSecret())
             .setOAuthRedirectUrl(DatabricksJdbcConstants.U2M_AUTH_REDIRECT_URL);
     if (!config.isAzure()) {
-      // Default scope is already being set for Azure in databricks-sdk.
-      config.setScopes(this.connectionContext.getOAuthScopesForU2M());
+      config.setScopes(connectionContext.getOAuthScopesForU2M());
     }
-    return new WorkspaceClient(config);
+    return config;
   }
 
-  public WorkspaceClient authenticateAccessToken() throws DatabricksParsingException {
-    DatabricksConfig config =
-        new DatabricksConfig()
-            .setAuthType(DatabricksJdbcConstants.ACCESS_TOKEN_AUTH_TYPE)
-            .setHost(this.connectionContext.getHostUrl())
-            .setToken(this.connectionContext.getToken());
-    return new WorkspaceClient(config);
+  public DatabricksConfig createAccessTokenConfig() throws DatabricksParsingException {
+    return new DatabricksConfig()
+        .setAuthType(DatabricksJdbcConstants.ACCESS_TOKEN_AUTH_TYPE)
+        .setHost(connectionContext.getHostUrl())
+        .setToken(connectionContext.getToken());
   }
 
-  public WorkspaceClient authenticateM2M() throws DatabricksParsingException {
-    DatabricksConfig config =
-        new DatabricksConfig()
-            .setAuthType(DatabricksJdbcConstants.M2M_AUTH_TYPE)
-            .setHost(this.connectionContext.getHostForOAuth())
-            .setClientId(this.connectionContext.getClientId())
-            .setClientSecret(this.connectionContext.getClientSecret());
-    return new WorkspaceClient(config);
+  public DatabricksConfig createM2MConfig() throws DatabricksParsingException {
+    return new DatabricksConfig()
+        .setAuthType(DatabricksJdbcConstants.M2M_AUTH_TYPE)
+        .setHost(connectionContext.getHostForOAuth())
+        .setClientId(connectionContext.getClientId())
+        .setClientSecret(connectionContext.getClientSecret());
   }
 }