Skip to content

Commit 4b80c3d

Browse files
jayantsing-dbjayantsing-db
authored
Do not send repeated DBSQL version query (#590)
often times, multiple connections are created concurrently. one can reasonably assume that the customer does not change DBSQL version while the application is running. this allows us to cache the version string and send single query.
1 parent d27d6c7 commit 4b80c3d

File tree

2 files changed

+107
-22
lines changed

2 files changed

+107
-22
lines changed

‎src/main/java/com/databricks/jdbc/common/util/DriverUtil.java‎

Lines changed: 53 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,8 @@
1212
import com.google.common.annotations.VisibleForTesting;
1313
import java.io.IOException;
1414
import java.sql.ResultSet;
15+
import java.util.concurrent.ConcurrentHashMap;
16+
import java.util.concurrent.ConcurrentMap;
1517

1618
/**
1719
* Utility class for operations related to the Databricks JDBC driver.
@@ -22,10 +24,15 @@
2224
public class DriverUtil {
2325

2426
private static final JdbcLogger LOGGER = JdbcLoggerFactory.getLogger(DriverUtil.class);
27+
public static final String DBSQL_VERSION_SQL = "SELECT current_version().dbsql_version";
2528
private static final String VERSION = "0.9.7-oss";
26-
private static final String DBSQL_VERSION_SQL = "SELECT current_version().dbsql_version";
27-
public static final int DBSQL_MIN_MAJOR_VERSION_FOR_SEA_SUPPORT = 2024;
28-
public static final int DBSQL_MIN_MINOR_VERSION_FOR_SEA_SUPPORT = 30;
29+
private static final int DBSQL_MIN_MAJOR_VERSION_FOR_SEA_SUPPORT = 2024;
30+
private static final int DBSQL_MIN_MINOR_VERSION_FOR_SEA_SUPPORT = 30;
31+
32+
/** Cached DBSQL version mapped by HTTP path to avoid repeated queries to the cluster. */
33+
private static final ConcurrentMap<String, String> cachedDBSQLVersions =
34+
new ConcurrentHashMap<>();
35+
2936
private static final String[] VERSION_PARTS = VERSION.split("[.-]");
3037

3138
public static String getVersion() {
@@ -46,7 +53,7 @@ public static void resolveMetadataClient(IDatabricksConnection connection)
4653
LOGGER.warn("Empty metadata client is being used.");
4754
connection.getSession().setEmptyMetadataClient();
4855
}
49-
ensureUpdatedDBRVersionInUse(connection);
56+
ensureUpdatedDBSQLVersionInUse(connection);
5057
}
5158

5259
public static void setUpLogging(IDatabricksConnectionContext connectionContext)
@@ -73,6 +80,13 @@ public static String getRootCauseMessage(Throwable e) {
7380
: e.getMessage();
7481
}
7582

