Merge pull request #221 from datafuselabs/fix/setSchema-conn

fix: connection setSchema to switch database
databendlabs · May 21, 2024 · c65c22d · c65c22d
2 parents daf247f + 47e21eb
commit c65c22d
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 1 deletion.
diff --git a/databend-jdbc/src/main/java/com/databend/jdbc/DatabendConnection.java b/databend-jdbc/src/main/java/com/databend/jdbc/DatabendConnection.java
@@ -77,9 +77,9 @@ public class DatabendConnection implements Connection, FileTransferAPI, Consumer
     DatabendConnection(DatabendDriverUri uri, OkHttpClient httpClient) throws SQLException {
         requireNonNull(uri, "uri is null");
         this.httpUri = uri.getUri();
-        this.setSchema(uri.getDatabase());
         this.httpClient = httpClient;
         this.driverUri = uri;
+        this.setSchema(uri.getDatabase());
         DatabendSession session = new DatabendSession.Builder().setHost(this.getURI()).setDatabase(this.getSchema()).build();
         this.setSession(session);
     }
@@ -470,6 +470,7 @@ public void setSchema(String schema)
             throws SQLException {
         checkOpen();
         this.schema.set(schema);
+        this.startQuery("use " + schema);
     }
 
     @Override

diff --git a/databend-jdbc/src/test/java/com/databend/jdbc/TestDatabendDriverUri.java b/databend-jdbc/src/test/java/com/databend/jdbc/TestDatabendDriverUri.java
@@ -4,6 +4,7 @@
 import org.testng.Assert;
 import org.testng.annotations.Test;
 
+import java.sql.DriverManager;
 import java.sql.SQLException;
 import java.util.Properties;
 
@@ -207,4 +208,19 @@ public void testFull() throws SQLException {
         Assert.assertEquals("null", uri.nullDisplay().toString());
         Assert.assertEquals(false, uri.getStrNullAsNull());
     }
+    @Test
+    public void TestSetSchema() throws SQLException {
+        String url = "jdbc:databend://databend:databend@localhost:8000/";
+        Properties p = new Properties();
+        DatabendConnection connection = (DatabendConnection) DriverManager.getConnection(url, p);
+        try {
+            connection.createStatement().execute("create or replace database test2");
+            connection.createStatement().execute("create or replace table test2.test2(id int)");
+        }catch (SQLException e){
+            throw new RuntimeException(e);
+        }
+        connection.setSchema("test2");
+        Assert.assertEquals(connection.getSchema(), "test2");
+        connection.createStatement().execute("insert into test2 values (1)");
+    }
 }