Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[client-v2] Added session roles for client and operations #1861

Draft
wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

[client-v2] Added session roles for client and operations #1861

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions clickhouse-http-client/src/main/java/com/clickhouse/client/http/ClickHouseHttpProto.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -51,4 +51,5 @@ public class ClickHouseHttpProto {
*/
public static final String QPARAM_QUERY_ID = "query_id";

public static final String QPARAM_ROLE = "role";
}
30 changes: 28 additions & 2 deletions client-v2/src/main/java/com/clickhouse/client/api/Client.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1617,10 +1617,10 @@ public CompletableFuture<Records> queryRecords(String sqlQuery, QuerySettings se
* @param sqlQuery - SQL query
* @return - complete list of records
*/
public List<GenericRecord> queryAll(String sqlQuery) {
public List<GenericRecord> queryAll(String sqlQuery, QuerySettings settings) {
try {
int operationTimeout = getOperationTimeout();
QuerySettings settings = new QuerySettings().setFormat(ClickHouseFormat.RowBinaryWithNamesAndTypes)
settings.setFormat(ClickHouseFormat.RowBinaryWithNamesAndTypes)
.waitEndOfQuery(true);
try (QueryResponse response = operationTimeout == 0 ? query(sqlQuery, settings).get() :
query(sqlQuery, settings).get(operationTimeout, TimeUnit.MILLISECONDS)) {
Expand All @@ -1643,6 +1643,10 @@ public List<GenericRecord> queryAll(String sqlQuery) {
}
}

public List<GenericRecord> queryAll(String sqlQuery) {
return queryAll(sqlQuery, new QuerySettings());
}

public <T> List<T> queryAll(String sqlQuery, Class<T> clazz, TableSchema schema) {
return queryAll(sqlQuery, clazz, schema, null);
}
Expand Down Expand Up @@ -1891,6 +1895,28 @@ public Set<String> getEndpoints() {
return Collections.unmodifiableSet(endpoints);
}

/**
* Sets list of DB roles that should be applied to each query.
*
* @param dbRoles
*/
public void setDBRoles(Collection<String> dbRoles) {
this.configuration.put(ClientSettings.SESSION_DB_ROLES, ClientSettings.commaSeparated(dbRoles));
this.unmodifiableDbRolesView =
Collections.unmodifiableCollection(ClientSettings.valuesFromCommaSeparated(
this.configuration.get(ClientSettings.SESSION_DB_ROLES)));
}

private Collection<String> unmodifiableDbRolesView = Collections.emptyList();

/**
* Returns list of DB roles that should be applied to each query.
*
* @return List of DB roles
*/
public Collection<String> getDBRoles() {
return unmodifiableDbRolesView;
}

private ClickHouseNode getNextAliveNode() {
return serverNodes.get(0);
Expand Down
11 changes: 9 additions & 2 deletions client-v2/src/main/java/com/clickhouse/client/api/ClientSettings.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.List;
import java.util.stream.Collectors;

Expand All @@ -18,14 +19,20 @@ public class ClientSettings {
public static String commaSeparated(Collection<?> values) {
StringBuilder sb = new StringBuilder();
for (Object value : values) {
sb.append(value.toString().replaceAll(",", "\\,")).append(",");
sb.append(value.toString().replaceAll(",", "\\\\,")).append(",");
}
sb.setLength(sb.length() - 1);
return sb.toString();
}

public static List<String> valuesFromCommaSeparated(String value) {
return Arrays.stream(value.split(",")).map(s -> s.replaceAll("\\\\,", ","))
if (value == null || value.isEmpty()) {
return Collections.emptyList();
}

return Arrays.stream(value.split("(?<!\\\\),")).map(s -> s.replaceAll("\\\\,", ","))
.collect(Collectors.toList());
}

public static final String SESSION_DB_ROLES = "session_db_roles";
}
7 changes: 6 additions & 1 deletion client-v2/src/main/java/com/clickhouse/client/api/command/CommandResponse.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@
import com.clickhouse.client.api.metrics.ServerMetrics;
import com.clickhouse.client.api.query.QueryResponse;

public class CommandResponse{
public class CommandResponse implements AutoCloseable {

private final QueryResponse response;

Expand Down Expand Up @@ -71,4 +71,9 @@ public long getWrittenBytes() {
public long getServerTime() {
return response.getServerTime();
}

@Override
public void close() throws Exception {
response.close();
}
}
10 changes: 7 additions & 3 deletions client-v2/src/main/java/com/clickhouse/client/api/data_formats/internal/MapBackedRecord.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -241,8 +241,7 @@ public ClickHouseGeoMultiPolygonValue getGeoMultiPolygon(String colName) {

@Override
public <T> List<T> getList(String colName) {
ClickHouseArrayValue<?> array = readValue(colName);
return null;
return getList(schema.nameToIndex(colName));
}


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

@Override
public <T> List<T> getList(int index) {
return readValue(index);
Object value = readValue(index);
if (value instanceof BinaryStreamReader.ArrayValue) {
return ((BinaryStreamReader.ArrayValue) value).asList();
} else {
throw new ClientException("Column is not of array type");
}
}

@Override
Expand Down
19 changes: 19 additions & 0 deletions client-v2/src/main/java/com/clickhouse/client/api/insert/InsertSettings.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -205,4 +205,23 @@ public InsertSettings serverSetting(String name, Collection<String> values) {
rawSettings.put(ClientSettings.SERVER_SETTING_PREFIX + name, ClientSettings.commaSeparated(values));
return this;
}

/**
* Sets DB roles for an operation. Roles that were set by {@link Client#setDBRoles(Collection)} will be overridden.
*
* @param dbRoles
*/
public InsertSettings setDBRoles(Collection<String> dbRoles) {
rawSettings.put(ClientSettings.SESSION_DB_ROLES, dbRoles);
return this;
}

/**
* Gets DB roles for an operation.
*
* @return list of DB roles
*/
public Collection<String> getDBRoles() {
return (Collection<String>) rawSettings.get(ClientSettings.SESSION_DB_ROLES);
}
}
44 changes: 26 additions & 18 deletions client-v2/src/main/java/com/clickhouse/client/api/internal/HttpAPIClientHelper.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -70,6 +70,8 @@
import java.nio.charset.StandardCharsets;
import java.security.NoSuchAlgorithmException;
import java.util.Base64;
import java.util.Collection;
import java.util.Collections;
import java.util.EnumSet;
import java.util.HashSet;
import java.util.Map;
Expand Down Expand Up @@ -425,26 +427,27 @@ private void addHeaders(HttpPost req, Map<String, String> chConfig, Map<String,
}
}
private void addQueryParams(URIBuilder req, Map<String, String> chConfig, Map<String, Object> requestConfig) {
if (requestConfig == null) {
requestConfig = Collections.emptyMap();
}

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

if (requestConfig != null) {
if (requestConfig.containsKey(ClickHouseHttpOption.WAIT_END_OF_QUERY.getKey())) {
req.addParameter(ClickHouseHttpOption.WAIT_END_OF_QUERY.getKey(),
requestConfig.get(ClickHouseHttpOption.WAIT_END_OF_QUERY.getKey()).toString());
}
if (requestConfig.containsKey(ClickHouseClientOption.QUERY_ID.getKey())) {
req.addParameter(ClickHouseHttpProto.QPARAM_QUERY_ID, requestConfig.get(ClickHouseClientOption.QUERY_ID.getKey()).toString());
}
if (requestConfig.containsKey("statement_params")) {
Map<String, Object> params = (Map<String, Object>) requestConfig.get("statement_params");
for (Map.Entry<String, Object> entry : params.entrySet()) {
req.addParameter("param_" + entry.getKey(), String.valueOf(entry.getValue()));
}
if (requestConfig.containsKey(ClickHouseHttpOption.WAIT_END_OF_QUERY.getKey())) {
req.addParameter(ClickHouseHttpOption.WAIT_END_OF_QUERY.getKey(),
requestConfig.get(ClickHouseHttpOption.WAIT_END_OF_QUERY.getKey()).toString());
}
if (requestConfig.containsKey(ClickHouseClientOption.QUERY_ID.getKey())) {
req.addParameter(ClickHouseHttpProto.QPARAM_QUERY_ID, requestConfig.get(ClickHouseClientOption.QUERY_ID.getKey()).toString());
}
if (requestConfig.containsKey("statement_params")) {
Map<String, Object> params = (Map<String, Object>) requestConfig.get("statement_params");
for (Map.Entry<String, Object> entry : params.entrySet()) {
req.addParameter("param_" + entry.getKey(), String.valueOf(entry.getValue()));
}
}

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

if (requestConfig != null) {
for (Map.Entry<String, Object> entry : requestConfig.entrySet()) {
if (entry.getKey().startsWith(ClientSettings.SERVER_SETTING_PREFIX)) {
req.addParameter(entry.getKey().substring(ClientSettings.SERVER_SETTING_PREFIX.length()), entry.getValue().toString());
}
Collection<String> sessionRoles = (Collection<String>) requestConfig.getOrDefault(ClientSettings.SESSION_DB_ROLES,
ClientSettings.valuesFromCommaSeparated(chConfiguration.getOrDefault(ClientSettings.SESSION_DB_ROLES, "")));
if (!sessionRoles.isEmpty()) {

sessionRoles.forEach(r -> req.addParameter(ClickHouseHttpProto.QPARAM_ROLE, r));
}

for (Map.Entry<String, Object> entry : requestConfig.entrySet()) {
if (entry.getKey().startsWith(ClientSettings.SERVER_SETTING_PREFIX)) {
req.addParameter(entry.getKey().substring(ClientSettings.SERVER_SETTING_PREFIX.length()), entry.getValue().toString());
}
}
}
Expand Down
19 changes: 19 additions & 0 deletions client-v2/src/main/java/com/clickhouse/client/api/query/QuerySettings.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -221,4 +221,23 @@ public QuerySettings serverSetting(String name, Collection<String> values) {
rawSettings.put(ClientSettings.SERVER_SETTING_PREFIX + name, ClientSettings.commaSeparated(values));
return this;
}

/**
* Sets DB roles for an operation. Roles that were set by {@link Client#setDBRoles(Collection)} will be overridden.
*
* @param dbRoles
*/
public QuerySettings setDBRoles(Collection<String> dbRoles) {
rawSettings.put(ClientSettings.SESSION_DB_ROLES, dbRoles);
return this;
}

/**
* Gets DB roles for an operation.
*
* @return list of DB roles
*/
public Collection<String> getDBRoles() {
return (Collection<String>) rawSettings.get(ClientSettings.SESSION_DB_ROLES);
}
}
14 changes: 14 additions & 0 deletions client-v2/src/test/java/com/clickhouse/client/SettingsTests.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,19 @@
package com.clickhouse.client;

import com.clickhouse.client.api.ClientSettings;
import org.testng.Assert;
import org.testng.annotations.Test;

import java.util.Arrays;
import java.util.List;

public class SettingsTests {

@Test
void testClientSettings() {
List<String> source = Arrays.asList("ROL1", "ROL2,☺", "Rol,3,3");
String listA = ClientSettings.commaSeparated(source);
List<String> listB = ClientSettings.valuesFromCommaSeparated(listA);
Assert.assertEquals(listB, source);
}
}
71 changes: 71 additions & 0 deletions client-v2/src/test/java/com/clickhouse/client/query/QueryTests.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,8 +12,10 @@
import com.clickhouse.client.ClickHouseResponse;
import com.clickhouse.client.api.Client;
import com.clickhouse.client.api.ClientException;
import com.clickhouse.client.api.ClientSettings;
import com.clickhouse.client.api.DataTypeUtils;
import com.clickhouse.client.api.ServerException;
import com.clickhouse.client.api.command.CommandResponse;
import com.clickhouse.client.api.data_formats.ClickHouseBinaryFormatReader;
import com.clickhouse.client.api.enums.Protocol;
import com.clickhouse.client.api.insert.InsertSettings;
Expand All @@ -26,6 +28,7 @@
import com.clickhouse.client.api.query.QueryResponse;
import com.clickhouse.client.api.query.QuerySettings;
import com.clickhouse.client.api.query.Records;
import com.clickhouse.client.http.config.HttpConnectionProvider;
import com.clickhouse.data.ClickHouseDataType;
import com.clickhouse.data.ClickHouseFormat;
import com.fasterxml.jackson.databind.JsonNode;
Expand All @@ -36,6 +39,7 @@
import org.testng.annotations.BeforeMethod;
import org.testng.annotations.DataProvider;
import org.testng.annotations.Test;
import org.testng.util.Strings;

import java.io.BufferedReader;
import java.io.BufferedWriter;
Expand All @@ -56,6 +60,7 @@
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
Expand Down Expand Up @@ -1570,6 +1575,72 @@ public static BigDecimal cropDecimal(BigDecimal value, int scale) {
return new BigDecimal(bi, scale);
}

@DataProvider(name = "sessionRoles")
private static Object[][] sessionRoles() {
return new Object[][]{
{new String[]{"ROL1", "ROL2"}},
{new String[]{"ROL1", "ROL2"}},
{new String[]{"ROL1", "ROL2"}},
{new String[]{"ROL1", "ROL2,☺"}},
{new String[]{"ROL1", "ROL2"}},
};
}

@Test(groups = {"integration"}, dataProvider = "sessionRoles", dataProviderClass = QueryTests.class)
public void testOperationCustomRoles(String[] roles) throws Exception {
final String password = UUID.randomUUID().toString();
final String rolesList = "\"" + Strings.join("\",\"", roles) + "\"";
try (CommandResponse resp = client.execute("DROP ROLE IF EXISTS " + rolesList).get()) {
}
try (CommandResponse resp = client.execute("CREATE ROLE " + rolesList).get()) {
}
try (CommandResponse resp = client.execute("DROP USER IF EXISTS some_user").get()) {
}
try (CommandResponse resp = client.execute("CREATE USER some_user IDENTIFIED WITH sha256_password BY '" + password + "'" ).get()) {
}
try (CommandResponse resp = client.execute("GRANT " + rolesList + " TO some_user").get()) {
}

try (Client userClient = newClient().setUsername("some_user").setPassword(password).build()) {
QuerySettings settings = new QuerySettings().setDBRoles(Arrays.asList(roles));
List<GenericRecord> resp = userClient.queryAll("SELECT currentRoles()", settings);
Set<String> roleSet = new HashSet<>(Arrays.asList(roles));
Set<String> currentRoles = new HashSet<String> (resp.get(0).getList(1));
Assert.assertEquals(currentRoles, roleSet, "Roles " + roleSet + " not found in " + currentRoles);
}
}

@DataProvider(name = "clientSessionRoles")
private static Object[][] clientSessionRoles() {
return new Object[][]{
{new String[]{"ROL1", "ROL2"}},
{new String[]{"ROL1", "ROL2,☺"}},
};
}
@Test(groups = {"integration"}, dataProvider = "clientSessionRoles", dataProviderClass = QueryTests.class)
public void testClientCustomRoles(String[] roles) throws Exception {
final String password = UUID.randomUUID().toString();
final String rolesList = "\"" + Strings.join("\",\"", roles) + "\"";
try (CommandResponse resp = client.execute("DROP ROLE IF EXISTS " + rolesList).get()) {
}
try (CommandResponse resp = client.execute("CREATE ROLE " + rolesList).get()) {
}
try (CommandResponse resp = client.execute("DROP USER IF EXISTS some_user").get()) {
}
try (CommandResponse resp = client.execute("CREATE USER some_user IDENTIFIED WITH sha256_password BY '" + password + "'" ).get()) {
}
try (CommandResponse resp = client.execute("GRANT " + rolesList + " TO some_user").get()) {
}

try (Client userClient = newClient().setUsername("some_user").setPassword(password).build()) {
userClient.setDBRoles(Arrays.asList(roles));
List<GenericRecord> resp = userClient.queryAll("SELECT currentRoles()");
Set<String> roleSet = new HashSet<>(Arrays.asList(roles));
Set<String> currentRoles = new HashSet<String> (resp.get(0).getList(1));
Assert.assertEquals(currentRoles, roleSet, "Roles " + roleSet + " not found in " + currentRoles);
}
}


protected Client.Builder newClient() {
ClickHouseNode node = getServer(ClickHouseProtocol.HTTP);
Expand Down
Loading