refactor

databendlabs · Jan 31, 2024 · 52a746a · 52a746a
1 parent 7b359f0
commit 52a746a
Showing 1 changed file with 12 additions and 21 deletions.
diff --git 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<String, String> 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<String, String> 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<String, String> 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<String, String> setAdditionalHeaders() {