83+
/**
84+
* Returns whether the driver is running against fake services based on request/response stubs.
85+
*/
86+
public static boolean isRunningAgainstFake() {
87+
return Boolean.parseBoolean(System.getProperty(IS_FAKE_SERVICE_TEST_PROP));
88+
}
89+
7690
private static Throwable getRootCause(Throwable throwable) {
7791
Throwable cause;
7892
while ((cause = throwable.getCause()) != null && cause != throwable) {
@@ -82,38 +96,58 @@ private static Throwable getRootCause(Throwable throwable) {
8296
}
8397

8498
@VisibleForTesting
85-
static void ensureUpdatedDBRVersionInUse(IDatabricksConnection connection)
99+
static void ensureUpdatedDBSQLVersionInUse(IDatabricksConnection connection)
86100
throws DatabricksValidationException {
87101
if (connection.getConnectionContext().getClientType() != DatabricksClientType.SQL_EXEC
88102
|| isRunningAgainstFake()) {
89103
// Check applicable only for SEA flow
90104
return;
91105
}
92-
String dbrVersion = getDBRVersion(connection);
93-
if (!doesDriverSupportSEA(dbrVersion)) {
106+
String dbsqlVersion = getDBSQLVersionCached(connection);
107+
if (!doesDriverSupportSEA(dbsqlVersion)) {
94108
String errorMessage =
95109
String.format(
96-
"Unsupported DBR version %s. Please update your compute to use the latest DBR version.",
97-
dbrVersion);
110+
"Unsupported DBSQL version %s. Please update your compute to use the latest DBSQL version.",
111+
dbsqlVersion);
98112
LOGGER.error(errorMessage);
99113
throw new DatabricksValidationException(errorMessage);
100114
}
101115
}
102116

103-
private static String getDBRVersion(IDatabricksConnection connection) {
117+
private static String getDBSQLVersionCached(IDatabricksConnection connection) {
118+
String httpPath = connection.getConnectionContext().getHttpPath();
119+
String version = cachedDBSQLVersions.get(httpPath);
120+
if (version != null) {
121+
LOGGER.debug("Using cached DBSQL Version for path %s: %s", httpPath, version);
122+
return version;
123+
}
124+
125+
synchronized (DriverUtil.class) {
126+
version = cachedDBSQLVersions.get(httpPath);
127+
if (version != null) {
128+
return version;
129+
}
130+
131+
version = queryDBSQLVersion(connection);
132+
cachedDBSQLVersions.put(httpPath, version);
133+
return version;
134+
}
135+
}
136+
137+
private static String queryDBSQLVersion(IDatabricksConnection connection) {
104138
try (ResultSet resultSet = connection.createStatement().executeQuery(DBSQL_VERSION_SQL)) {
105139
resultSet.next();
106-
String dbrVersion = resultSet.getString(1);
107-
LOGGER.debug("DBR Version in use: %s", dbrVersion);
108-
return dbrVersion;
140+
String dbsqlVersion = resultSet.getString(1);
141+
LOGGER.debug("DBSQL Version in use: %s", dbsqlVersion);
142+
return dbsqlVersion;
109143
} catch (Exception e) {
110144
LOGGER.info(
111-
"Error retrieving DBR version: {%s}. Defaulting to minimum supported version.", e);
112-
return getDefaultDBRVersion();
145+
"Error retrieving DBSQL version: {%s}. Defaulting to minimum supported version.", e);
146+
return getDefaultDBSQLVersion();
113147
}
114148
}
115149

116-
private static String getDefaultDBRVersion() {
150+
private static String getDefaultDBSQLVersion() {
117151
return DBSQL_MIN_MAJOR_VERSION_FOR_SEA_SUPPORT + "." + DBSQL_MIN_MINOR_VERSION_FOR_SEA_SUPPORT;
118152
}
119153

@@ -127,7 +161,8 @@ private static boolean doesDriverSupportSEA(String dbsqlVersion) {
127161
return majorVersion > DBSQL_MIN_MAJOR_VERSION_FOR_SEA_SUPPORT;
128162
}
129163

130-
public static boolean isRunningAgainstFake() {
131-
return Boolean.parseBoolean(System.getProperty(IS_FAKE_SERVICE_TEST_PROP));
164+
@VisibleForTesting
165+
static void clearDBSQLVersionCache() {
166+
cachedDBSQLVersions.clear();
132167
}
133168
}

‎src/test/java/com/databricks/jdbc/common/util/DriverUtilTest.java‎

Lines changed: 54 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,8 @@
11
package com.databricks.jdbc.common.util;
22

3+
import static com.databricks.jdbc.common.util.DriverUtil.DBSQL_VERSION_SQL;
34
import static org.junit.jupiter.api.Assertions.*;
4-
import static org.mockito.Mockito.when;
5+
import static org.mockito.Mockito.*;
56

67
import com.databricks.jdbc.api.IDatabricksConnection;
78
import com.databricks.jdbc.api.IDatabricksConnectionContext;
@@ -10,6 +11,7 @@
1011
import com.databricks.jdbc.common.DatabricksClientType;
1112
import com.databricks.jdbc.exception.DatabricksValidationException;
1213
import java.sql.SQLException;
14+
import org.junit.jupiter.api.BeforeEach;
1315
import org.junit.jupiter.api.Test;
1416
import org.junit.jupiter.api.extension.ExtendWith;
1517
import org.junit.jupiter.params.ParameterizedTest;
@@ -19,11 +21,18 @@
1921

2022
@ExtendWith(MockitoExtension.class)
2123
public class DriverUtilTest {
24+
private static final String TEST_HTTP_PATH = "/sql/1.0/warehouses/warehouse_id";
25+
private static final String TEST_HTTP_PATH_2 = "/sql/1.0/warehouses/warehouse_id_2";
2226
@Mock IDatabricksConnection connection;
2327
@Mock IDatabricksConnectionContext connectionContext;
2428
@Mock IDatabricksStatement statement;
2529
@Mock IDatabricksResultSet resultSet;
2630

31+
@BeforeEach
32+
void setUp() {
33+
DriverUtil.clearDBSQLVersionCache(); // Clear the cache before each test case
34+
}
35+
2736
@ParameterizedTest
2837
@CsvSource({
2938
"2023.99, true",
@@ -35,22 +44,63 @@ public class DriverUtilTest {
3544
void testDriverSupportInSEA(String dbsqlVersion, boolean throwsError) throws SQLException {
3645
when(connection.getConnectionContext()).thenReturn(connectionContext);
3746
when(connectionContext.getClientType()).thenReturn(DatabricksClientType.SQL_EXEC);
47+
when(connectionContext.getHttpPath()).thenReturn(TEST_HTTP_PATH);
3848
when(connection.createStatement()).thenReturn(statement);
3949
when(statement.executeQuery("SELECT current_version().dbsql_version")).thenReturn(resultSet);
4050
when(resultSet.getString(1)).thenReturn(dbsqlVersion);
51+
4152
if (throwsError) {
4253
assertThrows(
4354
DatabricksValidationException.class,
44-
() -> DriverUtil.ensureUpdatedDBRVersionInUse(connection));
55+
() -> DriverUtil.ensureUpdatedDBSQLVersionInUse(connection));
4556
} else {
46-
assertDoesNotThrow(() -> DriverUtil.ensureUpdatedDBRVersionInUse(connection));
57+
assertDoesNotThrow(() -> DriverUtil.ensureUpdatedDBSQLVersionInUse(connection));
4758
}
4859
}
4960

5061
@Test
5162
void testDriverSupportInThrift() {
5263
when(connection.getConnectionContext()).thenReturn(connectionContext);
5364
when(connectionContext.getClientType()).thenReturn(DatabricksClientType.THRIFT);
54-
assertDoesNotThrow(() -> DriverUtil.ensureUpdatedDBRVersionInUse(connection));
65+
66+
assertDoesNotThrow(() -> DriverUtil.ensureUpdatedDBSQLVersionInUse(connection));
67+
}
68+
69+
@Test
70+
void testCacheIsSeparateForDifferentHttpPaths() throws SQLException {
71+
when(connection.getConnectionContext()).thenReturn(connectionContext);
72+
when(connectionContext.getClientType()).thenReturn(DatabricksClientType.SQL_EXEC);
73+
when(connection.createStatement()).thenReturn(statement);
74+
when(statement.executeQuery(DBSQL_VERSION_SQL)).thenReturn(resultSet);
75+
76+
// First connection with TEST_HTTP_PATH
77+
when(connectionContext.getHttpPath()).thenReturn(TEST_HTTP_PATH);
78+
when(resultSet.getString(1)).thenReturn("2024.30");
79+
DriverUtil.ensureUpdatedDBSQLVersionInUse(connection);
80+
81+
// Second connection with TEST_HTTP_PATH_2
82+
when(connectionContext.getHttpPath()).thenReturn(TEST_HTTP_PATH_2);
83+
when(resultSet.getString(1)).thenReturn("2024.31");
84+
DriverUtil.ensureUpdatedDBSQLVersionInUse(connection);
85+
86+
// Verify that the statement was executed twice (once for each path)
87+
verify(statement, times(2)).executeQuery(DBSQL_VERSION_SQL);
88+
}
89+
90+
@Test
91+
void testCacheIsReusedForSameHttpPath() throws SQLException {
92+
when(connection.getConnectionContext()).thenReturn(connectionContext);
93+
when(connectionContext.getClientType()).thenReturn(DatabricksClientType.SQL_EXEC);
94+
when(connectionContext.getHttpPath()).thenReturn(TEST_HTTP_PATH);
95+
when(connection.createStatement()).thenReturn(statement);
96+
when(statement.executeQuery(DBSQL_VERSION_SQL)).thenReturn(resultSet);
97+
when(resultSet.getString(1)).thenReturn("2024.30");
98+
99+
// Call twice with the same HTTP path
100+
DriverUtil.ensureUpdatedDBSQLVersionInUse(connection);
101+
DriverUtil.ensureUpdatedDBSQLVersionInUse(connection);
102+
103+
// Verify that the statement was executed only once
104+
verify(statement, times(1)).executeQuery(DBSQL_VERSION_SQL);
55105
}
56106
}

0 commit comments

Comments
 (0)