[ES-1436915] Pass context to public classes and their dependencies

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

@@ -2,7 +2,6 @@
 
 import static com.databricks.jdbc.common.MetadataResultConstants.*;
 import static com.databricks.jdbc.dbclient.impl.common.CommandConstants.METADATA_STATEMENT_ID;
-import static com.databricks.jdbc.dbclient.impl.common.MetadataResultSetBuilder.getFunctionsResult;
 import static com.databricks.jdbc.dbclient.impl.sqlexec.ResultConstants.CLIENT_INFO_PROPERTIES_RESULT;
 
 import com.databricks.jdbc.api.impl.converters.ConverterHelper;
@@ -40,10 +39,12 @@ public class DatabricksDatabaseMetaData implements DatabaseMetaData {
       "CURDATE,CURRENT_DATE,CURRENT_TIME,CURRENT_TIMESTAMP,CURTIME,DAYNAME,DAYOFMONTH,DAYOFWEEK,DAYOFYEAR,HOUR,MINUTE,MONTH,MONTHNAME,NOW,QUARTER,SECOND,TIMESTAMPADD,TIMESTAMPDIFF,WEEK,YEAR";
   private final IDatabricksConnectionInternal connection;
   private final IDatabricksSession session;
+  private final MetadataResultSetBuilder metadataResultSetBuilder;
 
   public DatabricksDatabaseMetaData(IDatabricksConnectionInternal connection) {
     this.connection = connection;
     this.session = connection.getSession();
+    this.metadataResultSetBuilder = new MetadataResultSetBuilder(session.getConnectionContext());
   }
 
   @Override
@@ -944,7 +945,7 @@ public ResultSet getProcedureColumns(
             catalog, schemaPattern, procedureNamePattern, columnNamePattern));
     throwExceptionIfConnectionIsClosed();
 
