feat: support streaming load

Author: youngsofun

.github/workflows/test.yml

@@ -51,6 +51,6 @@ jobs:
           curl -u databend:databend  --request POST localhost:8000/v1/query --header 'Content-Type:application/json' --data-raw '{"sql":"select 1"}'
 
       - name: Run Maven clean deploy with release profile
-        run: mvn test -DexcludedGroups=CLUSTER,FLAKY
+        run: mvn test -DexcludedGroups=MULTI_HOST,FLAKY
         env:
           MAVEN_GPG_PASSPHRASE: ${{ secrets.GPG_PASSPHRASE }}

.github/workflows/test_cluster.yml

@@ -31,7 +31,7 @@ jobs:
       - uses: ./.github/actions/setup_databend_cluster
         timeout-minutes: 15
         with:
-          version: '1.2.754-nightly'
+          version: '1.2.797-nightly'
           target: 'x86_64-unknown-linux-gnu'
 
       - name: Test with conn to node 1

databend-client/pom.xml

@@ -18,25 +18,21 @@
     <properties>
         <air.main.basedir>${project.parent.basedir}</air.main.basedir>
         <project.build.targetJdk>8</project.build.targetJdk>
-        <jackson.version>2.15.2</jackson.version>
     </properties>
     <dependencies>
         <dependency>
             <groupId>com.fasterxml.jackson.core</groupId>
             <artifactId>jackson-annotations</artifactId>
-            <version>${jackson.version}</version>
         </dependency>
 
         <dependency>
             <groupId>com.fasterxml.jackson.core</groupId>
             <artifactId>jackson-core</artifactId>
-            <version>${jackson.version}</version>
         </dependency>
 
         <dependency>
             <groupId>com.fasterxml.jackson.core</groupId>
             <artifactId>jackson-databind</artifactId>
-            <version>${jackson.version}</version>
         </dependency>
 
         <dependency>
@@ -53,14 +49,13 @@
         <dependency>
             <groupId>com.google.guava</groupId>
             <artifactId>guava</artifactId>
-            <version>32.0.1-jre</version>
         </dependency>
 
         <dependency>
             <groupId>com.squareup.okhttp3</groupId>
             <artifactId>okhttp</artifactId>
         </dependency>
-        <!-- https://mvnrepository.com/artifact/com.squareup.okio/okio -->
+
         <dependency>
             <groupId>com.squareup.okio</groupId>
             <artifactId>okio</artifactId>
@@ -76,18 +71,8 @@
             <artifactId>okhttp-urlconnection</artifactId>
         </dependency>
 
-        <dependency>
-            <groupId>com.github.zafarkhaja</groupId>
-            <artifactId>java-semver</artifactId>
-        </dependency>
-
         <!-- for testing -->
 
-        <dependency>
-            <groupId>io.airlift</groupId>
-            <artifactId>json</artifactId>
-            <scope>test</scope>
-        </dependency>
         <dependency>
             <groupId>org.testng</groupId>
             <artifactId>testng</artifactId>

databend-client/src/main/java/com/databend/client/ClientSettings.java

