Read all bytes from response on non query statements (#216)

build.gradle

@@ -9,7 +9,7 @@ plugins {
 }
 
 group 'com.firebolt'
-version = '2.4.0'
+version = '2.4.1-SNAPSHOT'
 
 repositories {
     mavenCentral()

lombok.config

@@ -1,4 +1,4 @@
 lombok.anyConstructor.addConstructorProperties = true
 config.stopBubbling = true
 lombok.addLombokGeneratedAnnotation = true
-lombok.log.custom.declaration = com.firebolt.jdbc.log.FireboltLogger com.firebolt.jdbc.LoggerUtil.getLogger(NAME)
+lombok.log.custom.declaration = com.firebolt.jdbc.log.FireboltLogger com.firebolt.jdbc.util.LoggerUtil.getLogger(NAME)

src/integrationTest/java/integration/tests/DatabaseMetaDataTest.java

@@ -110,7 +110,7 @@ void shouldReturnColumns() throws SQLException {
 				contains(schemaName, schemaName, schemaName, schemaName, schemaName));
 		assertThat(result.get(REMARKS), contains(null, null, null, null, null));
 		assertThat(result.get(NULLABLE), contains("0", "1", "1", "1", "0"));
-		assertThat(result.get(DECIMAL_DIGITS), contains(null, null, null, null, null));
+		assertThat(result.get(DECIMAL_DIGITS), contains("0", "0", "0", "0", "0"));
 		assertThat(result.get(SQL_DATETIME_SUB), contains(null, null, null, null, null));
 		assertThat(result.get(NUM_PREC_RADIX), contains("10", "10", "10", "10", "10"));
 		assertThat(result.get(IS_GENERATEDCOLUMN), contains("NO", "NO", "NO", "NO", "NO"));
@@ -121,9 +121,9 @@ void shouldReturnColumns() throws SQLException {
 		assertThat(result.get(SCOPE_SCHEMA), contains(null, null, null, null, null));
 		assertThat(result.get(ORDINAL_POSITION), contains("1", "2", "3", "4", "5"));
 		assertThat(result.get(TYPE_NAME),
-				contains("bigint", "timestamp", "string", "boolean", "integer"));
+				contains("bigint", "timestamp", "text", "boolean", "integer"));
 		assertThat(result.get(DATA_TYPE), contains("-5", "93", "12", "16", "4"));
-		assertThat(result.get(COLUMN_SIZE), contains("20", "19", "0", "1", "11"));
+		assertThat(result.get(COLUMN_SIZE), contains("0", "6", "0", "1", "0"));
 	}
 
 }

src/main/java/com/firebolt/FireboltDriver.java

@@ -6,8 +6,8 @@
 
 import org.apache.commons.lang3.StringUtils;
 
-import com.firebolt.jdbc.PropertyUtil;
-import com.firebolt.jdbc.VersionUtil;
+import com.firebolt.jdbc.util.PropertyUtil;
+import com.firebolt.jdbc.util.VersionUtil;
 import com.firebolt.jdbc.connection.FireboltConnection;
 import com.firebolt.jdbc.exception.FireboltSQLFeatureNotSupportedException;
 

src/main/java/com/firebolt/jdbc/client/FireboltClient.java

@@ -15,7 +15,7 @@
 import org.apache.commons.lang3.tuple.Pair;
 
 import com.fasterxml.jackson.databind.ObjectMapper;
-import com.firebolt.jdbc.CloseableUtil;
+import com.firebolt.jdbc.util.CloseableUtil;
 import com.firebolt.jdbc.connection.FireboltConnection;
 import com.firebolt.jdbc.exception.FireboltException;
 import com.firebolt.jdbc.resultset.compress.LZ4InputStream;

src/main/java/com/firebolt/jdbc/client/UsageTrackerUtil.java

@@ -5,7 +5,7 @@
 
 import org.apache.commons.lang3.StringUtils;
 
-import com.firebolt.jdbc.VersionUtil;
+import com.firebolt.jdbc.util.VersionUtil;
 import com.google.common.collect.ImmutableMap;
 
 import lombok.CustomLog;

src/main/java/com/firebolt/jdbc/client/query/StatementClientImpl.java

@@ -10,8 +10,8 @@
 import org.apache.commons.lang3.tuple.Pair;
 
 import com.fasterxml.jackson.databind.ObjectMapper;
-import com.firebolt.jdbc.CloseableUtil;
-import com.firebolt.jdbc.PropertyUtil;
+import com.firebolt.jdbc.util.CloseableUtil;
+import com.firebolt.jdbc.util.PropertyUtil;
 import com.firebolt.jdbc.client.FireboltClient;
 import com.firebolt.jdbc.connection.FireboltConnection;
 import com.firebolt.jdbc.connection.FireboltConnectionTokens;