-    return MetadataResultSetBuilder.getResultSetWithGivenRowsAndColumns(
+    return metadataResultSetBuilder.getResultSetWithGivenRowsAndColumns(
         PROCEDURE_COLUMNS_COLUMNS,
         new ArrayList<>(),
         METADATA_STATEMENT_ID,
@@ -969,7 +970,7 @@ public ResultSet getTables(
   public ResultSet getSchemas() throws SQLException {
     LOGGER.debug("public ResultSet getSchemas()");
     if (session.getConnectionContext().getClientType() == DatabricksClientType.SEA) {
-      return MetadataResultSetBuilder.getSchemasResult(null);
+      return metadataResultSetBuilder.getSchemasResult(null);
     }
     return getSchemas(null /* catalog */, null /* schema pattern */);
   }
@@ -1011,7 +1012,7 @@ public ResultSet getColumnPrivileges(
             "public ResultSet getColumnPrivileges(String catalog = {%s}, String schema = {%s}, String table = {%s}, String columnNamePattern = {%s})",
             catalog, schema, table, columnNamePattern));
     throwExceptionIfConnectionIsClosed();
-    return MetadataResultSetBuilder.getResultSetWithGivenRowsAndColumns(
+    return metadataResultSetBuilder.getResultSetWithGivenRowsAndColumns(
         COLUMN_PRIVILEGES_COLUMNS,
         new ArrayList<>(),
         METADATA_STATEMENT_ID,
@@ -1026,7 +1027,7 @@ public ResultSet getTablePrivileges(String catalog, String schemaPattern, String
             "public ResultSet getTablePrivileges(String catalog = {%s}, String schemaPattern = {%s}, String tableNamePattern = {%s})",
             catalog, schemaPattern, tableNamePattern));
     throwExceptionIfConnectionIsClosed();
-    return MetadataResultSetBuilder.getResultSetWithGivenRowsAndColumns(
+    return metadataResultSetBuilder.getResultSetWithGivenRowsAndColumns(
         TABLE_PRIVILEGES_COLUMNS,
         new ArrayList<>(),
         METADATA_STATEMENT_ID,
@@ -1045,7 +1046,7 @@ public ResultSet getBestRowIdentifier(
       case 0:
       case 1:
       case 2:
-        return MetadataResultSetBuilder.getResultSetWithGivenRowsAndColumns(
+        return metadataResultSetBuilder.getResultSetWithGivenRowsAndColumns(
             BEST_ROW_IDENTIFIER_COLUMNS,
             new ArrayList<>(),
             METADATA_STATEMENT_ID,
@@ -1064,7 +1065,7 @@ public ResultSet getVersionColumns(String catalog, String schema, String table)
             "public ResultSet getVersionColumns(String catalog = {%s}, String schema = {%s}, String table = {%s})",
             catalog, schema, table));
     throwExceptionIfConnectionIsClosed();
-    return MetadataResultSetBuilder.getResultSetWithGivenRowsAndColumns(
+    return metadataResultSetBuilder.getResultSetWithGivenRowsAndColumns(
         VERSION_COLUMNS_COLUMNS,
         new ArrayList<>(),
         METADATA_STATEMENT_ID,
@@ -1155,7 +1156,7 @@ public ResultSet getIndexInfo(
             catalog, schema, table, unique, approximate));
     throwExceptionIfConnectionIsClosed();
 
-    return MetadataResultSetBuilder.getResultSetWithGivenRowsAndColumns(
+    return metadataResultSetBuilder.getResultSetWithGivenRowsAndColumns(
         INDEX_INFO_COLUMNS, new ArrayList<>(), METADATA_STATEMENT_ID, CommandName.GET_INDEX_INFO);
   }
 
@@ -1333,7 +1334,7 @@ public ResultSet getSuperTypes(String catalog, String schemaPattern, String type
             "public ResultSet getSuperTypes(String catalog = {%s}, String schemaPattern = {%s}, String typeNamePattern = {%s})",
             catalog, schemaPattern, typeNamePattern));
     throwExceptionIfConnectionIsClosed();
-    return MetadataResultSetBuilder.getResultSetWithGivenRowsAndColumns(
+    return metadataResultSetBuilder.getResultSetWithGivenRowsAndColumns(
         SUPER_TYPES_COLUMNS, new ArrayList<>(), METADATA_STATEMENT_ID, CommandName.GET_SUPER_TYPES);
   }
 
@@ -1346,7 +1347,7 @@ public ResultSet getSuperTables(String catalog, String schemaPattern, String tab
             catalog, schemaPattern, tableNamePattern));
     throwExceptionIfConnectionIsClosed();
 
-    return MetadataResultSetBuilder.getResultSetWithGivenRowsAndColumns(
+    return metadataResultSetBuilder.getResultSetWithGivenRowsAndColumns(
         SUPER_TABLES_COLUMNS,
         new ArrayList<>(),
         METADATA_STATEMENT_ID,
@@ -1361,7 +1362,7 @@ public ResultSet getAttributes(
         String.format(
             "public ResultSet getAttributes(String catalog = {%s}, String schemaPattern = {%s}, String typeNamePattern = {%s}, String attributeNamePattern = {%s})",
             catalog, schemaPattern, typeNamePattern, attributeNamePattern));
-    return MetadataResultSetBuilder.getResultSetWithGivenRowsAndColumns(
+    return metadataResultSetBuilder.getResultSetWithGivenRowsAndColumns(
         ATTRIBUTES_COLUMNS, new ArrayList<>(), METADATA_STATEMENT_ID, CommandName.GET_ATTRIBUTES);
   }
 
@@ -1478,7 +1479,7 @@ public ResultSet getFunctions(String catalog, String schemaPattern, String funct
           .listFunctions(session, catalog, schemaPattern, functionNamePattern);
     } catch (Exception e) {
       LOGGER.error("Unable to fetch functions, returning empty result set", e);
-      return getFunctionsResult(catalog, List.of());
+      return metadataResultSetBuilder.getFunctionsResult(catalog, List.of());
     }
   }
 
@@ -1492,7 +1493,7 @@ public ResultSet getFunctionColumns(
             catalog, schemaPattern, functionNamePattern, columnNamePattern));
     throwExceptionIfConnectionIsClosed();
 
-    return MetadataResultSetBuilder.getResultSetWithGivenRowsAndColumns(
+    return metadataResultSetBuilder.getResultSetWithGivenRowsAndColumns(
         FUNCTION_COLUMNS_COLUMNS,
         new ArrayList<>(),
         METADATA_STATEMENT_ID,
@@ -1509,7 +1510,7 @@ public ResultSet getPseudoColumns(
             catalog, schemaPattern, tableNamePattern, columnNamePattern));
     throwExceptionIfConnectionIsClosed();
 
-    return MetadataResultSetBuilder.getResultSetWithGivenRowsAndColumns(
+    return metadataResultSetBuilder.getResultSetWithGivenRowsAndColumns(
         PSEUDO_COLUMNS_COLUMNS,
         new ArrayList<>(),
         METADATA_STATEMENT_ID,

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

@@ -89,7 +89,10 @@ public DatabricksResultSet(
               resultData, resultManifest, statementId, session, parentStatement);
       this.resultSetMetaData =
           new DatabricksResultSetMetaData(
-              statementId, resultManifest, resultData.getExternalLinks() != null);
+              statementId,
+              resultManifest,
+              resultData.getExternalLinks() != null,
+              session.getConnectionContext());
       switch (resultManifest.getFormat()) {
         case ARROW_STREAM:
           this.resultSetType = ResultSetType.SEA_ARROW_ENABLED;
@@ -156,7 +159,8 @@ public DatabricksResultSet(
               resultsResp.getResultSetMetadata(),
               rowSize,
               executionResult.getChunkCount(),
-              arrowMetadata);
+              arrowMetadata,
+              session.getConnectionContext());
       switch (resultsResp.getResultSetMetadata().getResultFormat()) {
         case COLUMN_BASED_SET:
           this.resultSetType = ResultSetType.THRIFT_INLINE;

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

@@ -10,12 +10,13 @@
 import static com.databricks.jdbc.common.util.DatabricksTypeUtil.TIMESTAMP_NTZ;
 import static com.databricks.jdbc.common.util.DatabricksTypeUtil.VARIANT;
 import static com.databricks.jdbc.common.util.DatabricksTypeUtil.getBasePrecisionAndScale;
-import static com.databricks.jdbc.dbclient.impl.common.MetadataResultSetBuilder.stripTypeName;
 
+import com.databricks.jdbc.api.internal.IDatabricksConnectionContext;
 import com.databricks.jdbc.common.AccessType;
 import com.databricks.jdbc.common.Nullable;
 import com.databricks.jdbc.common.util.DatabricksTypeUtil;
 import com.databricks.jdbc.common.util.WrapperUtil;
+import com.databricks.jdbc.dbclient.impl.common.MetadataResultSetBuilder;
 import com.databricks.jdbc.dbclient.impl.common.StatementId;
 import com.databricks.jdbc.log.JdbcLogger;
 import com.databricks.jdbc.log.JdbcLoggerFactory;
@@ -38,6 +39,7 @@ public class DatabricksResultSetMetaData implements ResultSetMetaData {
   private static final JdbcLogger LOGGER =
       JdbcLoggerFactory.getLogger(DatabricksResultSetMetaData.class);
   private final StatementId statementId;
+  private IDatabricksConnectionContext ctx;
   private final ImmutableList<ImmutableDatabricksColumn> columns;
   private final ImmutableMap<String, Integer> columnNameIndex;
   private final long totalRows;
@@ -54,10 +56,15 @@ public class DatabricksResultSetMetaData implements ResultSetMetaData {
    *     used)
    */
   public DatabricksResultSetMetaData(
-      StatementId statementId, ResultManifest resultManifest, boolean usesExternalLinks) {
+      StatementId statementId,
+      ResultManifest resultManifest,
+      boolean usesExternalLinks,
+      IDatabricksConnectionContext ctx) {
+    this.ctx = ctx;
     this.statementId = statementId;
     Map<String, Integer> columnNameToIndexMap = new HashMap<>();
     ImmutableList.Builder<ImmutableDatabricksColumn> columnsBuilder = ImmutableList.builder();
+    MetadataResultSetBuilder metadataResultSetBuilder = new MetadataResultSetBuilder(ctx);
 
     int currIndex = 0;
     if (resultManifest.getIsVolumeOperation() != null && resultManifest.getIsVolumeOperation()) {
@@ -99,7 +106,7 @@ public DatabricksResultSetMetaData(
               .columnTypeClassName(DatabricksTypeUtil.getColumnTypeClassName(columnTypeName))
               .columnType(columnType)
               .columnTypeText(
-                  stripTypeName(
+                  metadataResultSetBuilder.stripTypeName(
                       columnInfo
                           .getTypeText())) // store base type eg. DECIMAL instead of DECIMAL(7,2)
               .typePrecision(precision)
@@ -137,7 +144,9 @@ public DatabricksResultSetMetaData(
       TGetResultSetMetadataResp resultManifest,
       long rows,
       long chunkCount,
-      List<String> arrowMetadata) {
+      List<String> arrowMetadata,
+      IDatabricksConnectionContext ctx) {
+    this.ctx = ctx;
     this.statementId = statementId;
     Map<String, Integer> columnNameToIndexMap = new HashMap<>();
     ImmutableList.Builder<ImmutableDatabricksColumn> columnsBuilder = ImmutableList.builder();
@@ -519,7 +528,7 @@ public Long getChunkCount() {
   }
 
   public int[] getPrecisionAndScale(ColumnInfo columnInfo, int columnType) {
-    int[] result = getBasePrecisionAndScale(columnType);
+    int[] result = getBasePrecisionAndScale(columnType, ctx);
     if (columnInfo.getTypePrecision() != null) {
       result[0] = Math.toIntExact(columnInfo.getTypePrecision()); // precision
       result[1] = Math.toIntExact(columnInfo.getTypeScale()); // scale
@@ -528,7 +537,7 @@ public int[] getPrecisionAndScale(ColumnInfo columnInfo, int columnType) {
   }
 
   public int[] getPrecisionAndScale(TColumnDesc columnInfo, int columnType) {
-    int[] result = getBasePrecisionAndScale(columnType);
+    int[] result = getBasePrecisionAndScale(columnType, ctx);
     if (columnInfo.getTypeDesc() != null && columnInfo.getTypeDesc().getTypesSize() > 0) {
       TTypeEntry tTypeEntry = columnInfo.getTypeDesc().getTypes().get(0);
       if (tTypeEntry.isSetPrimitiveEntry()

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

@@ -259,7 +259,7 @@ public IDatabricksConnectionContext getConnectionContext() {
 
   @Override
   public void setEmptyMetadataClient() {
-    databricksMetadataClient = new DatabricksEmptyMetadataClient();
+    databricksMetadataClient = new DatabricksEmptyMetadataClient(connectionContext);
   }
 
   @Override

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

@@ -198,9 +198,8 @@ public static String getColumnTypeClassName(ColumnInfoTypeName typeName) {
   /*
    * Returns default precision and scale based on column type. For string columns, returns the default string col length in precision.
    */
-  public static int[] getBasePrecisionAndScale(int columnType) {
+  public static int[] getBasePrecisionAndScale(int columnType, IDatabricksConnectionContext ctx) {
     if (columnType == Types.VARCHAR || columnType == Types.CHAR) {
-      IDatabricksConnectionContext ctx = DatabricksThreadContextHolder.getConnectionContext();
       return new int[] {ctx.getDefaultStringColumnLength(), 0};
     }
     return new int[] {