diff --git a/databend-jdbc/src/main/java/com/databend/jdbc/DatabendConnection.java b/databend-jdbc/src/main/java/com/databend/jdbc/DatabendConnection.java index 77466e53..21a853fd 100644 --- a/databend-jdbc/src/main/java/com/databend/jdbc/DatabendConnection.java +++ b/databend-jdbc/src/main/java/com/databend/jdbc/DatabendConnection.java @@ -520,22 +520,9 @@ public URI getURI() { } public void PingDatabendClientV1() throws IOException { - PaginationOptions options = getPaginationOptions(); - Map additionalHeaders = setAdditionalHeaders(); - additionalHeaders.put(X_Databend_Query_ID, UUID.randomUUID().toString()); - ClientSettings settings = new ClientSettings.Builder(). - setQueryTimeoutSecs(this.driverUri.getQueryTimeout()). - setConnectionTimeout(this.driverUri.getConnectionTimeout()). - setSocketTimeout(this.driverUri.getSocketTimeout()). - setSession(this.session.get()). - setHost(this.getURI().toString()). - setAdditionalHeaders(additionalHeaders). - setPaginationOptions(options).build(); + ClientSettings settings = makeClientSettings(); String query = "select 1"; HttpUrl url = HttpUrl.get(settings.getHost()); - if (url == null) { - throw new IllegalArgumentException("Invalid host: " + settings.getHost()); - } QueryRequest req = QueryRequest.builder().setSession(settings.getSession()).setStageAttachment(settings.getStageAttachment()).setPaginationOptions(settings.getPaginationOptions()).setSql(query).build(); String reqString = req.toString(); if (reqString == null || reqString.isEmpty()) { @@ -571,21 +558,26 @@ public void executePing(Request request) throws IOException { // TODO(zhihanz): session property push down DatabendClient startQuery(String sql) throws SQLException { + return new DatabendClientV1(httpClient, sql, makeClientSettings()); + } + + DatabendClient startQuery(String sql, StageAttachment attach) throws SQLException { PaginationOptions options = getPaginationOptions(); Map additionalHeaders = setAdditionalHeaders(); - additionalHeaders.put(X_Databend_Query_ID, UUID.randomUUID().toString()); ClientSettings s = new ClientSettings.Builder(). + setSession(this.session.get()). + setHost(this.getURI().toString()). setQueryTimeoutSecs(this.driverUri.getQueryTimeout()). setConnectionTimeout(this.driverUri.getConnectionTimeout()). setSocketTimeout(this.driverUri.getSocketTimeout()). - setSession(this.session.get()). - setHost(this.getURI().toString()). + setPaginationOptions(options). setAdditionalHeaders(additionalHeaders). - setPaginationOptions(options).build(); + setStageAttachment(attach). + build(); return new DatabendClientV1(httpClient, sql, s); } - DatabendClient startQuery(String sql, StageAttachment attach) throws SQLException { + private ClientSettings makeClientSettings() { PaginationOptions options = getPaginationOptions(); Map additionalHeaders = setAdditionalHeaders(); ClientSettings s = new ClientSettings.Builder(). @@ -596,9 +588,8 @@ DatabendClient startQuery(String sql, StageAttachment attach) throws SQLExceptio setSocketTimeout(this.driverUri.getSocketTimeout()). setPaginationOptions(options). setAdditionalHeaders(additionalHeaders). - setStageAttachment(attach). build(); - return new DatabendClientV1(httpClient, sql, s); + return s; } private Map setAdditionalHeaders() {