src/main/java/com/firebolt/jdbc/connection/FireboltConnection.java

@@ -14,7 +14,7 @@
 import org.apache.commons.lang3.tuple.Pair;
 
 import com.fasterxml.jackson.databind.ObjectMapper;
-import com.firebolt.jdbc.PropertyUtil;
+import com.firebolt.jdbc.util.PropertyUtil;
 import com.firebolt.jdbc.annotation.ExcludeFromJacocoGeneratedReport;
 import com.firebolt.jdbc.annotation.NotImplemented;
 import com.firebolt.jdbc.client.FireboltObjectMapper;

src/main/java/com/firebolt/jdbc/metadata/FireboltDatabaseMetadata.java

@@ -12,7 +12,7 @@
 import org.apache.commons.lang3.StringUtils;
 
 import com.firebolt.jdbc.QueryResult;
-import com.firebolt.jdbc.VersionUtil;
+import com.firebolt.jdbc.util.VersionUtil;
 import com.firebolt.jdbc.annotation.ExcludeFromJacocoGeneratedReport;
 import com.firebolt.jdbc.annotation.NotImplemented;
 import com.firebolt.jdbc.connection.FireboltConnection;

src/main/java/com/firebolt/jdbc/resultset/FireboltResultSet.java

@@ -15,7 +15,7 @@
 import org.apache.commons.lang3.exception.ExceptionUtils;
 import org.apache.commons.text.StringEscapeUtils;
 
-import com.firebolt.jdbc.LoggerUtil;
+import com.firebolt.jdbc.util.LoggerUtil;
 import com.firebolt.jdbc.QueryResult;
 import com.firebolt.jdbc.annotation.ExcludeFromJacocoGeneratedReport;
 import com.firebolt.jdbc.annotation.NotImplemented;
@@ -123,7 +123,7 @@ private BufferedReader createStreamReader(InputStream is, Integer bufferSize, bo
 		} else {
 			inputStreamReader = new InputStreamReader(is, StandardCharsets.UTF_8);
 		}
-		return bufferSize != null ? new BufferedReader(inputStreamReader, bufferSize)
+		return bufferSize != null && bufferSize != 0 ? new BufferedReader(inputStreamReader, bufferSize)
 				: new BufferedReader(inputStreamReader);
 	}
 

src/main/java/com/firebolt/jdbc/service/FireboltStatementService.java

@@ -3,12 +3,21 @@
 import static com.firebolt.jdbc.exception.ExceptionType.INVALID_REQUEST;
 
 import java.io.InputStream;
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import java.util.Optional;
 
 import com.firebolt.jdbc.client.query.StatementClient;
 import com.firebolt.jdbc.connection.settings.FireboltProperties;
 import com.firebolt.jdbc.exception.FireboltException;
+import com.firebolt.jdbc.resultset.FireboltResultSet;
+import com.firebolt.jdbc.statement.FireboltStatement;
 import com.firebolt.jdbc.statement.StatementInfoWrapper;
 
+import com.firebolt.jdbc.statement.StatementType;
+import com.firebolt.jdbc.statement.rawstatement.QueryRawStatement;
+import com.firebolt.jdbc.util.CloseableUtil;
+import com.firebolt.jdbc.util.InputStreamUtil;
 import lombok.CustomLog;
 import lombok.NonNull;
 import lombok.RequiredArgsConstructor;
@@ -17,24 +26,36 @@
 @CustomLog
 public class FireboltStatementService {
 
+	private static final String UNKNOWN_TABLE_NAME = "unknown";
 	private final StatementClient statementClient;
 	private final boolean systemEngine;
 
 	/**
 	 * Executes statement
 	 *
-	 * @param statementInfoWrapper the statement
-	 * @param connectionProperties the connection properties
+	 * @param statementInfoWrapper the statement info
+	 * @param properties the connection properties
 	 * @param queryTimeout         query timeout
 	 * @param maxRows              max rows
 	 * @param standardSql          indicates if standard sql should be used
+	 * @param statement           the statement
 	 * @return an InputStream with the result
 	 */
-	public InputStream execute(@NonNull StatementInfoWrapper statementInfoWrapper,
-			@NonNull FireboltProperties connectionProperties, int queryTimeout, int maxRows, boolean standardSql)
-			throws FireboltException {
-		return statementClient.postSqlStatement(statementInfoWrapper, connectionProperties, systemEngine, queryTimeout,
+	public Optional<ResultSet> execute(StatementInfoWrapper statementInfoWrapper,
+									   FireboltProperties properties, int queryTimeout, int maxRows, boolean standardSql,
+									   FireboltStatement statement)
+			throws SQLException {
+		InputStream is = statementClient.postSqlStatement(statementInfoWrapper, properties, systemEngine, queryTimeout,
 				maxRows, standardSql);
+		if (statementInfoWrapper.getType() == StatementType.QUERY) {
+			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).
+			InputStreamUtil.readAllBytes(is);
+			CloseableUtil.close(is);
+		}
+		return Optional.empty();
 	}
 
 	public void abortStatement(@NonNull String statementId, @NonNull FireboltProperties properties)
@@ -53,4 +74,12 @@ public void abortStatementHttpRequest(@NonNull String statementId) throws Firebo
 	public boolean isStatementRunning(String statementId) {
 		return statementClient.isStatementRunning(statementId);
 	}
+
+	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),
+				Optional.ofNullable(initialQuery.getDatabase()).orElse(properties.getDatabase()),
+				properties.getBufferSize(), properties.isCompress(), statement,
+				properties.isLogResultSet());
+	}
 }

