Merge pull request #1984 from ClickHouse/add-role-support

Client bug fix, add role support, placeholder cluster flag

client-v2/src/main/java/com/clickhouse/client/api/ClientConfigProperties.java

@@ -140,7 +140,10 @@ public static String commaSeparated(Collection<?> values) {
         for (Object value : values) {
             sb.append(value.toString().replaceAll(",", "\\\\,")).append(",");
         }
-        sb.setLength(sb.length() - 1);
+
+        if (sb.length() > 0) {
+            sb.setLength(sb.length() - 1);
+        }
         return sb.toString();
     }
 

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

@@ -23,6 +23,8 @@ public class ConnectionImpl implements Connection, JdbcV2Wrapper {
     protected final JdbcConfiguration config;
 
     private boolean closed = false;
+    protected boolean onCluster;//TODO: Placeholder for cluster support
+    protected String cluster;
     private String catalog;
     private String schema;
     private QuerySettings defaultQuerySettings;
@@ -31,6 +33,8 @@ public ConnectionImpl(String url, Properties info) {
         log.debug("Creating connection to {}", url);
         this.url = url;//Raw URL
         this.config = new JdbcConfiguration(url, info);
+        this.onCluster = false;
+        this.cluster = null;
         String clientName = "ClickHouse JDBC Driver V2/" + Driver.driverVersion;
 
         if (this.config.isDisableFrameworkDetection()) {

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

@@ -476,10 +476,14 @@ private static String encodeObject(Object x) throws SQLException {
             } else if (x instanceof Array) {
                 StringBuilder listString = new StringBuilder();
                 listString.append("[");
+                int i = 0;
                 for (Object item : (Object[])((Array) x).getArray()) {
-                    listString.append(encodeObject(item)).append(", ");
+                    if (i > 0) {
+                        listString.append(", ");
+                    }
+                    listString.append(encodeObject(item));
+                    i++;
                 }
-                listString.delete(listString.length() - 2, listString.length());
                 listString.append("]");
 
                 return listString.toString();

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

@@ -5,6 +5,7 @@
 import com.clickhouse.client.api.metrics.ServerMetrics;
 import com.clickhouse.client.api.query.QueryResponse;
 import com.clickhouse.client.api.query.QuerySettings;
+import com.clickhouse.jdbc.internal.JdbcUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -15,7 +16,9 @@
 import java.sql.SQLWarning;
 import java.sql.Statement;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
+import java.util.UUID;
 import java.util.concurrent.TimeUnit;
 
 public class StatementImpl implements Statement, JdbcV2Wrapper {
@@ -27,6 +30,7 @@ public class StatementImpl implements Statement, JdbcV2Wrapper {
     private ResultSetImpl currentResultSet;
     private OperationMetrics metrics;
     private List<String> batch;
+    private String lastQueryId;
 
     public StatementImpl(ConnectionImpl connection) {
         this.connection = connection;
@@ -121,6 +125,7 @@ public ResultSet executeQuery(String sql, QuerySettings settings) throws SQLExce
             ClickHouseBinaryFormatReader reader = connection.client.newBinaryFormatReader(response);
             currentResultSet = new ResultSetImpl(this, response, reader);
             metrics = response.getMetrics();
+            lastQueryId = response.getQueryId();
         } catch (Exception e) {
             throw new SQLException(e);
         }
@@ -153,6 +158,7 @@ public int executeUpdate(String sql, QuerySettings settings) throws SQLException
             }
             currentResultSet = null;
             metrics = response.getMetrics();
+            lastQueryId = response.getQueryId();
             response.close();
         } catch (Exception e) {
             throw new RuntimeException(e);
@@ -212,8 +218,17 @@ public void setQueryTimeout(int seconds) throws SQLException {
 
     @Override
     public void cancel() throws SQLException {
-        checkClosed();
-        throw new UnsupportedOperationException("Cancel is not supported.");
+        if (closed) {
+            return;
+        }
+
+        try {
+            connection.client.query(String.format("KILL QUERY%sWHERE query_id = '%s'",
+                    connection.onCluster ? " ON CLUSTER " + connection.cluster + " " : " ",
+                    lastQueryId), connection.getDefaultQuerySettings()).get();
+        } catch (Exception e) {
+            throw new SQLException(e);
+        }
     }
 
     @Override
@@ -239,25 +254,34 @@ public boolean execute(String sql) throws SQLException {
 
     public boolean execute(String sql, QuerySettings settings) throws SQLException {
         checkClosed();
-        List<String> statements = List.of(sql.split(";"));
-        boolean firstIsResult = false;
-
-        int index = 0;
-        for (String statement : statements) {
-            StatementType type = parseStatementType(statement);
-
-            if (type == StatementType.SELECT) {
-                executeQuery(statement, settings);
-                if (index == 0) {
-                    firstIsResult = true;
+        StatementType type = parseStatementType(sql);
+
+        if (type == StatementType.SELECT) {
+            executeQuery(sql, settings);
+            return true;
+        } else if(type == StatementType.SET) {
+            //SET ROLE
+            List<String> tokens = JdbcUtils.tokenizeSQL(sql);
+            if (JdbcUtils.containsIgnoresCase(tokens, "ROLE")) {
+                List<String> roles = new ArrayList<>();
+                int roleIndex = JdbcUtils.indexOfIgnoresCase(tokens, "ROLE");
+                if (roleIndex == 1) {
+                    for (int i = 2; i < tokens.size(); i++) {
+                        roles.add(tokens.get(i));
+                    }
+
+                    if (JdbcUtils.containsIgnoresCase(roles, "NONE")) {
+                        connection.client.setDBRoles(Collections.emptyList());
+                    } else {
+                        connection.client.setDBRoles(roles);
+                    }
                 }
-            } else {
-                executeUpdate(statement, settings);
             }
-
-            index++;
+            return false;
+        } else {
+            executeUpdate(sql, settings);
+            return false;
         }
-        return firstIsResult;
     }
 
     @Override

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

@@ -3,8 +3,12 @@
 import com.clickhouse.data.ClickHouseDataType;
 
 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;
 
 public class JdbcUtils {
     //Define a map to store the mapping between ClickHouse data types and SQL data types
@@ -52,4 +56,53 @@ public static String generateSqlTypeEnum(String columnName) {
         sql.append(Types.OTHER + ")");
         return sql.toString();
     }
+
+    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);
+            }
+        }
+
+        return tokens;
+    }
+
+    public static boolean isBlank(String str) {
+        return str == null || str.isEmpty() || str.trim().isEmpty();
+    }
+
+    public static boolean containsIgnoresCase(List<String> list, String str) {
+        if (list == null || list.isEmpty() || isBlank(str)) {
+            return false;
+        }
+
+        for (String s : list) {
+            if (s.equalsIgnoreCase(str)) {
+                return true;
+            }
+        }
+
+        return false;
+    }
+
+    public static int indexOfIgnoresCase(List<String> list, String str) {
+        if (list == null || list.isEmpty() || isBlank(str)) {
+            return -1;
+        }
+
+        for (int i = 0; i < list.size(); i++) {
+            if (list.get(i).equalsIgnoreCase(str)) {
+                return i;
+            }
+        }
+
+        return -1;
+    }
 }

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

@@ -4,6 +4,7 @@
 
 import com.clickhouse.client.BaseIntegrationTest;
 import com.clickhouse.client.ClickHouseProtocol;
+import com.clickhouse.client.api.query.GenericRecord;
 import com.clickhouse.logging.Logger;
 import com.clickhouse.logging.LoggerFactory;
 
@@ -47,4 +48,35 @@ protected boolean runQuery(String query) {
             return false;
         }
     }
+
+
+    protected boolean earlierThan(int major, int minor) {
+        String serverVersion = getServerVersion();
+        if (serverVersion == null) {
+            return false;
+        }
+
+        String[] parts = serverVersion.split("\\.");
+        if (parts.length < 2) {
+            return false;
+        }
+
+        try {
+            int serverMajor = Integer.parseInt(parts[0]);
+            int serverMinor = Integer.parseInt(parts[1]);
+            return serverMajor < major || (serverMajor == major && serverMinor < minor);
+        } catch (NumberFormatException e) {
+            return false;
+        }
+    }
+
+    protected String getServerVersion() {
+        try (ConnectionImpl connection = (ConnectionImpl) getJdbcConnection()) {
+            GenericRecord result = connection.client.queryAll("SELECT version() as server_version").stream()
+                    .findFirst().orElseThrow(() -> new SQLException("Failed to retrieve server version."));
+            return result.getString("server_version");
+        } catch (SQLException e) {
+            return null;
+        }
+    }
 }