Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Feature] Remove proxy settings from databricks config and use commons http client builder for connection manager #337

Merged
merged 8 commits into from
Sep 13, 2024
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
Loading