src/main/java/com/firebolt/jdbc/statement/FireboltStatement.java

@@ -7,18 +7,16 @@
 
 import org.apache.commons.lang3.StringUtils;
 
-import com.firebolt.jdbc.CloseableUtil;
-import com.firebolt.jdbc.PropertyUtil;
 import com.firebolt.jdbc.annotation.ExcludeFromJacocoGeneratedReport;
 import com.firebolt.jdbc.annotation.NotImplemented;
 import com.firebolt.jdbc.connection.FireboltConnection;
 import com.firebolt.jdbc.connection.settings.FireboltProperties;
 import com.firebolt.jdbc.exception.FireboltException;
 import com.firebolt.jdbc.exception.FireboltSQLFeatureNotSupportedException;
 import com.firebolt.jdbc.exception.FireboltUnsupportedOperationException;
-import com.firebolt.jdbc.resultset.FireboltResultSet;
 import com.firebolt.jdbc.service.FireboltStatementService;
-import com.firebolt.jdbc.statement.rawstatement.QueryRawStatement;
+import com.firebolt.jdbc.util.CloseableUtil;
+import com.firebolt.jdbc.util.PropertyUtil;
 
 import lombok.Builder;
 import lombok.CustomLog;
@@ -111,15 +109,13 @@ private Optional<ResultSet> execute(StatementInfoWrapper statementInfoWrapper, b
 					this.connection.addProperty(statementInfoWrapper.getParam());
 					log.debug("The property from the query {} was stored", runningStatementId);
 				} else {
-					inputStream = statementService.execute(statementInfoWrapper, this.sessionProperties,
-							this.queryTimeout, this.maxRows, isStandardSql);
-					if (statementInfoWrapper.getType() == StatementType.QUERY) {
-						resultSet = getResultSet(inputStream,
-								(QueryRawStatement) statementInfoWrapper.getInitialStatement());
+					Optional<ResultSet> currentRs = statementService.execute(statementInfoWrapper,
+							this.sessionProperties, this.queryTimeout, this.maxRows, isStandardSql, this);
+					if (currentRs.isPresent()) {
+						resultSet = currentRs.get();
 						currentUpdateCount = -1; // Always -1 when returning a ResultSet
 					} else {
 						currentUpdateCount = 0;
-						CloseableUtil.close(inputStream);
 					}
 					log.info("The query with the id {} was executed with success", runningStatementId);
 				}
@@ -152,14 +148,6 @@ private boolean isStatementNotCancelled(StatementInfoWrapper statementInfoWrappe
 		}
 	}
 
