[Optimization-3888][core] Optimize SQL execution steps

.gitignore

@@ -62,4 +62,5 @@ tmp/*
 extends/*
 /.run/
 
-.idea
+.idea
+*/tmp/*

dinky-client/dinky-client-1.14/src/main/java/org/dinky/executor/CustomTableEnvironmentImpl.java

@@ -19,6 +19,7 @@
 
 package org.dinky.executor;
 
+import org.dinky.data.exception.DinkyException;
 import org.dinky.data.model.LineageRel;
 import org.dinky.data.result.SqlExplainResult;
 import org.dinky.parser.CustomParserImpl;
@@ -60,6 +61,7 @@
 import org.apache.flink.table.operations.ModifyOperation;
 import org.apache.flink.table.operations.Operation;
 import org.apache.flink.table.operations.QueryOperation;
+import org.apache.flink.table.operations.ddl.CreateTableASOperation;
 import org.apache.flink.table.planner.delegation.DefaultExecutor;
 
 import java.lang.reflect.Field;
@@ -69,6 +71,10 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.stream.Collectors;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 import com.fasterxml.jackson.databind.node.ObjectNode;
 
@@ -82,6 +88,8 @@
  */
 @Slf4j
 public class CustomTableEnvironmentImpl extends AbstractCustomTableEnvironment {
+    private static final Logger log = LoggerFactory.getLogger(CustomTableEnvironmentImpl.class);
+
     private final CustomExtendedOperationExecutorImpl extendedExecutor = new CustomExtendedOperationExecutorImpl(this);
     private static final String UNSUPPORTED_QUERY_IN_EXECUTE_SQL_MSG =
             "Unsupported SQL query! executeSql() only accepts a single SQL statement of type "
@@ -91,6 +99,8 @@ public class CustomTableEnvironmentImpl extends AbstractCustomTableEnvironment {
                     + "CREATE VIEW, DROP VIEW, SHOW VIEWS, INSERT, DESCRIBE, LOAD MODULE, UNLOAD "
                     + "MODULE, USE MODULES, SHOW [FULL] MODULES.";
 
+    private List<ModifyOperation> modifyOperations = new ArrayList<>();
+
     public CustomTableEnvironmentImpl(
             CatalogManager catalogManager,
             ModuleManager moduleManager,
@@ -199,29 +209,59 @@ private static Executor lookupExecutor(
         }
     }
 
-    @Override
+    public List<ModifyOperation> getModifyOperations() {
+        return modifyOperations;
+    }
+
+    public void addModifyOperations(ModifyOperation modifyOperation) {
+        if (modifyOperation instanceof CreateTableASOperation) {
+            CreateTableASOperation ctasOperation = (CreateTableASOperation) modifyOperation;
+            executeInternal(ctasOperation.getCreateTableOperation());
+            modifyOperations.add(ctasOperation.getInsertOperation());
+        } else {
+            modifyOperations.add(modifyOperation);
+        }
+    }
+
+    public void addOperator(Transformation<?> transformation) {
+        getStreamExecutionEnvironment().addOperator(transformation);
+    }
+
+    public void clearModifyOperations() {
+        modifyOperations.clear();
+    }
+
+    public List<Transformation<?>> transOperatoinsToTransformation(List<ModifyOperation> modifyOperations) {
+        return getPlanner().translate(modifyOperations);
+    }
+
     public ObjectNode getStreamGraph(String statement) {
         List<Operation> operations = super.getParser().parse(statement);
         if (operations.size() != 1) {
             throw new TableException("Unsupported SQL query! explainSql() only accepts a single SQL query.");
-        } else {
-            List<ModifyOperation> modifyOperations = new ArrayList<>();
-            for (int i = 0; i < operations.size(); i++) {
-                if (operations.get(i) instanceof ModifyOperation) {
-                    modifyOperations.add((ModifyOperation) operations.get(i));
-                }
-            }
-            List<Transformation<?>> trans = getPlanner().translate(modifyOperations);
-            if (executor instanceof DefaultExecutor) {
-                StreamGraph streamGraph =
-                        ((DefaultExecutor) executor).getExecutionEnvironment().generateStreamGraph(trans);
-                JSONGenerator jsonGenerator = new JSONGenerator(streamGraph);
-                String json = jsonGenerator.getJSON();
-                return JsonUtils.parseObject(json);
-            } else {
-                throw new TableException("Unsupported SQL query! explainSql() need a single SQL to query.");
-            }
         }
+
+        List<ModifyOperation> modifyOperations = operations.stream()
+                .filter(ModifyOperation.class::isInstance)
+                .map(ModifyOperation.class::cast)
+                .collect(Collectors.toList());
+
+        StreamGraph streamGraph = transOperatoinsToStreamGraph(modifyOperations);
+        JSONGenerator jsonGenerator = new JSONGenerator(streamGraph);
+        return JsonUtils.parseObject(jsonGenerator.getJSON());
+    }
+
+    private StreamGraph transOperatoinsToStreamGraph(List<ModifyOperation> modifyOperations) {
+        List<Transformation<?>> trans = getPlanner().translate(modifyOperations);
+        final StreamExecutionEnvironment environment = getStreamExecutionEnvironment();
+        trans.forEach(environment::addOperator);
+
+        StreamGraph streamGraph = environment.getStreamGraph();
+        final Configuration configuration = getConfig().getConfiguration();
+        if (configuration.containsKey(PipelineOptions.NAME.key())) {
+            streamGraph.setJobName(configuration.getString(PipelineOptions.NAME));
+        }
+        return streamGraph;
     }
 
     @Override
@@ -281,31 +321,106 @@ public StreamGraph getStreamGraphFromInserts(List<String> statements) {
     }
 
     @Override
+    public Operation getOperationFromStatement(String statement) {
+        List<Operation> operations = getParser().parse(statement);
+        if (operations.isEmpty()) {
+            throw new TableException("No statement is parsed.");
+        }
+        if (operations.size() > 1) {
+            throw new TableException("Only single statement is supported.");
+        }
+        return operations.get(0);
+    }
+
+    public ModifyOperation getModifyOperationFromInsert(String statement) {
+        List<Operation> operations = getParser().parse(statement);
+        if (operations.isEmpty()) {
+            throw new TableException("No statement is parsed.");
+        }
+        if (operations.size() > 1) {
+            throw new TableException("Only single statement is supported.");
+        }
+        Operation operation = operations.get(0);
+        if (operation instanceof ModifyOperation) {
+            return (ModifyOperation) operation;
+        } else if (operation instanceof QueryOperation) {
+            log.info("Select statement is skipped.");
+            return null;
+        } else {
+            throw new TableException("Only insert or select statement is supported now.");
+        }
+    }
+
+    public StreamGraph getStreamGraph() {
+        return transOperatoinsToStreamGraph(modifyOperations);
+    }
+
+    public StreamGraph getStreamGraphFromModifyOperations(List<ModifyOperation> modifyOperations) {
+        List<Transformation<?>> trans = getPlanner().translate(modifyOperations);
+        for (Transformation<?> transformation : trans) {
+            getStreamExecutionEnvironment().addOperator(transformation);
+        }
+        StreamGraph streamGraph = getStreamExecutionEnvironment().getStreamGraph();
+        if (getConfig().getConfiguration().containsKey(PipelineOptions.NAME.key())) {
+            streamGraph.setJobName(getConfig().getConfiguration().getString(PipelineOptions.NAME));
+        }
+        return streamGraph;
+    }
+
     public SqlExplainResult explainSqlRecord(String statement, ExplainDetail... extraDetails) {
-        SqlExplainResult record = new SqlExplainResult();
         List<Operation> operations = getParser().parse(statement);
-        record.setParseTrue(true);
         if (operations.size() != 1) {
-            throw new TableException("Unsupported SQL query! explainSql() only accepts a single SQL query.");
+            throw new DinkyException("Unsupported SQL explain! explainSql() only accepts a single SQL.");
         }
-
+        SqlExplainResult record = new SqlExplainResult();
+        if (operations.isEmpty()) {
+            throw new DinkyException("No statement is explained.");
+        }
+        record.setParseTrue(true);
         Operation operation = operations.get(0);
         if (operation instanceof ModifyOperation) {
-            record.setType("Modify DML");
+            if (operation instanceof CreateTableASOperation) {
+                record.setExplain(operation.asSummaryString());
+                record.setType("CTAS");
+            } else {
+                record.setExplain(getPlanner().explain(operations, extraDetails));
+                record.setType("DML");
+            }
         } else if (operation instanceof ExplainOperation) {
-            record.setType("Explain DML");
+            record.setExplain(operation.asSummaryString());
+            record.setType("Explain");
         } else if (operation instanceof QueryOperation) {
-            record.setType("Query DML");
+            record.setExplain(getPlanner().explain(operations, extraDetails));
+            record.setType("DQL");
         } else {
             record.setExplain(operation.asSummaryString());
             record.setType("DDL");
         }
         record.setExplainTrue(true);
-        if ("DDL".equals(record.getType())) {
-            // record.setExplain("DDL语句不进行解释。");
-            return record;
+        return record;
+    }
+
+    public SqlExplainResult explainModifyOperations(
+            List<ModifyOperation> modifyOperations, ExplainDetail... extraDetails) {
+        SqlExplainResult record = new SqlExplainResult();
+        if (modifyOperations.isEmpty()) {
+            throw new DinkyException("No modify operation is explained.");
+        }
+        record.setParseTrue(true);
+        if (modifyOperations.size() == 1) {
+            Operation operation = modifyOperations.get(0);
+            if (operation instanceof CreateTableASOperation) {
+                record.setExplain(operation.asSummaryString());
+                record.setType("CTAS");
+            } else {
+                record.setExplain(getPlanner().explain(new ArrayList<>(modifyOperations), extraDetails));
+                record.setType("DML");
+            }
+        } else {
+            record.setExplain(getPlanner().explain(new ArrayList<>(modifyOperations), extraDetails));
+            record.setType("Statement Set");
         }
-        record.setExplain(getPlanner().explain(operations, extraDetails));
+        record.setExplainTrue(true);
         return record;
     }