From 261c84ebcb9526ab68d5989f6ad7f853aab33de2 Mon Sep 17 00:00:00 2001
From: wenmo <32723967+wenmo@users.noreply.github.com>
Date: Tue, 12 Dec 2023 23:53:33 +0800
Subject: [PATCH] [Fix-2628] [core] Fix set config is not effective

---
 .../executor/CustomTableEnvironmentImpl.java  | 10 ++++++
 ...Operation.java => CustomSetOperation.java} | 35 +++++++++++++++----
 .../trans/parse/SetSqlParseStrategy.java      |  6 ++--
 3 files changed, 42 insertions(+), 9 deletions(-)
 rename dinky-client/dinky-client-base/src/main/java/org/dinky/trans/ddl/{SetOperation.java => CustomSetOperation.java} (78%)

diff --git a/dinky-client/dinky-client-1.16/src/main/java/org/dinky/executor/CustomTableEnvironmentImpl.java b/dinky-client/dinky-client-1.16/src/main/java/org/dinky/executor/CustomTableEnvironmentImpl.java
index eea3517885..f1ad9ce279 100644
--- a/dinky-client/dinky-client-1.16/src/main/java/org/dinky/executor/CustomTableEnvironmentImpl.java
+++ b/dinky-client/dinky-client-1.16/src/main/java/org/dinky/executor/CustomTableEnvironmentImpl.java
@@ -22,6 +22,7 @@
 import org.dinky.assertion.Asserts;
 import org.dinky.data.model.LineageRel;
 import org.dinky.data.result.SqlExplainResult;
+import org.dinky.trans.ddl.CustomSetOperation;
 import org.dinky.utils.JsonUtils;
 import org.dinky.utils.LineageContext;
 
@@ -106,6 +107,15 @@ public boolean parseAndLoadConfiguration(String statement, Map<String, Object> s
             } else if (operation instanceof ResetOperation) {
                 callReset((ResetOperation) operation, getStreamExecutionEnvironment(), setMap);
                 return true;
+            } else if (operation instanceof CustomSetOperation) {
+                CustomSetOperation customSetOperation = (CustomSetOperation) operation;
+                if (customSetOperation.isValid()) {
+                    callSet(
+                            new SetOperation(customSetOperation.getKey(), customSetOperation.getValue()),
+                            getStreamExecutionEnvironment(),
+                            setMap);
+                }
+                return true;
             }
         }
         return false;
diff --git a/dinky-client/dinky-client-base/src/main/java/org/dinky/trans/ddl/SetOperation.java b/dinky-client/dinky-client-base/src/main/java/org/dinky/trans/ddl/CustomSetOperation.java
similarity index 78%
rename from dinky-client/dinky-client-base/src/main/java/org/dinky/trans/ddl/SetOperation.java
rename to dinky-client/dinky-client-base/src/main/java/org/dinky/trans/ddl/CustomSetOperation.java
index b1229f0b90..9202e12d6a 100644
--- a/dinky-client/dinky-client-base/src/main/java/org/dinky/trans/ddl/SetOperation.java
+++ b/dinky-client/dinky-client-base/src/main/java/org/dinky/trans/ddl/CustomSetOperation.java
@@ -40,12 +40,36 @@
  *
  * @since 2021/10/21 19:56
  */
-public class SetOperation extends AbstractOperation implements ExtendOperation {
+public class CustomSetOperation extends AbstractOperation implements ExtendOperation {
 
-    public SetOperation() {}
+    private String key;
+    private String value;
 
-    public SetOperation(String statement) {
+    public CustomSetOperation() {}
+
+    public CustomSetOperation(String statement) {
         super(statement);
+        parseConfig();
+    }
+
+    public String getKey() {
+        return key;
+    }
+
+    public String getValue() {
+        return value;
+    }
+
+    public boolean isValid() {
+        return Asserts.isAllNotNullString(key, value);
+    }
+
+    private void parseConfig() {
+        Map<String, List<String>> map = SetSqlParseStrategy.getInfo(statement);
+        if (Asserts.isNotNullMap(map) && map.size() == 2) {
+            key = StringUtils.join(map.get("SET"), ".");
+            value = StringUtils.join(map.get("="), ",");
+        }
     }
 
     @Override
@@ -58,10 +82,9 @@ public Optional<? extends TableResult> execute(CustomTableEnvironment tEnv) {
         } catch (ClassNotFoundException e) {
             logger.error("Class not found: org.apache.log4j.Logger");
         }
-        Map<String, List<String>> map = SetSqlParseStrategy.getInfo(statement);
-        if (Asserts.isNotNullMap(map) && map.size() == 2) {
+        if (Asserts.isAllNotNullString(key, value)) {
             Map<String, String> confMap = new HashMap<>();
-            confMap.put(StringUtils.join(map.get("SET"), "."), StringUtils.join(map.get("="), ","));
+            confMap.put(key, value);
             TableConfig config = tEnv.getConfig();
             config.addConfiguration(Configuration.fromMap(confMap));
             Configuration configuration = Configuration.fromMap(confMap);
diff --git a/dinky-client/dinky-client-base/src/main/java/org/dinky/trans/parse/SetSqlParseStrategy.java b/dinky-client/dinky-client-base/src/main/java/org/dinky/trans/parse/SetSqlParseStrategy.java
index cd5fa3f8c2..b868c7f28e 100644
--- a/dinky-client/dinky-client-base/src/main/java/org/dinky/trans/parse/SetSqlParseStrategy.java
+++ b/dinky-client/dinky-client-base/src/main/java/org/dinky/trans/parse/SetSqlParseStrategy.java
@@ -20,7 +20,7 @@
 package org.dinky.trans.parse;
 
 import org.dinky.parser.SqlSegment;
-import org.dinky.trans.ddl.SetOperation;
+import org.dinky.trans.ddl.CustomSetOperation;
 import org.dinky.utils.SqlSegmentUtil;
 
 import org.apache.flink.table.operations.Operation;
@@ -51,13 +51,13 @@ public static Map<String, List<String>> getInfo(String statement) {
         // SET(\s+(\S+)\s*=(.*))?
         List<SqlSegment> segments = new ArrayList<>();
         segments.add(new SqlSegment("(set)\\s+(.+)(\\s*=)", "[.]"));
-        segments.add(new SqlSegment("(=)\\s*(.*)( ENDOFSQL)", ","));
+        segments.add(new SqlSegment("(=)\\s*(.*)($)", ","));
         return SqlSegmentUtil.splitSql2Segment(segments, statement);
     }
 
     @Override
     public Operation convert(String statement) {
-        return new SetOperation(statement);
+        return new CustomSetOperation(statement);
     }
 
     @Override