Skip to content

Commit ab78e9b

Browse files
author
Paultagoras
committed
Client bug fix, add role support, placeholder cluster flag,
1 parent 87e8d8e commit ab78e9b

File tree

5 files changed

+124
-19
lines changed

5 files changed

+124
-19
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.isEmpty()) {
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: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ public class ConnectionImpl implements Connection, JdbcV2Wrapper {
2323
protected final JdbcConfiguration config;
2424

2525
private boolean closed = false;
26+
protected boolean onCluster = false;//TODO: Placeholder for cluster support
2627
private String catalog;
2728
private String schema;
2829
private QuerySettings defaultQuerySettings;

‎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 " : " ",
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 (tokens.contains("ROLE")) {
266+
List<String> roles = new ArrayList<>();
267+
int roleIndex = tokens.indexOf("ROLE");
268+
if (roleIndex == 1) {
269+
for (int i = 2; i < tokens.size(); i++) {
270+
roles.add(tokens.get(i));
271+
}
272+
273+
if (roles.contains("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: 25 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,25 @@ 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+
}
5580
}

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

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
package com.clickhouse.jdbc;
22

3+
import com.clickhouse.client.api.query.GenericRecord;
4+
import org.testng.Assert;
35
import org.testng.annotations.Test;
46

57
import java.sql.Connection;
@@ -8,6 +10,9 @@
810
import java.sql.SQLException;
911
import java.sql.Statement;
1012
import java.time.LocalDate;
13+
import java.util.Arrays;
14+
import java.util.List;
15+
import java.util.Properties;
1116

1217
import static org.testng.Assert.assertEquals;
1318
import static org.testng.Assert.assertFalse;
@@ -324,4 +329,51 @@ public void testExecuteQueryTimeout() throws Exception {
324329
}
325330
}
326331
}
332+
333+
334+
@Test
335+
private void testSettingRole() throws SQLException {
336+
List<String> roles = Arrays.asList("role1", "role2");
337+
338+
try (ConnectionImpl conn = (ConnectionImpl) getJdbcConnection()) {
339+
try (Statement stmt = conn.createStatement()) {
340+
stmt.execute("DROP ROLE IF EXISTS " + String.join(", ", roles));
341+
stmt.execute("DROP USER IF EXISTS some_user");
342+
stmt.execute("CREATE ROLE " + String.join(", ", roles));
343+
stmt.execute("CREATE USER some_user IDENTIFIED WITH no_password");
344+
stmt.execute("GRANT " + String.join(", ", roles) + " TO some_user");
345+
stmt.execute("SET DEFAULT ROLE NONE TO some_user");
346+
}
347+
}
348+
349+
Properties info = new Properties();
350+
info.setProperty("user", "some_user");
351+
info.setProperty("password", "");
352+
353+
try (ConnectionImpl conn = new ConnectionImpl(getEndpointString(), info)) {
354+
GenericRecord record = conn.client.queryAll("SELECT currentRoles()").get(0);
355+
assertEquals(record.getList(1).size(), 2);
356+
357+
try (Statement stmt = conn.createStatement()) {
358+
stmt.execute("SET ROLE role1");
359+
}
360+
361+
record = conn.client.queryAll("SELECT currentRoles()").get(0);
362+
assertEquals(record.getList(1).size(), 1);
363+
364+
try (Statement stmt = conn.createStatement()) {
365+
stmt.execute("SET ROLE role2");
366+
}
367+
368+
record = conn.client.queryAll("SELECT currentRoles()").get(0);
369+
assertEquals(record.getList(1).size(), 1);
370+
371+
try (Statement stmt = conn.createStatement()) {
372+
stmt.execute("SET ROLE NONE");
373+
}
374+
375+
record = conn.client.queryAll("SELECT currentRoles()").get(0);
376+
assertEquals(record.getList(1).size(), 2);
377+
}
378+
}
327379
}

0 commit comments

Comments
 (0)