From cba23b4c18811b86bae6b635c75b4c32b953d35e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=9D=A8=E6=B3=BD=E7=BF=B0?= <627617031@qq.com>
Date: Wed, 20 Dec 2023 17:17:22 +0800
Subject: [PATCH] Debugging a normal job returns multiple query result sets

---
 .../service/impl/DataBaseServiceImpl.java     | 40 ++++++++++---------
 .../metadata/driver/AbstractJdbcDriver.java   |  1 +
 2 files changed, 22 insertions(+), 19 deletions(-)

diff --git a/dinky-admin/src/main/java/org/dinky/service/impl/DataBaseServiceImpl.java b/dinky-admin/src/main/java/org/dinky/service/impl/DataBaseServiceImpl.java
index d499b221af..27e69a8498 100644
--- a/dinky-admin/src/main/java/org/dinky/service/impl/DataBaseServiceImpl.java
+++ b/dinky-admin/src/main/java/org/dinky/service/impl/DataBaseServiceImpl.java
@@ -27,6 +27,7 @@
 import org.dinky.data.dto.TaskDTO;
 import org.dinky.data.enums.ProcessStepType;
 import org.dinky.data.enums.Status;
+import org.dinky.data.exception.TaskNotDoneException;
 import org.dinky.data.model.Column;
 import org.dinky.data.model.DataBase;
 import org.dinky.data.model.QueryData;
@@ -44,10 +45,7 @@
 import org.apache.commons.lang3.StringUtils;
 
 import java.time.LocalDateTime;
-import java.util.ArrayList;
-import java.util.Collections;
-import java.util.List;
-import java.util.UUID;
+import java.util.*;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
@@ -275,7 +273,6 @@ public JobResult executeCommonSql(SqlDTO sqlDTO) {
             result.setEndTime(LocalDateTime.now());
             return result;
         }
-
         DataBase dataBase = getById(sqlDTO.getDatabaseId());
         if (Asserts.isNull(dataBase)) {
             result.setSuccess(false);
@@ -334,24 +331,29 @@ public JobResult StreamExecuteCommonSql(SqlDTO sqlDTO) {
             result.setEndTime(LocalDateTime.now());
             return result;
         }
+        List<JdbcSelectResult> jdbcSelectResults=new ArrayList<>();
         try (Driver driver = Driver.build(dataBase.getDriverConfig())) {
-
             Stream<JdbcSelectResult> jdbcSelectResultStream =
                     driver.StreamExecuteSql(sqlDTO.getStatement(), sqlDTO.getMaxRowNum());
-            List<JdbcSelectResult> jdbcSelectResults = jdbcSelectResultStream
-                    .takeWhile(res -> {
-                        if (!res.isSuccess()) {
-                            result.setSuccess(false);
-                            result.setError(res.getError());
-                            result.setEndTime(LocalDateTime.now());
-                            return false;
-                        }
-                        return true;
-                    })
-                    .collect(Collectors.toList());
-            if (result.getError() == null) {
-                result.setSuccess(true);
+            jdbcSelectResultStream.forEach(res->{
+                jdbcSelectResults.add(res);
+                if (!res.isSuccess()){
+                    throw new RuntimeException();
+                }
+            });
+            result.setResults(jdbcSelectResults);
+            result.setSuccess(true);
+            return result;
+        }
+        catch (RuntimeException e){
+            if (!jdbcSelectResults.isEmpty()){
+                result.setError(jdbcSelectResults.get(jdbcSelectResults.size()-1).getError());
+            }
+            else {
+                result.setError(e.getMessage());
             }
+            result.setSuccess(false);
+            result.setEndTime(LocalDateTime.now());
             result.setResults(jdbcSelectResults);
             return result;
         }
diff --git a/dinky-metadata/dinky-metadata-base/src/main/java/org/dinky/metadata/driver/AbstractJdbcDriver.java b/dinky-metadata/dinky-metadata-base/src/main/java/org/dinky/metadata/driver/AbstractJdbcDriver.java
index c922df4850..35e2a8d4d1 100644
--- a/dinky-metadata/dinky-metadata-base/src/main/java/org/dinky/metadata/driver/AbstractJdbcDriver.java
+++ b/dinky-metadata/dinky-metadata-base/src/main/java/org/dinky/metadata/driver/AbstractJdbcDriver.java
@@ -708,6 +708,7 @@ public Stream<JdbcSelectResult> StreamExecuteSql(String sql, Integer limit) {
                     return result;
                 }
             }
+            result.success();
             return result;
         });
     }