Skip to content

Commit

Permalink
FIR-32353: refactored FireboltStatementService.execute() (#389)
Browse files Browse the repository at this point in the history
  • Loading branch information
Alexander Radzin authored Apr 21, 2024
1 parent de8d4eb commit b03f605
Show file tree
Hide file tree
Showing 8 changed files with 79 additions and 102 deletions.
15 changes: 10 additions & 5 deletions src/main/java/com/firebolt/jdbc/resultset/FireboltResultSet.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
import com.firebolt.jdbc.QueryResult;
import com.firebolt.jdbc.annotation.ExcludeFromJacocoGeneratedReport;
import com.firebolt.jdbc.annotation.NotImplemented;
import com.firebolt.jdbc.connection.settings.FireboltProperties;
import com.firebolt.jdbc.exception.ExceptionType;
import com.firebolt.jdbc.exception.FireboltException;
import com.firebolt.jdbc.exception.FireboltSQLFeatureNotSupportedException;
Expand Down Expand Up @@ -94,11 +95,11 @@ public FireboltResultSet(InputStream is) throws SQLException {
this(is, null, null, null, false, null, false);
}

public FireboltResultSet(InputStream is, String tableName, String dbName, Integer bufferSize) throws SQLException {
FireboltResultSet(InputStream is, String tableName, String dbName, Integer bufferSize) throws SQLException {
this(is, tableName, dbName, bufferSize, false, null, false);
}

public FireboltResultSet(InputStream is, String tableName, String dbName, Integer bufferSize,
FireboltResultSet(InputStream is, String tableName, String dbName, Integer bufferSize,
FireboltStatement fireboltStatement) throws SQLException {
this(is, tableName, dbName, bufferSize, false, fireboltStatement, false);
}
Expand All @@ -114,13 +115,17 @@ private FireboltResultSet() {
maxFieldSize = 0; // 0 value means unlimited
}

public FireboltResultSet(InputStream is, String tableName, String dbName, Integer bufferSize, boolean isCompressed,
FireboltResultSet(InputStream is, String tableName, String dbName, Integer bufferSize, boolean isCompressed,
FireboltStatement statement, boolean logResultSet) throws SQLException {
this(is, tableName, dbName, bufferSize, 0, 0, isCompressed, statement, logResultSet);
this(is, tableName, dbName, bufferSize, statement == null ? 0 : statement.getMaxRows(), statement == null ? 0 : statement.getMaxFieldSize(), isCompressed, statement, logResultSet);
}

public FireboltResultSet(InputStream is, String tableName, String dbName, FireboltProperties properties, FireboltStatement statement) throws SQLException {
this(is, tableName, dbName, properties.getBufferSize(), statement == null ? 0 : statement.getMaxRows(), statement == null ? 0 : statement.getMaxFieldSize(), properties.isCompress(), statement, properties.isLogResultSet());
}

@SuppressWarnings("java:S107") //Number of parameters (8) > max (7). This is the price of the immutability
public FireboltResultSet(InputStream is, String tableName, String dbName, Integer bufferSize, int maxRows, int maxFieldSize, boolean isCompressed,
private FireboltResultSet(InputStream is, String tableName, String dbName, Integer bufferSize, int maxRows, int maxFieldSize, boolean isCompressed,
FireboltStatement statement, boolean logResultSet) throws SQLException {
log.debug("Creating resultSet...");
this.statement = statement;
Expand Down
20 changes: 9 additions & 11 deletions src/main/java/com/firebolt/jdbc/service/FireboltStatementService.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -31,21 +31,18 @@ public class FireboltStatementService {
*
* @param statementInfoWrapper the statement info
* @param properties the connection properties
* @param queryTimeout query timeout
* @param maxRows max rows
* @param maxFieldSize max field size
* @param standardSql indicates if standard sql should be used
* @param systemEngine indicates if system engine is used
* @param statement the statement
* @return an InputStream with the result
*/
@SuppressWarnings("java:S107") //Number of parameters (8) > max (7). Not nice but probably not so bad for internal class
public Optional<ResultSet> execute(StatementInfoWrapper statementInfoWrapper,
FireboltProperties properties, int queryTimeout, int maxRows, int maxFieldSize, boolean standardSql, boolean systemEngine, FireboltStatement statement)
FireboltProperties properties, boolean standardSql, FireboltStatement statement)
throws SQLException {
int queryTimeout = statement.getQueryTimeout();
boolean systemEngine = properties.isSystemEngine();
InputStream is = statementClient.executeSqlStatement(statementInfoWrapper, properties, systemEngine, queryTimeout, standardSql);
if (statementInfoWrapper.getType() == StatementType.QUERY) {
return Optional.of(createResultSet(is, (QueryRawStatement) statementInfoWrapper.getInitialStatement(), properties, statement, maxRows, maxFieldSize));
return Optional.of(createResultSet(is, (QueryRawStatement) statementInfoWrapper.getInitialStatement(), properties, statement));
} else {
// If the statement is not a query, read all bytes from the input stream and close it.
// This is needed otherwise the stream with the server will be closed after having received the first chunk of data (resulting in incomplete inserts).
Expand All @@ -63,11 +60,12 @@ public boolean isStatementRunning(String statementLabel) {
return statementClient.isStatementRunning(statementLabel);
}

private FireboltResultSet createResultSet(InputStream inputStream, QueryRawStatement initialQuery, FireboltProperties properties, FireboltStatement statement, int maxRows, int maxFieldSize)
private FireboltResultSet createResultSet(InputStream inputStream, QueryRawStatement initialQuery, FireboltProperties properties, FireboltStatement statement)
throws SQLException {
return new FireboltResultSet(inputStream, Optional.ofNullable(initialQuery.getTable()).orElse(UNKNOWN_TABLE_NAME),
return new FireboltResultSet(inputStream,
Optional.ofNullable(initialQuery.getTable()).orElse(UNKNOWN_TABLE_NAME),
Optional.ofNullable(initialQuery.getDatabase()).orElse(properties.getDatabase()),
properties.getBufferSize(), maxRows, maxFieldSize, properties.isCompress(), statement,
properties.isLogResultSet());
properties,
statement);
}
}
3 changes: 1 addition & 2 deletions src/main/java/com/firebolt/jdbc/statement/FireboltStatement.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -111,8 +111,7 @@ private Optional<ResultSet> execute(StatementInfoWrapper statementInfoWrapper, b
connection.addProperty(statementInfoWrapper.getParam());
log.debug("The property from the query {} was stored", runningStatementLabel);
} else {
Optional<ResultSet> currentRs = statementService.execute(statementInfoWrapper,
sessionProperties, queryTimeout, maxRows, maxFieldSize, isStandardSql, sessionProperties.isSystemEngine(), this);
Optional<ResultSet> currentRs = statementService.execute(statementInfoWrapper, sessionProperties, isStandardSql, this);
if (currentRs.isPresent()) {
resultSet = currentRs.get();
currentUpdateCount = -1; // Always -1 when returning a ResultSet
Expand Down
25 changes: 12 additions & 13 deletions src/test/java/com/firebolt/jdbc/connection/FireboltConnectionTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -190,7 +190,7 @@ void shouldPrepareStatementNoGeneratedKeys() throws SQLException {
}

private void shouldPrepareStatement(CheckedBiFunction<Connection, String, PreparedStatement> preparedStatementFactoryMethod) throws SQLException {
when(fireboltStatementService.execute(any(), any(), anyInt(), anyInt(), anyInt(), anyBoolean(),anyBoolean(), any()))
when(fireboltStatementService.execute(any(), any(), anyBoolean(), any()))
.thenReturn(Optional.empty());
try (FireboltConnection fireboltConnection = createConnection(URL, connectionProperties)) {
PreparedStatement statement = preparedStatementFactoryMethod.apply(fireboltConnection, "INSERT INTO cars(sales, name) VALUES (?, ?)");
Expand All @@ -199,7 +199,7 @@ private void shouldPrepareStatement(CheckedBiFunction<Connection, String, Prepar
statement.execute();
assertNotNull(fireboltConnection);
assertNotNull(statement);
verify(fireboltStatementService).execute(queryInfoWrapperArgumentCaptor.capture(), any(), anyInt(), anyInt(), anyInt(), anyBoolean(), anyBoolean(), any());
verify(fireboltStatementService).execute(queryInfoWrapperArgumentCaptor.capture(), any(), anyBoolean(), any());
assertEquals("INSERT INTO cars(sales, name) VALUES (500, 'Ford')",
queryInfoWrapperArgumentCaptor.getValue().getSql());
}
Expand Down Expand Up @@ -300,13 +300,13 @@ private <T> void notSupported(CheckedFunction<Connection, T> getter) throws SQLE
@Test
void shouldNotSetNewPropertyWhenConnectionIsNotValidWithTheNewProperty() throws SQLException {
try (FireboltConnection fireboltConnection = createConnection(URL, connectionProperties)) {
when(fireboltStatementService.execute(any(), any(), anyInt(), anyInt(), anyInt(), anyBoolean(), anyBoolean(), any()))
when(fireboltStatementService.execute(any(), any(), anyBoolean(), any()))
.thenThrow(new FireboltException(ExceptionType.TOO_MANY_REQUESTS));
assertThrows(FireboltException.class,
() -> fireboltConnection.addProperty(Map.entry("custom_1", "1")));

verify(fireboltStatementService).execute(queryInfoWrapperArgumentCaptor.capture(),
propertiesArgumentCaptor.capture(), anyInt(), anyInt(), anyInt(), anyBoolean(), anyBoolean(), any());
propertiesArgumentCaptor.capture(), anyBoolean(), any());
assertEquals("1", propertiesArgumentCaptor.getValue().getAdditionalProperties().get("custom_1"));
assertEquals("SELECT 1", queryInfoWrapperArgumentCaptor.getValue().getSql());
assertNull(fireboltConnection.getSessionProperties().getAdditionalProperties().get("custom_1"));
Expand All @@ -315,7 +315,7 @@ void shouldNotSetNewPropertyWhenConnectionIsNotValidWithTheNewProperty() throws

@Test
void shouldSetNewPropertyWhenConnectionIsValidWithTheNewProperty() throws SQLException {
when(fireboltStatementService.execute(any(), any(), anyInt(), anyInt(), anyInt(), anyBoolean(), anyBoolean(), any()))
when(fireboltStatementService.execute(any(), any(), anyBoolean(), any()))
.thenReturn(Optional.empty());

try (FireboltConnection fireboltConnection = createConnection(URL, connectionProperties)) {
Expand All @@ -324,7 +324,7 @@ void shouldSetNewPropertyWhenConnectionIsValidWithTheNewProperty() throws SQLExc
fireboltConnection.addProperty(newProperties);

verify(fireboltStatementService).execute(queryInfoWrapperArgumentCaptor.capture(),
propertiesArgumentCaptor.capture(), anyInt(), anyInt(), anyInt(), anyBoolean(), anyBoolean(), any());
propertiesArgumentCaptor.capture(), anyBoolean(), any());
assertEquals("1", propertiesArgumentCaptor.getValue().getAdditionalProperties().get("custom_1"));
assertEquals("1", fireboltConnection.getSessionProperties().getAdditionalProperties().get("custom_1"));
assertEquals(List.of("SELECT 1"), queryInfoWrapperArgumentCaptor.getAllValues().stream().map(StatementInfoWrapper::getSql).collect(toList()));
Expand All @@ -333,19 +333,19 @@ void shouldSetNewPropertyWhenConnectionIsValidWithTheNewProperty() throws SQLExc

@Test
void shouldValidateConnectionWhenCallingIsValid() throws SQLException {
when(fireboltStatementService.execute(any(), any(), anyInt(), anyInt(), anyInt(), anyBoolean(), anyBoolean(), any()))
when(fireboltStatementService.execute(any(), any(), anyBoolean(), any()))
.thenReturn(Optional.empty());
try (FireboltConnection fireboltConnection = createConnection(URL, connectionProperties)) {
fireboltConnection.isValid(500);
verify(fireboltStatementService).execute(queryInfoWrapperArgumentCaptor.capture(),
propertiesArgumentCaptor.capture(), anyInt(), anyInt(), anyInt(), anyBoolean(), anyBoolean(), any());
propertiesArgumentCaptor.capture(), anyBoolean(), any());
assertEquals(List.of("SELECT 1"), queryInfoWrapperArgumentCaptor.getAllValues().stream().map(StatementInfoWrapper::getSql).collect(toList()));
}
}

@Test
void shouldIgnore429WhenValidatingConnection() throws SQLException {
when(fireboltStatementService.execute(any(), any(), anyInt(), anyInt(), anyInt(), anyBoolean(), anyBoolean(), any()))
when(fireboltStatementService.execute(any(), any(), anyBoolean(), any()))
.thenThrow(new FireboltException(ExceptionType.TOO_MANY_REQUESTS));
try (FireboltConnection fireboltConnection = createConnection(URL, connectionProperties)) {
assertTrue(fireboltConnection.isValid(500));
Expand All @@ -354,7 +354,7 @@ void shouldIgnore429WhenValidatingConnection() throws SQLException {

@Test
void shouldReturnFalseWhenValidatingConnectionThrowsAnException() throws SQLException {
when(fireboltStatementService.execute(any(), any(), anyInt(), anyInt(), anyInt(), anyBoolean(), anyBoolean(), any()))
when(fireboltStatementService.execute(any(), any(), anyBoolean(), any()))
.thenThrow(new FireboltException(ExceptionType.ERROR));
try (FireboltConnection fireboltConnection = createConnection(URL, connectionProperties)) {
assertFalse(fireboltConnection.isValid(500));
Expand All @@ -378,7 +378,7 @@ void shouldReturnFalseWhenValidatingClosedConnection() throws SQLException {

@Test
void shouldExtractConnectorOverrides() throws SQLException {
when(fireboltStatementService.execute(any(), any(), anyInt(), anyInt(), anyInt(), anyBoolean(), anyBoolean(), any()))
when(fireboltStatementService.execute(any(), any(), anyBoolean(), any()))
.thenReturn(Optional.empty());
connectionProperties.put("user_clients", "ConnA:1.0.9,ConnB:2.8.0");
connectionProperties.put("user_drivers", "DriverA:2.0.9,DriverB:3.8.0");
Expand All @@ -387,8 +387,7 @@ void shouldExtractConnectorOverrides() throws SQLException {
PreparedStatement statement = fireboltConnection.prepareStatement("SELECT 1");
statement.execute();

verify(fireboltStatementService).execute(any(), propertiesArgumentCaptor.capture(), anyInt(), anyInt(), anyInt(),
anyBoolean(), anyBoolean(), any());
verify(fireboltStatementService).execute(any(), propertiesArgumentCaptor.capture(), anyBoolean(), any());
assertNull(propertiesArgumentCaptor.getValue().getAdditionalProperties().get("user_clients"));
assertNull(propertiesArgumentCaptor.getValue().getAdditionalProperties().get("user_drivers"));
assertNull(fireboltConnection.getSessionProperties().getAdditionalProperties().get("user_clients"));
Expand Down
2 changes: 2 additions & 0 deletions src/test/java/com/firebolt/jdbc/resultset/FireboltResultSetTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -679,6 +679,8 @@ void shouldCloseStatementWhenCloseOnCompletion() throws SQLException {
void shouldNotCloseStatementWhenNotCloseOnCompletion() throws SQLException {
when(fireboltStatement.isCloseOnCompletion()).thenReturn(false);
inputStream = getInputStreamWithCommonResponseExample();
when(fireboltStatement.getMaxRows()).thenReturn(1024);
when(fireboltStatement.getMaxFieldSize()).thenReturn(0);
resultSet = new FireboltResultSet(inputStream, "any_name", "array_db", 65535, fireboltStatement);
resultSet.close();
verifyNoMoreInteractions(fireboltStatement);
Expand Down
29 changes: 17 additions & 12 deletions src/test/java/com/firebolt/jdbc/service/FireboltStatementServiceTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -43,9 +43,9 @@ void shouldExecuteQueryAndCreateResultSet() throws SQLException {
StatementInfoWrapper statementInfoWrapper = StatementUtil.parseToStatementInfoWrappers("SELECT 1").get(0);
FireboltProperties fireboltProperties = fireboltProperties("firebolt1", false);
FireboltStatementService fireboltStatementService = new FireboltStatementService(statementClient);

fireboltStatementService.execute(statementInfoWrapper, fireboltProperties, 10, -1, 0, true, false,
mock(FireboltStatement.class));
FireboltStatement statement = mock(FireboltStatement.class);
when(statement.getQueryTimeout()).thenReturn(10);
fireboltStatementService.execute(statementInfoWrapper, fireboltProperties, true, statement);
verify(statementClient).executeSqlStatement(statementInfoWrapper, fireboltProperties, false, 10, true);
Assertions.assertEquals(1, mocked.constructed().size());
}
Expand All @@ -57,8 +57,9 @@ void shouldExecuteQueryWithLocalHostFormatParameters() throws SQLException {
StatementInfoWrapper statementInfoWrapper = StatementUtil.parseToStatementInfoWrappers("SELECT 1").get(0);
FireboltProperties fireboltProperties = fireboltProperties("localhost", false);
FireboltStatementService fireboltStatementService = new FireboltStatementService(statementClient);
fireboltStatementService.execute(statementInfoWrapper, fireboltProperties, -1, 10, 0, true, false,
mock(FireboltStatement.class));
FireboltStatement statement = mock(FireboltStatement.class);
when(statement.getQueryTimeout()).thenReturn(-1);
fireboltStatementService.execute(statementInfoWrapper, fireboltProperties, true, statement);
Assertions.assertEquals(1, mocked.constructed().size());
verify(statementClient).executeSqlStatement(statementInfoWrapper, fireboltProperties, false, -1, true);
}
Expand All @@ -85,10 +86,11 @@ void shouldThrowExceptionWhenTryingToCancelQueryWithASystemEngine() {
void shouldExecuteQueryWithParametersForSystemEngine() throws SQLException {
try (MockedConstruction<FireboltResultSet> mocked = Mockito.mockConstruction(FireboltResultSet.class)) {
StatementInfoWrapper statementInfoWrapper = StatementUtil.parseToStatementInfoWrappers("SELECT 1").get(0);
FireboltProperties fireboltProperties = fireboltProperties("firebolt1", false);
FireboltProperties fireboltProperties = fireboltProperties("firebolt1", true);
FireboltStatementService fireboltStatementService = new FireboltStatementService(statementClient);
fireboltStatementService.execute(statementInfoWrapper, fireboltProperties, 10, 10, 0, true, true,
mock(FireboltStatement.class));
FireboltStatement statement = mock(FireboltStatement.class);
when(statement.getQueryTimeout()).thenReturn(10);
fireboltStatementService.execute(statementInfoWrapper, fireboltProperties, true, statement);
Assertions.assertEquals(1, mocked.constructed().size());
verify(statementClient).executeSqlStatement(statementInfoWrapper, fireboltProperties, true, 10, true);
}
Expand All @@ -99,11 +101,12 @@ void shouldIncludeNonStandardSqlQueryParamForNonStandardSql() throws SQLExceptio
try (MockedConstruction<FireboltResultSet> mocked = Mockito.mockConstruction(FireboltResultSet.class)) {

StatementInfoWrapper statementInfoWrapper = StatementUtil.parseToStatementInfoWrappers("SELECT 1").get(0);
FireboltProperties fireboltProperties = fireboltProperties("localhost", false);
FireboltProperties fireboltProperties = fireboltProperties("localhost", true);

FireboltStatementService fireboltStatementService = new FireboltStatementService(statementClient);
fireboltStatementService.execute(statementInfoWrapper, fireboltProperties, -1, 0, 0, false, true,
mock(FireboltStatement.class));
FireboltStatement statement = mock(FireboltStatement.class);
when(statement.getQueryTimeout()).thenReturn(-1);
fireboltStatementService.execute(statementInfoWrapper, fireboltProperties, false, statement);
Assertions.assertEquals(1, mocked.constructed().size());
verify(statementClient).executeSqlStatement(statementInfoWrapper, fireboltProperties, true, -1, false);
}
Expand All @@ -116,8 +119,10 @@ void shouldBeEmptyWithNonQueryStatement() throws SQLException {
FireboltProperties fireboltProperties = fireboltProperties("localhost", false);

FireboltStatementService fireboltStatementService = new FireboltStatementService(statementClient);
FireboltStatement statement = mock(FireboltStatement.class);
when(statement.getQueryTimeout()).thenReturn(-1);
Assertions.assertEquals(Optional.empty(), fireboltStatementService.execute(statementInfoWrapper,
fireboltProperties, -1, 10, 0, true, false, mock(FireboltStatement.class)));
fireboltProperties, true, statement));
verify(statementClient).executeSqlStatement(statementInfoWrapper, fireboltProperties, false, -1, true);
}

Expand Down
Loading

0 comments on commit b03f605

Please sign in to comment.