diff --git a/databend-jdbc/src/main/java/com/databend/jdbc/ConnectionProperties.java b/databend-jdbc/src/main/java/com/databend/jdbc/ConnectionProperties.java index 40eddb03..e38b0bfb 100644 --- a/databend-jdbc/src/main/java/com/databend/jdbc/ConnectionProperties.java +++ b/databend-jdbc/src/main/java/com/databend/jdbc/ConnectionProperties.java @@ -31,6 +31,7 @@ public final class ConnectionProperties { public static final ConnectionProperty PRESIGNED_URL_DISABLED = new PresignedUrlDisabled(); public static final ConnectionProperty COPY_PURGE = new CopyPurge(); + public static final ConnectionProperty STAGE_NAME = new StageName(); public static final ConnectionProperty NULL_DISPLAY = new NullDisplay(); public static final ConnectionProperty BINARY_FORMAT = new BinaryFormat(); public static final ConnectionProperty WAIT_TIME_SECS = new WaitTimeSecs(); @@ -130,6 +131,13 @@ public Tenant() { } } + private static class StageName + extends AbstractConnectionProperty { + public StageName() { + super("stage_name", Optional.of("~"), NOT_REQUIRED, ALLOWED, STRING_CONVERTER); + } + } + private static class AccessToken extends AbstractConnectionProperty { public AccessToken() { diff --git a/databend-jdbc/src/main/java/com/databend/jdbc/DatabendConnection.java b/databend-jdbc/src/main/java/com/databend/jdbc/DatabendConnection.java index 66289228..0b4490e1 100644 --- a/databend-jdbc/src/main/java/com/databend/jdbc/DatabendConnection.java +++ b/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(); } diff --git a/databend-jdbc/src/main/java/com/databend/jdbc/DatabendDriverUri.java b/databend-jdbc/src/main/java/com/databend/jdbc/DatabendDriverUri.java index 199d466f..f80fcae9 100644 --- a/databend-jdbc/src/main/java/com/databend/jdbc/DatabendDriverUri.java +++ b/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; } diff --git a/databend-jdbc/src/main/java/com/databend/jdbc/DatabendPreparedStatement.java b/databend-jdbc/src/main/java/com/databend/jdbc/DatabendPreparedStatement.java index 4b83af54..fc9731ca 100644 --- a/databend-jdbc/src/main/java/com/databend/jdbc/DatabendPreparedStatement.java +++ b/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 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; diff --git a/databend-jdbc/src/test/java/com/databend/jdbc/TestDatabendDriverUri.java b/databend-jdbc/src/test/java/com/databend/jdbc/TestDatabendDriverUri.java index 6c5bd491..2adacf8b 100644 --- a/databend-jdbc/src/test/java/com/databend/jdbc/TestDatabendDriverUri.java +++ b/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"}) diff --git a/docs/Connection.md b/docs/Connection.md index a60f47c9..df3c50d0 100644 --- a/docs/Connection.md +++ b/docs/Connection.md @@ -91,3 +91,4 @@ String url="jdbc:databend://databend:secret@0.0.0.0: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 |