Merge pull request #167 from youngsofun/dev

fix: remove checks in updateClientSession.
databendlabs · Mar 5, 2024 · f74c32a · f74c32a
2 parents 8ccd928 + 3631c8b
commit f74c32a
Showing 1 changed file with 4 additions and 13 deletions.
diff --git a/databend-jdbc/src/main/java/com/databend/jdbc/DatabendStatement.java b/databend-jdbc/src/main/java/com/databend/jdbc/DatabendStatement.java
@@ -1,9 +1,6 @@
 package com.databend.jdbc;
 
-import com.databend.client.DatabendClient;
-import com.databend.client.QueryAffect;
-import com.databend.client.QueryResults;
-import com.databend.client.StageAttachment;
+import com.databend.client.*;
 import com.databend.jdbc.annotation.NotImplemented;
 
 import java.sql.Connection;
@@ -168,19 +165,13 @@ private void updateClientSession(QueryResults q) {
         if (q == null) {
             return;
         }
-        if (q.getAffect() == null) {
-            return;
-        }
-        if (q.getSession() == null) {
-            return;
-        }
-
-        if (q.getAffect().getClass() != QueryAffect.UseDB.class && q.getAffect().getClass() != QueryAffect.ChangeSettings.class) {
+        DatabendSession session = q.getSession();
+        if (session == null) {
             return;
         }
         // current query has result on update client session
         DatabendConnection connection = this.connection.get();
-        connection.setSession(q.getSession());
+        connection.setSession(session);
     }
 
     final boolean internalExecute(String sql, StageAttachment attachment) throws SQLException {