Skip to content

Commit e41db81

Browse files
jayantsing-dbjayantsing-db
authored
[ES-1436915] Pass context to public classes and their dependencies (#806)
For public classes and interfaces, we can't make assumptions about the threading environment in which they will be used. Specifically, host applications may utilize these classes, such as ResultSet, in a thread that is different from the one that created the JDBC connection resource. In the current implementation of certain APIs, along with the execution flow of DatabricksResultSet, there is an expectation that the thread will automatically be populated with an internal context object. This assumption is problematic, given the scenario described above, where threads may not have the expected context object set by default. To address this issue, this pull request explicitly passes the context object to the areas where it was previously missing, correcting the faulty implementation. Key changes in this PR: - The metadata result set builder class has been refactored from a static utility class to an instance-based class, with the context object now stored as a member variable. - The context object is now passed explicitly to the DatabricksResultSet class, ensuring that it functions correctly in multithreaded environments. These changes eliminate the previous assumption about the thread's context
1 parent 2fba0bb commit e41db81

14 files changed

+281
-179
lines changed

‎src/main/java/com/databricks/jdbc/api/impl/DatabricksDatabaseMetaData.java‎

Lines changed: 15 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,6 @@
22

33
import static com.databricks.jdbc.common.MetadataResultConstants.*;
44
import static com.databricks.jdbc.dbclient.impl.common.CommandConstants.METADATA_STATEMENT_ID;
5-
import static com.databricks.jdbc.dbclient.impl.common.MetadataResultSetBuilder.getFunctionsResult;
65
import static com.databricks.jdbc.dbclient.impl.sqlexec.ResultConstants.CLIENT_INFO_PROPERTIES_RESULT;
76

87
import com.databricks.jdbc.api.impl.converters.ConverterHelper;
@@ -40,10 +39,12 @@ public class DatabricksDatabaseMetaData implements DatabaseMetaData {
4039
"CURDATE,CURRENT_DATE,CURRENT_TIME,CURRENT_TIMESTAMP,CURTIME,DAYNAME,DAYOFMONTH,DAYOFWEEK,DAYOFYEAR,HOUR,MINUTE,MONTH,MONTHNAME,NOW,QUARTER,SECOND,TIMESTAMPADD,TIMESTAMPDIFF,WEEK,YEAR";
4140
private final IDatabricksConnectionInternal connection;
4241
private final IDatabricksSession session;
42+
private final MetadataResultSetBuilder metadataResultSetBuilder;
4343

4444
public DatabricksDatabaseMetaData(IDatabricksConnectionInternal connection) {
4545
this.connection = connection;
4646
this.session = connection.getSession();
47+
this.metadataResultSetBuilder = new MetadataResultSetBuilder(session.getConnectionContext());
4748
}
4849

4950
@Override
@@ -944,7 +945,7 @@ public ResultSet getProcedureColumns(
944945
catalog, schemaPattern, procedureNamePattern, columnNamePattern));
945946
throwExceptionIfConnectionIsClosed();
946947

947-
return MetadataResultSetBuilder.getResultSetWithGivenRowsAndColumns(
948+
return metadataResultSetBuilder.getResultSetWithGivenRowsAndColumns(
948949
PROCEDURE_COLUMNS_COLUMNS,
949950
new ArrayList<>(),
950951
METADATA_STATEMENT_ID,
@@ -969,7 +970,7 @@ public ResultSet getTables(
969970
public ResultSet getSchemas() throws SQLException {
970971
LOGGER.debug("public ResultSet getSchemas()");
971972
if (session.getConnectionContext().getClientType() == DatabricksClientType.SEA) {
972-
return MetadataResultSetBuilder.getSchemasResult(null);
973+
return metadataResultSetBuilder.getSchemasResult(null);
973974
}
974975
return getSchemas(null /* catalog */, null /* schema pattern */);
975976
}
@@ -1011,7 +1012,7 @@ public ResultSet getColumnPrivileges(
10111012
"public ResultSet getColumnPrivileges(String catalog = {%s}, String schema = {%s}, String table = {%s}, String columnNamePattern = {%s})",
10121013
catalog, schema, table, columnNamePattern));
10131014
throwExceptionIfConnectionIsClosed();
1014-
return MetadataResultSetBuilder.getResultSetWithGivenRowsAndColumns(
1015+
return metadataResultSetBuilder.getResultSetWithGivenRowsAndColumns(
10151016
COLUMN_PRIVILEGES_COLUMNS,
10161017
new ArrayList<>(),
10171018
METADATA_STATEMENT_ID,
@@ -1026,7 +1027,7 @@ public ResultSet getTablePrivileges(String catalog, String schemaPattern, String
10261027
"public ResultSet getTablePrivileges(String catalog = {%s}, String schemaPattern = {%s}, String tableNamePattern = {%s})",
10271028
catalog, schemaPattern, tableNamePattern));
10281029
throwExceptionIfConnectionIsClosed();
1029-
return MetadataResultSetBuilder.getResultSetWithGivenRowsAndColumns(
1030+
return metadataResultSetBuilder.getResultSetWithGivenRowsAndColumns(
10301031
TABLE_PRIVILEGES_COLUMNS,
10311032
new ArrayList<>(),
10321033
METADATA_STATEMENT_ID,
@@ -1045,7 +1046,7 @@ public ResultSet getBestRowIdentifier(
10451046
case 0:
10461047
case 1:
10471048
case 2:
1048-
return MetadataResultSetBuilder.getResultSetWithGivenRowsAndColumns(
1049+
return metadataResultSetBuilder.getResultSetWithGivenRowsAndColumns(
10491050
BEST_ROW_IDENTIFIER_COLUMNS,
10501051
new ArrayList<>(),
10511052
METADATA_STATEMENT_ID,
@@ -1064,7 +1065,7 @@ public ResultSet getVersionColumns(String catalog, String schema, String table)
10641065
"public ResultSet getVersionColumns(String catalog = {%s}, String schema = {%s}, String table = {%s})",
10651066
catalog, schema, table));
10661067
throwExceptionIfConnectionIsClosed();
1067-
return MetadataResultSetBuilder.getResultSetWithGivenRowsAndColumns(
1068+
return metadataResultSetBuilder.getResultSetWithGivenRowsAndColumns(
10681069
VERSION_COLUMNS_COLUMNS,
10691070
new ArrayList<>(),
10701071
METADATA_STATEMENT_ID,
@@ -1155,7 +1156,7 @@ public ResultSet getIndexInfo(
11551156
catalog, schema, table, unique, approximate));
11561157
throwExceptionIfConnectionIsClosed();
11571158

1158-
return MetadataResultSetBuilder.getResultSetWithGivenRowsAndColumns(
1159+
return metadataResultSetBuilder.getResultSetWithGivenRowsAndColumns(
11591160
INDEX_INFO_COLUMNS, new ArrayList<>(), METADATA_STATEMENT_ID, CommandName.GET_INDEX_INFO);
11601161
}
11611162

@@ -1333,7 +1334,7 @@ public ResultSet getSuperTypes(String catalog, String schemaPattern, String type
13331334
"public ResultSet getSuperTypes(String catalog = {%s}, String schemaPattern = {%s}, String typeNamePattern = {%s})",
13341335
catalog, schemaPattern, typeNamePattern));
13351336
throwExceptionIfConnectionIsClosed();
1336-
return MetadataResultSetBuilder.getResultSetWithGivenRowsAndColumns(
1337+
return metadataResultSetBuilder.getResultSetWithGivenRowsAndColumns(
13371338
SUPER_TYPES_COLUMNS, new ArrayList<>(), METADATA_STATEMENT_ID, CommandName.GET_SUPER_TYPES);
13381339
}
13391340

@@ -1346,7 +1347,7 @@ public ResultSet getSuperTables(String catalog, String schemaPattern, String tab
13461347
catalog, schemaPattern, tableNamePattern));
13471348
throwExceptionIfConnectionIsClosed();
13481349

1349-
return MetadataResultSetBuilder.getResultSetWithGivenRowsAndColumns(
1350+
return metadataResultSetBuilder.getResultSetWithGivenRowsAndColumns(
13501351
SUPER_TABLES_COLUMNS,
13511352
new ArrayList<>(),
13521353
METADATA_STATEMENT_ID,
@@ -1361,7 +1362,7 @@ public ResultSet getAttributes(
13611362
String.format(
13621363
"public ResultSet getAttributes(String catalog = {%s}, String schemaPattern = {%s}, String typeNamePattern = {%s}, String attributeNamePattern = {%s})",
13631364
catalog, schemaPattern, typeNamePattern, attributeNamePattern));
1364-
return MetadataResultSetBuilder.getResultSetWithGivenRowsAndColumns(
1365+
return metadataResultSetBuilder.getResultSetWithGivenRowsAndColumns(
13651366
ATTRIBUTES_COLUMNS, new ArrayList<>(), METADATA_STATEMENT_ID, CommandName.GET_ATTRIBUTES);
13661367
}
13671368

@@ -1478,7 +1479,7 @@ public ResultSet getFunctions(String catalog, String schemaPattern, String funct
14781479
.listFunctions(session, catalog, schemaPattern, functionNamePattern);
14791480
} catch (Exception e) {
14801481
LOGGER.error("Unable to fetch functions, returning empty result set", e);
1481-
return getFunctionsResult(catalog, List.of());
1482+
return metadataResultSetBuilder.getFunctionsResult(catalog, List.of());
14821483
}
14831484
}
14841485

@@ -1492,7 +1493,7 @@ public ResultSet getFunctionColumns(
14921493
catalog, schemaPattern, functionNamePattern, columnNamePattern));
14931494
throwExceptionIfConnectionIsClosed();
14941495

1495-
return MetadataResultSetBuilder.getResultSetWithGivenRowsAndColumns(
1496+
return metadataResultSetBuilder.getResultSetWithGivenRowsAndColumns(
14961497
FUNCTION_COLUMNS_COLUMNS,
14971498
new ArrayList<>(),
14981499
METADATA_STATEMENT_ID,
@@ -1509,7 +1510,7 @@ public ResultSet getPseudoColumns(
15091510
catalog, schemaPattern, tableNamePattern, columnNamePattern));
15101511
throwExceptionIfConnectionIsClosed();
15111512

1512-
return MetadataResultSetBuilder.getResultSetWithGivenRowsAndColumns(
1513+
return metadataResultSetBuilder.getResultSetWithGivenRowsAndColumns(
15131514
PSEUDO_COLUMNS_COLUMNS,
15141515
new ArrayList<>(),
15151516
METADATA_STATEMENT_ID,

‎src/main/java/com/databricks/jdbc/api/impl/DatabricksResultSet.java‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -89,7 +89,10 @@ public DatabricksResultSet(
8989
resultData, resultManifest, statementId, session, parentStatement);
9090
this.resultSetMetaData =
9191
new DatabricksResultSetMetaData(
92-
statementId, resultManifest, resultData.getExternalLinks() != null);
92+
statementId,
93+
resultManifest,
94+
resultData.getExternalLinks() != null,
95+
session.getConnectionContext());
9396
switch (resultManifest.getFormat()) {
9497
case ARROW_STREAM:
9598
this.resultSetType = ResultSetType.SEA_ARROW_ENABLED;
@@ -156,7 +159,8 @@ public DatabricksResultSet(
156159
resultsResp.getResultSetMetadata(),
157160
rowSize,
158161
executionResult.getChunkCount(),
159-
arrowMetadata);
162+
arrowMetadata,
163+
session.getConnectionContext());
160164
switch (resultsResp.getResultSetMetadata().getResultFormat()) {
161165
case COLUMN_BASED_SET:
162166
this.resultSetType = ResultSetType.THRIFT_INLINE;

‎src/main/java/com/databricks/jdbc/api/impl/DatabricksResultSetMetaData.java‎

Lines changed: 15 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -10,12 +10,13 @@
1010
import static com.databricks.jdbc.common.util.DatabricksTypeUtil.TIMESTAMP_NTZ;
1111
import static com.databricks.jdbc.common.util.DatabricksTypeUtil.VARIANT;
1212
import static com.databricks.jdbc.common.util.DatabricksTypeUtil.getBasePrecisionAndScale;
13-
import static com.databricks.jdbc.dbclient.impl.common.MetadataResultSetBuilder.stripTypeName;
1413

14+
import com.databricks.jdbc.api.internal.IDatabricksConnectionContext;
1515
import com.databricks.jdbc.common.AccessType;
1616
import com.databricks.jdbc.common.Nullable;
1717
import com.databricks.jdbc.common.util.DatabricksTypeUtil;
1818
import com.databricks.jdbc.common.util.WrapperUtil;
19+
import com.databricks.jdbc.dbclient.impl.common.MetadataResultSetBuilder;
1920
import com.databricks.jdbc.dbclient.impl.common.StatementId;
2021
import com.databricks.jdbc.log.JdbcLogger;
2122
import com.databricks.jdbc.log.JdbcLoggerFactory;
@@ -38,6 +39,7 @@ public class DatabricksResultSetMetaData implements ResultSetMetaData {
3839
private static final JdbcLogger LOGGER =
3940
JdbcLoggerFactory.getLogger(DatabricksResultSetMetaData.class);
4041
private final StatementId statementId;
42+
private IDatabricksConnectionContext ctx;
4143
private final ImmutableList<ImmutableDatabricksColumn> columns;
4244
private final ImmutableMap<String, Integer> columnNameIndex;
4345
private final long totalRows;
@@ -54,10 +56,15 @@ public class DatabricksResultSetMetaData implements ResultSetMetaData {
5456
* used)
5557
*/
5658
public DatabricksResultSetMetaData(
57-
StatementId statementId, ResultManifest resultManifest, boolean usesExternalLinks) {
59+
StatementId statementId,
60+
ResultManifest resultManifest,
61+
boolean usesExternalLinks,
62+
IDatabricksConnectionContext ctx) {
63+
this.ctx = ctx;
5864
this.statementId = statementId;
5965
Map<String, Integer> columnNameToIndexMap = new HashMap<>();
6066
ImmutableList.Builder<ImmutableDatabricksColumn> columnsBuilder = ImmutableList.builder();
67+
MetadataResultSetBuilder metadataResultSetBuilder = new MetadataResultSetBuilder(ctx);
6168

6269
int currIndex = 0;
6370
if (resultManifest.getIsVolumeOperation() != null && resultManifest.getIsVolumeOperation()) {
@@ -99,7 +106,7 @@ public DatabricksResultSetMetaData(
99106
.columnTypeClassName(DatabricksTypeUtil.getColumnTypeClassName(columnTypeName))
100107
.columnType(columnType)
101108
.columnTypeText(
102-
stripTypeName(
109+
metadataResultSetBuilder.stripTypeName(
103110
columnInfo
104111
.getTypeText())) // store base type eg. DECIMAL instead of DECIMAL(7,2)
105112
.typePrecision(precision)
@@ -137,7 +144,9 @@ public DatabricksResultSetMetaData(
137144
TGetResultSetMetadataResp resultManifest,
138145
long rows,
139146
long chunkCount,
140-
List<String> arrowMetadata) {
147+
List<String> arrowMetadata,
148+
IDatabricksConnectionContext ctx) {
149+
this.ctx = ctx;
141150
this.statementId = statementId;
142151
Map<String, Integer> columnNameToIndexMap = new HashMap<>();
143152
ImmutableList.Builder<ImmutableDatabricksColumn> columnsBuilder = ImmutableList.builder();
@@ -519,7 +528,7 @@ public Long getChunkCount() {
519528
}
520529

521530
public int[] getPrecisionAndScale(ColumnInfo columnInfo, int columnType) {
522-
int[] result = getBasePrecisionAndScale(columnType);
531+
int[] result = getBasePrecisionAndScale(columnType, ctx);
523532
if (columnInfo.getTypePrecision() != null) {
524533
result[0] = Math.toIntExact(columnInfo.getTypePrecision()); // precision
525534
result[1] = Math.toIntExact(columnInfo.getTypeScale()); // scale
@@ -528,7 +537,7 @@ public int[] getPrecisionAndScale(ColumnInfo columnInfo, int columnType) {
528537
}
529538

530539
public int[] getPrecisionAndScale(TColumnDesc columnInfo, int columnType) {
531-
int[] result = getBasePrecisionAndScale(columnType);
540+
int[] result = getBasePrecisionAndScale(columnType, ctx);
532541
if (columnInfo.getTypeDesc() != null && columnInfo.getTypeDesc().getTypesSize() > 0) {
533542
TTypeEntry tTypeEntry = columnInfo.getTypeDesc().getTypes().get(0);
534543
if (tTypeEntry.isSetPrimitiveEntry()

‎src/main/java/com/databricks/jdbc/api/impl/DatabricksSession.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -259,7 +259,7 @@ public IDatabricksConnectionContext getConnectionContext() {
259259

260260
@Override
261261
public void setEmptyMetadataClient() {
262-
databricksMetadataClient = new DatabricksEmptyMetadataClient();
262+
databricksMetadataClient = new DatabricksEmptyMetadataClient(connectionContext);
263263
}
264264

265265
@Override

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

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -198,9 +198,8 @@ public static String getColumnTypeClassName(ColumnInfoTypeName typeName) {
198198
/*
199199
* Returns default precision and scale based on column type. For string columns, returns the default string col length in precision.
200200
*/
201-
public static int[] getBasePrecisionAndScale(int columnType) {
201+
public static int[] getBasePrecisionAndScale(int columnType, IDatabricksConnectionContext ctx) {
202202
if (columnType == Types.VARCHAR || columnType == Types.CHAR) {
203-
IDatabricksConnectionContext ctx = DatabricksThreadContextHolder.getConnectionContext();
204203
return new int[] {ctx.getDefaultStringColumnLength(), 0};
205204
}
206205
return new int[] {

0 commit comments

Comments
 (0)