Skip to content

Commit f451d85

Browse files
hantmachantmac
authored
Merge pull request #391 from databendlabs/feat/add-auto-presign
feat: support auto presign
2 parents 54b10c0 + e963e3a commit f451d85

5 files changed

Lines changed: 78 additions & 6 deletions

File tree

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

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@ public final class ConnectionProperties {
3636
public static final ConnectionProperty<Integer> SOCKET_TIMEOUT = new SocketTimeout();
3737

3838
public static final ConnectionProperty<Boolean> PRESIGNED_URL_DISABLED = new PresignedUrlDisabled();
39+
public static final ConnectionProperty<String> PRESIGN = new Presign();
3940
public static final ConnectionProperty<Boolean> COPY_PURGE = new CopyPurge();
4041
public static final ConnectionProperty<String> NULL_DISPLAY = new NullDisplay();
4142
public static final ConnectionProperty<String> BINARY_FORMAT = new BinaryFormat();
@@ -59,6 +60,7 @@ public final class ConnectionProperties {
5960
.add(DATABASE)
6061
.add(ACCESS_TOKEN)
6162
.add(PRESIGNED_URL_DISABLED)
63+
.add(PRESIGN)
6264
.add(QUERY_TIMEOUT)
6365
.add(CONNECTION_TIMEOUT)
6466
.add(SOCKET_TIMEOUT)
@@ -194,6 +196,13 @@ public PresignedUrlDisabled() {
194196
}
195197
}
196198

199+
private static class Presign
200+
extends AbstractConnectionProperty<String> {
201+
public Presign() {
202+
super("presign", NOT_REQUIRED, ALLOWED, STRING_CONVERTER);
203+
}
204+
}
205+
197206
private static class CopyPurge extends AbstractConnectionProperty<Boolean> {
198207
public CopyPurge() {
199208
super("copy_purge", Optional.of("true"), NOT_REQUIRED, ALLOWED, BOOLEAN_CONVERTER);

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

Lines changed: 48 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -92,6 +92,7 @@ public class DatabendConnection implements Connection, DatabendConnectionExtensi
9292
private final AtomicReference<String> lastNodeID = new AtomicReference<>();
9393
private Semver serverVersion = null;
9494
private Capability serverCapability = null;
95+
private boolean presignDisabled;
9596

9697
private static volatile ExecutorService heartbeatScheduler = null;
9798
private final HeartbeatManager heartbeatManager = new HeartbeatManager();
@@ -135,6 +136,7 @@ private void initializeFileHandler() {
135136

136137
initializeFileHandler();
137138
this.login();
139+
this.checkPresign();
138140
}
139141

140142
Semver getServerVersion() {
@@ -172,6 +174,49 @@ private void login() throws SQLException {
172174
}
173175
}
174176

177+
private void checkPresign() {
178+
String presign = this.driverUri.getPresign();
179+
if (presign == null || presign.isEmpty()) {
180+
// fallback to legacy presigned_url_disabled boolean
181+
this.presignDisabled = this.driverUri.presignedUrlDisabled();
182+
return;
183+
}
184+
switch (presign.toLowerCase(Locale.US)) {
185+
case "auto":
186+
String host = this.httpUri.getHost();
187+
if (host != null && (host.endsWith(".databend.com")
188+
|| host.endsWith(".databend.cn")
189+
|| host.endsWith(".tidbcloud.com"))) {
190+
this.presignDisabled = false;
191+
} else {
192+
this.presignDisabled = true;
193+
}
194+
break;
195+
case "detect":
196+
try {
197+
PresignContext.newPresignContext(this, PresignContext.PresignMethod.UPLOAD, "~", ".databend-jdbc/check");
198+
this.presignDisabled = false;
199+
} catch (Exception e) {
200+
logger.warning("presign off: detect failed: " + e.getMessage());
201+
this.presignDisabled = true;
202+
}
203+
break;
204+
case "on":
205+
this.presignDisabled = false;
206+
break;
207+
case "off":
208+
this.presignDisabled = true;
209+
break;
210+
default:
211+
logger.warning("Unknown presign value: " + presign + ", falling back to presigned_url_disabled");
212+
this.presignDisabled = this.driverUri.presignedUrlDisabled();
213+
break;
214+
}
215+
if (this.debug()) {
216+
logger.info("presign disabled: " + this.presignDisabled + " (presign=" + presign + ", host=" + this.httpUri.getHost() + ")");
217+
}
218+
}
219+
175220
private static String randRouteHint() {
176221
String charset = "abcdef0123456789";
177222
Random rand = new Random();
@@ -636,7 +681,7 @@ public boolean isWrapperFor(Class<?> aClass)
636681
}
637682

638683
boolean presignedUrlDisabled() {
639-
return this.driverUri.presignedUrlDisabled();
684+
return this.presignDisabled;
640685
}
641686

642687
boolean copyPurge() {
@@ -858,8 +903,8 @@ public void uploadStream(String stageName, String destPrefix, InputStream inputS
858903
// Update the file size to the compressed size
859904
fileSize = byteArrayOutputStream.size();
860905
}
861-
if (this.driverUri.presignedUrlDisabled()) {
862-
DatabendPresignClient cli = new DatabendPresignClientV1(httpClient, this.httpUri.toString());
906+
if (this.presignDisabled) {
907+
DatabendPresignClient cli = new DatabendPresignClientV1(httpClient, this.httpUri.toString(), this.driverUri.getWarehouse());
863908
cli.presignUpload(null, dataStream, s, p + "/", destFileName, fileSize, true);
864909
} else {
865910
// logger.log(Level.FINE, "presign to @" + s + "/" + dest);

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

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,7 @@ final class DatabendDriverUri {
5454
private final String binaryFormat;
5555
private final String database;
5656
private final boolean presignedUrlDisabled;
57+
private final String presign;
5758
private final Integer connectionTimeout;
5859
private final Integer queryTimeout;
5960
private final Integer socketTimeout;
@@ -80,6 +81,7 @@ private DatabendDriverUri(String url, Properties driverProperties)
8081
this.uri = canonicalizeUri(rawUri, this.useSecureConnection, this.sslmode);
8182
this.database = DATABASE.getValue(properties).orElse("default");
8283
this.presignedUrlDisabled = PRESIGNED_URL_DISABLED.getRequiredValue(properties);
84+
this.presign = PRESIGN.getValue(properties).orElse("");
8385
this.copyPurge = COPY_PURGE.getValue(properties).orElse(true);
8486
this.nullDisplay = NULL_DISPLAY.getValue(properties).orElse("\\N");
8587
this.binaryFormat = BINARY_FORMAT.getValue(properties).orElse("");
@@ -309,6 +311,10 @@ public Boolean presignedUrlDisabled() {
309311
return presignedUrlDisabled;
310312
}
311313

314+
public String getPresign() {
315+
return presign;
316+
}
317+
312318
public Boolean copyPurge() {
313319
return copyPurge;
314320
}

‎databend-jdbc/src/main/java/com/databend/jdbc/cloud/DatabendPresignClientV1.java‎

Lines changed: 14 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
import java.util.logging.Level;
2929
import java.util.logging.Logger;
3030

31+
import static com.databend.client.ClientSettings.DatabendWarehouseHeader;
3132
import static com.databend.client.ClientSettings.X_DATABEND_RELATIVE_PATH;
3233
import static com.databend.client.ClientSettings.X_DATABEND_STAGE_NAME;
3334
import static java.lang.String.format;
@@ -43,9 +44,15 @@ public class DatabendPresignClientV1
4344
private static final Duration RetryTimeout = Duration.ofMinutes(5);
4445
private final OkHttpClient client;
4546
private final String uri;
47+
private final String warehouse;
4648
private static final Logger logger = Logger.getLogger(DatabendPresignClientV1.class.getPackage().getName());
4749

4850
public DatabendPresignClientV1(OkHttpClient client, String uri)
51+
{
52+
this(client, uri, null);
53+
}
54+
55+
public DatabendPresignClientV1(OkHttpClient client, String uri, String warehouse)
4956
{
5057
Logger.getLogger(OkHttpClient.class.getName()).setLevel(Level.FINEST);
5158
OkHttpClient.Builder builder = client.newBuilder();
@@ -90,6 +97,7 @@ public DatabendPresignClientV1(OkHttpClient client, String uri)
9097
return response;
9198
}).build();
9299
this.uri = uri;
100+
this.warehouse = warehouse;
93101
}
94102

95103
private void uploadFromStream(InputStream inputStream, String stageName, String relativePath, String name, long fileSize)
@@ -100,10 +108,13 @@ private void uploadFromStream(InputStream inputStream, String stageName, String
100108
.setType(MultipartBody.FORM)
101109
.addFormDataPart("upload", name, new InputStreamRequestBody(null, inputStream, fileSize))
102110
.build();
103-
Headers headers = new Headers.Builder()
111+
Headers.Builder headersBuilder = new Headers.Builder()
104112
.add(X_DATABEND_STAGE_NAME, stageName)
105-
.add(X_DATABEND_RELATIVE_PATH, relativePath)
106-
.build();
113+
.add(X_DATABEND_RELATIVE_PATH, relativePath);
114+
if (this.warehouse != null && !this.warehouse.isEmpty()) {
115+
headersBuilder.add(DatabendWarehouseHeader, this.warehouse);
116+
}
117+
Headers headers = headersBuilder.build();
107118

108119
HttpUrl url = HttpUrl.get(this.uri);
109120
url = new HttpUrl.Builder()

‎docs/Connection.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -85,6 +85,7 @@ String url="jdbc:databend://databend:secret@0.0.0.0:8000/hello_databend";
8585
| sslmode | SSL mode | disable | jdbc:databend://0.0.0.0:8000/hello_databend?sslmode=enable |
8686
| copy_purge | If True, the command will purge the files in the stage after they are loaded successfully into the table | false | jdbc:databend://0.0.0.0:8000/hello_databend?copy_purge=true |
8787
| presigned_url_disabled | whether use presigned url to upload data, generally if you use local disk as your storage layer, it should be set as true | false | jdbc:databend://0.0.0.0:8000/hello_databend?presigned_url_disabled=true |
88+
| presign | Controls presign mode for data upload. Values: `auto` (enable for *.databend.com, *.databend.cn, *.tidbcloud.com hosts, disable otherwise), `detect` (probe the server to determine support), `on` (always enable), `off` (always disable). When set, takes precedence over presigned_url_disabled | none | jdbc:databend://0.0.0.0:8000/hello_databend?presign=auto |
8889
| wait_time_secs | Restful query api blocking time, if the query is not finished, the api will block for wait_time_secs seconds | 10 | jdbc:databend://0.0.0.0:8000/hello_databend?wait_time_secs=10 |
8990
| max_rows_per_page | the maximum rows per page in response data body | 100000 | jdbc:databend://0.0.0.0:8000/default?max_rows_per_page=100000 |
9091
| null_display | null value display | \N | jdbc:databend://0.0.0.0:8000/hello_databend?null_display=null |

0 commit comments

Comments
 (0)