-	private FireboltResultSet getResultSet(InputStream inputStream, QueryRawStatement initialQuery)
-			throws SQLException {
-		return new FireboltResultSet(inputStream, Optional.ofNullable(initialQuery.getTable()).orElse("unknown"),
-				Optional.ofNullable(initialQuery.getDatabase()).orElse(this.sessionProperties.getDatabase()),
-				this.sessionProperties.getBufferSize(), this.sessionProperties.isCompress(), this,
-				this.sessionProperties.isLogResultSet());
-	}
-
 	private void closeAllResults() {
 		synchronized (this) {
 			if (this.firstUnclosedStatementResult != null) {
@@ -496,14 +484,15 @@ public void clearBatch() throws SQLException {
 		// Batch are not supported by the driver
 		throw new FireboltUnsupportedOperationException();
 	}
+
 	@Override
 	@NotImplemented
 	@ExcludeFromJacocoGeneratedReport
 	public int[] executeBatch() throws SQLException {
 		// Batch are not supported by the driver
 		throw new FireboltUnsupportedOperationException();
 	}
-	
+
 	@Override
 	@NotImplemented
 	@ExcludeFromJacocoGeneratedReport

src/main/java/com/firebolt/jdbc/CloseableUtil.java → src/main/java/com/firebolt/jdbc/util/CloseableUtil.java

@@ -1,4 +1,4 @@
-package com.firebolt.jdbc;
+package com.firebolt.jdbc.util;
 
 import java.io.Closeable;
 import java.io.IOException;

src/main/java/com/firebolt/jdbc/util/InputStreamUtil.java

@@ -0,0 +1,30 @@
+package com.firebolt.jdbc.util;
+
+import lombok.CustomLog;
+import lombok.experimental.UtilityClass;
+
+import javax.annotation.Nullable;
+import java.io.IOException;
+import java.io.InputStream;
+
+@UtilityClass
+@CustomLog
+public class InputStreamUtil {
+
+    /**
+     * Read all bytes from the input stream if the stream is not null
+     *
+     * @param is input stream
+     */
+    public void readAllBytes(@Nullable InputStream is) {
+        if (is != null) {
+            while (true) {
+                try {
+                    if (is.read() == -1) break;
+                } catch (IOException e) {
+                    log.warn("Could not read entire input stream for non query statement", e);
+                }
+            }
+        }
+    }
+}

src/main/java/com/firebolt/jdbc/LoggerUtil.java → src/main/java/com/firebolt/jdbc/util/LoggerUtil.java

@@ -1,4 +1,4 @@
-package com.firebolt.jdbc;
+package com.firebolt.jdbc.util;
 
 import java.io.*;
 import java.nio.charset.StandardCharsets;

src/main/java/com/firebolt/jdbc/PropertyUtil.java → src/main/java/com/firebolt/jdbc/util/PropertyUtil.java

@@ -1,4 +1,4 @@
-package com.firebolt.jdbc;
+package com.firebolt.jdbc.util;
 
 import java.sql.DriverPropertyInfo;
 import java.util.ArrayList;

src/main/java/com/firebolt/jdbc/VersionUtil.java → src/main/java/com/firebolt/jdbc/util/VersionUtil.java

@@ -1,4 +1,4 @@
-package com.firebolt.jdbc;
+package com.firebolt.jdbc.util;
 
 import java.io.IOException;
 import java.util.Properties;

src/test/java/com/firebolt/jdbc/client/UsageTrackerUtilTest.java

@@ -17,7 +17,7 @@
 import org.mockito.MockedStatic;
 import org.mockito.Mockito;
 
-import com.firebolt.jdbc.VersionUtil;
+import com.firebolt.jdbc.util.VersionUtil;
 
 @SetSystemProperty(key = "java.version", value = "8.0.1")
 @SetSystemProperty(key = "os.version", value = "10.1")

src/test/java/com/firebolt/jdbc/client/account/FireboltAccountClientTest.java

@@ -20,7 +20,7 @@
 
 import com.fasterxml.jackson.databind.DeserializationFeature;
 import com.fasterxml.jackson.databind.ObjectMapper;
-import com.firebolt.jdbc.VersionUtil;
+import com.firebolt.jdbc.util.VersionUtil;
 import com.firebolt.jdbc.client.account.response.FireboltAccountResponse;
 import com.firebolt.jdbc.client.account.response.FireboltDefaultDatabaseEngineResponse;
 import com.firebolt.jdbc.client.account.response.FireboltEngineIdResponse;

src/test/java/com/firebolt/jdbc/client/authentication/FireboltAuthenticationClientTest.java

@@ -15,7 +15,7 @@
 
 import com.fasterxml.jackson.databind.DeserializationFeature;
 import com.fasterxml.jackson.databind.ObjectMapper;
-import com.firebolt.jdbc.VersionUtil;
+import com.firebolt.jdbc.util.VersionUtil;
 import com.firebolt.jdbc.client.authentication.response.FireboltAuthenticationResponse;
 import com.firebolt.jdbc.connection.FireboltConnection;
 import com.firebolt.jdbc.exception.FireboltException;

src/test/java/com/firebolt/jdbc/client/query/StatementClientImplTest.java

@@ -24,7 +24,7 @@
 import org.mockito.junit.jupiter.MockitoExtension;
 
 import com.fasterxml.jackson.databind.ObjectMapper;
-import com.firebolt.jdbc.VersionUtil;
+import com.firebolt.jdbc.util.VersionUtil;
 import com.firebolt.jdbc.connection.FireboltConnection;
 import com.firebolt.jdbc.connection.FireboltConnectionTokens;
 import com.firebolt.jdbc.connection.settings.FireboltProperties;