Skip to content

Commit

Permalink
throw special exception when service account does not have associated…
Browse files Browse the repository at this point in the history 
… user
  • Loading branch information
@alexradzin
alexradzin committed Nov 15, 2023
1 parent 1567b93 commit 44dd30e
Show file tree
Hide file tree
Showing 3 changed files with 70 additions and 34 deletions.
52 changes: 29 additions & 23 deletions src/main/java/com/firebolt/jdbc/client/FireboltClient.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;

import static java.lang.String.format;
import static java.net.HttpURLConnection.HTTP_UNAUTHORIZED;
import static java.net.HttpURLConnection.HTTP_UNAVAILABLE;

Expand Down Expand Up @@ -136,40 +137,45 @@ protected void validateResponse(String host, Response response, Boolean isCompre
int statusCode = response.code();
if (!isCallSuccessful(statusCode)) {
if (statusCode == HTTP_UNAVAILABLE) {
throw new FireboltException(
String.format("Could not query Firebolt at %s. The engine is not running.", host), statusCode);
throw new FireboltException(format("Could not query Firebolt at %s. The engine is not running.", host), statusCode);
}
String errorResponseMessage;
try {
String errorMessageFromServer = extractErrorMessage(response, isCompress);
errorResponseMessage = String.format(
"Server failed to execute query with the following error:%n%s%ninternal error:%n%s",
errorMessageFromServer, this.getInternalErrorWithHeadersText(response));
if (statusCode == HTTP_UNAUTHORIZED) {
this.getConnection().removeExpiredTokens();
throw new FireboltException(String.format(
"Could not query Firebolt at %s. The operation is not authorized or the token is expired and has been cleared from the cache.%n%s",
host, errorResponseMessage), statusCode, errorMessageFromServer);
}
throw new FireboltException(errorResponseMessage, statusCode, errorMessageFromServer);
} catch (IOException e) {
log.warn("Could not parse response containing the error message from Firebolt", e);
errorResponseMessage = String.format("Server failed to execute query%ninternal error:%n%s",
this.getInternalErrorWithHeadersText(response));
throw new FireboltException(errorResponseMessage, statusCode, e);
String errorMessageFromServer = extractErrorMessage(response, isCompress);
validateResponse(host, statusCode, errorMessageFromServer);
String errorResponseMessage = format(
"Server failed to execute query with the following error:%n%s%ninternal error:%n%s",
errorMessageFromServer, getInternalErrorWithHeadersText(response));
if (statusCode == HTTP_UNAUTHORIZED) {
getConnection().removeExpiredTokens();
throw new FireboltException(format(
"Could not query Firebolt at %s. The operation is not authorized or the token is expired and has been cleared from the cache.%n%s",
host, errorResponseMessage), statusCode, errorMessageFromServer);
}
throw new FireboltException(errorResponseMessage, statusCode, errorMessageFromServer);
}
}

protected void validateResponse(String host, int statusCode, String errorMessageFromServer) throws FireboltException {
// empty implementation
}

protected String getResponseAsString(Response response) throws FireboltException, IOException {
if (response.body() == null) {
throw new FireboltException("Cannot get resource: the response from the server is empty");
}
return response.body().string();
}

private String extractErrorMessage(Response response, boolean isCompress) throws IOException {
byte[] entityBytes = response.body() != null ? response.body().bytes() : null;
private String extractErrorMessage(Response response, boolean isCompress) throws FireboltException {
byte[] entityBytes;
try {
entityBytes = response.body() != null ? response.body().bytes() : null;
} catch (IOException e) {
log.warn("Could not parse response containing the error message from Firebolt", e);
String errorResponseMessage = format("Server failed to execute query%ninternal error:%n%s",
this.getInternalErrorWithHeadersText(response));
throw new FireboltException(errorResponseMessage, response.code(), e);
}

if (entityBytes == null) {
return null;
}
Expand All @@ -185,7 +191,7 @@ private String extractErrorMessage(Response response, boolean isCompress) throws
return new String(entityBytes, StandardCharsets.UTF_8);
}

private boolean isCallSuccessful(int statusCode) {
protected boolean isCallSuccessful(int statusCode) {
return statusCode >= 200 && statusCode <= 299; // Call is considered successful when the status code is 2XX
}

Expand Down
11 changes: 11 additions & 0 deletions src/main/java/com/firebolt/jdbc/client/query/StatementClientImpl.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,8 @@
import static com.firebolt.jdbc.connection.settings.FireboltQueryParameterKey.OUTPUT_FORMAT;
import static com.firebolt.jdbc.exception.ExceptionType.UNAUTHORIZED;
import static java.lang.String.format;
import static java.net.HttpURLConnection.HTTP_INTERNAL_ERROR;
import static java.net.HttpURLConnection.HTTP_UNAUTHORIZED;

@CustomLog
public class StatementClientImpl extends FireboltClient implements StatementClient {
Expand Down Expand Up @@ -232,4 +234,13 @@ private Optional<Pair<String, String>> getResponseFormatParameter(boolean isQuer
private Map<String, String> getCancelParameters(String statementId) {
return Map.of(FireboltQueryParameterKey.QUERY_ID.getKey(), statementId);
}

@Override
protected void validateResponse(String host, int statusCode, String errorMessageFromServer) throws FireboltException {
if (statusCode == HTTP_INTERNAL_ERROR && errorMessageFromServer != null && errorMessageFromServer.contains("HTTP status code: 401")) {
throw new FireboltException(
format("Could not query Firebolt at %s. Please associate user with your service account.", host),
HTTP_UNAUTHORIZED, errorMessageFromServer);
}
}
}
41 changes: 30 additions & 11 deletions src/test/java/com/firebolt/jdbc/client/query/StatementClientImplTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -76,8 +76,8 @@ private Entry<String, String> shouldPostSqlQueryForSystemEngine(boolean systemEn
.thenReturn(Optional.of("token"));
StatementClient statementClient = new StatementClientImpl(okHttpClient, mock(ObjectMapper.class), connection,
"ConnA:1.0.9", "ConnB:2.0.9");
injectMockedResponse(okHttpClient, 200);
Call call = getMockedCallWithResponse(200);
injectMockedResponse(okHttpClient, 200, "");
Call call = getMockedCallWithResponse(200, "");
when(okHttpClient.newCall(any())).thenReturn(call);
StatementInfoWrapper statementInfoWrapper = StatementUtil.parseToStatementInfoWrappers("show databases").get(0);
statementClient.executeSqlStatement(statementInfoWrapper, fireboltProperties, fireboltProperties.isSystemEngine(), 15, true);
Expand All @@ -100,8 +100,8 @@ void shouldCancelSqlQuery() throws FireboltException, IOException {
.host("firebolt1").port(555).build();
StatementClient statementClient = new StatementClientImpl(okHttpClient, mock(ObjectMapper.class), connection,
"", "");
injectMockedResponse(okHttpClient, 200);
Call call = getMockedCallWithResponse(200);
injectMockedResponse(okHttpClient, 200, "");
Call call = getMockedCallWithResponse(200, "");
when(okHttpClient.newCall(any())).thenReturn(call);
statementClient.abortStatement("12345", fireboltProperties);
verify(okHttpClient).newCall(requestArgumentCaptor.capture());
Expand All @@ -115,8 +115,8 @@ void shouldRetryOnUnauthorized() throws IOException, FireboltException {
.host("firebolt1").port(555).build();
when(connection.getAccessToken()).thenReturn(Optional.of("oldToken"))
.thenReturn(Optional.of("newToken"));
Call okCall = getMockedCallWithResponse(200);
Call unauthorizedCall = getMockedCallWithResponse(401);
Call okCall = getMockedCallWithResponse(200, "");
Call unauthorizedCall = getMockedCallWithResponse(401, "");
when(okHttpClient.newCall(any())).thenReturn(unauthorizedCall).thenReturn(okCall);
StatementClient statementClient = new StatementClientImpl(okHttpClient, mock(ObjectMapper.class), connection,
"ConnA:1.0.9", "ConnB:2.0.9");
Expand All @@ -132,8 +132,8 @@ void shouldRetryOnUnauthorized() throws IOException, FireboltException {
void shouldNotRetryNoMoreThanOnceOnUnauthorized() throws IOException, FireboltException {
FireboltProperties fireboltProperties = FireboltProperties.builder().database("db1").compress(true)
.host("firebolt1").port(555).build();
Call okCall = getMockedCallWithResponse(200);
Call unauthorizedCall = getMockedCallWithResponse(401);
Call okCall = getMockedCallWithResponse(200, "");
Call unauthorizedCall = getMockedCallWithResponse(401, "");
when(okHttpClient.newCall(any())).thenReturn(unauthorizedCall).thenReturn(unauthorizedCall).thenReturn(okCall);
StatementClient statementClient = new StatementClientImpl(okHttpClient, mock(ObjectMapper.class), connection,
"ConnA:1.0.9", "ConnB:2.0.9");
Expand All @@ -144,6 +144,22 @@ void shouldNotRetryNoMoreThanOnceOnUnauthorized() throws IOException, FireboltEx
verify(connection, times(2)).removeExpiredTokens();
}

@Test
void shouldThrowUnauthorizedExceptionWhenNoAssociatedUser() throws IOException, FireboltException {
FireboltProperties fireboltProperties = FireboltProperties.builder().database("db1").compress(true)
.host("firebolt1").port(555).build();
when(connection.getAccessToken()).thenReturn(Optional.of("token"));
Call unauthorizedCall = getMockedCallWithResponse(500, "HTTP status code: 401 Unauthorized, body: {\"error\":\"Authentication token is invalid\",\"code\":16,\"message\":\"Authentication token is invalid\",\"details\":[{\"@type\":\"type.googleapis.com/google.rpc.DebugInfo\",\"stack_entries\":[],\"detail\":\"failed to get user_id from fawkes: entity not found\"}]}");

when(okHttpClient.newCall(any())).thenReturn(unauthorizedCall);
StatementClient statementClient = new StatementClientImpl(okHttpClient, mock(ObjectMapper.class), connection,
"ConnA:1.0.9", "ConnB:2.0.9");
StatementInfoWrapper statementInfoWrapper = StatementUtil.parseToStatementInfoWrappers("show databases").get(0);
FireboltException exception = assertThrows(FireboltException.class, () -> statementClient.executeSqlStatement(statementInfoWrapper, fireboltProperties, false, 5, true));
assertEquals(ExceptionType.UNAUTHORIZED, exception.getType());
assertEquals("Could not query Firebolt at firebolt1. Please associate user with your service account.", exception.getMessage());
}

@ParameterizedTest
@CsvSource({
"java.io.IOException, ERROR",
Expand All @@ -163,10 +179,13 @@ <T extends Exception> void shouldThrowIOException(Class<T> exceptionClass, Excep
assertEquals(exceptionClass, ex.getCause().getClass());
}

private Call getMockedCallWithResponse(int statusCode) throws IOException {
private Call getMockedCallWithResponse(int statusCode, String content) throws IOException {
Call call = mock(Call.class);
Response response = mock(Response.class);
lenient().when(response.code()).thenReturn(statusCode);
ResponseBody body = mock(ResponseBody.class);
lenient().when(response.body()).thenReturn(body);
lenient().when(body.bytes()).thenReturn(content.getBytes());
lenient().when(call.execute()).thenReturn(response);
return call;
}
Expand All @@ -177,14 +196,14 @@ private Map<String, String> extractHeadersMap(Request request) {
return headers;
}


private void injectMockedResponse(OkHttpClient httpClient, int code) throws IOException {
private void injectMockedResponse(OkHttpClient httpClient, int code, String content) throws IOException {
Response response = mock(Response.class);
Call call = mock(Call.class);
lenient().when(httpClient.newCall(any())).thenReturn(call);
lenient().when(call.execute()).thenReturn(response);
ResponseBody body = mock(ResponseBody.class);
lenient().when(response.body()).thenReturn(body);
lenient().when(body.bytes()).thenReturn(content.getBytes());
lenient().when(response.code()).thenReturn(code);
}

Expand Down

0 comments on commit 44dd30e

Please sign in to comment.