[Enhancement] Transaction stream load supports to set warehouse header

Signed-off-by: PengFei Li <[email protected]>
StarRocks · Apr 28, 2024 · 1eb5e3a · 1eb5e3a
1 parent 5c2a334
commit 1eb5e3a
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 6 deletions.
diff --git a/.../src/main/java/com/starrocks/connector/flink/examples/datastream/WriteMultipleTables.java b/.../src/main/java/com/starrocks/connector/flink/examples/datastream/WriteMultipleTables.java
@@ -71,8 +71,8 @@ public static void main(String[] args) throws Exception {
         //
         // 2. replace the connector options with your cluster configurations
         MultipleParameterTool params = MultipleParameterTool.fromArgs(args);
-        String jdbcUrl = params.get("jdbcUrl", "jdbc:mysql://127.0.0.1:11903");
-        String loadUrl = params.get("loadUrl", "127.0.0.1:11901");
+        String jdbcUrl = params.get("jdbcUrl", "jdbc:mysql://127.0.0.1:9030");
+        String loadUrl = params.get("loadUrl", "127.0.0.1:8030");
         String userName = params.get("userName", "root");
         String password = params.get("password", "");
 

diff --git a/...stream-load-sdk/src/main/java/com/starrocks/data/load/stream/TransactionStreamLoader.java b/...stream-load-sdk/src/main/java/com/starrocks/data/load/stream/TransactionStreamLoader.java
@@ -72,6 +72,10 @@ protected void initTxHeaders(StreamLoadProperties properties) {
         } else {
             beginHeaders.put("timeout", timeout);
         }
+        String warehouse = properties.getHeaders().get("warehouse");
+        if (warehouse != null) {
+            beginHeaders.put("warehouse", warehouse);
+        }
         this.beginTxnHeader = beginHeaders.entrySet().stream()
                 .map(entry -> new BasicHeader(entry.getKey(), entry.getValue()))
                 .toArray(Header[]::new);

diff --git a/...oad-sdk/src/main/java/com/starrocks/data/load/stream/properties/StreamLoadProperties.java b/...oad-sdk/src/main/java/com/starrocks/data/load/stream/properties/StreamLoadProperties.java
@@ -356,17 +356,17 @@ public Builder expectDelayTime(long expectDelayTime) {
         }
 
         public Builder connectTimeout(int connectTimeout) {
-            if (connectTimeout < 100 || connectTimeout > 60000) {
-                throw new IllegalArgumentException("connectTimeout `" + connectTimeout + "ms` set failed, must range in [100, 60000]");
+            if (connectTimeout < 100) {
+                throw new IllegalArgumentException("connectTimeout `" + connectTimeout + "ms` set failed, must be larger than 100ms");
             }
             this.connectTimeout = connectTimeout;
             return this;
         }
 
         public Builder waitForContinueTimeoutMs(int waitForContinueTimeoutMs) {
-            if (waitForContinueTimeoutMs < DEFAULT_WAIT_FOR_CONTINUE || waitForContinueTimeoutMs > 60000) {
+            if (waitForContinueTimeoutMs < DEFAULT_WAIT_FOR_CONTINUE) {
                 throw new IllegalArgumentException("waitForContinueTimeoutMs `" + waitForContinueTimeoutMs +
-                        "ms` set failed, must be in range in [3000, 60000]");
+                        "ms` set failed, must be be larger than 3000ms");
             }
             this.waitForContinueTimeoutMs = waitForContinueTimeoutMs;
             return this;