@@ -31,6 +31,8 @@ public class ClientSettings {
     public static final String X_DATABEND_STICKY_NODE = "X-DATABEND-STICKY-NODE";
     public static final String DatabendWarehouseHeader = "X-DATABEND-WAREHOUSE";
     public static final String DatabendTenantHeader = "X-DATABEND-TENANT";
+    public static final String DatabendSQLHeader = "X-DATABEND-SQL";
+    public static final String DatabendQueryContextHeader = "X-DATABEND-QUERY-CONTEXT";
     private final String host;
     private final DatabendSession session;
     private final Integer queryTimeoutSecs;
@@ -40,14 +42,14 @@ public class ClientSettings {
     private final PaginationOptions paginationOptions;
 
     private final StageAttachment stageAttachment;
-    private Map<String, String> additionalHeaders;
+    private final Map<String, String> additionalHeaders;
 
     private final int retryAttempts;
     // TODO(zhihanz) timezone and locale info
 
     //ClientSettings for test case use
     public ClientSettings(String host) {
-        this(host, DatabendSession.createDefault(), DEFAULT_QUERY_TIMEOUT, DEFAULT_CONNECTION_TIMEOUT, DEFAULT_SOCKET_TIMEOUT, PaginationOptions.defaultPaginationOptions(), new HashMap<String, String>(), null, DEFAULT_RETRY_ATTEMPTS);
+        this(host, DatabendSession.createDefault(), DEFAULT_QUERY_TIMEOUT, DEFAULT_CONNECTION_TIMEOUT, DEFAULT_SOCKET_TIMEOUT, PaginationOptions.defaultPaginationOptions(), new HashMap<>(), null, DEFAULT_RETRY_ATTEMPTS);
     }
 
     public ClientSettings(String host, String database) {

databend-client/src/main/java/com/databend/client/DatabendClient.java

@@ -29,8 +29,6 @@ public interface DatabendClient extends Closeable {
 
     String getNodeID();
 
-    String getServerVersion();
-
 
     Map<String, String> getAdditionalHeaders();
 

databend-client/src/main/java/com/databend/client/DatabendClientFactory.java

@@ -16,7 +16,6 @@
 
 import com.databend.client.errors.CloudErrors;
 import okhttp3.*;
-import okio.Buffer;
 
 import javax.annotation.concurrent.ThreadSafe;
 import java.io.IOException;
@@ -50,10 +49,6 @@ public class DatabendClientV1
     public static final MediaType MEDIA_TYPE_JSON = MediaType.parse("application/json; charset=utf-8");
     public static final JsonCodec<QueryResults> QUERY_RESULTS_CODEC = jsonCodec(QueryResults.class);
     public static final JsonCodec<DiscoveryResponseCodec.DiscoveryResponse> DISCOVERY_RESULT_CODEC = jsonCodec(DiscoveryResponseCodec.DiscoveryResponse.class);
-    public static final String succeededState = "succeeded";
-    public static final String failedState = "failed";
-    public static final String runningState = "running";
-
 
     public static final String QUERY_PATH = "/v1/query";
     public static final String DISCOVERY_PATH = "/v1/discovery_nodes";
@@ -66,15 +61,14 @@ public class DatabendClientV1
     private final PaginationOptions paginationOptions;
     // request with retry timeout
     private final Integer requestTimeoutSecs;
-    private final Map<String, String> additonalHeaders;
-    private String serverVersion;
+    private final Map<String, String> additionalHeaders;
     // client session
     private final AtomicReference<DatabendSession> databendSession;
     private String nodeID;
     private final AtomicReference<QueryResults> currentResults = new AtomicReference<>(null);
     private static final Logger logger = Logger.getLogger(DatabendClientV1.class.getPackage().getName());
 
-    private Consumer<DatabendSession> on_session_state_update;
+    private final Consumer<DatabendSession> on_session_state_update;
 
     public DatabendClientV1(OkHttpClient httpClient, String sql, ClientSettings settings,
             Consumer<DatabendSession> on_session_state_update,
@@ -89,7 +83,7 @@ public DatabendClientV1(OkHttpClient httpClient, String sql, ClientSettings sett
         this.host = settings.getHost();
         this.paginationOptions = settings.getPaginationOptions();
         this.requestTimeoutSecs = settings.getQueryTimeoutSecs();
-        this.additonalHeaders = settings.getAdditionalHeaders();
+        this.additionalHeaders = settings.getAdditionalHeaders();
         this.maxRetryAttempts = settings.getRetryAttempts();
         this.databendSession = new AtomicReference<>(settings.getSession());
         this.nodeID = last_node_id.get();
@@ -126,7 +120,7 @@ public static Request.Builder prepareRequest(HttpUrl url, Map<String, String> ad
     }
 
     private Request buildQueryRequest(String query, ClientSettings settings) {
-        HttpUrl url = HttpUrl.get(settings.getHost());
+        HttpUrl url = HttpUrl.parse(settings.getHost());
         if (url == null) {
             // TODO(zhihanz) use custom exception
             throw new IllegalArgumentException("Invalid host: " + settings.getHost());
@@ -139,7 +133,7 @@ private Request buildQueryRequest(String query, ClientSettings settings) {
         }
 
         url = url.newBuilder().encodedPath(QUERY_PATH).build();
-        Request.Builder builder = prepareRequest(url, this.additonalHeaders);
+        Request.Builder builder = prepareRequest(url, this.additionalHeaders);
         DatabendSession session = databendSession.get();
         if (session != null && session.getNeedSticky()) {
             builder.addHeader(ClientSettings.X_DATABEND_STICKY_NODE, nodeID);
@@ -149,10 +143,6 @@ private Request buildQueryRequest(String query, ClientSettings settings) {
 
     private static Request buildDiscoveryRequest(ClientSettings settings) {
         HttpUrl url = HttpUrl.get(settings.getHost());
-        if (url == null) {
-            // TODO(zhihanz) use custom exception
-            throw new IllegalArgumentException("Invalid host: " + settings.getHost());
-        }
         String discoveryPath = DISCOVERY_PATH;
         // intentionally use unsupported discovery path for testing
         if (settings.getAdditionalHeaders().get("~mock.unsupported.discovery") != null && BOOLEAN_TRUE_STR.equals(settings.getAdditionalHeaders().get("~mock.unsupported.discovery"))) {
@@ -313,19 +303,6 @@ private boolean executeInternal(Request request, OptionalLong materializedJsonSi
         }
     }
 
-    private String requestBodyToString(Request request) {
-        try {
-            final Request copy = request.newBuilder().build();
-            final Buffer buffer = new Buffer();
-            if (copy.body() != null) {
-                copy.body().writeTo(buffer);
-            }
-            return buffer.readUtf8();
-        } catch (final IOException e) {
-            return "did not work";
-        }
-    }
-
     @Override
     public boolean execute(Request request) {
         return executeInternal(request, OptionalLong.empty());
@@ -340,20 +317,13 @@ private void processResponse(Headers headers, QueryResults results) {
                 this.on_session_state_update.accept(session);
             }
         }
-        if (results.getQueryId() != null && this.additonalHeaders.get(ClientSettings.X_Databend_Query_ID) == null) {
-            this.additonalHeaders.put(ClientSettings.X_Databend_Query_ID, results.getQueryId());
+        if (results.getQueryId() != null && this.additionalHeaders.get(ClientSettings.X_Databend_Query_ID) == null) {
+            this.additionalHeaders.put(ClientSettings.X_Databend_Query_ID, results.getQueryId());
         }
         if (headers != null) {
-            String serverVersionString = headers.get(ClientSettings.X_DATABEND_VERSION);
-            if (serverVersionString != null) {
-                try {
-                    serverVersion = serverVersionString;
-                } catch (Exception ignored) {
-                }
-            }
             String route_hint = headers.get(ClientSettings.X_DATABEND_ROUTE_HINT);
             if (route_hint != null) {
-                this.additonalHeaders.put(ClientSettings.X_DATABEND_ROUTE_HINT, route_hint);
+                this.additionalHeaders.put(ClientSettings.X_DATABEND_ROUTE_HINT, route_hint);
             }
         }
         currentResults.set(results);
@@ -375,7 +345,7 @@ public boolean advance() {
         String nextUriPath = this.currentResults.get().getNextUri().toString();
         HttpUrl url = HttpUrl.get(this.host);
         url = url.newBuilder().encodedPath(nextUriPath).build();
-        Request.Builder builder = prepareRequest(url, this.additonalHeaders);
+        Request.Builder builder = prepareRequest(url, this.additionalHeaders);
         builder.addHeader(ClientSettings.X_DATABEND_STICKY_NODE, this.nodeID);
         Request request = builder.get().build();
         return executeInternal(request, OptionalLong.of(MAX_MATERIALIZED_JSON_RESPONSE_SIZE));
@@ -388,7 +358,7 @@ public boolean hasNext() {
 
     @Override
     public Map<String, String> getAdditionalHeaders() {
-        return additonalHeaders;
+        return additionalHeaders;
     }
 
     @Override
@@ -406,12 +376,6 @@ public String getNodeID() {
         return this.nodeID;
     }
 
-    @Override
-    public String getServerVersion() {
-        return this.serverVersion;
-    }
-
-
     @Override
     public void close() {
         closeQuery();
@@ -432,7 +396,7 @@ private void closeQuery() {
         String path = uri.toString();
         HttpUrl url = HttpUrl.get(this.host);
         url = url.newBuilder().encodedPath(path).build();
-        Request r = prepareRequest(url, this.additonalHeaders).get().build();
+        Request r = prepareRequest(url, this.additionalHeaders).get().build();
         try {
             httpClient.newCall(r).execute().close();
         } catch (IOException ignored) {