[PECOBLR-414]Prepared Statement getMetaData to work without executing Query

NEXT_CHANGELOG.md

@@ -3,7 +3,7 @@
 ## [Unreleased]
 
 ### Added
-- 
+- Support to fetch metadata in PreparedStatement for SELECT queries before executing the query.
 
 ### Updated
 - 

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

@@ -8,6 +8,7 @@
 
 import com.databricks.jdbc.common.StatementType;
 import com.databricks.jdbc.common.util.DatabricksTypeUtil;
+import com.databricks.jdbc.dbclient.impl.common.StatementId;
 import com.databricks.jdbc.exception.*;
 import com.databricks.jdbc.log.JdbcLogger;
 import com.databricks.jdbc.log.JdbcLoggerFactory;
@@ -395,7 +396,14 @@ public ResultSetMetaData getMetaData() throws SQLException {
     LOGGER.debug("public ResultSetMetaData getMetaData()");
     checkIfClosed();
     if (resultSet == null) {
-      return null;
+
+      if (DatabricksStatement.isSelectQuery(sql)) {
+        LOGGER.info(
+            "Fetching metadata before executing the query, some values may not be available");
+        return getMetaDataFromDescribeQuery();
+      } else {
+        return null;
+      }
     }
     return resultSet.getMetaData();
   }
@@ -809,10 +817,57 @@ private DatabricksResultSet interpolateIfRequiredAndExecute(StatementType statem
         this.interpolateParameters
             ? interpolateSQL(sql, this.databricksParameterMetaData.getParameterBindings())
             : sql;
+
     Map<Integer, ImmutableSqlParameter> paramMap =
         this.interpolateParameters
             ? new HashMap<>()
             : this.databricksParameterMetaData.getParameterBindings();
     return executeInternal(interpolatedSql, paramMap, statementType);
   }
+
+  /**
+   * Executes a DESCRIBE QUERY command to retrieve metadata about the SQL query.
+   *
+   * <p>This method is used when the result set is null
+   *
+   * @return a {@link ResultSetMetaData} object containing the metadata of the query.
+   * @throws DatabricksSQLException if there is an error executing the DESCRIBE QUERY command
+   */
+  private ResultSetMetaData getMetaDataFromDescribeQuery() throws DatabricksSQLException {
+
+    try (DatabricksResultSet metadataResultSet = executeDescribeQueryCommand()) {
+      ArrayList<String> columnNames = new ArrayList<>();
+      ArrayList<String> columnDataTypes = new ArrayList<>();
+
+      while (metadataResultSet.next()) {
+        columnNames.add(metadataResultSet.getString(1));
+        columnDataTypes.add(metadataResultSet.getString(2));
+      }
+
+      return new DatabricksResultSetMetaData(
+          StatementId.deserialize(metadataResultSet.getStatementId()),
+          columnNames,
+          columnDataTypes,
+          this.connection.getConnectionContext());
+    } catch (SQLException e) {
+      String errorMessage = "Failed to get query metadata";
+      LOGGER.error(e, errorMessage);
+      throw new DatabricksSQLException(
+          errorMessage, e, DatabricksDriverErrorCode.EXECUTE_STATEMENT_FAILED);
+    }
+  }
+
+  private DatabricksResultSet executeDescribeQueryCommand() throws SQLException {
+    String interpolatedSql =
+        this.interpolateParameters
+            ? interpolateSQL(sql, this.databricksParameterMetaData.getParameterBindings())
+            : sql;
+
+    interpolatedSql = "DESCRIBE QUERY " + interpolatedSql;
+    Map<Integer, ImmutableSqlParameter> paramMap =
+        this.interpolateParameters
+            ? new HashMap<>()
+            : this.databricksParameterMetaData.getParameterBindings();
+    return executeInternal(interpolatedSql, paramMap, StatementType.QUERY);
+  }
 }

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

@@ -30,9 +30,13 @@
 import java.sql.ResultSetMetaData;
 import java.sql.SQLException;
 import java.sql.Types;
