fixed warnings + more unit tests

firebolt-db · Dec 4, 2023 · 63644b0 · 63644b0
1 parent 00bd5cd
commit 63644b0
Show file tree

Hide file tree

Showing 2 changed files with 62 additions and 2 deletions.
diff --git 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
@@ -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)