Add comment parsing for shouldReturnResultSet and countParameters (#1207)

## Description

Closes #1201
Closes #1254

Goes through each character in the sql string while keeping track of the
current state
(normal, comment, string literal, identifier). Provides an interface for
consumers to iterate over the non-comment characters.

## Testing

See unit tests. I was also testing with this script:

<details>

<summary>Main.java</summary>

```java
import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.Statement;
import java.util.List;

public class Main {

    public static void main(String[] args) throws Exception {
        String url = System.getenv("DATABRICKS_URL");
        if (url == null) {
            System.err.println("DATABRICKS_URL environment variable is not set");
            System.exit(1);
        }

        Class.forName("com.databricks.client.jdbc.Driver");

        List<String> examples;
        examples = List.of(
                // "/* This is a comment */",
                "SELECT 1; /* This is also a comment */",
                """
                                SELECT /* This is a comment
                                that spans multiple lines */ 1;
                                """,
                "SELECT /* Comments are not limited to Latin characters: 评论 😊 */ 1;",
                "SELECT /* Comments /* can be */ nested */ 1;",
                "SELECT /* Quotes in '/*' comments \"/*\" are not special */ */ */ 1;",
                "/* A prefixed comment */ SELECT 1;",
                """
                                SELECT '/* This is not a comment */';
                                /* This is not a comment */
                                """
        );
        for (int i = 0; i < examples.size(); i++) {
            String example = examples.get(i);
            System.out.println("==============================");
            try {
                Connection conn = DriverManager.getConnection(url);
                Statement stmt = conn.createStatement();
                ResultSet rs = stmt.executeQuery(example);
                while (rs.next()) {
                    System.out.println("Result " + i + ": " + rs.getObject(1));
                }
            } catch (Exception e) {
                System.out.println(e.getMessage());
            }
            System.out.println("==============================");
            System.out.println("");
            System.out.println("");
        }

        System.out.println("");
        System.out.println("");

        System.out.println("==============================");
        try {
            Connection conn = DriverManager.getConnection(url);
            Statement stmt = conn.createStatement();
            ResultSet rs = stmt.executeQuery("""
            /* 
            hi */
            SELECT  1;
            """);
            while (rs.next()) {
                System.out.println("Multi line block: " + rs.getObject(1));
            }
        } catch (Exception e) {
            System.out.println("Multi line block: " + e.getMessage());
        }
        System.out.println("==============================");

        System.out.println("");
        System.out.println("");

        System.out.println("==============================");
        try {
            Connection conn = DriverManager.getConnection(url);
            Statement stmt = conn.createStatement();
            ResultSet rs = stmt.executeQuery("""
            /* /* */ */
            SELECT 2;
            """);
            while (rs.next()) {
                System.out.println("Nested: " + rs.getObject(1));
            }
        } catch (Exception e) {
            System.out.println("Nested: " + e.getMessage());
        }
        System.out.println("==============================");

        System.out.println("");
        System.out.println("");

        System.out.println("==============================");
        try {
            Connection conn = DriverManager.getConnection(url);
            PreparedStatement pstmt = conn.prepareStatement("""
            /* ? /* ? */ ? */
            SELECT ? /* ? /* ? */ ? */
            /* ? /* ? */ ? */;
            """);
            pstmt.setString(1, "hello");
            ResultSet rs = pstmt.executeQuery();
            while (rs.next()) {
                System.out.println("Nested param: " + rs.getObject(1));
            }
        } catch (Exception e) {
            System.out.println("Nested param: " + e.getMessage());
        }
        System.out.println("==============================");

    }
}
```

</details>

## Additional Notes to the Reviewer

---------

Signed-off-by: rileythomp <[email protected]>

Author: rileythomp

NEXT_CHANGELOG.md

@@ -12,6 +12,7 @@
 - PECOBLR-1121 Arrow patch to circumvent Arrow issues with JDK 16+.
 
 ### Fixed
+- Fixed `DatabricksParameterMetaData.countParameters` and `DatabricksStatement.trimCommentsAndWhitespaces` with a `SqlCommentParser` utility class.
 - Fixed `rollback()` to throw `SQLException` when called in auto-commit mode (no active transaction), aligning with JDBC spec. Previously it silently sent a ROLLBACK command to the server.
 - Fixed `fetchAutoCommitStateFromServer()` to accept both `"1"`/`"0"` and `"true"`/`"false"` responses from `SET AUTOCOMMIT` query, since different server implementations return different formats.
 - Fixed socket leak in SDK HTTP client that prevented CRaC checkpointing. The SDK's connection pool was not shut down on `connection.close()`, leaving TCP sockets open.

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

@@ -3,6 +3,7 @@
 import static com.databricks.jdbc.common.DatabricksJdbcConstants.EMPTY_STRING;
 
 import com.databricks.jdbc.common.util.DatabricksTypeUtil;
+import com.databricks.jdbc.common.util.SqlCommentParser;
 import com.databricks.jdbc.common.util.WrapperUtil;
 import com.databricks.jdbc.log.JdbcLogger;
 import com.databricks.jdbc.log.JdbcLoggerFactory;
@@ -136,8 +137,7 @@ private ImmutableSqlParameter getObject(int param) {
   }
 
   /**
-   * Counts the number of parameter markers (?) in the SQL statement. This is currently a hacky
-   * implementation that may need improvement for complex SQL.
+   * Counts the number of parameter markers (?) in the SQL statement.
    *
    * @param sql The SQL statement to analyze
    * @return The number of parameter markers in the SQL statement
@@ -147,67 +147,14 @@ private int countParameters(String sql) {
       return 0;
     }
 
-    int count = 0;
-    boolean inSingleQuote = false;
-    boolean inDoubleQuote = false;
-    boolean inLineComment = false;
-    boolean inBlockComment = false;
-
-    for (int i = 0; i < sql.length(); i++) {
-      char c = sql.charAt(i);
-      char next = (i < sql.length() - 1) ? sql.charAt(i + 1) : '\0';
-
-      // Handle comments
-      if (!inSingleQuote && !inDoubleQuote) {
-        if (!inBlockComment && !inLineComment && c == '-' && next == '-') {
-          inLineComment = true;
-          i++; // Skip next dash
-          continue;
-        } else if (!inBlockComment && !inLineComment && c == '/' && next == '*') {
-          inBlockComment = true;
-          i++; // Skip next asterisk
-          continue;
-        } else if (inLineComment && (c == '\n' || c == '\r')) {
-          inLineComment = false;
-        } else if (inBlockComment && c == '*' && next == '/') {
-          inBlockComment = false;
-          i++; // Skip next slash
-          continue;
-        }
-      }
-
-      // Skip if in comment
-      if (inLineComment || inBlockComment) {
-        continue;
-      }
-
-      // Handle quotes with escaped quotes
-      if (c == '\'') {
-        if (!inDoubleQuote) {
-          // Check for escaped single quote
-          if (inSingleQuote && i < sql.length() - 1 && sql.charAt(i + 1) == '\'') {
-            i++; // Skip the escaped quote
-          } else {
-            inSingleQuote = !inSingleQuote;
+    int[] count = {0};
+    SqlCommentParser.forEachNonCommentChar(
+        sql,
+        (state, c) -> {
+          if (state == SqlCommentParser.State.NORMAL && c == '?') {
+            count[0]++;
           }
-        }
-      } else if (c == '"') {
-        if (!inSingleQuote) {
-          // Check for escaped double quote
-          if (inDoubleQuote && i < sql.length() - 1 && sql.charAt(i + 1) == '"') {
-            i++; // Skip the escaped quote
-          } else {
-            inDoubleQuote = !inDoubleQuote;
-          }
-        }
-      }
-
-      // Count parameter markers only when not inside quotes or comments
-      if (c == '?' && !inSingleQuote && !inDoubleQuote) {
-        count++;
-      }
-    }
-
-    return count;
+        });
+    return count[0];
   }
 }

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

@@ -717,12 +717,7 @@ static String trimCommentsAndWhitespaces(String query) {
           "Query cannot be null or empty", DatabricksDriverErrorCode.INPUT_VALIDATION_ERROR);
     }
 
-    // Trim and remove comments and whitespaces.
-    String trimmedQuery = query.trim().replaceAll("(?m)--.*$", "");
-    trimmedQuery = trimmedQuery.replaceAll("/\\*.*?\\*/", "");
-    trimmedQuery = trimmedQuery.replaceAll("\\s+", " ").trim();
-
-    return trimmedQuery;
+    return SqlCommentParser.stripCommentsAndWhitespaces(query);
   }
 
   @VisibleForTesting

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