+import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
+import java.util.stream.Collectors;
 
 public class DatabricksResultSetMetaData implements ResultSetMetaData {
 
@@ -376,6 +380,82 @@ public DatabricksResultSetMetaData(
     this.isCloudFetchUsed = false;
   }
 
+  /**
+   * Constructs a {@code DatabricksResultSetMetaData} object for metadata result set obtained from
+   * DESCRIBE QUERY Works for both SEA and Thrift flows as result set obtained from DESCRIBE QUERY
+   * is already parsed.
+   *
+   * @param statementId the unique identifier of the SQL statement execution
+   * @param columnNames names of each column
+   * @param columnDataTypes types of each column
+   * @param ctx connection context
+   */
+  public DatabricksResultSetMetaData(
+      StatementId statementId,
+      List<String> columnNames,
+      List<String> columnDataTypes,
+      IDatabricksConnectionContext ctx) {
+    this.ctx = ctx;
+    ImmutableList.Builder<ImmutableDatabricksColumn> columnsBuilder = ImmutableList.builder();
+    Map<String, Integer> columnNameToIndexMap = new HashMap<>();
+    MetadataResultSetBuilder metadataResultSetBuilder = new MetadataResultSetBuilder(ctx);
+
+    // Capitalize all the columnDataTypes
+    columnDataTypes =
+        columnDataTypes.stream()
+            .map(String::toUpperCase)
+            .collect(Collectors.toCollection(ArrayList::new));
+
+    for (int i = 0; i < columnNames.size(); i++) {
+      String columnName = columnNames.get(i);
+      String columnTypeText = columnDataTypes.get(i);
+
+      ColumnInfoTypeName columnTypeName;
+      if (columnTypeText.equalsIgnoreCase(TIMESTAMP_NTZ)) {
+        columnTypeName = ColumnInfoTypeName.TIMESTAMP;
+        columnTypeText = TIMESTAMP;
+      } else if (columnTypeText.equalsIgnoreCase(VARIANT)) {
+        columnTypeName = ColumnInfoTypeName.STRING;
+        columnTypeText = VARIANT;
+      } else {
+        columnTypeName =
+            ColumnInfoTypeName.valueOf(metadataResultSetBuilder.stripBaseTypeName(columnTypeText));
+      }
+
+      int columnType = DatabricksTypeUtil.getColumnType(columnTypeName);
+      int[] precisionAndScale = getPrecisionAndScale(columnTypeText, columnType);
+      int precision = precisionAndScale[0];
+      int scale = precisionAndScale[1];
+
+      ImmutableDatabricksColumn.Builder columnBuilder = getColumnBuilder();
+      columnBuilder
+          .columnName(columnName)
+          .columnTypeClassName(DatabricksTypeUtil.getColumnTypeClassName(columnTypeName))
+          .columnType(columnType)
+          .columnTypeText(
+              metadataResultSetBuilder.stripBaseTypeName(
+                  columnTypeText)) // store base type eg. DECIMAL instead of DECIMAL(7,2), ARRAY
+          // instead of ARRAY<STRING>
+          .typePrecision(precision)
+          .typeScale(scale)
+          .displaySize(DatabricksTypeUtil.getDisplaySize(columnTypeName, precision, scale))
+          .isSearchable(true) // set all columns to be searchable in execute query result set
+          .schemaName(
+              null) // set schema and table name to null, as server do not return these fields.
+          .tableName(null)
+          .isSigned(DatabricksTypeUtil.isSigned(columnTypeName));
+      columnsBuilder.add(columnBuilder.build());
+      // Keep index starting from 1, to be consistent with JDBC convention
+      columnNameToIndexMap.putIfAbsent(columnName, i + 1);
+    }
+    this.statementId = statementId;
+    this.isCloudFetchUsed = false;
+    this.totalRows = -1;
+    this.columns = columnsBuilder.build();
+    this.columnNameIndex = ImmutableMap.copyOf(columnNameToIndexMap);
+    ;
+  }
+
   @Override
   public int getColumnCount() throws SQLException {
     return columns.size();
@@ -527,6 +607,18 @@ public Long getChunkCount() {
     return chunkCount;
   }
 
+  public int[] getPrecisionAndScale(String columnTypeText, int columnType) {
+    int[] result = getBasePrecisionAndScale(columnType, ctx);
+    Pattern pattern = Pattern.compile("decimal\\((\\d+),\\s*(\\d+)\\)", Pattern.CASE_INSENSITIVE);
+    Matcher matcher = pattern.matcher(columnTypeText);
+
+    if (matcher.matches()) {
+      result[0] = Integer.parseInt(matcher.group(1));
+      result[1] = Integer.parseInt(matcher.group(2));
+    }
+    return result;
+  }
+
   public int[] getPrecisionAndScale(ColumnInfo columnInfo, int columnType) {
     int[] result = getBasePrecisionAndScale(columnType, ctx);
     if (columnInfo.getTypePrecision() != null) {

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

@@ -34,7 +34,7 @@ public class DatabricksStatement implements IDatabricksStatement, IDatabricksSta
   private final ExecutorService executor = MoreExecutors.newDirectExecutorService();
 
   private int timeoutInSeconds;
-  private final DatabricksConnection connection;
+  protected final DatabricksConnection connection;
   DatabricksResultSet resultSet;
   private StatementId statementId;
   private boolean isClosed;
@@ -641,8 +641,7 @@ public boolean isSimpleIdentifier(String identifier) throws SQLException {
     return IDatabricksStatement.super.isSimpleIdentifier(identifier);
   }
 
-  @VisibleForTesting
-  static boolean shouldReturnResultSet(String query) {
+  static String trimCommentsAndWhitespaces(String query) {
     if (query == null || query.trim().isEmpty()) {
       throw new DatabricksDriverException(
           "Query cannot be null or empty", DatabricksDriverErrorCode.INPUT_VALIDATION_ERROR);
@@ -653,6 +652,13 @@ static boolean shouldReturnResultSet(String query) {
     trimmedQuery = trimmedQuery.replaceAll("/\\*.*?\\*/", "");
     trimmedQuery = trimmedQuery.replaceAll("\\s+", " ").trim();
 
+    return trimmedQuery;
+  }
+
+  @VisibleForTesting
+  static boolean shouldReturnResultSet(String query) {
+    String trimmedQuery = trimCommentsAndWhitespaces(query);
+
     // Check if the query matches any of the patterns that return a ResultSet
     return SELECT_PATTERN.matcher(trimmedQuery).find()
         || SHOW_PATTERN.matcher(trimmedQuery).find()
@@ -675,6 +681,11 @@ static boolean shouldReturnResultSet(String query) {
     // Otherwise, it should not return a ResultSet
   }
 
+  static boolean isSelectQuery(String query) {
+    String trimmedQuery = trimCommentsAndWhitespaces(query);
+    return SELECT_PATTERN.matcher(trimmedQuery).find();
+  }
+
   DatabricksResultSet executeInternal(
       String sql,
       Map<Integer, ImmutableSqlParameter> params,

src/main/java/com/databricks/jdbc/dbclient/impl/common/MetadataResultSetBuilder.java

@@ -449,6 +449,28 @@ public String stripTypeName(String typeName) {
     if (typeArgumentIndex != -1) {
       return typeName.substring(0, typeArgumentIndex);
     }
+
+    return typeName;
+  }
+
+  public String stripBaseTypeName(String typeName) {
+    if (typeName == null) {
+      return null;
+    }
+    // Checking '<' first and then '(' to handle cases like MAP<STRING,INT>(50)
+
+    // Checks for ARRAY<STRING> -> ARRAY
+    int typeArgumentIndex = typeName.indexOf('<');
+    if (typeArgumentIndex != -1) {
+      return typeName.substring(0, typeArgumentIndex);
+    }
+
+    // Checks for DECIMAL(10,2) -> DECIMAL
+    typeArgumentIndex = typeName.indexOf('(');
+    if (typeArgumentIndex != -1) {
+      return typeName.substring(0, typeArgumentIndex);
+    }
+
     return typeName;
   }
 

src/test/java/com/databricks/jdbc/api/impl/DatabricksPreparedStatementTest.java

@@ -245,6 +245,25 @@ public void testExecuteBatchStatement() throws Exception {
     assertTrue(statement.isClosed());
   }
 
+  @Test
+  public void testGetMetaData_NoResultSet_NonSelectQuery_ReturnNull() throws Exception {
+    IDatabricksConnectionContext connectionContext =
+        DatabricksConnectionContext.parse(JDBC_URL, new Properties());
+    DatabricksConnection connection = new DatabricksConnection(connectionContext, client);
+    DatabricksPreparedStatement statement =
+        new DatabricksPreparedStatement(connection, BATCH_STATEMENT);
+    // Setting to execute a batch of 4 statements
+    for (int i = 1; i <= 4; i++) {
+      statement.setLong(1, 100);
+      statement.setShort(2, (short) 10);
+      statement.setByte(3, (byte) 15);
+      statement.setString(4, "value");
+      statement.addBatch();
+    }
+
+    assertNull(statement.getMetaData());
+  }
+
   @Test
   public void testExecuteBatchStatementThrowsError() throws Exception {
     IDatabricksConnectionContext connectionContext =

src/test/java/com/databricks/jdbc/api/impl/DatabricksResultSetMetaDataTest.java

@@ -20,6 +20,7 @@
 import java.sql.SQLException;
 import java.sql.Types;
 import java.util.*;
+import java.util.stream.Collectors;
 import java.util.stream.Stream;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.Assertions;
@@ -179,6 +180,48 @@ public void testDatabricksResultSetMetaDataInitialization() throws SQLException
     assertEquals(ResultSetMetaData.columnNullable, metaData.isNullable(3));
   }
 
+  @Test
+  public void testDatabricksResultSetMetaDataInitialization_DescribeQuery() throws SQLException {
+
+    // [columnName, columnType, expectedTypeName, expectedIntegerType, expectedPrecision,
+    // expectedScale]
+    Object[][] columnData = {
+      {"col_int", "int", "INT", Types.INTEGER, 10, 0},
+      {"col_string", "string", "STRING", Types.VARCHAR, 255, 0},
+      {"col_decimal", "decimal(10,2)", "DECIMAL", Types.DECIMAL, 10, 2},
+      {"col_date", "date", "DATE", Types.DATE, 10, 0},
+      {"col_timestamp", "timestamp", "TIMESTAMP", Types.TIMESTAMP, 29, 9},
+      {"col_timestamp_ntz", "timestamp_ntz", "TIMESTAMP", Types.TIMESTAMP, 29, 9},
+      {"col_bool", "boolean", "BOOLEAN", Types.BOOLEAN, 1, 0},
+      {"col_binary", "binary", "BINARY", Types.BINARY, 1, 0},
+      {"col_struct", "struct<col_int:int,col_string:string>", "STRUCT", Types.STRUCT, 255, 0},
+      {"col_array", "array<int>", "ARRAY", Types.ARRAY, 255, 0},
+      {"col_map", "map<string,string>", "MAP", Types.VARCHAR, 255, 0},
+      {"col_variant", "variant", "VARIANT", Types.VARCHAR, 255, 0}
+    };
+
+    List<String> columnNames =
+        Arrays.stream(columnData).map(row -> (String) row[0]).collect(Collectors.toList());
+
+    List<String> columnTypes =
+        Arrays.stream(columnData).map(row -> (String) row[1]).collect(Collectors.toList());
+
+    DatabricksResultSetMetaData metaData =
+        new DatabricksResultSetMetaData(STATEMENT_ID, columnNames, columnTypes, connectionContext);
+
+    assertEquals(columnNames.size(), metaData.getColumnCount());
+    // With Describe query we can't determine total rows
+    assertEquals(-1, metaData.getTotalRows());
+
+    for (int i = 0; i < columnData.length; i++) {
+      assertEquals(columnData[i][0], metaData.getColumnName(i + 1));
+      assertEquals(columnData[i][2], metaData.getColumnTypeName(i + 1));
+      assertEquals(columnData[i][3], metaData.getColumnType(i + 1));
+      assertEquals(columnData[i][4], metaData.getPrecision(i + 1));
+      assertEquals(columnData[i][5], metaData.getScale(i + 1));
+    }
+  }
+
   @Test
   public void testColumnsForVolumeOperation() throws SQLException {
     ResultManifest resultManifest = getResultManifest().setIsVolumeOperation(true);

src/test/java/com/databricks/jdbc/api/impl/DatabricksStatementTest.java

@@ -608,6 +608,16 @@ public void testShouldReturnResultSet_CommentSurroundingQuery() {
     assertTrue(DatabricksStatement.shouldReturnResultSet(query));
   }
 
+  @Test
+  public void testIsSelectQuery() {
+    String query =
+        "-- Single-line comment\n/* Multi-line comment */ SELECT * FROM table; /* Another comment */ -- End comment";
+    assertTrue(DatabricksStatement.isSelectQuery(query));
+
+    query = "REMOVE some_data FROM table;";
+    assertFalse(DatabricksStatement.isSelectQuery(query));
+  }
+
   private DatabricksConnection getTestConnection() throws DatabricksSQLException {
     IDatabricksConnectionContext connectionContext =
         DatabricksConnectionContext.parse(JDBC_URL, new Properties());