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 8cae6aacd..27ec12691 100644 --- a/src/main/java/com/firebolt/jdbc/client/query/StatementClientImpl.java +++ b/src/main/java/com/firebolt/jdbc/client/query/StatementClientImpl.java @@ -117,7 +117,7 @@ private InputStream executeSqlStatementWithRetryOnUnauthorized(String label, @No private InputStream postSqlStatement(@NonNull FireboltProperties connectionProperties, String formattedStatement, String uri, String label) throws FireboltException, IOException { - Request post = createPostRequest(uri, label, formattedStatement, getConnection().getAccessToken().orElse(null)); //, statementInfoWrapper.getId()); + Request post = createPostRequest(uri, label, formattedStatement, getConnection().getAccessToken().orElse(null)); Response response = this.execute(post, connectionProperties.getHost(), connectionProperties.isCompress()); InputStream is = Optional.ofNullable(response.body()).map(ResponseBody::byteStream).orElse(null); if (is == null) { 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 3982c0d06..6663cdac4 100644 --- a/src/test/java/com/firebolt/jdbc/client/query/StatementClientImplTest.java +++ b/src/test/java/com/firebolt/jdbc/client/query/StatementClientImplTest.java @@ -22,7 +22,6 @@ import org.mockito.ArgumentCaptor; import org.mockito.Captor; import org.mockito.Mock; -import org.mockito.Mockito; import org.mockito.junit.jupiter.MockitoExtension; import java.io.IOException; @@ -124,6 +123,67 @@ void shouldCancelSqlQuery() throws SQLException, IOException { requestArgumentCaptor.getValue().url().uri().toString()); } + @Test + void shouldIgnoreIfStatementIsNotFoundInDbWhenCancelSqlQuery() throws SQLException, IOException { + FireboltProperties fireboltProperties = FireboltProperties.builder().database("db1").compress(true) + .host("firebolt1").port(555).build(); + String id = "12345"; + StatementClient statementClient = new StatementClientImpl(okHttpClient, mock(ObjectMapper.class), connection, "", ""); + PreparedStatement ps = mock(PreparedStatement.class); + ResultSet rs = mock(ResultSet.class); + when(connection.prepareStatement(anyString())).thenReturn(ps); + when(ps.executeQuery()).thenReturn(rs); + when(rs.next()).thenReturn(true); + when(rs.getString(1)).thenReturn(id); + injectMockedResponse(okHttpClient, 200, ""); + Call call = getMockedCallWithResponse(400, ""); // BAD REQUEST + when(okHttpClient.newCall(any())).thenReturn(call); + when(okHttpClient.dispatcher()).thenReturn(mock(Dispatcher.class)); + statementClient.abortStatement(id, fireboltProperties); + verify(okHttpClient).newCall(requestArgumentCaptor.capture()); + assertEquals("http://firebolt1:555/cancel?query_id=12345", + requestArgumentCaptor.getValue().url().uri().toString()); + } + + @Test + void cannotGetStatementIdWhenCancelling() throws SQLException, IOException { + FireboltProperties fireboltProperties = FireboltProperties.builder().database("db1").compress(true) + .host("firebolt1").port(555).build(); + String id = "12345"; + StatementClient statementClient = new StatementClientImpl(okHttpClient, mock(ObjectMapper.class), connection, "", ""); + PreparedStatement ps = mock(PreparedStatement.class); + ResultSet rs = mock(ResultSet.class); + when(connection.prepareStatement(anyString())).thenReturn(ps); + when(ps.executeQuery()).thenReturn(rs); + when(rs.next()).thenReturn(true); + when(rs.getString(1)).thenReturn(null); + when(okHttpClient.dispatcher()).thenReturn(mock(Dispatcher.class)); + assertEquals("Cannot retrieve id for statement with label " + id, assertThrows(FireboltException.class, () -> statementClient.abortStatement(id, fireboltProperties)).getMessage()); + } + + @ParameterizedTest + @CsvSource({ + "401,The operation is not authorized", + "500, Server failed to execute query" + }) + void shouldFailToCancelSqlQuery(int httpStatus, String errorMessage) throws SQLException, IOException { + FireboltProperties fireboltProperties = FireboltProperties.builder().database("db1").compress(true) + .host("firebolt1").port(555).build(); + String id = "12345"; + StatementClient statementClient = new StatementClientImpl(okHttpClient, mock(ObjectMapper.class), connection, "", ""); + PreparedStatement ps = mock(PreparedStatement.class); + ResultSet rs = mock(ResultSet.class); + when(connection.prepareStatement(anyString())).thenReturn(ps); + when(ps.executeQuery()).thenReturn(rs); + when(rs.next()).thenReturn(true); + when(rs.getString(1)).thenReturn(id); + Call call = getMockedCallWithResponse(httpStatus, ""); + when(okHttpClient.newCall(any())).thenReturn(call); + when(okHttpClient.dispatcher()).thenReturn(mock(Dispatcher.class)); + FireboltException e = assertThrows(FireboltException.class, () -> statementClient.abortStatement(id, fireboltProperties)); + assertTrue(e.getMessage().contains(errorMessage)); + } + @Test void shouldRetryOnUnauthorized() throws IOException, SQLException { FireboltProperties fireboltProperties = FireboltProperties.builder().database("db1").compress(true)