Skip to content

Commit 8e10e18

Browse files
chernserchernser
committed
added session roles for client and operations
1 parent a499238 commit 8e10e18

File tree

10 files changed

+200
-26
lines changed

10 files changed

+200
-26
lines changed

‎clickhouse-http-client/src/main/java/com/clickhouse/client/http/ClickHouseHttpProto.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -51,4 +51,5 @@ public class ClickHouseHttpProto {
5151
*/
5252
public static final String QPARAM_QUERY_ID = "query_id";
5353

54+
public static final String QPARAM_ROLE = "role";
5455
}

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

Lines changed: 28 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1617,10 +1617,10 @@ public CompletableFuture<Records> queryRecords(String sqlQuery, QuerySettings se
16171617
* @param sqlQuery - SQL query
16181618
* @return - complete list of records
16191619
*/
1620-
public List<GenericRecord> queryAll(String sqlQuery) {
1620+
public List<GenericRecord> queryAll(String sqlQuery, QuerySettings settings) {
16211621
try {
16221622
int operationTimeout = getOperationTimeout();
1623-
QuerySettings settings = new QuerySettings().setFormat(ClickHouseFormat.RowBinaryWithNamesAndTypes)
1623+
settings.setFormat(ClickHouseFormat.RowBinaryWithNamesAndTypes)
16241624
.waitEndOfQuery(true);
16251625
try (QueryResponse response = operationTimeout == 0 ? query(sqlQuery, settings).get() :
16261626
query(sqlQuery, settings).get(operationTimeout, TimeUnit.MILLISECONDS)) {
@@ -1643,6 +1643,10 @@ public List<GenericRecord> queryAll(String sqlQuery) {
16431643
}
16441644
}
16451645

1646+
public List<GenericRecord> queryAll(String sqlQuery) {
1647+
return queryAll(sqlQuery, new QuerySettings());
1648+
}
1649+
16461650
public <T> List<T> queryAll(String sqlQuery, Class<T> clazz, TableSchema schema) {
16471651
return queryAll(sqlQuery, clazz, schema, null);
16481652
}
@@ -1891,6 +1895,28 @@ public Set<String> getEndpoints() {
18911895
return Collections.unmodifiableSet(endpoints);
18921896
}
18931897

1898+
/**
1899+
* Sets list of DB roles that should be applied to each query.
1900+
*
1901+
* @param dbRoles
1902+
*/
1903+
public void setDBRoles(Collection<String> dbRoles) {
1904+
this.configuration.put(ClientSettings.SESSION_DB_ROLES, ClientSettings.commaSeparated(dbRoles));
1905+
this.unmodifiableDbRolesView =
1906+
Collections.unmodifiableCollection(ClientSettings.valuesFromCommaSeparated(
1907+
this.configuration.get(ClientSettings.SESSION_DB_ROLES)));
1908+
}
1909+
1910+
private Collection<String> unmodifiableDbRolesView = Collections.emptyList();
1911+
1912+
/**
1913+
* Returns list of DB roles that should be applied to each query.
1914+
*
1915+
* @return List of DB roles
1916+
*/
1917+
public Collection<String> getDBRoles() {
1918+
return unmodifiableDbRolesView;
1919+
}
18941920

18951921
private ClickHouseNode getNextAliveNode() {
18961922
return serverNodes.get(0);

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

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
import java.util.Arrays;
44
import java.util.Collection;
5+
import java.util.Collections;
56
import java.util.List;
67
import java.util.stream.Collectors;
78

@@ -18,14 +19,20 @@ public class ClientSettings {
1819
public static String commaSeparated(Collection<?> values) {
1920
StringBuilder sb = new StringBuilder();
2021
for (Object value : values) {
21-
sb.append(value.toString().replaceAll(",", "\\,")).append(",");
22+
sb.append(value.toString().replaceAll(",", "\\\\,")).append(",");
2223
}
2324
sb.setLength(sb.length() - 1);
2425
return sb.toString();
2526
}
2627

2728
public static List<String> valuesFromCommaSeparated(String value) {
28-
return Arrays.stream(value.split(",")).map(s -> s.replaceAll("\\\\,", ","))
29+
if (value == null || value.isEmpty()) {
30+
return Collections.emptyList();
31+
}
32+
33+
return Arrays.stream(value.split("(?<!\\\\),")).map(s -> s.replaceAll("\\\\,", ","))
2934
.collect(Collectors.toList());
3035
}
36+
37+
public static final String SESSION_DB_ROLES = "session_db_roles";
3138
}

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

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
import com.clickhouse.client.api.metrics.ServerMetrics;
66
import com.clickhouse.client.api.query.QueryResponse;
77

8-
public class CommandResponse{
8+
public class CommandResponse implements AutoCloseable {
99

1010
private final QueryResponse response;
1111

@@ -71,4 +71,9 @@ public long getWrittenBytes() {
7171
public long getServerTime() {
7272
return response.getServerTime();
7373
}
74+
75+
@Override
76+
public void close() throws Exception {
77+
response.close();
78+
}
7479
}

‎client-v2/src/main/java/com/clickhouse/client/api/data_formats/internal/MapBackedRecord.java‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -241,8 +241,7 @@ public ClickHouseGeoMultiPolygonValue getGeoMultiPolygon(String colName) {
241241

242242
@Override
243243
public <T> List<T> getList(String colName) {
244-
ClickHouseArrayValue<?> array = readValue(colName);
245-
return null;
244+
return getList(schema.nameToIndex(colName));
246245
}
247246

248247

@@ -387,7 +386,12 @@ public ClickHouseGeoMultiPolygonValue getGeoMultiPolygon(int index) {
387386

388387
@Override
389388
public <T> List<T> getList(int index) {
390-
return readValue(index);
389+
Object value = readValue(index);
390+
if (value instanceof BinaryStreamReader.ArrayValue) {
391+
return ((BinaryStreamReader.ArrayValue) value).asList();
392+
} else {
393+
throw new ClientException("Column is not of array type");
394+
}
391395
}
392396

393397
@Override

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

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -205,4 +205,23 @@ public InsertSettings serverSetting(String name, Collection<String> values) {
205205
rawSettings.put(ClientSettings.SERVER_SETTING_PREFIX + name, ClientSettings.commaSeparated(values));
206206
return this;
207207
}
208+
209+
/**
210+
* Sets DB roles for an operation. Roles that were set by {@link Client#setDBRoles(Collection)} will be overridden.
211+
*
212+
* @param dbRoles
213+
*/
214+
public InsertSettings setDBRoles(Collection<String> dbRoles) {
215+
rawSettings.put(ClientSettings.SESSION_DB_ROLES, dbRoles);
216+
return this;
217+
}
218+
219+
/**
220+
* Gets DB roles for an operation.
221+
*
222+
* @return list of DB roles
223+
*/
224+
public Collection<String> getDBRoles() {
225+
return (Collection<String>) rawSettings.get(ClientSettings.SESSION_DB_ROLES);
226+
}
208227
}

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

Lines changed: 26 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,8 @@
7070
import java.nio.charset.StandardCharsets;
7171
import java.security.NoSuchAlgorithmException;
7272
import java.util.Base64;
73+
import java.util.Collection;
74+
import java.util.Collections;
7375
import java.util.EnumSet;
7476
import java.util.HashSet;
7577
import java.util.Map;
@@ -425,26 +427,27 @@ private void addHeaders(HttpPost req, Map<String, String> chConfig, Map<String,
425427
}
426428
}
427429
private void addQueryParams(URIBuilder req, Map<String, String> chConfig, Map<String, Object> requestConfig) {
430+
if (requestConfig == null) {
431+
requestConfig = Collections.emptyMap();
432+
}
428433

429434
for (Map.Entry<String, String> entry : chConfig.entrySet()) {
430435
if (entry.getKey().startsWith(ClientSettings.SERVER_SETTING_PREFIX)) {
431436
req.addParameter(entry.getKey().substring(ClientSettings.SERVER_SETTING_PREFIX.length()), entry.getValue());
432437
}
433438
}
434439

435-
if (requestConfig != null) {
436-
if (requestConfig.containsKey(ClickHouseHttpOption.WAIT_END_OF_QUERY.getKey())) {
437-
req.addParameter(ClickHouseHttpOption.WAIT_END_OF_QUERY.getKey(),
438-
requestConfig.get(ClickHouseHttpOption.WAIT_END_OF_QUERY.getKey()).toString());
439-
}
440-
if (requestConfig.containsKey(ClickHouseClientOption.QUERY_ID.getKey())) {
441-
req.addParameter(ClickHouseHttpProto.QPARAM_QUERY_ID, requestConfig.get(ClickHouseClientOption.QUERY_ID.getKey()).toString());
442-
}
443-
if (requestConfig.containsKey("statement_params")) {
444-
Map<String, Object> params = (Map<String, Object>) requestConfig.get("statement_params");
445-
for (Map.Entry<String, Object> entry : params.entrySet()) {
446-
req.addParameter("param_" + entry.getKey(), String.valueOf(entry.getValue()));
447-
}
440+
if (requestConfig.containsKey(ClickHouseHttpOption.WAIT_END_OF_QUERY.getKey())) {
441+
req.addParameter(ClickHouseHttpOption.WAIT_END_OF_QUERY.getKey(),
442+
requestConfig.get(ClickHouseHttpOption.WAIT_END_OF_QUERY.getKey()).toString());
443+
}
444+
if (requestConfig.containsKey(ClickHouseClientOption.QUERY_ID.getKey())) {
445+
req.addParameter(ClickHouseHttpProto.QPARAM_QUERY_ID, requestConfig.get(ClickHouseClientOption.QUERY_ID.getKey()).toString());
446+
}
447+
if (requestConfig.containsKey("statement_params")) {
448+
Map<String, Object> params = (Map<String, Object>) requestConfig.get("statement_params");
449+
for (Map.Entry<String, Object> entry : params.entrySet()) {
450+
req.addParameter("param_" + entry.getKey(), String.valueOf(entry.getValue()));
448451
}
449452
}
450453

@@ -467,11 +470,16 @@ private void addQueryParams(URIBuilder req, Map<String, String> chConfig, Map<St
467470
}
468471
}
469472

470-
if (requestConfig != null) {
471-
for (Map.Entry<String, Object> entry : requestConfig.entrySet()) {
472-
if (entry.getKey().startsWith(ClientSettings.SERVER_SETTING_PREFIX)) {
473-
req.addParameter(entry.getKey().substring(ClientSettings.SERVER_SETTING_PREFIX.length()), entry.getValue().toString());
474-
}
473+
Collection<String> sessionRoles = (Collection<String>) requestConfig.getOrDefault(ClientSettings.SESSION_DB_ROLES,
474+
ClientSettings.valuesFromCommaSeparated(chConfiguration.getOrDefault(ClientSettings.SESSION_DB_ROLES, "")));
475+
if (!sessionRoles.isEmpty()) {
476+
477+
sessionRoles.forEach(r -> req.addParameter(ClickHouseHttpProto.QPARAM_ROLE, r));
478+
}
479+
480+
for (Map.Entry<String, Object> entry : requestConfig.entrySet()) {
481+
if (entry.getKey().startsWith(ClientSettings.SERVER_SETTING_PREFIX)) {
482+
req.addParameter(entry.getKey().substring(ClientSettings.SERVER_SETTING_PREFIX.length()), entry.getValue().toString());
475483
}
476484
}
477485
}

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

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -221,4 +221,23 @@ public QuerySettings serverSetting(String name, Collection<String> values) {
221221
rawSettings.put(ClientSettings.SERVER_SETTING_PREFIX + name, ClientSettings.commaSeparated(values));
222222
return this;
223223
}
224+
225+
/**
226+
* Sets DB roles for an operation. Roles that were set by {@link Client#setDBRoles(Collection)} will be overridden.
227+
*
228+
* @param dbRoles
229+
*/
230+
public QuerySettings setDBRoles(Collection<String> dbRoles) {
231+
rawSettings.put(ClientSettings.SESSION_DB_ROLES, dbRoles);
232+
return this;
233+
}
234+
235+
/**
236+
* Gets DB roles for an operation.
237+
*
238+
* @return list of DB roles
239+
*/
240+
public Collection<String> getDBRoles() {
241+
return (Collection<String>) rawSettings.get(ClientSettings.SESSION_DB_ROLES);
242+
}
224243
}

‎client-v2/src/test/java/com/clickhouse/client/SettingsTests.java‎

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

3+
import com.clickhouse.client.api.ClientSettings;
4+
import org.testng.Assert;
5+
import org.testng.annotations.Test;
6+
7+
import java.util.Arrays;
8+
import java.util.List;
9+
310
public class SettingsTests {
411

12+
@Test
13+
void testClientSettings() {
14+
List<String> source = Arrays.asList("ROL1", "ROL2,☺", "Rol,3,3");
15+
String listA = ClientSettings.commaSeparated(source);
16+
List<String> listB = ClientSettings.valuesFromCommaSeparated(listA);
17+
Assert.assertEquals(listB, source);
18+
}
519
}

‎client-v2/src/test/java/com/clickhouse/client/query/QueryTests.java‎

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,8 +12,10 @@
1212
import com.clickhouse.client.ClickHouseResponse;
1313
import com.clickhouse.client.api.Client;
1414
import com.clickhouse.client.api.ClientException;
15+
import com.clickhouse.client.api.ClientSettings;
1516
import com.clickhouse.client.api.DataTypeUtils;
1617
import com.clickhouse.client.api.ServerException;
18+
import com.clickhouse.client.api.command.CommandResponse;
1719
import com.clickhouse.client.api.data_formats.ClickHouseBinaryFormatReader;
1820
import com.clickhouse.client.api.enums.Protocol;
1921
import com.clickhouse.client.api.insert.InsertSettings;
@@ -26,6 +28,7 @@
2628
import com.clickhouse.client.api.query.QueryResponse;
2729
import com.clickhouse.client.api.query.QuerySettings;
2830
import com.clickhouse.client.api.query.Records;
31+
import com.clickhouse.client.http.config.HttpConnectionProvider;
2932
import com.clickhouse.data.ClickHouseDataType;
3033
import com.clickhouse.data.ClickHouseFormat;
3134
import com.fasterxml.jackson.databind.JsonNode;
@@ -36,6 +39,7 @@
3639
import org.testng.annotations.BeforeMethod;
3740
import org.testng.annotations.DataProvider;
3841
import org.testng.annotations.Test;
42+
import org.testng.util.Strings;
3943

4044
import java.io.BufferedReader;
4145
import java.io.BufferedWriter;
@@ -56,6 +60,7 @@
5660
import java.util.ArrayList;
5761
import java.util.Arrays;
5862
import java.util.Collections;
63+
import java.util.Comparator;
5964
import java.util.HashMap;
6065
import java.util.HashSet;
6166
import java.util.Iterator;
@@ -1570,6 +1575,72 @@ public static BigDecimal cropDecimal(BigDecimal value, int scale) {
15701575
return new BigDecimal(bi, scale);
15711576
}
15721577

1578+
@DataProvider(name = "sessionRoles")
1579+
private static Object[][] sessionRoles() {
1580+
return new Object[][]{
1581+
{new String[]{"ROL1", "ROL2"}},
1582+
{new String[]{"ROL1", "ROL2"}},
1583+
{new String[]{"ROL1", "ROL2"}},
1584+
{new String[]{"ROL1", "ROL2,☺"}},
1585+
{new String[]{"ROL1", "ROL2"}},
1586+
};
1587+
}
1588+
1589+
@Test(groups = {"integration"}, dataProvider = "sessionRoles", dataProviderClass = QueryTests.class)
1590+
public void testOperationCustomRoles(String[] roles) throws Exception {
1591+
final String password = UUID.randomUUID().toString();
1592+
final String rolesList = "\"" + Strings.join("\",\"", roles) + "\"";
1593+
try (CommandResponse resp = client.execute("DROP ROLE IF EXISTS " + rolesList).get()) {
1594+
}
1595+
try (CommandResponse resp = client.execute("CREATE ROLE " + rolesList).get()) {
1596+
}
1597+
try (CommandResponse resp = client.execute("DROP USER IF EXISTS some_user").get()) {
1598+
}
1599+
try (CommandResponse resp = client.execute("CREATE USER some_user IDENTIFIED WITH sha256_password BY '" + password + "'" ).get()) {
1600+
}
1601+
try (CommandResponse resp = client.execute("GRANT " + rolesList + " TO some_user").get()) {
1602+
}
1603+
1604+
try (Client userClient = newClient().setUsername("some_user").setPassword(password).build()) {
1605+
QuerySettings settings = new QuerySettings().setDBRoles(Arrays.asList(roles));
1606+
List<GenericRecord> resp = userClient.queryAll("SELECT currentRoles()", settings);
1607+
Set<String> roleSet = new HashSet<>(Arrays.asList(roles));
1608+
Set<String> currentRoles = new HashSet<String> (resp.get(0).getList(1));
1609+
Assert.assertEquals(currentRoles, roleSet, "Roles " + roleSet + " not found in " + currentRoles);
1610+
}
1611+
}
1612+
1613+
@DataProvider(name = "clientSessionRoles")
1614+
private static Object[][] clientSessionRoles() {
1615+
return new Object[][]{
1616+
{new String[]{"ROL1", "ROL2"}},
1617+
{new String[]{"ROL1", "ROL2,☺"}},
1618+
};
1619+
}
1620+
@Test(groups = {"integration"}, dataProvider = "clientSessionRoles", dataProviderClass = QueryTests.class)
1621+
public void testClientCustomRoles(String[] roles) throws Exception {
1622+
final String password = UUID.randomUUID().toString();
1623+
final String rolesList = "\"" + Strings.join("\",\"", roles) + "\"";
1624+
try (CommandResponse resp = client.execute("DROP ROLE IF EXISTS " + rolesList).get()) {
1625+
}
1626+
try (CommandResponse resp = client.execute("CREATE ROLE " + rolesList).get()) {
1627+
}
1628+
try (CommandResponse resp = client.execute("DROP USER IF EXISTS some_user").get()) {
1629+
}
1630+
try (CommandResponse resp = client.execute("CREATE USER some_user IDENTIFIED WITH sha256_password BY '" + password + "'" ).get()) {
1631+
}
1632+
try (CommandResponse resp = client.execute("GRANT " + rolesList + " TO some_user").get()) {
1633+
}
1634+
1635+
try (Client userClient = newClient().setUsername("some_user").setPassword(password).build()) {
1636+
userClient.setDBRoles(Arrays.asList(roles));
1637+
List<GenericRecord> resp = userClient.queryAll("SELECT currentRoles()");
1638+
Set<String> roleSet = new HashSet<>(Arrays.asList(roles));
1639+
Set<String> currentRoles = new HashSet<String> (resp.get(0).getList(1));
1640+
Assert.assertEquals(currentRoles, roleSet, "Roles " + roleSet + " not found in " + currentRoles);
1641+
}
1642+
}
1643+
15731644

15741645
protected Client.Builder newClient() {
15751646
ClickHouseNode node = getServer(ClickHouseProtocol.HTTP);

0 commit comments

Comments
 (0)