Skip to content

Commit 6c413e7

Browse files
PaultagorasPaultagoras
authored
Merge pull request #1984 from ClickHouse/add-role-support
Client bug fix, add role support, placeholder cluster flag
2 parents 87e8d8e + 72e0d71 commit 6c413e7

File tree

8 files changed

+264
-37
lines changed

8 files changed

+264
-37
lines changed

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

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -140,7 +140,10 @@ public static String commaSeparated(Collection<?> values) {
140140
for (Object value : values) {
141141
sb.append(value.toString().replaceAll(",", "\\\\,")).append(",");
142142
}
143-
sb.setLength(sb.length() - 1);
143+
144+
if (sb.length() > 0) {
145+
sb.setLength(sb.length() - 1);
146+
}
144147
return sb.toString();
145148
}
146149

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,8 @@ public class ConnectionImpl implements Connection, JdbcV2Wrapper {
2323
protected final JdbcConfiguration config;
2424

2525
private boolean closed = false;
26+
protected boolean onCluster;//TODO: Placeholder for cluster support
27+
protected String cluster;
2628
private String catalog;
2729
private String schema;
2830
private QuerySettings defaultQuerySettings;
@@ -31,6 +33,8 @@ public ConnectionImpl(String url, Properties info) {
3133
log.debug("Creating connection to {}", url);
3234
this.url = url;//Raw URL
3335
this.config = new JdbcConfiguration(url, info);
36+
this.onCluster = false;
37+
this.cluster = null;
3438
String clientName = "ClickHouse JDBC Driver V2/" + Driver.driverVersion;
3539

3640
if (this.config.isDisableFrameworkDetection()) {

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

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -476,10 +476,14 @@ private static String encodeObject(Object x) throws SQLException {
476476
} else if (x instanceof Array) {
477477
StringBuilder listString = new StringBuilder();
478478
listString.append("[");
479+
int i = 0;
479480
for (Object item : (Object[])((Array) x).getArray()) {
480-
listString.append(encodeObject(item)).append(", ");
481+
if (i > 0) {
482+
listString.append(", ");
483+
}
484+
listString.append(encodeObject(item));
485+
i++;
481486
}
482-
listString.delete(listString.length() - 2, listString.length());
483487
listString.append("]");
484488

485489
return listString.toString();

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

Lines changed: 42 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
import com.clickhouse.client.api.metrics.ServerMetrics;
66
import com.clickhouse.client.api.query.QueryResponse;
77
import com.clickhouse.client.api.query.QuerySettings;
8+
import com.clickhouse.jdbc.internal.JdbcUtils;
89
import org.slf4j.Logger;
910
import org.slf4j.LoggerFactory;
1011

@@ -15,7 +16,9 @@
1516
import java.sql.SQLWarning;
1617
import java.sql.Statement;
1718
import java.util.ArrayList;
19+
import java.util.Collections;
1820
import java.util.List;
21+
import java.util.UUID;
1922
import java.util.concurrent.TimeUnit;
2023

2124
public class StatementImpl implements Statement, JdbcV2Wrapper {
@@ -27,6 +30,7 @@ public class StatementImpl implements Statement, JdbcV2Wrapper {
2730
private ResultSetImpl currentResultSet;
2831
private OperationMetrics metrics;
2932
private List<String> batch;
33+
private String lastQueryId;
3034

3135
public StatementImpl(ConnectionImpl connection) {
3236
this.connection = connection;
@@ -121,6 +125,7 @@ public ResultSet executeQuery(String sql, QuerySettings settings) throws SQLExce
121125
ClickHouseBinaryFormatReader reader = connection.client.newBinaryFormatReader(response);
122126
currentResultSet = new ResultSetImpl(this, response, reader);
123127
metrics = response.getMetrics();
128+
lastQueryId = response.getQueryId();
124129
} catch (Exception e) {
125130
throw new SQLException(e);
126131
}
@@ -153,6 +158,7 @@ public int executeUpdate(String sql, QuerySettings settings) throws SQLException
153158
}
154159
currentResultSet = null;
155160
metrics = response.getMetrics();
161+
lastQueryId = response.getQueryId();
156162
response.close();
157163
} catch (Exception e) {
158164
throw new RuntimeException(e);
@@ -212,8 +218,17 @@ public void setQueryTimeout(int seconds) throws SQLException {
212218

213219
@Override
214220
public void cancel() throws SQLException {
215-
checkClosed();
216-
throw new UnsupportedOperationException("Cancel is not supported.");
221+
if (closed) {
222+
return;
223+
}
224+
225+
try {
226+
connection.client.query(String.format("KILL QUERY%sWHERE query_id = '%s'",
227+
connection.onCluster ? " ON CLUSTER " + connection.cluster + " " : " ",
228+
lastQueryId), connection.getDefaultQuerySettings()).get();
229+
} catch (Exception e) {
230+
throw new SQLException(e);
231+
}
217232
}
218233

219234
@Override
@@ -239,25 +254,34 @@ public boolean execute(String sql) throws SQLException {
239254

240255
public boolean execute(String sql, QuerySettings settings) throws SQLException {
241256
checkClosed();
242-
List<String> statements = List.of(sql.split(";"));
243-
boolean firstIsResult = false;
244-
245-
int index = 0;
246-
for (String statement : statements) {
247-
StatementType type = parseStatementType(statement);
248-
249-
if (type == StatementType.SELECT) {
250-
executeQuery(statement, settings);
251-
if (index == 0) {
252-
firstIsResult = true;
257+
StatementType type = parseStatementType(sql);
258+
259+
if (type == StatementType.SELECT) {
260+
executeQuery(sql, settings);
261+
return true;
262+
} else if(type == StatementType.SET) {
263+
//SET ROLE
264+
List<String> tokens = JdbcUtils.tokenizeSQL(sql);
265+
if (JdbcUtils.containsIgnoresCase(tokens, "ROLE")) {
266+
List<String> roles = new ArrayList<>();
267+
int roleIndex = JdbcUtils.indexOfIgnoresCase(tokens, "ROLE");
268+
if (roleIndex == 1) {
269+
for (int i = 2; i < tokens.size(); i++) {
270+
roles.add(tokens.get(i));
271+
}
272+
273+
if (JdbcUtils.containsIgnoresCase(roles, "NONE")) {
274+
connection.client.setDBRoles(Collections.emptyList());
275+
} else {
276+
connection.client.setDBRoles(roles);
277+
}
253278
}
254-
} else {
255-
executeUpdate(statement, settings);
256279
}
257-
258-
index++;
280+
return false;
281+
} else {
282+
executeUpdate(sql, settings);
283+
return false;
259284
}
260-
return firstIsResult;
261285
}
262286

263287
@Override

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

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,12 @@
33
import com.clickhouse.data.ClickHouseDataType;
44

55
import java.sql.Types;
6+
import java.util.ArrayList;
67
import java.util.HashMap;
8+
import java.util.List;
79
import java.util.Map;
10+
import java.util.regex.Matcher;
11+
import java.util.regex.Pattern;
812

913
public class JdbcUtils {
1014
//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) {
5256
sql.append(Types.OTHER + ")");
5357
return sql.toString();
5458
}
59+
60+
public static List<String> tokenizeSQL(String sql) {
61+
List<String> tokens = new ArrayList<>();
62+
Matcher m = Pattern.compile("([^\"]\\S*|\".+?\")\\s*").matcher(sql);
63+
while (m.find()) {
64+
String token = m.group(1).replace("\"", "").trim();
65+
if (!token.isEmpty() && token.charAt(token.length() - 1) == ',') {
66+
token = token.substring(0, token.length() - 1);
67+
}
68+
69+
if (!isBlank(token)) {
70+
tokens.add(token);
71+
}
72+
}
73+
74+
return tokens;
75+
}
76+
77+
public static boolean isBlank(String str) {
78+
return str == null || str.isEmpty() || str.trim().isEmpty();
79+
}
80+
81+
public static boolean containsIgnoresCase(List<String> list, String str) {
82+
if (list == null || list.isEmpty() || isBlank(str)) {
83+
return false;
84+
}
85+
86+
for (String s : list) {
87+
if (s.equalsIgnoreCase(str)) {
88+
return true;
89+
}
90+
}
91+
92+
return false;
93+
}
94+
95+
public static int indexOfIgnoresCase(List<String> list, String str) {
96+
if (list == null || list.isEmpty() || isBlank(str)) {
97+
return -1;
98+
}
99+
100+
for (int i = 0; i < list.size(); i++) {
101+
if (list.get(i).equalsIgnoreCase(str)) {
102+
return i;
103+
}
104+
}
105+
106+
return -1;
107+
}
55108
}

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

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44

55
import com.clickhouse.client.BaseIntegrationTest;
66
import com.clickhouse.client.ClickHouseProtocol;
7+
import com.clickhouse.client.api.query.GenericRecord;
78
import com.clickhouse.logging.Logger;
89
import com.clickhouse.logging.LoggerFactory;
910

@@ -47,4 +48,35 @@ protected boolean runQuery(String query) {
4748
return false;
4849
}
4950
}
51+
52+
53+
protected boolean earlierThan(int major, int minor) {
54+
String serverVersion = getServerVersion();
55+
if (serverVersion == null) {
56+
return false;
57+
}
58+
59+
String[] parts = serverVersion.split("\\.");
60+
if (parts.length < 2) {
61+
return false;
62+
}
63+
64+
try {
65+
int serverMajor = Integer.parseInt(parts[0]);
66+
int serverMinor = Integer.parseInt(parts[1]);
67+
return serverMajor < major || (serverMajor == major && serverMinor < minor);
68+
} catch (NumberFormatException e) {
69+
return false;
70+
}
71+
}
72+
73+
protected String getServerVersion() {
74+
try (ConnectionImpl connection = (ConnectionImpl) getJdbcConnection()) {
75+
GenericRecord result = connection.client.queryAll("SELECT version() as server_version").stream()
76+
.findFirst().orElseThrow(() -> new SQLException("Failed to retrieve server version."));
77+
return result.getString("server_version");
78+
} catch (SQLException e) {
79+
return null;
80+
}
81+
}
5082
}

0 commit comments

Comments
 (0)