Merge pull request #2048 from ClickHouse/adjust-roles-parsing

Adjust roles parsing

Author: Paultagoras

jdbc-v2/src/main/java/com/clickhouse/jdbc/StatementImpl.java

@@ -309,7 +309,11 @@ private boolean execute(String sql, QuerySettings settings) throws SQLException
                 int roleIndex = JdbcUtils.indexOfIgnoresCase(tokens, "ROLE");
                 if (roleIndex == 1) {
                     for (int i = 2; i < tokens.size(); i++) {
-                        roles.add(tokens.get(i));
+                        String token = tokens.get(i);
+                        String[] roleTokens = token.split(",");
+                        for (String roleToken : roleTokens) {
+                            roles.add(roleToken.replace("\"", ""));//Remove double quotes
+                        }
                     }
 
                     if (JdbcUtils.containsIgnoresCase(roles, "NONE")) {

jdbc-v2/src/main/java/com/clickhouse/jdbc/internal/JdbcUtils.java

@@ -63,18 +63,36 @@ public static String generateSqlTypeEnum(String columnName) {
 
     public static List<String> tokenizeSQL(String sql) {
         List<String> tokens = new ArrayList<>();
-        Matcher m = Pattern.compile("([^\"]\\S*|\".+?\")\\s*").matcher(sql);
-        while (m.find()) {
-            String token = m.group(1).replace("\"", "").trim();
-            if (!token.isEmpty() && token.charAt(token.length() - 1) == ',') {
-                token = token.substring(0, token.length() - 1);
-            }
 
-            if (!isBlank(token)) {
-                tokens.add(token);
+        // Remove SQL comments
+        String withoutComments = sql
+                .replaceAll("--.*?$", "") // Remove single-line comments
+                .replaceAll("/\\*.*?\\*/", ""); // Remove multi-line comments
+
+        StringBuilder currentToken = new StringBuilder();
+        boolean insideQuotes = false;
+
+        for (int i = 0; i < withoutComments.length(); i++) {
+            char c = withoutComments.charAt(i);
+
+            if (c == '"') {
+                insideQuotes = !insideQuotes; // Toggle the insideQuotes flag
+                currentToken.append(c); // Include the quote in the token
+            } else if (Character.isWhitespace(c) && !insideQuotes) {
+                if (currentToken.length() > 0) {
+                    tokens.add(currentToken.toString());
+                    currentToken.setLength(0); // Clear the current token
+                }
+            } else {
+                currentToken.append(c);
             }
         }
 
+        // Add the last token if it exists
+        if (currentToken.length() > 0) {
+            tokens.add(currentToken.toString());
+        }
+
         return tokens;
     }
 

jdbc-v2/src/test/java/com/clickhouse/jdbc/StatementTest.java

@@ -345,7 +345,7 @@ private void testSettingRole() throws SQLException {
             return;
         }
 
-        List<String> roles = Arrays.asList("role1", "role2");
+        List<String> roles = Arrays.asList("role1", "role2", "role3");
 
         try (ConnectionImpl conn = (ConnectionImpl) getJdbcConnection()) {
             try (Statement stmt = conn.createStatement()) {
@@ -388,6 +388,25 @@ record = conn.client.queryAll("SELECT currentRoles()").get(0);
 
             record = conn.client.queryAll("SELECT currentRoles()").get(0);
             assertEquals(record.getList(1).size(), 0);
+
+            try (Statement stmt = conn.createStatement()) {
+                stmt.execute("SET ROLE \"role1\",\"role2\"");
+            }
+
+            record = conn.client.queryAll("SELECT currentRoles()").get(0);
+            assertEquals(record.getList(1).size(), 2);
+            assertEquals(record.getList(1).get(0), "role1");
+            assertEquals(record.getList(1).get(1), "role2");
+
+            try (Statement stmt = conn.createStatement()) {
+                stmt.execute("SET ROLE \"role1\",\"role2\",\"role3\"");
+            }
+
+            record = conn.client.queryAll("SELECT currentRoles()").get(0);
+            assertEquals(record.getList(1).size(), 3);
+            assertEquals(record.getList(1).get(0), "role1");
+            assertEquals(record.getList(1).get(1), "role2");
+            assertEquals(record.getList(1).get(2), "role3");
         }
     }
 

jdbc-v2/src/test/java/com/clickhouse/jdbc/internal/JdbcUtilsTest.java

@@ -46,6 +46,6 @@ public void testTokenizeSQL() {
         assertEquals(tokens3.get(2), "FROM");
         assertEquals(tokens3.get(3), "table");
         assertEquals(tokens3.get(4), "WHERE");
-        assertEquals(tokens3.get(5), "id = 1 AND name = 'John' OR age = 30");
+        assertEquals(tokens3.get(5).replace("\"", ""), "id = 1 AND name = 'John' OR age = 30");
     }
 }