@@ -0,0 +1,147 @@
+package com.databricks.jdbc.common.util;
+
+/** Utility class for parsing out comments from SQL strings */
+public class SqlCommentParser {
+
+  public enum State {
+    NORMAL,
+    IN_SINGLE_QUOTE,
+    IN_DOUBLE_QUOTE,
+    IN_BACKTICK,
+    IN_LINE_COMMENT,
+    IN_BLOCK_COMMENT
+  }
+
+  @FunctionalInterface
+  public interface SqlCharConsumer {
+    void accept(State state, char c);
+  }
+
+  /**
+   * Iterates over each character in the SQL string while keeping track of comment, string literal,
+   * and identifier state. Each character that is not part of a comment calls the consumer with the
+   * current state and character. Emits a ' ' character after each comment to avoid token fusion.
+   *
+   * @param sql the SQL string to parse
+   * @param consumer called for each visible character with its parsing state
+   */
+  public static void forEachNonCommentChar(String sql, SqlCharConsumer consumer) {
+    if (sql == null || sql.isEmpty()) {
+      return;
+    }
+
+    State state = State.NORMAL;
+    int blockCommentDepth = 0;
+
+    for (int i = 0; i < sql.length(); i++) {
+      char c = sql.charAt(i);
+      char next = (i + 1 < sql.length()) ? sql.charAt(i + 1) : '\0';
+
+      switch (state) {
+        case NORMAL:
+          if (c == '-' && next == '-') {
+            state = State.IN_LINE_COMMENT;
+            i++; // skip second '-'
+          } else if (c == '/' && next == '*') {
+            state = State.IN_BLOCK_COMMENT;
+            blockCommentDepth = 1;
+            i++; // skip '*'
+          } else if (c == '\'') {
+            state = State.IN_SINGLE_QUOTE;
+            consumer.accept(state, c);
+          } else if (c == '"') {
+            state = State.IN_DOUBLE_QUOTE;
+            consumer.accept(state, c);
+          } else if (c == '`') {
+            state = State.IN_BACKTICK;
+            consumer.accept(state, c);
+          } else {
+            consumer.accept(state, c);
+          }
+          break;
+
+        case IN_SINGLE_QUOTE:
+          consumer.accept(state, c);
+          if (c == '\'' && next == '\'') {
+            consumer.accept(state, next);
+            i++; // skip escaped quote
+          } else if (c == '\'') {
+            state = State.NORMAL;
+          }
+          break;
+
+        case IN_DOUBLE_QUOTE:
+          consumer.accept(state, c);
+          if (c == '"' && next == '"') {
+            consumer.accept(state, next);
+            i++; // skip escaped quote
+          } else if (c == '"') {
+            state = State.NORMAL;
+          }
+          break;
+
+        case IN_BACKTICK:
+          consumer.accept(state, c);
+          if (c == '`' && next == '`') {
+            consumer.accept(state, next);
+            i++; // skip escaped backtick
+          } else if (c == '`') {
+            state = State.NORMAL;
+          }
+          break;
+
+        case IN_LINE_COMMENT:
+          if (c == '\n' || c == '\r') {
+            state = State.NORMAL;
+            consumer.accept(state, ' ');
+            if (c == '\r' && next == '\n') {
+              i++; // Treat \r\n as a single line ending
+            }
+          }
+          // else: skip character (part of the comment)
+          break;
+
+        case IN_BLOCK_COMMENT:
+          if (c == '/' && next == '*') {
+            blockCommentDepth++;
+            i++; // skip '*'
+          } else if (c == '*' && next == '/') {
+            blockCommentDepth--;
+            i++; // skip '/'
+            if (blockCommentDepth == 0) {
+              state = State.NORMAL;
+              consumer.accept(state, ' ');
+            }
+          }
+          // else: skip character (part of the comment)
+          break;
+      }
+    }
+  }
+
+  /**
+   * Removes all comments and extra whitespace from the SQL string.
+   *
+   * @param sql the SQL string to remove comments and extra whitespace from
+   * @return the SQL string with all comments and extra whitespace removed
+   */
+  public static String stripCommentsAndWhitespaces(String sql) {
+    if (sql == null || sql.isEmpty()) {
+      return sql;
+    }
+
+    StringBuilder result = new StringBuilder(sql.length());
+    boolean[] lastWasSpace = {false};
+    forEachNonCommentChar(
+        sql,
+        (state, c) -> {
+          boolean isNormalWhitespace = state == State.NORMAL && Character.isWhitespace(c);
+          if (isNormalWhitespace && lastWasSpace[0]) {
+            return;
+          }
+          lastWasSpace[0] = isNormalWhitespace;
+          result.append(c);
+        });
+    return result.toString().trim();
+  }
+}

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

