Skip to content

Commit d76de89

Browse files
vikrantpuppalavikrantpuppala
authored
[PECO-1486] Support proxy (#292)
This PR adds support for proxy to the JDBC driver
1 parent b5d7e05 commit d76de89

File tree

14 files changed

+182
-136
lines changed

14 files changed

+182
-136
lines changed

‎pom.xml‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,7 @@
3636

3737
<spotless.version>2.43.0</spotless.version>
3838
<commons-io.version>2.13.0</commons-io.version>
39-
<databricks-sdk.version>0.17.0</databricks-sdk.version>
39+
<databricks-sdk.version>0.26.0</databricks-sdk.version>
4040
<maven-surefire-plugin.version>3.1.2</maven-surefire-plugin.version>
4141
<sql-logic-test.version>0.3</sql-logic-test.version>
4242
<lz4-compression.version>1.8.0</lz4-compression.version>
@@ -67,6 +67,11 @@
6767
<artifactId>arrow-vector</artifactId>
6868
<version>${arrow.version}</version>
6969
</dependency>
70+
<dependency>
71+
<groupId>org.slf4j</groupId>
72+
<artifactId>slf4j-reload4j</artifactId>
73+
<version>2.0.13</version>
74+
</dependency>
7075
<dependency>
7176
<groupId>com.diffplug.spotless</groupId>
7277
<artifactId>spotless-maven-plugin</artifactId>

‎src/main/java/com/databricks/jdbc/client/http/DatabricksHttpClient.java‎

Lines changed: 39 additions & 59 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,7 @@
11
package com.databricks.jdbc.client.http;
22

33
import static com.databricks.jdbc.client.http.RetryHandler.*;
4-
import static com.databricks.jdbc.driver.DatabricksJdbcConstants.FAKE_SERVICE_URI_PROP_SUFFIX;
5-
import static com.databricks.jdbc.driver.DatabricksJdbcConstants.IS_FAKE_SERVICE_TEST_PROP;
4+
import static com.databricks.jdbc.driver.DatabricksJdbcConstants.*;
65
import static io.netty.util.NetUtil.LOCALHOST;
76

87
import com.databricks.jdbc.client.DatabricksHttpException;
@@ -11,7 +10,10 @@
1110
import com.databricks.jdbc.commons.LogLevel;
1211
import com.databricks.jdbc.commons.util.LoggingUtil;
1312
import com.databricks.jdbc.driver.IDatabricksConnectionContext;
13+
import com.databricks.sdk.core.DatabricksConfig;
14+
import com.databricks.sdk.core.ProxyConfig;
1415
import com.databricks.sdk.core.UserAgent;
16+
import com.databricks.sdk.core.utils.ProxyUtils;
1517
import com.google.common.annotations.VisibleForTesting;
1618
import java.io.IOException;
1719
import java.util.Objects;
@@ -21,19 +23,13 @@
2123
import org.apache.http.HttpException;
2224
import org.apache.http.HttpHost;
2325
import org.apache.http.HttpResponse;
24-
import org.apache.http.HttpResponseInterceptor;
25-
import org.apache.http.auth.AuthScope;
26-
import org.apache.http.auth.UsernamePasswordCredentials;
27-
import org.apache.http.client.CredentialsProvider;
2826
import org.apache.http.client.config.RequestConfig;
2927
import org.apache.http.client.methods.CloseableHttpResponse;
3028
import org.apache.http.client.methods.HttpUriRequest;
3129
import org.apache.http.conn.UnsupportedSchemeException;
3230
import org.apache.http.conn.routing.HttpRoute;
33-
import org.apache.http.impl.client.BasicCredentialsProvider;
3431
import org.apache.http.impl.client.CloseableHttpClient;
3532
import org.apache.http.impl.client.HttpClientBuilder;
36-
import org.apache.http.impl.client.ProxyAuthenticationStrategy;
3733
import org.apache.http.impl.conn.DefaultSchemePortResolver;
3834
import org.apache.http.impl.conn.PoolingHttpClientConnectionManager;
3935
import org.apache.http.protocol.HttpContext;
@@ -138,16 +134,12 @@ private CloseableHttpClient makeClosableHttpClient(
138134
.setConnectionManager(connectionManager)
139135
.setUserAgent(getUserAgent())
140136
.setDefaultRequestConfig(makeRequestConfig())
141-
.setRetryHandler(
142-
(exception, executionCount, context) ->
143-
handleRetry(exception, executionCount, context))
144-
.addInterceptorFirst(
145-
(HttpResponseInterceptor)
146-
(httpResponse, httpContext) ->
147-
handleResponseInterceptor(httpResponse, httpContext));
148-
149-
configureProxy(connectionContext, builder);
150-
137+
.setRetryHandler(this::handleRetry)
138+
.addInterceptorFirst(this::handleResponseInterceptor);
139+
setupProxy(connectionContext, builder);
140+
if (Boolean.parseBoolean(System.getProperty(IS_FAKE_SERVICE_TEST_PROP))) {
141+
setFakeServiceRouteInHttpClient(builder);
142+
}
151143
return builder.build();
152144
}
153145

@@ -240,51 +232,39 @@ private void initializeRetryCounts(HttpContext httpContext) {
240232
}
241233
}
242234

