feat: add stageName param in DSN

databend-jdbc/src/main/java/com/databend/jdbc/ConnectionProperties.java

@@ -31,6 +31,7 @@ public final class ConnectionProperties {
 
     public static final ConnectionProperty<Boolean> PRESIGNED_URL_DISABLED = new PresignedUrlDisabled();
     public static final ConnectionProperty<Boolean> COPY_PURGE = new CopyPurge();
+    public static final ConnectionProperty<String> STAGE_NAME = new StageName();
     public static final ConnectionProperty<String> NULL_DISPLAY = new NullDisplay();
     public static final ConnectionProperty<String> BINARY_FORMAT = new BinaryFormat();
     public static final ConnectionProperty<Integer> WAIT_TIME_SECS = new WaitTimeSecs();
@@ -130,6 +131,13 @@ public Tenant() {
         }
     }
 
+    private static class StageName
+            extends AbstractConnectionProperty<String> {
+        public StageName() {
+            super("stage_name", Optional.of("~"), NOT_REQUIRED, ALLOWED, STRING_CONVERTER);
+        }
+    }
+
     private static class AccessToken
             extends AbstractConnectionProperty<String> {
         public AccessToken() {

databend-jdbc/src/main/java/com/databend/jdbc/DatabendConnection.java

@@ -499,6 +499,10 @@ public boolean copyPurge() {
         return this.driverUri.copyPurge();
     }
 
+    public String stageName() {
+        return this.driverUri.stageName();
+    }
+
     public String warehouse() {
         return this.driverUri.getWarehouse();
     }

databend-jdbc/src/main/java/com/databend/jdbc/DatabendDriverUri.java

@@ -46,6 +46,7 @@ public final class DatabendDriverUri {
     private final String sslmode;
     private final String tenant;
     private final boolean copyPurge;
+    private final String stageName;
     private final String nullDisplay;
     private final String binaryFormat;
     private final String database;
@@ -74,6 +75,7 @@ private DatabendDriverUri(String url, Properties driverProperties)
         this.database = DATABASE.getValue(properties).orElse("default");
         this.presignedUrlDisabled = PRESIGNED_URL_DISABLED.getRequiredValue(properties);
         this.copyPurge = COPY_PURGE.getValue(properties).orElse(true);
+        this.stageName = STAGE_NAME.getValue(properties).orElse("~");
         this.nullDisplay = NULL_DISPLAY.getValue(properties).orElse("\\N");
         this.binaryFormat = BINARY_FORMAT.getValue(properties).orElse("");
         this.waitTimeSecs = WAIT_TIME_SECS.getRequiredValue(properties);
@@ -258,6 +260,10 @@ public Boolean copyPurge() {
         return copyPurge;
     }
 
+    public String stageName() {
+        return stageName;
+    }
+
     public String getWarehouse() {
         return warehouse;
     }

databend-jdbc/src/main/java/com/databend/jdbc/DatabendPreparedStatement.java

@@ -168,8 +168,8 @@ private DatabendCopyParams uploadBatchesForCopyInto() throws SQLException {
                     LocalDateTime.now().getSecond(),
                     uuid);
             String fileName = saved.getName();
-            c.uploadStream(null, stagePrefix, fis, fileName, saved.length(), false);
-            String stageName = "~";
+            String stageName = c.stageName();
+            c.uploadStream(stageName, stagePrefix, fis, fileName, saved.length(), false);
             Map<String, String> copyOptions = new HashMap<>();
             copyOptions.put("PURGE", String.valueOf(c.copyPurge()));
             copyOptions.put("NULL_DISPLAY", String.valueOf(c.nullDisplay()));
@@ -209,9 +209,10 @@ private StageAttachment uploadBatches() throws SQLException {
                     LocalDateTime.now().getSecond(),
                     uuid);
             String fileName = saved.getName();
+            String stageName = c.stageName();
             // upload to stage
-            c.uploadStream(null, stagePrefix, fis, fileName, saved.length(), false);
-            String stagePath = "@~/" + stagePrefix + fileName;
+            c.uploadStream(stageName, stagePrefix, fis, fileName, saved.length(), false);
+            String stagePath = "@" + stageName + "/" + stagePrefix + fileName;
             StageAttachment attachment = buildStateAttachment(c, stagePath);
             return attachment;
         } catch (Exception e) {
@@ -524,7 +525,7 @@ public void setString(int i, String s)
             String finalS1 = s;
             batchInsertUtils.ifPresent(insertUtils -> insertUtils.setPlaceHolderValue(i, finalS1));
         } else {
-            if (s.contains("'")){
+            if (s.contains("'")) {
                 s = s.replace("'", "\\\'");
             }
             String finalS = s;

databend-jdbc/src/test/java/com/databend/jdbc/TestDatabendDriverUri.java

@@ -130,7 +130,7 @@ public void testUserDatabaseProp() throws SQLException {
         props.setProperty("SSL", "true");
         props.setProperty("binary_format", "base64");
         props.setProperty("sslmode", "enable");
-        DatabendDriverUri uri = DatabendDriverUri.create("jdbc:databend://u1@localhost:33101/db1?password=p1&database=db2&query_timeout=120&connection_timeout=15&socket_timeout=15", props);
+        DatabendDriverUri uri = DatabendDriverUri.create("jdbc:databend://u1@localhost:33101/db1?password=p1&database=db2&stage_name=test_stage&query_timeout=120&connection_timeout=15&socket_timeout=15", props);
 
         Assert.assertEquals(uri.getProperties().get("user"), "u1");
         Assert.assertEquals(uri.getProperties().get("password"), "p1");
@@ -149,6 +149,7 @@ public void testUserDatabaseProp() throws SQLException {
         Assert.assertEquals("\\N", uri.nullDisplay().toString());
         Assert.assertEquals("base64", uri.binaryFormat().toString());
         Assert.assertEquals("enable", uri.getSslmode().toString());
+        Assert.assertEquals("test_stage", uri.stageName().toString());
     }
 
     @Test(groups = {"unit"})
@@ -173,6 +174,7 @@ public void testUserDatabasePropFull() throws SQLException {
         Assert.assertTrue(uri.presignedUrlDisabled().booleanValue());
         Assert.assertTrue(uri.copyPurge().booleanValue());
         Assert.assertEquals("", uri.binaryFormat().toString());
+        Assert.assertEquals("~", uri.stageName().toString());
     }
 
     @Test(groups = {"unit"})

docs/Connection.md

@@ -91,3 +91,4 @@ String url="jdbc:databend://databend:[email protected]:8000/hello_databend";
 | null_display           | null value display                                                                                                        | \N      | jdbc:databend://0.0.0.0:8000/hello_databend?null_display=null           |
 | binary_format          | binary format, support hex and base64                                                                                     | hex     | jdbc:databend://0.0.0.0:8000/default?binary_format=hex                  |
 | use_verify             | whether verify the server before establishing the connection                                                              | true    | jdbc:databend://0.0.0.0:8000/default?use_verify=true                    |
+| stage_name             | User specify stage name                                                                                                   | ~       | jdbc:databend://0.0.0.0:8000/default?stage_name=own_stage               |