@@ -249,4 +249,25 @@ public void testParameterCountExceedsBindingsThrowsException() {
             .getMessage()
             .contains("Number of parameter bindings (3) exceeds parameter count (2)"));
   }
+
+  @Test
+  public void testParameterCountWithNestedBlockComment() throws SQLException {
+    String sql = "SELECT /* outer ? /* inner ? */ outer ? */ * FROM table WHERE id = ?";
+    DatabricksParameterMetaData metadata = new DatabricksParameterMetaData(sql);
+    assertEquals(1, metadata.getParameterCount());
+  }
+
+  @Test
+  public void testParameterCountWithTripleNestedBlockComment() throws SQLException {
+    String sql = "SELECT /* a ? /* b? /* c ? */ b ? */ a ? */ * FROM table WHERE id = ?";
+    DatabricksParameterMetaData metadata = new DatabricksParameterMetaData(sql);
+    assertEquals(1, metadata.getParameterCount());
+  }
+
+  @Test
+  public void testParameterCountWithApostropheInColumnName() throws SQLException {
+    String sql = "select foo, sum(`x'y`) as sum from table where foo = ? group by foo;";
+    DatabricksParameterMetaData metadata = new DatabricksParameterMetaData(sql);
+    assertEquals(1, metadata.getParameterCount());
+  }
 }

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

@@ -700,6 +700,18 @@ public void testShouldReturnResultSet_CallStatement() {
             "CALL send_notifications(12);", Collections.emptyList()));
   }
 
+  @Test
+  public void testShouldReturnResultSet_MultiLineBlockComment() {
+    String query = "/*\nMulti-line comment\n*/ SELECT * FROM table;";
+    assertTrue(DatabricksStatement.shouldReturnResultSet(query, Collections.emptyList()));
+  }
+
+  @Test
+  public void testShouldReturnResultSet_NestedBlockComment() {
+    String query = "/* /* Nested block comment */ */ SELECT * FROM table;";
+    assertTrue(DatabricksStatement.shouldReturnResultSet(query, Collections.emptyList()));
+  }
+
   @Test
   public void testShouldReturnResultSet_StartWithBegin() {
     assertTrue(DatabricksStatement.shouldReturnResultSet("BEGIN", Collections.emptyList()));