Skip to content

Commit

Permalink
[Feature] Remove proxy settings from databricks config and use common…
Browse files Browse the repository at this point in the history
…s http client builder for connection manager (#337)

## Changes
<!-- Summary of your changes that are easy to understand -->
Add the ability to provide non proxy hosts via databricks config as well

## Tests
<!-- How is this tested? -->
Local testing
  • Loading branch information
vikrantpuppala authored Sep 13, 2024
1 parent ab39312 commit b692350
Show file tree
Hide file tree
Showing 4 changed files with 21 additions and 132 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -132,24 +132,6 @@ public class DatabricksConfig {
@ConfigAttribute(env = "DATABRICKS_RATE_LIMIT")
private Integer rateLimit;

@ConfigAttribute(env = "PROXY_HOST")
private String proxyHost;

@ConfigAttribute(env = "PROXY_PORT")
private Integer proxyPort;

@ConfigAttribute(env = "PROXY_USERNAME")
private String proxyUsername;

@ConfigAttribute(env = "PROXY_PASSWORD")
private String proxyPassword;

@ConfigAttribute(env = "PROXY_AUTH_TYPE")
private ProxyConfig.ProxyAuthType proxyAuthType;

@ConfigAttribute(env = "USE_SYSTEM_PROPERTIES_HTTP")
private Boolean useSystemPropertiesHttp;

private volatile boolean resolved;
private HeaderFactory headerFactory;

Expand Down Expand Up @@ -530,60 +512,6 @@ public DatabricksConfig setHttpClient(HttpClient httpClient) {
return this;
}

public String getProxyHost() {
return proxyHost;
}

public DatabricksConfig setProxyHost(String proxyHost) {
this.proxyHost = proxyHost;
return this;
}

public Integer getProxyPort() {
return proxyPort;
}

public DatabricksConfig setProxyPort(Integer proxyPort) {
this.proxyPort = proxyPort;
return this;
}

public String getProxyUsername() {
return proxyUsername;
}

public DatabricksConfig setProxyUsername(String proxyUsername) {
this.proxyUsername = proxyUsername;
return this;
}

public String getProxyPassword() {
return proxyPassword;
}

public DatabricksConfig setProxyPassword(String proxyPassword) {
this.proxyPassword = proxyPassword;
return this;
}

public ProxyConfig.ProxyAuthType getProxyAuthType() {
return proxyAuthType;
}

public DatabricksConfig setProxyAuthType(ProxyConfig.ProxyAuthType proxyAuthType) {
this.proxyAuthType = proxyAuthType;
return this;
}

public Boolean getUseSystemPropertiesHttp() {
return useSystemPropertiesHttp;
}

public DatabricksConfig setUseSystemPropertiesHttp(Boolean useSystemPropertiesHttp) {
this.useSystemPropertiesHttp = useSystemPropertiesHttp;
return this;
}

public boolean isAzure() {
if (azureWorkspaceResourceId != null) {
return true;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -25,15 +25,6 @@ public enum ProxyAuthType {

public ProxyConfig() {}

public ProxyConfig(DatabricksConfig config) {
this.host = config.getProxyHost();
this.port = config.getProxyPort();
this.username = config.getProxyUsername();
this.password = config.getProxyPassword();
this.proxyAuthType = config.getProxyAuthType();
this.useSystemProperties = config.getUseSystemPropertiesHttp();
}

public String getHost() {
return host;
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,7 @@ public static class Builder {
private Integer timeoutSeconds;
private ProxyConfig proxyConfig;
private SSLConnectionSocketFactory sslSocketFactory;
private PoolingHttpClientConnectionManager connectionManager;

/**
* @param databricksConfig The DatabricksConfig to use for the HttpClient. If the
Expand Down Expand Up @@ -86,19 +87,28 @@ public Builder withSslSocketFactory(SSLConnectionSocketFactory sslSocketFactory)
return this;
}

/**
* @param connectionManager the PoolingHttpClientConnectionManager to use for the HttpClient.
* @return This builder.
*/
public Builder withConnectionManager(PoolingHttpClientConnectionManager connectionManager) {
this.connectionManager = connectionManager;
return this;
}

/** Builds a new instance of CommonsHttpClient with the configured parameters. */
public CommonsHttpClient build() {
return new CommonsHttpClient(this);
}
}

private static final Logger LOG = LoggerFactory.getLogger(CommonsHttpClient.class);
private final PoolingHttpClientConnectionManager connectionManager =
new PoolingHttpClientConnectionManager();
private final CloseableHttpClient hc;
private int timeout;

private CommonsHttpClient(Builder builder) {
HttpClientBuilder httpClientBuilder =
HttpClientBuilder.create().setDefaultRequestConfig(makeRequestConfig());
int timeoutSeconds = 300;
if (builder.databricksConfig != null
&& builder.databricksConfig.getHttpTimeoutSeconds() != null) {
Expand All @@ -108,45 +118,23 @@ private CommonsHttpClient(Builder builder) {
timeoutSeconds = builder.timeoutSeconds;
}
timeout = timeoutSeconds * 1000;
connectionManager.setMaxTotal(100);
HttpClientBuilder httpClientBuilder =
HttpClientBuilder.create()
.setConnectionManager(connectionManager)
.setDefaultRequestConfig(makeRequestConfig());
if (builder.proxyConfig != null) {
ProxyUtils.setupProxy(builder.proxyConfig, httpClientBuilder);
}
if (builder.sslSocketFactory != null) {
httpClientBuilder.setSSLSocketFactory(builder.sslSocketFactory);
}
if (builder.connectionManager != null) {
httpClientBuilder.setConnectionManager(builder.connectionManager);
} else {
PoolingHttpClientConnectionManager connectionManager =
new PoolingHttpClientConnectionManager();
connectionManager.setMaxTotal(100);
httpClientBuilder.setConnectionManager(connectionManager);
}
hc = httpClientBuilder.build();
}

// These constructors have been deprecate in favour of a builder pattern.
// They will be removed in a future release.
@Deprecated
public CommonsHttpClient(int timeoutSeconds) {
timeout = timeoutSeconds * 1000;
connectionManager.setMaxTotal(100);
hc = makeClosableHttpClient();
}

@Deprecated
public CommonsHttpClient(DatabricksConfig databricksConfig) {
this(
databricksConfig.getHttpTimeoutSeconds() == null
? 300
: databricksConfig.getHttpTimeoutSeconds(),
new ProxyConfig(databricksConfig));
}

@Deprecated
public CommonsHttpClient(int timeoutSeconds, ProxyConfig proxyConfig) {
timeout = timeoutSeconds * 1000;
connectionManager.setMaxTotal(100);
hc = makeClosableHttpClient(proxyConfig);
}

private RequestConfig makeRequestConfig() {
return RequestConfig.custom()
.setConnectionRequestTimeout(timeout)
Expand All @@ -155,24 +143,6 @@ private RequestConfig makeRequestConfig() {
.build();
}

@Deprecated
private CloseableHttpClient makeClosableHttpClient() {
return HttpClientBuilder.create()
.setConnectionManager(connectionManager)
.setDefaultRequestConfig(makeRequestConfig())
.build();
}

@Deprecated
private CloseableHttpClient makeClosableHttpClient(ProxyConfig proxyConfig) {
HttpClientBuilder builder =
HttpClientBuilder.create()
.setConnectionManager(connectionManager)
.setDefaultRequestConfig(makeRequestConfig());
ProxyUtils.setupProxy(proxyConfig, builder);
return builder.build();
}

@Override
public Response execute(Request in) throws IOException {
HttpUriRequest request = transformRequest(in);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -129,7 +129,7 @@ public void testDiscoveryEndpoint() throws IOException {
new DatabricksConfig()
.setHost(server.getUrl())
.setDiscoveryUrl(discoveryUrl)
.setHttpClient(new CommonsHttpClient(30))
.setHttpClient(new CommonsHttpClient.Builder().withTimeoutSeconds(30).build())
.getOidcEndpoints();

assertEquals(
Expand Down

0 comments on commit b692350

Please sign in to comment.