Merge pull request #1732 from ClickHouse/fix_inputstream_client_v2

chernser · web-flow · commit f49976215f37 · 2024-07-22T12:15:32.000-07:00
[client-v2] Migrate to use InputStream do read data
diff --git a/clickhouse-http-client/src/main/java/com/clickhouse/client/http/ApacheHttpConnectionImpl.java b/clickhouse-http-client/src/main/java/com/clickhouse/client/http/ApacheHttpConnectionImpl.java
@@ -115,7 +115,8 @@ private ClickHouseHttpResponse buildResponse(ClickHouseConfig config, CloseableH
         // {"read_rows":"0","read_bytes":"0","written_rows":"0","written_bytes":"0","total_rows_to_read":"0"}
         String displayName = getResponseHeader(response, "X-ClickHouse-Server-Display-Name", server.getHost());
         String queryId = getResponseHeader(response, "X-ClickHouse-Query-Id", "");
-        String summary = getResponseHeader(response, "X-ClickHouse-Summary", "{}");
+        Header hSum = response.getLastHeader("X-ClickHouse-Summary");
+        String summary = hSum == null ? "{}" : hSum.getValue(); // getResponseHeader(response, "X-ClickHouse-Summary", "{}");
 
         ClickHouseFormat format = config.getFormat();
         TimeZone timeZone = config.getServerTimeZone();
diff --git a/clickhouse-http-client/src/main/java/com/clickhouse/client/http/config/ClickHouseHttpOption.java b/clickhouse-http-client/src/main/java/com/clickhouse/client/http/config/ClickHouseHttpOption.java
@@ -58,7 +58,7 @@ public enum ClickHouseHttpOption implements ClickHouseOption {
     // "Enables or disables X-ClickHouse-Progress HTTP response headers in
     // clickhouse-server responses."),
     // SEND_PROGRESS_INTERVAL("http_headers_progress_interval_ms", 3000, ""),
-    // WAIT_END_OF_QUERY("wait_end_of_query", false, ""),
+     WAIT_END_OF_QUERY("wait_end_of_query", false, ""),
 
     /**
      * Whether to remember last set role and send them in every next requests as query parameters.
diff --git a/client-v2/src/main/java/com/clickhouse/client/api/Client.java b/client-v2/src/main/java/com/clickhouse/client/api/Client.java
@@ -30,6 +30,7 @@
 import com.clickhouse.client.api.query.Records;
 import com.clickhouse.client.config.ClickHouseClientOption;
 import com.clickhouse.client.config.ClickHouseDefaults;
+import com.clickhouse.client.http.config.ClickHouseHttpOption;
 import com.clickhouse.data.ClickHouseColumn;
 import com.clickhouse.data.ClickHouseDataStreamFactory;
 import com.clickhouse.data.ClickHouseFormat;
@@ -820,6 +821,7 @@ public CompletableFuture<Records> queryRecords(String sqlQuery, QuerySettings se
             settings = new QuerySettings();
         }
         settings.setFormat(ClickHouseFormat.RowBinaryWithNamesAndTypes);
+        settings.waitEndOfQuery(true); // we rely on the summery
         ClientStatisticsHolder clientStats = new ClientStatisticsHolder();
         clientStats.start("query");
         ClickHouseClient client = ClientV1AdaptorHelper.createClient(configuration);
@@ -865,9 +867,9 @@ public CompletableFuture<Records> queryRecords(String sqlQuery, QuerySettings se
     public List<GenericRecord> queryAll(String sqlQuery) {
         try {
             int operationTimeout = getOperationTimeout();
-
-            try (QueryResponse response = operationTimeout == 0 ? query(sqlQuery).get() :
-                    query(sqlQuery).get(operationTimeout, TimeUnit.MILLISECONDS)) {
+            QuerySettings settings = new QuerySettings().waitEndOfQuery(true);
+            try (QueryResponse response = operationTimeout == 0 ? query(sqlQuery, settings).get() :
+                    query(sqlQuery, settings).get(operationTimeout, TimeUnit.MILLISECONDS)) {
                 List<GenericRecord> records = new ArrayList<>();
                 if (response.getResultRows() > 0) {
                     ClickHouseBinaryFormatReader reader = new RowBinaryWithNamesAndTypesFormatReader(response.getInputStream());
diff --git a/client-v2/src/main/java/com/clickhouse/client/api/data_formats/NativeFormatReader.java b/client-v2/src/main/java/com/clickhouse/client/api/data_formats/NativeFormatReader.java
@@ -1,12 +1,10 @@
 package com.clickhouse.client.api.data_formats;
 
-import com.clickhouse.client.api.ClientException;
 import com.clickhouse.client.api.data_formats.internal.AbstractBinaryFormatReader;
+import com.clickhouse.client.api.data_formats.internal.BinaryStreamReader;
 import com.clickhouse.client.api.query.QuerySettings;
 import com.clickhouse.data.ClickHouseColumn;
-import com.clickhouse.data.ClickHouseDataType;
 
-import java.io.EOFException;
 import java.io.IOException;
 import java.io.InputStream;
 import java.util.ArrayList;
@@ -43,14 +41,15 @@ protected void readRecord(Map<String, Object> record) throws IOException {
     }
 
     private void readBlock() throws IOException {
-        int nColumns = chInputStream.readVarInt();
-        int nRows = chInputStream.readVarInt();
+        int nColumns = BinaryStreamReader.readVarInt(input);
+        int nRows = BinaryStreamReader.readVarInt(input);
 
         List<String> names = new ArrayList<>(nColumns);
         List<String> types = new ArrayList<>(nColumns);
         currentBlock = new Block(names, types, nRows);
         for (int i = 0; i < nColumns; i++) {
-            ClickHouseColumn column = ClickHouseColumn.of(chInputStream.readUnicodeString(), chInputStream.readUnicodeString());
+            ClickHouseColumn column = ClickHouseColumn.of(BinaryStreamReader.readString(input),
+                    BinaryStreamReader.readString(input));
             names.add(column.getColumnName());
             types.add(column.getDataType().name());
 
diff --git a/client-v2/src/main/java/com/clickhouse/client/api/data_formats/RowBinaryFormatReader.java b/client-v2/src/main/java/com/clickhouse/client/api/data_formats/RowBinaryFormatReader.java
@@ -1,12 +1,10 @@
 package com.clickhouse.client.api.data_formats;
 
-import com.clickhouse.client.api.ClientException;
 import com.clickhouse.client.api.data_formats.internal.AbstractBinaryFormatReader;
 import com.clickhouse.client.api.metadata.TableSchema;
 import com.clickhouse.client.api.query.QuerySettings;
 import com.clickhouse.data.ClickHouseColumn;
 
-import java.io.EOFException;
 import java.io.IOException;
 import java.io.InputStream;
 import java.util.Map;
diff --git a/client-v2/src/main/java/com/clickhouse/client/api/data_formats/RowBinaryWithNamesAndTypesFormatReader.java b/client-v2/src/main/java/com/clickhouse/client/api/data_formats/RowBinaryWithNamesAndTypesFormatReader.java
@@ -2,18 +2,17 @@
 
 import com.clickhouse.client.api.ClientException;
 import com.clickhouse.client.api.data_formats.internal.AbstractBinaryFormatReader;
+import com.clickhouse.client.api.data_formats.internal.BinaryStreamReader;
 import com.clickhouse.client.api.metadata.TableSchema;
 import com.clickhouse.client.api.query.QuerySettings;
 import com.clickhouse.data.ClickHouseColumn;
 
-import java.io.EOFException;
 import java.io.IOException;
 import java.io.InputStream;
 import java.util.ArrayList;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
-import java.util.NoSuchElementException;
 
 public class RowBinaryWithNamesAndTypesFormatReader extends AbstractBinaryFormatReader implements Iterator<Map<String, Object>> {
 
@@ -28,18 +27,18 @@ public RowBinaryWithNamesAndTypesFormatReader(InputStream inputStream, QuerySett
 
     private void readSchema() {
         try {
-            if (inputStream.available() < 1) {
+            if (input.available() < 1) {
                 return;
             }
             TableSchema headerSchema = new TableSchema();
             List<String> columns = new ArrayList<>();
-            int nCol = chInputStream.readVarInt();
+            int nCol = BinaryStreamReader.readVarInt(input);
             for (int i = 0; i < nCol; i++) {
-                columns.add(chInputStream.readUnicodeString());
+                columns.add(BinaryStreamReader.readString(input));
             }
 
             for (int i = 0; i < nCol; i++) {
-                headerSchema.addColumn(columns.get(i), chInputStream.readUnicodeString());
+                headerSchema.addColumn(columns.get(i), BinaryStreamReader.readString(input));
             }
 
             setSchema(headerSchema);
diff --git a/client-v2/src/main/java/com/clickhouse/client/api/data_formats/RowBinaryWithNamesFormatReader.java b/client-v2/src/main/java/com/clickhouse/client/api/data_formats/RowBinaryWithNamesFormatReader.java
@@ -1,12 +1,11 @@
 package com.clickhouse.client.api.data_formats;
 
-import com.clickhouse.client.api.ClientException;
 import com.clickhouse.client.api.data_formats.internal.AbstractBinaryFormatReader;
+import com.clickhouse.client.api.data_formats.internal.BinaryStreamReader;
 import com.clickhouse.client.api.metadata.TableSchema;
 import com.clickhouse.client.api.query.QuerySettings;
 import com.clickhouse.data.ClickHouseColumn;
 
-import java.io.EOFException;
 import java.io.IOException;
 import java.io.InputStream;
 import java.util.ArrayList;
@@ -30,9 +29,9 @@ public RowBinaryWithNamesFormatReader(InputStream inputStream, QuerySettings que
     public void readRecord(Map<String, Object> record) throws IOException {
         if (columns == null) {
             columns = new ArrayList<>();
-            int nCol = chInputStream.readVarInt();
+            int nCol = BinaryStreamReader.readVarInt(input);
             for (int i = 0; i < nCol; i++) {
-                columns.add(chInputStream.readUnicodeString());
+                columns.add(BinaryStreamReader.readString(input));
             }
 
             columns = Collections.unmodifiableList(columns);
diff --git a/client-v2/src/main/java/com/clickhouse/client/api/data_formats/internal/AbstractBinaryFormatReader.java b/client-v2/src/main/java/com/clickhouse/client/api/data_formats/internal/AbstractBinaryFormatReader.java
@@ -6,7 +6,6 @@
 import com.clickhouse.client.api.query.NullValueException;
 import com.clickhouse.client.api.query.QuerySettings;
 import com.clickhouse.data.ClickHouseColumn;
-import com.clickhouse.data.ClickHouseInputStream;
 import com.clickhouse.data.value.ClickHouseArrayValue;
 import com.clickhouse.data.value.ClickHouseGeoMultiPolygonValue;
 import com.clickhouse.data.value.ClickHouseGeoPointValue;
@@ -41,9 +40,7 @@ public abstract class AbstractBinaryFormatReader implements ClickHouseBinaryForm
 
     private static final Logger LOG = LoggerFactory.getLogger(AbstractBinaryFormatReader.class);
 
-    protected InputStream inputStream;
-
-    protected ClickHouseInputStream chInputStream;
+    protected InputStream input;
 
     protected Map<String, Object> settings;
 
@@ -54,15 +51,12 @@ public abstract class AbstractBinaryFormatReader implements ClickHouseBinaryForm
     protected volatile boolean hasNext = true;
 
     protected AbstractBinaryFormatReader(InputStream inputStream, QuerySettings querySettings, TableSchema schema) {
-        this.inputStream = inputStream;
-        this.chInputStream = inputStream instanceof ClickHouseInputStream ?
-                (ClickHouseInputStream) inputStream : ClickHouseInputStream.of(inputStream);
+        this.input = inputStream;
         this.settings = querySettings == null ? Collections.emptyMap() : new HashMap<>(querySettings.getAllSettings());
-        this.binaryStreamReader = new BinaryStreamReader(chInputStream, LOG);
+        this.binaryStreamReader = new BinaryStreamReader(inputStream, LOG);
         setSchema(schema);
     }
 
-
     protected Map<String, Object> currentRecord = new ConcurrentHashMap<>();
 
     protected abstract void readRecord(Map<String, Object> record) throws IOException;
@@ -85,7 +79,7 @@ public <T> T readValue(String colName) {
     public boolean hasNext() {
         if (hasNext) {
             try {
-                hasNext = chInputStream.available() > 0;
+                hasNext = input.available() > 0;
                 return hasNext;
             } catch (IOException e) {
                 hasNext = false;
diff --git a/client-v2/src/main/java/com/clickhouse/client/api/data_formats/internal/BinaryStreamReader.java b/client-v2/src/main/java/com/clickhouse/client/api/data_formats/internal/BinaryStreamReader.java
diff --git a/client-v2/src/main/java/com/clickhouse/client/api/query/QueryResponse.java b/client-v2/src/main/java/com/clickhouse/client/api/query/QueryResponse.java
diff --git a/client-v2/src/main/java/com/clickhouse/client/api/query/QuerySettings.java b/client-v2/src/main/java/com/clickhouse/client/api/query/QuerySettings.java
diff --git a/client-v2/src/test/java/com/clickhouse/client/query/QueryTests.java b/client-v2/src/test/java/com/clickhouse/client/query/QueryTests.java
