Skip to content

Commit 3b831cd

Browse files
chernserchernser
authored
Merge pull request #1861 from ClickHouse/v2_set_roles
[client-v2] Added session roles for client and operations
2 parents 54cb400 + 76d0e7c commit 3b831cd

File tree

9 files changed

+208
-25
lines changed

9 files changed

+208
-25
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
@@ -53,4 +53,5 @@ public class ClickHouseHttpProto {
5353
*/
5454
public static final String QPARAM_QUERY_ID = "query_id";
5555

56+
public static final String QPARAM_ROLE = "role";
5657
}

‎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
@@ -1676,10 +1676,10 @@ public CompletableFuture<Records> queryRecords(String sqlQuery, QuerySettings se
16761676
* @param sqlQuery - SQL query
16771677
* @return - complete list of records
16781678
*/
1679-
public List<GenericRecord> queryAll(String sqlQuery) {
1679+
public List<GenericRecord> queryAll(String sqlQuery, QuerySettings settings) {
16801680
try {
16811681
int operationTimeout = getOperationTimeout();
1682-
QuerySettings settings = new QuerySettings().setFormat(ClickHouseFormat.RowBinaryWithNamesAndTypes)
1682+
settings.setFormat(ClickHouseFormat.RowBinaryWithNamesAndTypes)
16831683
.waitEndOfQuery(true);
16841684
try (QueryResponse response = operationTimeout == 0 ? query(sqlQuery, settings).get() :
16851685
query(sqlQuery, settings).get(operationTimeout, TimeUnit.MILLISECONDS)) {
@@ -1702,6 +1702,10 @@ public List<GenericRecord> queryAll(String sqlQuery) {
17021702
}
17031703
}
17041704

1705+
public List<GenericRecord> queryAll(String sqlQuery) {
1706+
return queryAll(sqlQuery, new QuerySettings());
1707+
}
1708+
17051709
public <T> List<T> queryAll(String sqlQuery, Class<T> clazz, TableSchema schema) {
17061710
return queryAll(sqlQuery, clazz, schema, null);
17071711
}
@@ -1952,6 +1956,28 @@ public Set<String> getEndpoints() {
19521956
return Collections.unmodifiableSet(endpoints);
19531957
}
19541958

1959+
/**
1960+
* Sets list of DB roles that should be applied to each query.
1961+
*
1962+
* @param dbRoles
1963+
*/
1964+
public void setDBRoles(Collection<String> dbRoles) {
1965+
this.configuration.put(ClientSettings.SESSION_DB_ROLES, ClientSettings.commaSeparated(dbRoles));
1966+
this.unmodifiableDbRolesView =
1967+
Collections.unmodifiableCollection(ClientSettings.valuesFromCommaSeparated(
1968+
this.configuration.get(ClientSettings.SESSION_DB_ROLES)));
1969+
}
1970+
1971+
private Collection<String> unmodifiableDbRolesView = Collections.emptyList();
1972+
1973+
/**
1974+
* Returns list of DB roles that should be applied to each query.
1975+
*
1976+
* @return List of DB roles
1977+
*/
1978+
public Collection<String> getDBRoles() {
1979+
return unmodifiableDbRolesView;
1980+
}
19551981

19561982
private ClickHouseNode getNextAliveNode() {
19571983
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/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;
@@ -434,26 +436,27 @@ private void addHeaders(HttpPost req, Map<String, String> chConfig, Map<String,
434436
}
435437
}
436438
private void addQueryParams(URIBuilder req, Map<String, String> chConfig, Map<String, Object> requestConfig) {
439+
if (requestConfig == null) {
440+
requestConfig = Collections.emptyMap();
441+
}
437442

438443
for (Map.Entry<String, String> entry : chConfig.entrySet()) {
439444
if (entry.getKey().startsWith(ClientSettings.SERVER_SETTING_PREFIX)) {
440445
req.addParameter(entry.getKey().substring(ClientSettings.SERVER_SETTING_PREFIX.length()), entry.getValue());
441446
}
442447
}
443448

444-
if (requestConfig != null) {
445-
if (requestConfig.containsKey(ClickHouseHttpOption.WAIT_END_OF_QUERY.getKey())) {
446-
req.addParameter(ClickHouseHttpOption.WAIT_END_OF_QUERY.getKey(),
447-
requestConfig.get(ClickHouseHttpOption.WAIT_END_OF_QUERY.getKey()).toString());
448-
}
449-
if (requestConfig.containsKey(ClickHouseClientOption.QUERY_ID.getKey())) {
450-
req.addParameter(ClickHouseHttpProto.QPARAM_QUERY_ID, requestConfig.get(ClickHouseClientOption.QUERY_ID.getKey()).toString());
451-
}
452-
if (requestConfig.containsKey("statement_params")) {
453-
Map<String, Object> params = (Map<String, Object>) requestConfig.get("statement_params");
454-
for (Map.Entry<String, Object> entry : params.entrySet()) {
455-
req.addParameter("param_" + entry.getKey(), String.valueOf(entry.getValue()));
456-
}
449+
if (requestConfig.containsKey(ClickHouseHttpOption.WAIT_END_OF_QUERY.getKey())) {
450+
req.addParameter(ClickHouseHttpOption.WAIT_END_OF_QUERY.getKey(),
451+
requestConfig.get(ClickHouseHttpOption.WAIT_END_OF_QUERY.getKey()).toString());
452+
}
453+
if (requestConfig.containsKey(ClickHouseClientOption.QUERY_ID.getKey())) {
454+
req.addParameter(ClickHouseHttpProto.QPARAM_QUERY_ID, requestConfig.get(ClickHouseClientOption.QUERY_ID.getKey()).toString());
455+
}
456+
if (requestConfig.containsKey("statement_params")) {
457+
Map<String, Object> params = (Map<String, Object>) requestConfig.get("statement_params");
458+
for (Map.Entry<String, Object> entry : params.entrySet()) {
459+
req.addParameter("param_" + entry.getKey(), String.valueOf(entry.getValue()));
457460
}
458461
}
459462

@@ -476,11 +479,16 @@ private void addQueryParams(URIBuilder req, Map<String, String> chConfig, Map<St
476479
}
477480
}
478481

479-
if (requestConfig != null) {
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());
483-
}
482+
Collection<String> sessionRoles = (Collection<String>) requestConfig.getOrDefault(ClientSettings.SESSION_DB_ROLES,
483+
ClientSettings.valuesFromCommaSeparated(chConfiguration.getOrDefault(ClientSettings.SESSION_DB_ROLES, "")));
484+
if (!sessionRoles.isEmpty()) {
485+
486+
sessionRoles.forEach(r -> req.addParameter(ClickHouseHttpProto.QPARAM_ROLE, r));
487+
}
488+
489+
for (Map.Entry<String, Object> entry : requestConfig.entrySet()) {
490+
if (entry.getKey().startsWith(ClientSettings.SERVER_SETTING_PREFIX)) {
491+
req.addParameter(entry.getKey().substring(ClientSettings.SERVER_SETTING_PREFIX.length()), entry.getValue().toString());
484492
}
485493
}
486494
}

‎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: 85 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.data_formats.internal.BinaryStreamReader;
1921
import com.clickhouse.client.api.enums.Protocol;
@@ -27,8 +29,10 @@
2729
import com.clickhouse.client.api.query.QueryResponse;
2830
import com.clickhouse.client.api.query.QuerySettings;
2931
import com.clickhouse.client.api.query.Records;
32+
import com.clickhouse.client.http.config.HttpConnectionProvider;
3033
import com.clickhouse.data.ClickHouseDataType;
3134
import com.clickhouse.data.ClickHouseFormat;
35+
import com.clickhouse.data.ClickHouseVersion;
3236
import com.fasterxml.jackson.databind.JsonNode;
3337
import com.fasterxml.jackson.databind.MappingIterator;
3438
import com.fasterxml.jackson.databind.ObjectMapper;
@@ -37,6 +41,7 @@
3741
import org.testng.annotations.BeforeMethod;
3842
import org.testng.annotations.DataProvider;
3943
import org.testng.annotations.Test;
44+
import org.testng.util.Strings;
4045

4146
import java.io.BufferedReader;
4247
import java.io.BufferedWriter;
@@ -58,6 +63,7 @@
5863
import java.util.ArrayList;
5964
import java.util.Arrays;
6065
import java.util.Collections;
66+
import java.util.Comparator;
6167
import java.util.HashMap;
6268
import java.util.HashSet;
6369
import java.util.Iterator;
@@ -1657,6 +1663,85 @@ public static BigDecimal cropDecimal(BigDecimal value, int scale) {
16571663
return new BigDecimal(bi, scale);
16581664
}
16591665

1666+
@DataProvider(name = "sessionRoles")
1667+
private static Object[][] sessionRoles() {
1668+
return new Object[][]{
1669+
{new String[]{"ROL1", "ROL2"}},
1670+
{new String[]{"ROL1", "ROL2"}},
1671+
{new String[]{"ROL1", "ROL2"}},
1672+
{new String[]{"ROL1", "ROL2,☺"}},
1673+
{new String[]{"ROL1", "ROL2"}},
1674+
};
1675+
}
1676+
1677+
@Test(groups = {"integration"}, dataProvider = "sessionRoles", dataProviderClass = QueryTests.class)
1678+
public void testOperationCustomRoles(String[] roles) throws Exception {
1679+
List<GenericRecord> serverVersion = client.queryAll("SELECT version()");
1680+
if (ClickHouseVersion.of(serverVersion.get(0).getString(1)).check("(,24.3]")) {
1681+
System.out.println("Test is skipped: feature is supported since 24.4");
1682+
return;
1683+
}
1684+
1685+
final String password = UUID.randomUUID().toString();
1686+
final String rolesList = "\"" + Strings.join("\",\"", roles) + "\"";
1687+
try (CommandResponse resp = client.execute("DROP ROLE IF EXISTS " + rolesList).get()) {
1688+
}
1689+
try (CommandResponse resp = client.execute("CREATE ROLE " + rolesList).get()) {
1690+
}
1691+
try (CommandResponse resp = client.execute("DROP USER IF EXISTS some_user").get()) {
1692+
}
1693+
try (CommandResponse resp = client.execute("CREATE USER some_user IDENTIFIED WITH sha256_password BY '" + password + "'" ).get()) {
1694+
}
1695+
try (CommandResponse resp = client.execute("GRANT " + rolesList + " TO some_user").get()) {
1696+
}
1697+
1698+
1699+
try (Client userClient = newClient().setUsername("some_user").setPassword(password).build()) {
1700+
QuerySettings settings = new QuerySettings().setDBRoles(Arrays.asList(roles));
1701+
List<GenericRecord> resp = userClient.queryAll("SELECT currentRoles()", settings);
1702+
Set<String> roleSet = new HashSet<>(Arrays.asList(roles));
1703+
Set<String> currentRoles = new HashSet<String> (resp.get(0).getList(1));
1704+
Assert.assertEquals(currentRoles, roleSet, "Roles " + roleSet + " not found in " + currentRoles);
1705+
}
1706+
}
1707+
1708+
@DataProvider(name = "clientSessionRoles")
1709+
private static Object[][] clientSessionRoles() {
1710+
return new Object[][]{
1711+
{new String[]{"ROL1", "ROL2"}},
1712+
{new String[]{"ROL1", "ROL2,☺"}},
1713+
};
1714+
}
1715+
@Test(groups = {"integration"}, dataProvider = "clientSessionRoles", dataProviderClass = QueryTests.class)
1716+
public void testClientCustomRoles(String[] roles) throws Exception {
1717+
List<GenericRecord> serverVersion = client.queryAll("SELECT version()");
1718+
if (ClickHouseVersion.of(serverVersion.get(0).getString(1)).check("(,24.3]")) {
1719+
System.out.println("Test is skipped: feature is supported since 24.4");
1720+
return;
1721+
}
1722+
1723+
final String password = UUID.randomUUID().toString();
1724+
final String rolesList = "\"" + Strings.join("\",\"", roles) + "\"";
1725+
try (CommandResponse resp = client.execute("DROP ROLE IF EXISTS " + rolesList).get()) {
1726+
}
1727+
try (CommandResponse resp = client.execute("CREATE ROLE " + rolesList).get()) {
1728+
}
1729+
try (CommandResponse resp = client.execute("DROP USER IF EXISTS some_user").get()) {
1730+
}
1731+
try (CommandResponse resp = client.execute("CREATE USER some_user IDENTIFIED WITH sha256_password BY '" + password + "'" ).get()) {
1732+
}
1733+
try (CommandResponse resp = client.execute("GRANT " + rolesList + " TO some_user").get()) {
1734+
}
1735+
1736+
try (Client userClient = newClient().setUsername("some_user").setPassword(password).build()) {
1737+
userClient.setDBRoles(Arrays.asList(roles));
1738+
List<GenericRecord> resp = userClient.queryAll("SELECT currentRoles()");
1739+
Set<String> roleSet = new HashSet<>(Arrays.asList(roles));
1740+
Set<String> currentRoles = new HashSet<String> (resp.get(0).getList(1));
1741+
Assert.assertEquals(currentRoles, roleSet, "Roles " + roleSet + " not found in " + currentRoles);
1742+
}
1743+
}
1744+
16601745

16611746
protected Client.Builder newClient() {
16621747
ClickHouseNode node = getServer(ClickHouseProtocol.HTTP);

0 commit comments

Comments
 (0)