243-
private void configureProxy(
235+
@VisibleForTesting
236+
public static void setupProxy(
244237
IDatabricksConnectionContext connectionContext, HttpClientBuilder builder) {
245-
if (connectionContext.getUseSystemProxy()) {
246-
builder.useSystemProperties();
247-
}
248-
// Override system proxy if proxy details are explicitly provided
249-
// If cloud fetch proxy is provided use that, else use the regular proxy
238+
String proxyHost = null;
239+
Integer proxyPort = null;
240+
String proxyUser = null;
241+
String proxyPassword = null;
242+
ProxyConfig.ProxyAuthType proxyAuth = connectionContext.getProxyAuthType();
243+
// System proxy is handled by the SDK.
244+
// If proxy details are explicitly provided use those for the connection.
250245
if (connectionContext.getUseCloudFetchProxy()) {
251-
setProxyDetailsInHttpClient(
252-
builder,
253-
connectionContext.getCloudFetchProxyHost(),
254-
connectionContext.getCloudFetchProxyPort(),
255-
connectionContext.getUseCloudFetchProxyAuth(),
256-
connectionContext.getCloudFetchProxyUser(),
257-
connectionContext.getCloudFetchProxyPassword());
246+
proxyHost = connectionContext.getCloudFetchProxyHost();
247+
proxyPort = connectionContext.getCloudFetchProxyPort();
248+
proxyUser = connectionContext.getCloudFetchProxyUser();
249+
proxyPassword = connectionContext.getCloudFetchProxyPassword();
250+
proxyAuth = connectionContext.getCloudFetchProxyAuthType();
258251
} else if (connectionContext.getUseProxy()) {
259-
setProxyDetailsInHttpClient(
260-
builder,
261-
connectionContext.getProxyHost(),
262-
connectionContext.getProxyPort(),
263-
connectionContext.getUseProxyAuth(),
264-
connectionContext.getProxyUser(),
265-
connectionContext.getProxyPassword());
266-
} else if (Boolean.parseBoolean(System.getProperty(IS_FAKE_SERVICE_TEST_PROP))) {
267-
setFakeServiceRouteInHttpClient(builder);
252+
proxyHost = connectionContext.getProxyHost();
253+
proxyPort = connectionContext.getProxyPort();
254+
proxyUser = connectionContext.getProxyUser();
255+
proxyPassword = connectionContext.getProxyPassword();
256+
proxyAuth = connectionContext.getProxyAuthType();
268257
}
269-
}
270-
271-
@VisibleForTesting
272-
public static void setProxyDetailsInHttpClient(
273-
HttpClientBuilder builder,
274-
String proxyHost,
275-
int proxyPort,
276-
Boolean useProxyAuth,
277-
String proxyUser,
278-
String proxyPassword) {
279-
builder.setProxy(new HttpHost(proxyHost, proxyPort));
280-
if (useProxyAuth) {
281-
CredentialsProvider credsProvider = new BasicCredentialsProvider();
282-
credsProvider.setCredentials(
283-
new AuthScope(proxyHost, proxyPort),
284-
new UsernamePasswordCredentials(proxyUser, proxyPassword));
285-
builder
286-
.setDefaultCredentialsProvider(credsProvider)
287-
.setProxyAuthenticationStrategy(new ProxyAuthenticationStrategy());
258+
if (proxyHost != null || connectionContext.getUseSystemProxy()) {
259+
ProxyConfig proxyConfig =
260+
new ProxyConfig(new DatabricksConfig())
261+
.setUseSystemProperties(connectionContext.getUseSystemProxy())
262+
.setHost(proxyHost)
263+
.setPort(proxyPort)
264+
.setUsername(proxyUser)
265+
.setPassword(proxyPassword)
266+
.setProxyAuthType(proxyAuth);
267+
ProxyUtils.setupProxy(proxyConfig, builder);
288268
}
289269
}
290270

‎src/main/java/com/databricks/jdbc/client/impl/helper/ClientUtils.java‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
package com.databricks.jdbc.client.impl.helper;
2+
3+
import com.databricks.jdbc.core.DatabricksParsingException;
4+
import com.databricks.jdbc.driver.IDatabricksConnectionContext;
5+
import com.databricks.sdk.core.DatabricksConfig;
6+
7+
public class ClientUtils {
8+
public static DatabricksConfig generateDatabricksConfig(
9+
IDatabricksConnectionContext connectionContext) throws DatabricksParsingException {
10+
DatabricksConfig databricksConfig =
11+
new DatabricksConfig()
12+
.setHost(connectionContext.getHostUrl())
13+
.setToken(connectionContext.getToken())
14+
.setUseSystemPropertiesHttp(connectionContext.getUseSystemProxy());
15+
// Setup proxy settings
16+
if (connectionContext.getUseProxy()) {
17+
databricksConfig
18+
.setProxyHost(connectionContext.getProxyHost())
19+
.setProxyPort(connectionContext.getProxyPort());
20+
}
21+
databricksConfig
22+
.setProxyAuthType(connectionContext.getProxyAuthType())
23+
.setProxyUsername(connectionContext.getProxyUser())
24+
.setProxyPassword(connectionContext.getProxyPassword());
25+
return databricksConfig;
26+
}
27+
}

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

Lines changed: 5 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55

66
import com.databricks.jdbc.client.DatabricksClient;
77
import com.databricks.jdbc.client.StatementType;
8+
import com.databricks.jdbc.client.impl.helper.ClientUtils;
89
import com.databricks.jdbc.client.sqlexec.*;
910
import com.databricks.jdbc.client.sqlexec.CloseStatementRequest;
1011
import com.databricks.jdbc.client.sqlexec.CreateSessionRequest;
@@ -25,6 +26,7 @@
2526
import com.databricks.sdk.core.ApiClient;
2627
import com.databricks.sdk.core.DatabricksConfig;
2728
import com.databricks.sdk.service.sql.*;
29+
import com.google.common.annotations.VisibleForTesting;
2830
import java.sql.SQLException;
2931
import java.time.Instant;
3032
import java.util.Collection;
@@ -36,8 +38,6 @@
3638
/** Implementation of DatabricksClient interface using Databricks Java SDK. */
3739
public class DatabricksSdkClient implements DatabricksClient {
3840
private static final String SYNC_TIMEOUT_VALUE = "10s";
39-
private static final String ASYNC_TIMEOUT_VALUE = "0s";
40-
4141
private final IDatabricksConnectionContext connectionContext;
4242
private final DatabricksConfig databricksConfig;
4343
private final WorkspaceClient workspaceClient;
@@ -51,17 +51,12 @@ private static Map<String, String> getHeaders() {
5151
public DatabricksSdkClient(IDatabricksConnectionContext connectionContext)
5252
throws DatabricksParsingException {
5353
this.connectionContext = connectionContext;
54-
// TODO: [PECO-1486] pass on proxy settings to SDK once changes are merged in SDK
55-
// Handle more auth types
56-
this.databricksConfig =
57-
new DatabricksConfig()
58-
.setHost(connectionContext.getHostUrl())
59-
.setToken(connectionContext.getToken());
60-
54+
this.databricksConfig = ClientUtils.generateDatabricksConfig(connectionContext);
6155
OAuthAuthenticator authenticator = new OAuthAuthenticator(connectionContext);
62-
this.workspaceClient = authenticator.getWorkspaceClient();
56+
this.workspaceClient = authenticator.getWorkspaceClient(this.databricksConfig);
6357
}
6458

59+
@VisibleForTesting
6560
public DatabricksSdkClient(
6661
IDatabricksConnectionContext connectionContext,
6762
StatementExecutionService statementExecutionService,

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

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
import com.databricks.jdbc.client.DatabricksHttpException;
88
import com.databricks.jdbc.client.StatementType;
99
import com.databricks.jdbc.client.http.DatabricksHttpClient;
10+
import com.databricks.jdbc.client.impl.helper.ClientUtils;
1011
import com.databricks.jdbc.client.impl.thrift.generated.*;
1112
import com.databricks.jdbc.commons.CommandName;
1213
import com.databricks.jdbc.commons.LogLevel;
@@ -33,7 +34,9 @@ public class DatabricksThriftAccessor {
3334
public DatabricksThriftAccessor(IDatabricksConnectionContext connectionContext)
3435
throws DatabricksParsingException {
3536
enableDirectResults = connectionContext.getDirectResultMode();
36-
this.databricksConfig = new OAuthAuthenticator(connectionContext).getDatabricksConfig();
37+
this.databricksConfig = ClientUtils.generateDatabricksConfig(connectionContext);
38+
OAuthAuthenticator authenticator = new OAuthAuthenticator(connectionContext);
39+
authenticator.setupDatabricksConfig(databricksConfig);
3740
this.databricksConfig.resolve();
3841
Map<String, String> authHeaders = databricksConfig.authenticate();
3942
String endPointUrl = connectionContext.getEndpointURL();

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

Lines changed: 29 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -13,53 +13,59 @@ public OAuthAuthenticator(IDatabricksConnectionContext connectionContext) {
1313
this.connectionContext = connectionContext;
1414
}
1515

16-
public WorkspaceClient getWorkspaceClient() throws DatabricksParsingException {
17-
return new WorkspaceClient(getDatabricksConfig());
16+
public WorkspaceClient getWorkspaceClient(DatabricksConfig databricksConfig)
17+
throws DatabricksParsingException {
18+
setupDatabricksConfig(databricksConfig);
19+
return new WorkspaceClient(databricksConfig);
1820
}
1921

20-
public DatabricksConfig getDatabricksConfig() throws DatabricksParsingException {
22+
public void setupDatabricksConfig(DatabricksConfig databricksConfig)
23+
throws DatabricksParsingException {
2124
if (this.connectionContext.getAuthMech().equals(IDatabricksConnectionContext.AuthMech.PAT)) {
22-
return createAccessTokenConfig();
25+
setupAccessTokenConfig(databricksConfig);
2326
}
2427
// TODO(Madhav): Revisit these to set JDBC values
2528
else if (this.connectionContext
2629
.getAuthMech()
2730
.equals(IDatabricksConnectionContext.AuthMech.OAUTH)) {
2831
switch (this.connectionContext.getAuthFlow()) {
2932
case TOKEN_PASSTHROUGH:
30-
return createAccessTokenConfig();
33+
setupAccessTokenConfig(databricksConfig);
34+
break;
3135
case CLIENT_CREDENTIALS:
32-
return createM2MConfig();
36+
setupM2MConfig(databricksConfig);
37+
break;
3338
case BROWSER_BASED_AUTHENTICATION:
34-
return createU2MConfig();
39+
setupU2MConfig(databricksConfig);
40+
break;
3541
}
42+
} else {
43+
setupAccessTokenConfig(databricksConfig);
3644
}
37-
return createAccessTokenConfig();
3845
}
3946

40-
public DatabricksConfig createU2MConfig() throws DatabricksParsingException {
41-
DatabricksConfig config =
42-
new DatabricksConfig()
43-
.setAuthType(DatabricksJdbcConstants.U2M_AUTH_TYPE)
44-
.setHost(connectionContext.getHostForOAuth())
45-
.setClientId(connectionContext.getClientId())
46-
.setClientSecret(connectionContext.getClientSecret())
47-
.setOAuthRedirectUrl(DatabricksJdbcConstants.U2M_AUTH_REDIRECT_URL);
48-
if (!config.isAzure()) {
49-
config.setScopes(connectionContext.getOAuthScopesForU2M());
47+
public void setupU2MConfig(DatabricksConfig databricksConfig) throws DatabricksParsingException {
48+
databricksConfig
49+
.setAuthType(DatabricksJdbcConstants.U2M_AUTH_TYPE)
50+
.setHost(connectionContext.getHostForOAuth())
51+
.setClientId(connectionContext.getClientId())
52+
.setClientSecret(connectionContext.getClientSecret())
53+
.setOAuthRedirectUrl(DatabricksJdbcConstants.U2M_AUTH_REDIRECT_URL);
54+
if (!databricksConfig.isAzure()) {
55+
databricksConfig.setScopes(connectionContext.getOAuthScopesForU2M());
5056
}
51-
return config;
5257
}
5358

54-
public DatabricksConfig createAccessTokenConfig() throws DatabricksParsingException {
55-
return new DatabricksConfig()
59+
public void setupAccessTokenConfig(DatabricksConfig databricksConfig)
60+
throws DatabricksParsingException {
61+
databricksConfig
5662
.setAuthType(DatabricksJdbcConstants.ACCESS_TOKEN_AUTH_TYPE)
5763
.setHost(connectionContext.getHostUrl())
5864
.setToken(connectionContext.getToken());
5965
}
6066

61-
public DatabricksConfig createM2MConfig() throws DatabricksParsingException {
62-
return new DatabricksConfig()
67+
public void setupM2MConfig(DatabricksConfig databricksConfig) throws DatabricksParsingException {
68+
databricksConfig
6369
.setAuthType(DatabricksJdbcConstants.M2M_AUTH_TYPE)
6470
.setHost(connectionContext.getHostForOAuth())
6571
.setClientId(connectionContext.getClientId())

‎src/main/java/com/databricks/jdbc/driver/DatabricksConnectionContext.java‎

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
import com.databricks.jdbc.core.types.ComputeResource;
1313
import com.databricks.jdbc.core.types.Warehouse;
1414
import com.databricks.jdbc.telemetry.DatabricksMetrics;
15+
import com.databricks.sdk.core.ProxyConfig;
1516
import com.google.common.annotations.VisibleForTesting;
1617
import com.google.common.collect.ImmutableMap;
1718
import java.util.*;
@@ -330,7 +331,7 @@ public String getClientUserAgent() {
330331
: DatabricksJdbcConstants.USER_AGENT_THRIFT_CLIENT;
331332
return nullOrEmptyString(customerUserAgent)
332333
? clientAgent
333-
: clientAgent + " " + customerUserAgent;
334+
: clientAgent + USER_AGENT_DELIMITER + customerUserAgent;
334335
}
335336

336337
// TODO: Make use of compression type
@@ -431,8 +432,9 @@ public Boolean getUseProxy() {
431432
}
432433

433434
@Override
434-
public Boolean getUseProxyAuth() {
435-
return Objects.equals(getParameter(USE_PROXY_AUTH), "1");
435+
public ProxyConfig.ProxyAuthType getProxyAuthType() {
436+
int proxyAuthTypeOrdinal = Integer.parseInt(getParameter(PROXY_AUTH, "0"));
437+
return ProxyConfig.ProxyAuthType.values()[proxyAuthTypeOrdinal];
436438
}
437439

438440
@Override
@@ -466,8 +468,9 @@ public String getCloudFetchProxyPassword() {
466468
}
467469

468470
@Override
469-
public Boolean getUseCloudFetchProxyAuth() {
470-
return Objects.equals(getParameter(USE_CF_PROXY_AUTH), "1");
471+
public ProxyConfig.ProxyAuthType getCloudFetchProxyAuthType() {
472+
int proxyAuthTypeOrdinal = Integer.parseInt(getParameter(CF_PROXY_AUTH, "0"));
473+
return ProxyConfig.ProxyAuthType.values()[proxyAuthTypeOrdinal];
471474
}
472475

473476
@Override

0 commit comments

Comments
 (0)