Fix leak + interpolation

Author: samikshya-db

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

@@ -4,11 +4,11 @@
 import static com.databricks.jdbc.common.util.DatabricksTypeUtil.getDatabricksTypeFromSQLType;
 import static com.databricks.jdbc.common.util.DatabricksTypeUtil.inferDatabricksType;
 import static com.databricks.jdbc.common.util.SQLInterpolator.interpolateSQL;
+import static com.databricks.jdbc.common.util.SQLInterpolator.surroundPlaceholdersWithQuotes;
 import static com.databricks.jdbc.common.util.ValidationUtil.throwErrorIfNull;
 
 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;
@@ -45,6 +45,18 @@ public DatabricksPreparedStatement(DatabricksConnection connection, String sql)
     this.databricksBatchParameterMetaData = new ArrayList<>();
   }
 
+  DatabricksPreparedStatement(
+      DatabricksConnection connection,
+      String sql,
+      boolean interpolateParameters,
+      DatabricksParameterMetaData databricksParameterMetaData) {
+    super(connection);
+    this.sql = sql;
+    this.interpolateParameters = interpolateParameters;
+    this.databricksParameterMetaData = databricksParameterMetaData;
+    this.databricksBatchParameterMetaData = new ArrayList<>();
+  }
+
   @Override
   public ResultSet executeQuery() throws SQLException {
     LOGGER.debug("public ResultSet executeQuery()");
@@ -834,18 +846,20 @@ private DatabricksResultSet interpolateIfRequiredAndExecute(StatementType statem
    * @throws DatabricksSQLException if there is an error executing the DESCRIBE QUERY command
    */
   private ResultSetMetaData getMetaDataFromDescribeQuery() throws DatabricksSQLException {
-
-    try (DatabricksResultSet metadataResultSet = executeDescribeQueryCommand()) {
+    String describeQuerySQL = "DESCRIBE QUERY " + surroundPlaceholdersWithQuotes(sql);
+    try (DatabricksPreparedStatement preparedStatement =
+            new DatabricksPreparedStatement(
+                connection, describeQuerySQL, interpolateParameters, databricksParameterMetaData);
+        ResultSet metadataResultSet = preparedStatement.executeQuery(); ) {
       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()),
+          preparedStatement.getStatementId(),
           columnNames,
           columnDataTypes,
           this.connection.getConnectionContext());
@@ -856,18 +870,4 @@ private ResultSetMetaData getMetaDataFromDescribeQuery() throws DatabricksSQLExc
           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/common/util/SQLInterpolator.java

@@ -67,4 +67,20 @@ public static String interpolateSQL(String sql, Map<Integer, ImmutableSqlParamet
     }
     return sb.toString();
   }
+
+  /**
+   * Surrounds unquoted placeholders (?) with single quotes, preserving already quoted ones. This is
+   * crucial for DESCRIBE QUERY commands as unquoted placeholders will cause a parse_syntax_error.
+   */
+  public static String surroundPlaceholdersWithQuotes(String sql) {
+    if (sql == null || sql.isEmpty()) {
+      return sql;
+    }
+    // First replace '?' with a temporary marker
+    String temp = sql.replace("'?'", "##QUOTED##");
+    // Then replace remaining ? with '?'
+    String result = temp.replace("?", "'?'");
+    // Finally restore the original quoted placeholders
+    return result.replace("##QUOTED##", "'?'");
+  }
 }

src/test/java/com/databricks/jdbc/common/util/SQLInterpolatorTest.java

@@ -9,7 +9,11 @@
 import com.databricks.jdbc.exception.DatabricksValidationException;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.stream.Stream;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.Arguments;
+import org.junit.jupiter.params.provider.MethodSource;
 
 public class SQLInterpolatorTest {
 
@@ -87,4 +91,52 @@ public void testBinaryType() throws DatabricksValidationException {
     String expected = "INSERT INTO sales (id, data) VALUES (101, X'0102030405')";
     assertEquals(expected, SQLInterpolator.interpolateSQL(sql, params));
   }
+
+  private static Stream<Arguments> providePlaceholderQuotingTestCases() {
+    return Stream.of(
+        // Basic placeholder quoting
+        Arguments.of(
+            "SELECT * FROM table WHERE id = ?",
+            "SELECT * FROM table WHERE id = '?'",
+            "Basic placeholder quoting"),
+
+        // Multiple placeholders
+        Arguments.of(
+            "SELECT * FROM table WHERE id = ? AND name = ?",
+            "SELECT * FROM table WHERE id = '?' AND name = '?'",
+            "Multiple placeholders"),
+
+        // Already quoted placeholders
+        Arguments.of(
+            "SELECT * FROM table WHERE id = '?' AND name = ?",
+            "SELECT * FROM table WHERE id = '?' AND name = '?'",
+            "Already quoted placeholders"),
+
+        // Mixed quoted and unquoted placeholders
+        Arguments.of(
+            "SELECT * FROM table WHERE id = '?' AND name = ? AND age = '?'",
+            "SELECT * FROM table WHERE id = '?' AND name = '?' AND age = '?'",
+            "Mixed quoted and unquoted placeholders"),
+
+        // Null input
+        Arguments.of(null, null, "Null input"),
+
+        // Empty input
+        Arguments.of("", "", "Empty input"),
+
+        // No placeholders
+        Arguments.of("SELECT * FROM table", "SELECT * FROM table", "No placeholders"),
+
+        // Complex query with multiple conditions
+        Arguments.of(
+            "SELECT * FROM table WHERE id = ? AND (name = ? OR age = ?) AND status = ?",
+            "SELECT * FROM table WHERE id = '?' AND (name = '?' OR age = '?') AND status = '?'",
+            "Complex query with multiple conditions"));
+  }
+
+  @ParameterizedTest(name = "{2}")
+  @MethodSource("providePlaceholderQuotingTestCases")
+  public void testSurroundPlaceholdersWithQuotes(String input, String expected, String testName) {
+    assertEquals(expected, SQLInterpolator.surroundPlaceholdersWithQuotes(input), testName);
+  }
 }