Add safeguard for multi-batch execution (#205)

(cherry picked from commit 8c6b9c5)
(cherry picked from commit 7b6c6d0)

Author: ferenc-csaky

flink-sql-runner/src/main/java/com/datasqrl/flinkrunner/SqlExecutor.java

@@ -47,9 +47,12 @@
 @RequiredArgsConstructor(access = AccessLevel.PACKAGE)
 class SqlExecutor {
 
-  private static final Pattern SET_STATEMENT_PATTERN =
+  private static final Pattern SET_PATTERN =
       Pattern.compile("SET\\s+'(\\S+)'\\s*=\\s*'(.+)';?", Pattern.CASE_INSENSITIVE);
 
+  private static final Pattern STATEMENT_SET_PATTERN =
+      Pattern.compile("EXECUTE\\s+STATEMENT\\s+SET", Pattern.CASE_INSENSITIVE | Pattern.DOTALL);
+
   private final TableEnvironment tEnv;
 
   SqlExecutor(Configuration config, @Nullable String udfPath) {
@@ -115,8 +118,11 @@ void setupSystemFunctions() {
   TableResult executeScript(String script) {
     var statements = SqlUtils.parseStatements(script);
     TableResult tableResult = null;
-    for (String statement : statements) {
-      tableResult = executeStatement(statement);
+
+    var it = statements.iterator();
+    while (it.hasNext()) {
+      var statement = it.next();
+      tableResult = executeStatement(statement, it.hasNext());
     }
 
     return tableResult;
@@ -141,24 +147,29 @@ TableResult executeCompiledPlan(String planJson) {
     }
   }
 
-  private TableResult executeStatement(String statement) {
+  private TableResult executeStatement(String statement, boolean intermediate) {
     TableResult tableResult = null;
     try {
-      var setMatcher = SET_STATEMENT_PATTERN.matcher(statement.trim());
+      var setMatcher = SET_PATTERN.matcher(statement.trim());
+      var statementSetMatcher = STATEMENT_SET_PATTERN.matcher(statement.trim());
 
       if (setMatcher.matches()) {
         // Handle SET statements
         var key = setMatcher.group(1);
         var value = setMatcher.group(2);
         tEnv.getConfig().getConfiguration().setString(key, value);
         log.info("Set configuration: {} = {}", key, value);
+
       } else {
         log.info("Executing statement:\n{}", statement);
         tableResult = tEnv.executeSql(statement);
+        if (statementSetMatcher.find() && intermediate) {
+          log.debug("Make sure to wait intermediate statement set to finish...");
+          tableResult.await();
+        }
       }
     } catch (Exception e) {
-      e.addSuppressed(new RuntimeException("Error while executing stmt: " + statement));
-      throw e;
+      throw new RuntimeException("Error while executing stmt: " + statement, e);
     }
     return tableResult;
   }