From 44dd30e25a6de965ec8af2bbdc6f852539a075eb Mon Sep 17 00:00:00 2001 From: alexradzin Date: Wed, 15 Nov 2023 12:36:54 +0200 Subject: [PATCH] throw special exception when service account does not have associated user --- .../firebolt/jdbc/client/FireboltClient.java | 52 +++++++++++-------- .../client/query/StatementClientImpl.java | 11 ++++ .../client/query/StatementClientImplTest.java | 41 +++++++++++---- 3 files changed, 70 insertions(+), 34 deletions(-) diff --git a/src/main/java/com/firebolt/jdbc/client/FireboltClient.java b/src/main/java/com/firebolt/jdbc/client/FireboltClient.java index a21c9481b..ceb8953c6 100644 --- a/src/main/java/com/firebolt/jdbc/client/FireboltClient.java +++ b/src/main/java/com/firebolt/jdbc/client/FireboltClient.java @@ -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; @@ -136,31 +137,27 @@ 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"); @@ -168,8 +165,17 @@ protected String getResponseAsString(Response response) throws FireboltException 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; } @@ -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 } diff --git a/src/main/java/com/firebolt/jdbc/client/query/StatementClientImpl.java b/src/main/java/com/firebolt/jdbc/client/query/StatementClientImpl.java index c58bc9765..8385bf16b 100644 --- a/src/main/java/com/firebolt/jdbc/client/query/StatementClientImpl.java +++ b/src/main/java/com/firebolt/jdbc/client/query/StatementClientImpl.java @@ -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 { @@ -232,4 +234,13 @@ private Optional> getResponseFormatParameter(boolean isQuer private Map 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); + } + } } diff --git a/src/test/java/com/firebolt/jdbc/client/query/StatementClientImplTest.java b/src/test/java/com/firebolt/jdbc/client/query/StatementClientImplTest.java index 903538225..5ba9c34e2 100644 --- a/src/test/java/com/firebolt/jdbc/client/query/StatementClientImplTest.java +++ b/src/test/java/com/firebolt/jdbc/client/query/StatementClientImplTest.java @@ -76,8 +76,8 @@ private Entry 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); @@ -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()); @@ -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"); @@ -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"); @@ -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", @@ -163,10 +179,13 @@ void shouldThrowIOException(Class 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; } @@ -177,14 +196,14 @@ private Map 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); }