Skip to content

Commit 881478a

Browse files
bhu1rdbbhu1rdb
committed
Merge branch 'main' into flushmetrics
2 parents eb2b57f + ee3e52c commit 881478a

File tree

12 files changed

+87
-623
lines changed

12 files changed

+87
-623
lines changed

‎.github/workflows/bugCatcher.yml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,7 @@ jobs:
4747
echo "DATABRICKS_USER=${DATABRICKS_USER}" >> $GITHUB_ENV
4848
echo "DATABRICKS_JDBC_M2M_URL=${DATABRICKS_JDBC_M2M_URL}" >> $GITHUB_ENV
4949
- name: Run Integration Tests
50-
run: mvn -B test -Dtest=*e2e/*
50+
run: mvn -B test -Dtest=*e2e/OAuthTests*
5151
env:
5252
DATABRICKS_HOST: ${{ env.DATABRICKS_HOST }}
5353
DATABRICKS_HTTP_PATH: ${{ env.DATABRICKS_HTTP_PATH }}

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

Lines changed: 60 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -16,10 +16,14 @@
1616
import com.databricks.sdk.core.utils.ProxyUtils;
1717
import com.google.common.annotations.VisibleForTesting;
1818
import java.io.IOException;
19+
import java.security.KeyManagementException;
20+
import java.security.KeyStoreException;
21+
import java.security.NoSuchAlgorithmException;
1922
import java.util.Objects;
2023
import java.util.Set;
2124
import java.util.concurrent.ConcurrentHashMap;
2225
import java.util.concurrent.TimeUnit;
26+
import javax.net.ssl.SSLContext;
2327
import org.apache.http.HttpException;
2428
import org.apache.http.HttpHost;
2529
import org.apache.http.HttpResponse;
@@ -28,11 +32,13 @@
2832
import org.apache.http.client.methods.HttpUriRequest;
2933
import org.apache.http.conn.UnsupportedSchemeException;
3034
import org.apache.http.conn.routing.HttpRoute;
35+
import org.apache.http.conn.ssl.NoopHostnameVerifier;
3136
import org.apache.http.impl.client.CloseableHttpClient;
3237
import org.apache.http.impl.client.HttpClientBuilder;
3338
import org.apache.http.impl.conn.DefaultSchemePortResolver;
3439
import org.apache.http.impl.conn.PoolingHttpClientConnectionManager;
3540
import org.apache.http.protocol.HttpContext;
41+
import org.apache.http.ssl.SSLContextBuilder;
3642

3743
/** Http client implementation to be used for executing http requests. */
3844
public class DatabricksHttpClient implements IDatabricksHttpClient {
@@ -72,6 +78,7 @@ public class DatabricksHttpClient implements IDatabricksHttpClient {
7278
private static boolean shouldRetryRateLimitError;
7379
private static int rateLimitRetryTimeout;
7480
protected static int idleHttpConnectionExpiry;
81+
private CloseableHttpClient httpDisabledSSLClient;
7582

7683
private DatabricksHttpClient(IDatabricksConnectionContext connectionContext) {
7784
initializeConnectionManager();
@@ -81,6 +88,7 @@ private DatabricksHttpClient(IDatabricksConnectionContext connectionContext) {
8188
shouldRetryRateLimitError = connectionContext.shouldRetryRateLimitError();
8289
rateLimitRetryTimeout = connectionContext.getRateLimitRetryTimeout();
8390
httpClient = makeClosableHttpClient(connectionContext);
91+
httpDisabledSSLClient = makeClosableDisabledSslHttpClient();
8492
idleHttpConnectionExpiry = connectionContext.getIdleHttpConnectionExpiry();
8593
}
8694

@@ -143,6 +151,26 @@ private CloseableHttpClient makeClosableHttpClient(
143151
return builder.build();
144152
}
145153

154+
private CloseableHttpClient makeClosableDisabledSslHttpClient() {
155+
try {
156+
// Create SSL context that trusts all certificates
157+
SSLContext sslContext =
158+
new SSLContextBuilder().loadTrustMaterial(null, (chain, authType) -> true).build();
159+
160+
// Create HttpClient with the SSL context
161+
return HttpClientBuilder.create()
162+
.setSSLContext(sslContext)
163+
.setSSLHostnameVerifier(new NoopHostnameVerifier())
164+
.build();
165+
} catch (NoSuchAlgorithmException | KeyManagementException | KeyStoreException e) {
166+
LoggingUtil.log(
167+
LogLevel.DEBUG,
168+
String.format(
169+
"Error in creating HttpClient with the SSL context [{%s}]", e.getMessage()));
170+
}
171+
return null;
172+
}
173+
146174
private boolean handleRetry(IOException exception, int executionCount, HttpContext context) {
147175
int errCode = getErrorCode(exception);
148176
if (!isErrorCodeRetryable(errCode)) {
@@ -322,20 +350,22 @@ public CloseableHttpResponse execute(HttpUriRequest request) throws DatabricksHt
322350
try {
323351
return httpClient.execute(request);
324352
} catch (IOException e) {
325-
Throwable cause = e;
326-
while (cause != null) {
327-
if (cause instanceof DatabricksRetryHandlerException) {
328-
throw new DatabricksHttpException(cause.getMessage(), cause);
329-
}
330-
cause = cause.getCause();
331-
}
332-
String errorMsg =
333-
String.format(
334-
"Caught error while executing http request: [%s]. Error Message: [%s]",
335-
RequestSanitizer.sanitizeRequest(request), e);
336-
LoggingUtil.log(LogLevel.ERROR, errorMsg);
337-
throw new DatabricksHttpException(errorMsg, e);
353+
throwHttpException(e, request, LogLevel.ERROR);
354+
}
355+
return null;
356+
}
357+
358+
public CloseableHttpResponse executeWithoutSSL(HttpUriRequest request)
359+
throws DatabricksHttpException {
360+
LoggingUtil.log(
361+
LogLevel.DEBUG,
362+
String.format("Executing HTTP request [{%s}]", RequestSanitizer.sanitizeRequest(request)));
363+
try {
364+
return httpDisabledSSLClient.execute(request);
365+
} catch (Exception e) {
366+
throwHttpException(e, request, LogLevel.DEBUG);
338367
}
368+
return null;
339369
}
340370

341371
@Override
@@ -385,4 +415,21 @@ public static synchronized void removeInstance(IDatabricksConnectionContext cont
385415
}
386416
}
387417
}
418+
419+
private static void throwHttpException(Exception e, HttpUriRequest request, LogLevel logLevel)
420+
throws DatabricksHttpException {
421+
Throwable cause = e;
422+
while (cause != null) {
423+
if (cause instanceof DatabricksRetryHandlerException) {
424+
throw new DatabricksHttpException(cause.getMessage(), cause);
425+
}
426+
cause = cause.getCause();
427+
}
428+
String errorMsg =
429+
String.format(
430+
"Caught error while executing http request: [%s]. Error Message: [%s]",
431+
RequestSanitizer.sanitizeRequest(request), e);
432+
LoggingUtil.log(logLevel, errorMsg);
433+
throw new DatabricksHttpException(errorMsg, e);
434+
}
388435
}

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -129,6 +129,7 @@ public static ColumnInfoTypeName getTypeFromTypeDesc(TTypeDesc typeDesc) {
129129
case BOOLEAN_TYPE:
130130
return ColumnInfoTypeName.BOOLEAN;
131131
case TINYINT_TYPE:
132+
return ColumnInfoTypeName.BYTE;
132133
case SMALLINT_TYPE:
133134
return ColumnInfoTypeName.SHORT;
134135
case INT_TYPE:

‎src/main/java/com/databricks/jdbc/commons/util/DefaultHttpClientUtil.java‎

Lines changed: 0 additions & 38 deletions
This file was deleted.

‎src/main/java/com/databricks/jdbc/commons/util/ValidationUtil.java‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -11,9 +11,8 @@ public class ValidationUtil {
1111

1212
public static void checkIfNonNegative(int number, String fieldName)
1313
throws DatabricksSQLException {
14-
// Todo : Add appropriate exception
1514
if (number < 0) {
16-
throw new DatabricksSQLException(
15+
throw new DatabricksValidationException(
1716
String.format("Invalid input for %s, : %d", fieldName, number));
1817
}
1918
}

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

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -357,11 +357,15 @@ public void setTimestamp(int parameterIndex, Timestamp x, Calendar cal) throws S
357357
LoggingUtil.log(
358358
LogLevel.DEBUG, "public void setTimestamp(int parameterIndex, Timestamp x, Calendar cal)");
359359
checkIfClosed();
360-
TimeZone originalTimeZone = TimeZone.getDefault();
361-
TimeZone.setDefault(cal.getTimeZone());
362-
x = new Timestamp(x.getTime());
363-
TimeZone.setDefault(originalTimeZone);
364-
setObject(parameterIndex, x, DatabricksTypeUtil.TIMESTAMP);
360+
if (cal != null) {
361+
TimeZone originalTimeZone = TimeZone.getDefault();
362+
TimeZone.setDefault(cal.getTimeZone());
363+
x = new Timestamp(x.getTime());
364+
TimeZone.setDefault(originalTimeZone);
365+
setObject(parameterIndex, x, DatabricksTypeUtil.TIMESTAMP);
366+
} else {
367+
setTimestamp(parameterIndex, x);
368+
}
365369
}
366370

367371
@Override

‎src/main/java/com/databricks/jdbc/telemetry/DatabricksMetrics.java‎

Lines changed: 5 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,8 @@
11
package com.databricks.jdbc.telemetry;
22

33
import com.databricks.jdbc.client.DatabricksHttpException;
4+
import com.databricks.jdbc.client.http.DatabricksHttpClient;
45
import com.databricks.jdbc.commons.LogLevel;
5-
import com.databricks.jdbc.commons.util.DefaultHttpClientUtil;
66
import com.databricks.jdbc.commons.util.LoggingUtil;
77
import com.databricks.jdbc.core.DatabricksSQLException;
88
import com.databricks.jdbc.driver.IDatabricksConnectionContext;
@@ -17,7 +17,6 @@
1717
import org.apache.http.client.methods.HttpPost;
1818
import org.apache.http.client.methods.HttpUriRequest;
1919
import org.apache.http.client.utils.URIBuilder;
20-
import org.apache.http.impl.client.DefaultHttpClient;
2120
import org.apache.http.util.EntityUtils;
2221

2322
public class DatabricksMetrics implements AutoCloseable {
@@ -32,7 +31,7 @@ public class DatabricksMetrics implements AutoCloseable {
3231
private final String METRICS_TYPE = "metrics_type";
3332
private Boolean hasInitialExportOccurred = false;
3433
private String workspaceId = null;
35-
private DefaultHttpClient telemetryClient = null;
34+
private final DatabricksHttpClient telemetryClient;
3635

3736
private void setWorkspaceId(String workspaceId) {
3837
this.workspaceId = workspaceId;
@@ -69,13 +68,7 @@ public DatabricksMetrics(IDatabricksConnectionContext context) throws Databricks
6968
}
7069
String resourceId = context.getComputeResource().getWorkspaceId();
7170
setWorkspaceId(resourceId);
72-
// TODO: Bhuvan: Replace telemetry client with DatabricksHttpClient when non-SSL variation is
73-
// implemented
74-
try {
75-
telemetryClient = DefaultHttpClientUtil.getDefaultHttpClient();
76-
} catch (Exception e) {
77-
throw new DatabricksSQLException("Failed to create telemetry client");
78-
}
71+
this.telemetryClient = DatabricksHttpClient.getInstance(context);
7972
scheduleExportMetrics();
8073
}
8174

@@ -102,7 +95,8 @@ public String sendRequest(Map<String, Double> map, MetricsType metricsType) thro
10295
uriBuilder.addParameter(METRICS_TYPE, metricsType.name().equals("GAUGE") ? "1" : "0");
10396
HttpUriRequest request = new HttpPost(uriBuilder.build());
10497
// TODO (Bhuvan): Add authentication headers
105-
CloseableHttpResponse response = telemetryClient.execute(request);
98+
// TODO (Bhuvan): execute request using SSL
99+
CloseableHttpResponse response = telemetryClient.executeWithoutSSL(request);
106100

107101
// Error handling
108102
if (response == null) {

‎src/test/java/com/databricks/jdbc/client/impl/thrift/commons/DatabricksThriftHelperTest.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -93,7 +93,7 @@ private static Stream<Arguments> thriftDirectResultSets() {
9393
private static Stream<Arguments> typeIdAndColumnInfoType() {
9494
return Stream.of(
9595
Arguments.of(TTypeId.BOOLEAN_TYPE, ColumnInfoTypeName.BOOLEAN),
96-
Arguments.of(TTypeId.TINYINT_TYPE, ColumnInfoTypeName.SHORT),
96+
Arguments.of(TTypeId.TINYINT_TYPE, ColumnInfoTypeName.BYTE),
9797
Arguments.of(TTypeId.SMALLINT_TYPE, ColumnInfoTypeName.SHORT),
9898
Arguments.of(TTypeId.INT_TYPE, ColumnInfoTypeName.INT),
9999
Arguments.of(TTypeId.BIGINT_TYPE, ColumnInfoTypeName.LONG),

‎src/test/java/com/databricks/jdbc/core/DatabricksPreparedStatementTest.java‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -267,6 +267,15 @@ public void testSetTimestampWithCalendar() throws DatabricksSQLException {
267267
1, new Timestamp(System.currentTimeMillis()), Calendar.getInstance()));
268268
}
269269

270+
@Test
271+
public void testSetTimestampWithNullCalendar() throws DatabricksSQLException {
272+
setupMocks();
273+
DatabricksPreparedStatement preparedStatement =
274+
new DatabricksPreparedStatement(connection, STATEMENT);
275+
assertDoesNotThrow(
276+
() -> preparedStatement.setTimestamp(1, new Timestamp(System.currentTimeMillis()), null));
277+
}
278+
270279
@Test
271280
public void testSetAsciiStream() throws DatabricksSQLException {
272281
setupMocks();

‎src/test/java/com/databricks/jdbc/integration/IntegrationTestUtil.java‎

Lines changed: 0 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -72,11 +72,6 @@ public static String getDatabricksDogfoodHost() {
7272
return System.getenv("DATABRICKS_DOGFOOD_HOST");
7373
}
7474

75-
public static String getDatabricksBenchmarkingHost() {
76-
// includes port
77-
return System.getenv("DATABRICKS_BENCHMARKING_HOST");
78-
}
79-
8075
public static String getDatabricksToken() {
8176
return System.getenv("DATABRICKS_TOKEN");
8277
}
@@ -89,10 +84,6 @@ public static String getDatabricksBenchfoodToken() {
8984
return System.getenv("DATABRICKS_BENCHFOOD_TOKEN");
9085
}
9186

92-
public static String getDatabricksBenchmarkingToken() {
93-
return System.getenv("DATABRICKS_BENCHMARKING_TOKEN");
94-
}
95-
9687
public static String getDatabricksHTTPPath() {
9788
return System.getenv("DATABRICKS_HTTP_PATH");
9889
}
@@ -105,10 +96,6 @@ public static String getDatabricksDogfoodHTTPPath() {
10596
return System.getenv("DATABRICKS_DOGFOOD_HTTP_PATH");
10697
}
10798

108-
public static String getDatabricksBenchmarkingHTTPPath() {
109-
return System.getenv("DATABRICKS_BENCHMARKING_HTTP_PATH");
110-
}
111-
11299
public static String getDatabricksCatalog() {
113100
return isFakeServiceTest
114101
? FakeServiceConfigLoader.getProperty(TEST_CATALOG)
@@ -172,11 +159,6 @@ public static Connection getBenchfoodJDBCConnection() throws SQLException {
172159
getBenchfoodJDBCUrl(), getDatabricksUser(), getDatabricksBenchfoodToken());
173160
}
174161

175-
public static Connection getBenchmarkingJDBCConnection() throws SQLException {
176-
return DriverManager.getConnection(
177-
getBenchmarkingJDBCUrl(), getDatabricksUser(), getDatabricksBenchmarkingToken());
178-
}
179-
180162
public static void resetJDBCConnection() {
181163
JDBCConnection = null;
182164
}
@@ -224,15 +206,6 @@ public static String getDogfoodJDBCUrl() {
224206
return String.format(template, host, httpPath);
225207
}
226208

227-
public static String getBenchmarkingJDBCUrl() {
228-
String template =
229-
"jdbc:databricks://%s/default;transportMode=http;ssl=1;AuthMech=3;httpPath=%s";
230-
String host = getDatabricksBenchmarkingHost();
231-
String httpPath = getDatabricksBenchmarkingHTTPPath();
232-
233-
return String.format(template, host, httpPath);
234-
}
235-
236209
public static boolean executeSQL(String sql) {
237210
try {
238211
if (JDBCConnection == null) JDBCConnection = getValidJDBCConnection();

0 commit comments

Comments
 (0)