From 81874651d09716ba957c5e6727b7235af0a6a8c3 Mon Sep 17 00:00:00 2001
From: wenmo <32723967+wenmo@users.noreply.github.com>
Date: Wed, 1 Nov 2023 22:53:50 +0800
Subject: [PATCH] [Feature-1709][core] Support use FlinkSQL CTE to preview data

---
 .../src/main/java/org/dinky/data/result/ResultBuilder.java      | 1 +
 dinky-core/src/main/java/org/dinky/explainer/Explainer.java     | 1 +
 dinky-executor/src/main/java/org/dinky/parser/SqlType.java      | 2 ++
 dinky-web/src/pages/DataStudio/BottomContainer/Result/index.tsx | 2 +-
 4 files changed, 5 insertions(+), 1 deletion(-)

diff --git a/dinky-core/src/main/java/org/dinky/data/result/ResultBuilder.java b/dinky-core/src/main/java/org/dinky/data/result/ResultBuilder.java
index 6178fa6be3..34c6157430 100644
--- a/dinky-core/src/main/java/org/dinky/data/result/ResultBuilder.java
+++ b/dinky-core/src/main/java/org/dinky/data/result/ResultBuilder.java
@@ -39,6 +39,7 @@ static ResultBuilder build(
             String timeZone) {
         switch (operationType) {
             case SELECT:
+            case WITH:
                 return new SelectResultBuilder(id, maxRowNum, isChangeLog, isAutoCancel, timeZone);
             case SHOW:
             case DESC:
diff --git a/dinky-core/src/main/java/org/dinky/explainer/Explainer.java b/dinky-core/src/main/java/org/dinky/explainer/Explainer.java
index 3e92249088..e14cc74296 100644
--- a/dinky-core/src/main/java/org/dinky/explainer/Explainer.java
+++ b/dinky-core/src/main/java/org/dinky/explainer/Explainer.java
@@ -126,6 +126,7 @@ public JobParam pretreatStatements(String[] statements) {
                 statementList.add(statement);
             } else if (operationType.equals(SqlType.INSERT)
                     || operationType.equals(SqlType.SELECT)
+                    || operationType.equals(SqlType.WITH)
                     || operationType.equals(SqlType.SHOW)
                     || operationType.equals(SqlType.DESCRIBE)
                     || operationType.equals(SqlType.DESC)
diff --git a/dinky-executor/src/main/java/org/dinky/parser/SqlType.java b/dinky-executor/src/main/java/org/dinky/parser/SqlType.java
index 1e41a590ed..8a56cfa81b 100644
--- a/dinky-executor/src/main/java/org/dinky/parser/SqlType.java
+++ b/dinky-executor/src/main/java/org/dinky/parser/SqlType.java
@@ -63,6 +63,8 @@ public enum SqlType {
 
     CTAS("CTAS", "^CREATE\\s.*AS\\sSELECT.*$"),
 
+    WITH("WITH", "^WITH.*"),
+
     UNKNOWN("UNKNOWN", "^UNKNOWN.*");
 
     private String type;
diff --git a/dinky-web/src/pages/DataStudio/BottomContainer/Result/index.tsx b/dinky-web/src/pages/DataStudio/BottomContainer/Result/index.tsx
index f122fc9095..56d105d5bf 100644
--- a/dinky-web/src/pages/DataStudio/BottomContainer/Result/index.tsx
+++ b/dinky-web/src/pages/DataStudio/BottomContainer/Result/index.tsx
@@ -147,7 +147,7 @@ const Result = (props: any) => {
             id: current.id
           });
           const historyData = res.data;
-          if ('2' == historyData.status) {
+          if (historyData && '2' == historyData.status) {
             const historyId = historyData.id;
             const tableData = await handleGetOption('api/studio/getJobData', 'Get Data', {
               jobId: historyId