Merge pull request #299 from youngsofun/warehouse

feat: set warehouse header according to settings
databendlabs · Jan 9, 2025 · 2128bd3 · 2128bd3
2 parents 50cdcad + af58e42
commit 2128bd3
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 15 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
@@ -786,9 +786,22 @@ private ClientSettings.Builder makeClientSettings(String queryID, String host) {
 
     private Map<String, String> setAdditionalHeaders() {
         Map<String, String> additionalHeaders = new HashMap<>();
-        if (!this.driverUri.getWarehouse().isEmpty()) {
-            additionalHeaders.put(DatabendWarehouseHeader, this.driverUri.getWarehouse());
+
+        DatabendSession session = this.getSession();
+        String warehouse = null;
+        if (session != null ) {
+            Map<String, String> settings = session.getSettings();
+            if (settings != null) {
+                warehouse = settings.get("warehouse");
+            }
+        }
+        if (warehouse == null && !this.driverUri.getWarehouse().isEmpty()) {
+            warehouse = this.driverUri.getWarehouse();
         }
+        if (warehouse!=null) {
+            additionalHeaders.put(DatabendWarehouseHeader, warehouse);
+        }
+
         if (!this.driverUri.getTenant().isEmpty()) {
             additionalHeaders.put(DatabendTenantHeader, this.driverUri.getTenant());
         }

diff --git a/databend-jdbc/src/main/java/com/databend/jdbc/DatabendStatement.java b/databend-jdbc/src/main/java/com/databend/jdbc/DatabendStatement.java
@@ -165,19 +165,6 @@ private void clearCurrentResults() {
         currentResult.set(null);
     }
 
-    private void updateClientSession(QueryResults q) {
-        if (q == null) {
-            return;
-        }
-        DatabendSession session = q.getSession();
-        if (session == null) {
-            return;
-        }
-        // current query has result on update client session
-        DatabendConnection connection = this.connection.get();
-        connection.setSession(session);
-    }
-
     final boolean internalExecute(String sql, StageAttachment attachment) throws SQLException {
         clearCurrentResults();
         checkOpen();