Skip to content

Commit 4713763

Browse files
chernserchernser
committed
added method to update bearer token
1 parent 0418232 commit 4713763

File tree

5 files changed

+77
-72
lines changed

5 files changed

+77
-72
lines changed

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

Lines changed: 17 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -117,7 +117,6 @@
117117
* ...
118118
* }
119119
* }
120-
*
121120
* }
122121
*
123122
*
@@ -131,6 +130,9 @@ public class Client implements AutoCloseable {
131130

132131
private final Set<String> endpoints;
133132
private final Map<String, String> configuration;
133+
134+
private final Map<String, String> readOnlyConfig;
135+
134136
private final List<ClickHouseNode> serverNodes = new ArrayList<>();
135137

136138
// POJO serializer mapping (class -> (schema -> (format -> serializer)))
@@ -154,10 +156,10 @@ public class Client implements AutoCloseable {
154156
private final ColumnToMethodMatchingStrategy columnToMethodMatchingStrategy;
155157

156158
private Client(Set<String> endpoints, Map<String,String> configuration, boolean useNewImplementation,
157-
ExecutorService sharedOperationExecutor, ColumnToMethodMatchingStrategy columnToMethodMatchingStrategy,
158-
Supplier<String> bearerTokenSupplier) {
159+
ExecutorService sharedOperationExecutor, ColumnToMethodMatchingStrategy columnToMethodMatchingStrategy) {
159160
this.endpoints = endpoints;
160161
this.configuration = configuration;
162+
this.readOnlyConfig = Collections.unmodifiableMap(this.configuration);
161163
this.endpoints.forEach(endpoint -> {
162164
this.serverNodes.add(ClickHouseNode.of(endpoint, this.configuration));
163165
});
@@ -172,7 +174,7 @@ private Client(Set<String> endpoints, Map<String,String> configuration, boolean
172174
}
173175
this.useNewImplementation = useNewImplementation;
174176
if (useNewImplementation) {
175-
this.httpClientHelper = new HttpAPIClientHelper(configuration, bearerTokenSupplier);
177+
this.httpClientHelper = new HttpAPIClientHelper(configuration);
176178
LOG.info("Using new http client implementation");
177179
} else {
178180
this.oldClient = ClientV1AdaptorHelper.createClient(configuration);
@@ -226,8 +228,6 @@ public static class Builder {
226228
private ExecutorService sharedOperationExecutor = null;
227229
private ColumnToMethodMatchingStrategy columnToMethodMatchingStrategy;
228230

229-
private Supplier<String> bearerTokenSupplier = null;
230-
231231
public Builder() {
232232
this.endpoints = new HashSet<>();
233233
this.configuration = new HashMap<String, String>();
@@ -855,7 +855,7 @@ public Builder allowBinaryReaderToReuseBuffers(boolean reuse) {
855855
* @return same instance of the builder
856856
*/
857857
public Builder httpHeader(String key, String value) {
858-
this.configuration.put(ClientConfigProperties.HTTP_HEADER_PREFIX + key.toUpperCase(Locale.US), value);
858+
this.configuration.put(ClientConfigProperties.httpHeader(key), value);
859859
return this;
860860
}
861861

@@ -866,7 +866,7 @@ public Builder httpHeader(String key, String value) {
866866
* @return same instance of the builder
867867
*/
868868
public Builder httpHeader(String key, Collection<String> values) {
869-
this.configuration.put(ClientConfigProperties.HTTP_HEADER_PREFIX + key.toUpperCase(Locale.US), ClientConfigProperties.commaSeparated(values));
869+
this.configuration.put(ClientConfigProperties.httpHeader(key), ClientConfigProperties.commaSeparated(values));
870870
return this;
871871
}
872872

@@ -965,21 +965,8 @@ public Builder setOptions(Map<String, String> options) {
965965
* @return same instance of the builder
966966
*/
967967
public Builder useBearerTokenAuth(String bearerToken) {
968-
this.httpHeader("Authorization", "Bearer " + bearerToken);
969-
return this;
970-
}
971-
972-
/**
973-
* Specifies a supplier for a bearer tokens. It is useful when token should be refreshed.
974-
* Supplier is called each time before sending a request.
975-
* Supplier should return encoded token.
976-
* This configuration cannot be used with {@link #useBearerTokenAuth(String)}.
977-
*
978-
* @param tokenSupplier - token supplier
979-
* @return
980-
*/
981-
public Builder useBearerTokenAuth(Supplier<String> tokenSupplier) {
982-
this.bearerTokenSupplier = tokenSupplier;
968+
// Most JWT libraries (https://jwt.io/libraries?language=Java) compact tokens in proper way
969+
this.httpHeader(HttpHeaders.AUTHORIZATION, "Bearer " + bearerToken);
983970
return this;
984971
}
985972

@@ -994,8 +981,7 @@ public Client build() {
994981
if (!this.configuration.containsKey("access_token") &&
995982
(!this.configuration.containsKey("user") || !this.configuration.containsKey("password")) &&
996983
!MapUtils.getFlag(this.configuration, "ssl_authentication", false) &&
997-
!this.configuration.containsKey(ClientConfigProperties.HTTP_HEADER_PREFIX + "Authorization") &&
998-
this.bearerTokenSupplier == null) {
984+
!this.configuration.containsKey(ClientConfigProperties.httpHeader(HttpHeaders.AUTHORIZATION))) {
999985
throw new IllegalArgumentException("Username and password (or access token or SSL authentication or pre-define Authorization header) are required");
1000986
}
1001987

@@ -1004,11 +990,6 @@ public Client build() {
1004990
throw new IllegalArgumentException("Only one of password, access token or SSL authentication can be used per client.");
1005991
}
1006992

1007-
if (this.configuration.containsKey(ClientConfigProperties.HTTP_HEADER_PREFIX + "Authorization") &&
1008-
this.bearerTokenSupplier != null) {
1009-
throw new IllegalArgumentException("Bearer token supplier cannot be used with a predefined Authorization header");
1010-
}
1011-
1012993
if (this.configuration.containsKey("ssl_authentication") &&
1013994
!this.configuration.containsKey(ClientConfigProperties.SSL_CERTIFICATE.getKey())) {
1014995
throw new IllegalArgumentException("SSL authentication requires a client certificate");
@@ -1047,15 +1028,8 @@ public Client build() {
10471028
throw new IllegalArgumentException("Nor server timezone nor specific timezone is set");
10481029
}
10491030

1050-
// check for only new implementation configuration
1051-
if (!this.useNewImplementation) {
1052-
if (this.bearerTokenSupplier != null) {
1053-
throw new IllegalArgumentException("Bearer token supplier cannot be used with old implementation");
1054-
}
1055-
}
1056-
10571031
return new Client(this.endpoints, this.configuration, this.useNewImplementation, this.sharedOperationExecutor,
1058-
this.columnToMethodMatchingStrategy, this.bearerTokenSupplier);
1032+
this.columnToMethodMatchingStrategy);
10591033
}
10601034

10611035
private static final int DEFAULT_NETWORK_BUFFER_SIZE = 300_000;
@@ -2147,7 +2121,7 @@ public String toString() {
21472121
* @return - configuration options
21482122
*/
21492123
public Map<String, String> getConfiguration() {
2150-
return Collections.unmodifiableMap(configuration);
2124+
return readOnlyConfig;
21512125
}
21522126

21532127
/** Returns operation timeout in seconds */
@@ -2194,6 +2168,10 @@ public Collection<String> getDBRoles() {
21942168
return unmodifiableDbRolesView;
21952169
}
21962170

2171+
public void updateBearerToken(String bearer) {
2172+
this.configuration.put(ClientConfigProperties.httpHeader(HttpHeaders.AUTHORIZATION), "Bearer " + bearer);
2173+
}
2174+
21972175
private ClickHouseNode getNextAliveNode() {
21982176
return serverNodes.get(0);
21992177
}

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
import java.util.Collection;
55
import java.util.Collections;
66
import java.util.List;
7+
import java.util.Locale;
78
import java.util.stream.Collectors;
89

910
/**
@@ -157,6 +158,10 @@ public static String serverSetting(String key) {
157158
return SERVER_SETTING_PREFIX + key;
158159
}
159160

161+
public static String httpHeader(String key) {
162+
return HTTP_HEADER_PREFIX + key.toUpperCase(Locale.US);
163+
}
164+
160165
public static String commaSeparated(Collection<?> values) {
161166
StringBuilder sb = new StringBuilder();
162167
for (Object value : values) {

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

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,12 +7,34 @@ public class ServerException extends RuntimeException {
77
public static final int TABLE_NOT_FOUND = 60;
88

99
private final int code;
10+
11+
private final int transportProtocolCode;
12+
1013
public ServerException(int code, String message) {
14+
this(code, message, 500);
15+
}
16+
17+
public ServerException(int code, String message, int transportProtocolCode) {
1118
super(message);
1219
this.code = code;
20+
this.transportProtocolCode = transportProtocolCode;
1321
}
1422

23+
/**
24+
* Returns CH server error code. May return 0 if code is unknown.
25+
* @return - error code from server response
26+
*/
1527
public int getCode() {
1628
return code;
1729
}
30+
31+
/**
32+
* Returns error code of underlying transport protocol. For example, HTTP status.
33+
* By default, will return {@code 500 } what is derived from HTTP Server Internal Error.
34+
*
35+
* @return - transport status code
36+
*/
37+
public int getTransportProtocolCode() {
38+
return transportProtocolCode;
39+
}
1840
}

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

Lines changed: 8 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -94,12 +94,9 @@ public class HttpAPIClientHelper {
9494

9595
private String httpClientUserAgentPart;
9696

97-
private final Supplier<String> bearerTokenSupplier;
98-
99-
public HttpAPIClientHelper(Map<String, String> configuration, Supplier<String> bearerTokenSupplier) {
97+
public HttpAPIClientHelper(Map<String, String> configuration) {
10098
this.chConfiguration = configuration;
10199
this.httpClient = createHttpClient();
102-
this.bearerTokenSupplier = bearerTokenSupplier;
103100

104101
this.httpClientUserAgentPart = this.httpClient.getClass().getPackage().getImplementationTitle() + "/" + this.httpClient.getClass().getPackage().getImplementationVersion();
105102

@@ -339,10 +336,13 @@ public Exception readError(ClassicHttpResponse httpResponse) {
339336

340337
String msg = msgBuilder.toString().replaceAll("\\s+", " ").replaceAll("\\\\n", " ")
341338
.replaceAll("\\\\/", "/");
342-
return new ServerException(serverCode, msg);
339+
if (msg.trim().isEmpty()) {
340+
msg = String.format(ERROR_CODE_PREFIX_PATTERN, serverCode) + " <Unreadable error message> (transport error: " + httpResponse.getCode() + ")";
341+
}
342+
return new ServerException(serverCode, msg, httpResponse.getCode());
343343
} catch (Exception e) {
344344
LOG.error("Failed to read error message", e);
345-
return new ServerException(serverCode, String.format(ERROR_CODE_PREFIX_PATTERN, serverCode) + " <Unreadable error message>");
345+
return new ServerException(serverCode, String.format(ERROR_CODE_PREFIX_PATTERN, serverCode) + " <Unreadable error message> (transport error: " + httpResponse.getCode() + ")", httpResponse.getCode());
346346
}
347347
}
348348

@@ -427,8 +427,6 @@ private void addHeaders(HttpPost req, Map<String, String> chConfig, Map<String,
427427
if (MapUtils.getFlag(chConfig, "ssl_authentication", false)) {
428428
req.addHeader(ClickHouseHttpProto.HEADER_DB_USER, chConfig.get(ClientConfigProperties.USER.getKey()));
429429
req.addHeader(ClickHouseHttpProto.HEADER_SSL_CERT_AUTH, "on");
430-
} else if (bearerTokenSupplier != null) {
431-
req.addHeader(HttpHeaders.AUTHORIZATION, "Bearer " + bearerTokenSupplier.get());
432430
} else if (chConfig.getOrDefault(ClientConfigProperties.HTTP_USE_BASIC_AUTH.getKey(), "true").equalsIgnoreCase("true")) {
433431
req.addHeader(HttpHeaders.AUTHORIZATION, "Basic " + Base64.getEncoder().encodeToString(
434432
(chConfig.get(ClientConfigProperties.USER.getKey()) + ":" + chConfig.get(ClientConfigProperties.PASSWORD.getKey())).getBytes(StandardCharsets.UTF_8)));
@@ -456,12 +454,12 @@ private void addHeaders(HttpPost req, Map<String, String> chConfig, Map<String,
456454

457455
for (Map.Entry<String, String> entry : chConfig.entrySet()) {
458456
if (entry.getKey().startsWith(ClientConfigProperties.HTTP_HEADER_PREFIX)) {
459-
req.addHeader(entry.getKey().substring(ClientConfigProperties.HTTP_HEADER_PREFIX.length()), entry.getValue());
457+
req.setHeader(entry.getKey().substring(ClientConfigProperties.HTTP_HEADER_PREFIX.length()), entry.getValue());
460458
}
461459
}
462460
for (Map.Entry<String, Object> entry : requestConfig.entrySet()) {
463461
if (entry.getKey().startsWith(ClientConfigProperties.HTTP_HEADER_PREFIX)) {
464-
req.addHeader(entry.getKey().substring(ClientConfigProperties.HTTP_HEADER_PREFIX.length()), entry.getValue().toString());
462+
req.setHeader(entry.getKey().substring(ClientConfigProperties.HTTP_HEADER_PREFIX.length()), entry.getValue().toString());
465463
}
466464
}
467465

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

Lines changed: 25 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -809,6 +809,7 @@ public void testBearerTokenAuth() throws Exception {
809809
.reduce((s1, s2) -> s1 + "." + s2).get();
810810
try (Client client = new Client.Builder().addEndpoint(Protocol.HTTP, "localhost", mockServer.port(), false)
811811
.useBearerTokenAuth(jwtToken1)
812+
.compressServerResponse(false)
812813
.build()) {
813814

814815
mockServer.addStubMapping(WireMock.post(WireMock.anyUrl())
@@ -828,37 +829,38 @@ public void testBearerTokenAuth() throws Exception {
828829
new String[]{"header2", "payload2", "signature2"})
829830
.map(s -> Base64.getEncoder().encodeToString(s.getBytes(StandardCharsets.UTF_8)))
830831
.reduce((s1, s2) -> s1 + "." + s2).get();
831-
final AtomicInteger callCount = new AtomicInteger(0);
832-
Supplier<String> tokenSupplier = () -> {
833-
callCount.incrementAndGet();
834-
return jwtToken2;
835-
};
836-
832+
833+
mockServer.resetAll();
837834
mockServer.addStubMapping(WireMock.post(WireMock.anyUrl())
838-
.withHeader("Authorization", WireMock.equalTo("Bearer " + jwtToken2))
835+
.withHeader("Authorization", WireMock.equalTo("Bearer " + jwtToken1))
839836
.willReturn(WireMock.aResponse()
840-
.withHeader("X-ClickHouse-Summary",
841-
"{ \"read_bytes\": \"10\", \"read_rows\": \"1\"}")).build());
837+
.withStatus(HttpStatus.SC_UNAUTHORIZED))
838+
.build());
842839

843840
try (Client client = new Client.Builder().addEndpoint(Protocol.HTTP, "localhost", mockServer.port(), false)
844-
.useBearerTokenAuth(tokenSupplier)
841+
.useBearerTokenAuth(jwtToken1)
842+
.compressServerResponse(false)
845843
.build()) {
846844

847-
for (int i = 0; i < 3; i++ ) {
848-
849-
try (QueryResponse response = client.query("SELECT 1").get(1, TimeUnit.SECONDS)) {
850-
Assert.assertEquals(response.getReadBytes(), 10);
851-
} catch (Exception e) {
852-
Assert.fail("Unexpected exception", e);
853-
}
845+
try {
846+
client.execute("SELECT 1").get();
847+
fail("Exception expected");
848+
} catch (ServerException e) {
849+
Assert.assertEquals(e.getTransportProtocolCode(), HttpStatus.SC_UNAUTHORIZED);
854850
}
855-
}
856851

857-
assertEquals(callCount.get(), 3);
852+
mockServer.resetAll();
853+
mockServer.addStubMapping(WireMock.post(WireMock.anyUrl())
854+
.withHeader("Authorization", WireMock.equalTo("Bearer " + jwtToken2))
855+
.willReturn(WireMock.aResponse()
856+
.withHeader("X-ClickHouse-Summary",
857+
"{ \"read_bytes\": \"10\", \"read_rows\": \"1\"}"))
858858

859-
assertThrows(IllegalArgumentException.class, () -> {
860-
new Client.Builder().useBearerTokenAuth("token")
861-
.useBearerTokenAuth(() -> "token2").build();
862-
});
859+
.build());
860+
861+
client.updateBearerToken(jwtToken2);
862+
863+
client.execute("SELECT 1").get();
864+
}
863865
}
864866
}

0 commit comments

Comments
 (0)