Introduce QueryType enum for Analysis

trinodb · Oct 8, 2021 · abb23d6 · abb23d6
1 parent b539167
commit abb23d6
Show file tree

Hide file tree

Showing 8 changed files with 44 additions and 14 deletions.
diff --git a/core/trino-main/src/main/java/io/trino/sql/analyzer/Analysis.java b/core/trino-main/src/main/java/io/trino/sql/analyzer/Analysis.java
@@ -96,6 +96,7 @@
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static com.google.common.collect.ImmutableMap.toImmutableMap;
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
+import static io.trino.sql.analyzer.QueryType.DESCRIBE;
 import static java.lang.String.format;
 import static java.util.Collections.emptyList;
 import static java.util.Collections.unmodifiableMap;
@@ -202,8 +203,7 @@ public class Analysis
     private Optional<TableHandle> analyzeTarget = Optional.empty();
     private Optional<List<ColumnMetadata>> updatedColumns = Optional.empty();
 
-    // for describe input and describe output
-    private final boolean isDescribe;
+    private final QueryType queryType;
 
     // for recursive view detection
     private final Deque<Table> tablesForView = new ArrayDeque<>();
@@ -213,11 +213,11 @@ public class Analysis
     private final Multimap<Field, SourceColumn> originColumnDetails = ArrayListMultimap.create();
     private final Multimap<NodeRef<Expression>, Field> fieldLineage = ArrayListMultimap.create();
 
-    public Analysis(@Nullable Statement root, Map<NodeRef<Parameter>, Expression> parameters, boolean isDescribe)
+    public Analysis(@Nullable Statement root, Map<NodeRef<Parameter>, Expression> parameters, QueryType queryType)
     {
         this.root = root;
         this.parameters = ImmutableMap.copyOf(requireNonNull(parameters, "parameters is null"));
-        this.isDescribe = isDescribe;
+        this.queryType = requireNonNull(queryType, "queryType is null");
     }
 
     public Statement getStatement()
@@ -840,9 +840,14 @@ public Map<NodeRef<Parameter>, Expression> getParameters()
         return parameters;
     }
 
+    public QueryType getQueryType()
+    {
+        return queryType;
+    }
+
     public boolean isDescribe()
     {
-        return isDescribe;
+        return queryType == DESCRIBE;
     }
 
     public void setJoinUsing(Join node, JoinUsingAnalysis analysis)

diff --git a/core/trino-main/src/main/java/io/trino/sql/analyzer/Analyzer.java b/core/trino-main/src/main/java/io/trino/sql/analyzer/Analyzer.java
@@ -38,6 +38,7 @@
 import static io.trino.sql.analyzer.ExpressionTreeUtils.extractAggregateFunctions;
 import static io.trino.sql.analyzer.ExpressionTreeUtils.extractExpressions;
 import static io.trino.sql.analyzer.ExpressionTreeUtils.extractWindowExpressions;
+import static io.trino.sql.analyzer.QueryType.OTHERS;
 import static io.trino.sql.analyzer.SemanticExceptions.semanticException;
 import static java.util.Objects.requireNonNull;
 
@@ -80,13 +81,13 @@ public Analyzer(
 
     public Analysis analyze(Statement statement)
     {
-        return analyze(statement, false);
+        return analyze(statement, OTHERS);
     }
 
-    public Analysis analyze(Statement statement, boolean isDescribe)
+    public Analysis analyze(Statement statement, QueryType queryType)
     {
         Statement rewrittenStatement = StatementRewrite.rewrite(session, metadata, sqlParser, queryExplainer, statement, parameters, parameterLookup, groupProvider, accessControl, warningCollector, statsCalculator);
-        Analysis analysis = new Analysis(rewrittenStatement, parameterLookup, isDescribe);
+        Analysis analysis = new Analysis(rewrittenStatement, parameterLookup, queryType);
         StatementAnalyzer analyzer = new StatementAnalyzer(analysis, metadata, sqlParser, groupProvider, accessControl, session, warningCollector, CorrelationSupport.ALLOWED);
         analyzer.analyze(rewrittenStatement, Optional.empty());
 

diff --git a/core/trino-main/src/main/java/io/trino/sql/analyzer/ExpressionAnalyzer.java b/core/trino-main/src/main/java/io/trino/sql/analyzer/ExpressionAnalyzer.java
@@ -2552,9 +2552,9 @@ public static ExpressionAnalysis analyzeExpressions(
             Iterable<Expression> expressions,
             Map<NodeRef<Parameter>, Expression> parameters,
             WarningCollector warningCollector,
-            boolean isDescribe)
+            QueryType queryType)
     {
-        Analysis analysis = new Analysis(null, parameters, isDescribe);
+        Analysis analysis = new Analysis(null, parameters, queryType);
         ExpressionAnalyzer analyzer = create(analysis, session, metadata, sqlParser, groupProvider, accessControl, types, warningCollector);
         for (Expression expression : expressions) {
             analyzer.analyze(

diff --git a/core/trino-main/src/main/java/io/trino/sql/analyzer/QueryType.java b/core/trino-main/src/main/java/io/trino/sql/analyzer/QueryType.java
@@ -0,0 +1,21 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.trino.sql.analyzer;
+
+public enum QueryType
+{
+    DESCRIBE,
+    OTHERS,
+    /**/;
+}
diff --git a/core/trino-main/src/main/java/io/trino/sql/analyzer/StatementAnalyzer.java b/core/trino-main/src/main/java/io/trino/sql/analyzer/StatementAnalyzer.java
@@ -1906,7 +1906,7 @@ protected Scope visitSampledRelation(SampledRelation relation, Optional<Scope> s
                     ImmutableList.of(samplePercentage),
                     analysis.getParameters(),
                     WarningCollector.NOOP,
-                    analysis.isDescribe())
+                    analysis.getQueryType())
                     .getExpressionTypes();
 
             Type samplePercentageType = expressionTypes.get(NodeRef.of(samplePercentage));

diff --git a/core/trino-main/src/main/java/io/trino/sql/planner/TypeAnalyzer.java b/core/trino-main/src/main/java/io/trino/sql/planner/TypeAnalyzer.java
@@ -30,6 +30,7 @@
 import java.util.Map;
 
 import static io.trino.sql.analyzer.ExpressionAnalyzer.analyzeExpressions;
+import static io.trino.sql.analyzer.QueryType.OTHERS;
 
 /**
  * This class is to facilitate obtaining the type of an expression and its subexpressions
@@ -59,7 +60,7 @@ public Map<NodeRef<Expression>, Type> getTypes(Session session, TypeProvider inp
                 expressions,
                 ImmutableMap.of(),
                 WarningCollector.NOOP,
-                false)
+                OTHERS)
                 .getExpressionTypes();
     }
 

diff --git a/core/trino-main/src/main/java/io/trino/sql/rewrite/DescribeInputRewrite.java b/core/trino-main/src/main/java/io/trino/sql/rewrite/DescribeInputRewrite.java
@@ -53,6 +53,7 @@
 import static io.trino.sql.QueryUtil.selectList;
 import static io.trino.sql.QueryUtil.simpleQuery;
 import static io.trino.sql.QueryUtil.values;
+import static io.trino.sql.analyzer.QueryType.DESCRIBE;
 import static io.trino.type.TypeUtils.getDisplayLabel;
 import static io.trino.type.UnknownType.UNKNOWN;
 import static java.util.Objects.requireNonNull;
@@ -123,7 +124,7 @@ protected Node visitDescribeInput(DescribeInput node, Void context)
 
             // create  analysis for the query we are describing.
             Analyzer analyzer = new Analyzer(session, metadata, parser, groupProvider, accessControl, queryExplainer, parameters, parameterLookup, warningCollector, statsCalculator);
-            Analysis analysis = analyzer.analyze(statement, true);
+            Analysis analysis = analyzer.analyze(statement, DESCRIBE);
 
             // get all parameters in query
             List<Parameter> parameters = getParameters(statement);

diff --git a/core/trino-main/src/main/java/io/trino/sql/rewrite/DescribeOutputRewrite.java b/core/trino-main/src/main/java/io/trino/sql/rewrite/DescribeOutputRewrite.java
@@ -53,6 +53,7 @@
 import static io.trino.sql.QueryUtil.selectList;
 import static io.trino.sql.QueryUtil.simpleQuery;
 import static io.trino.sql.QueryUtil.values;
+import static io.trino.sql.analyzer.QueryType.DESCRIBE;
 import static io.trino.type.TypeUtils.getDisplayLabel;
 import static java.util.Objects.requireNonNull;
 
@@ -121,7 +122,7 @@ protected Node visitDescribeOutput(DescribeOutput node, Void context)
             Statement statement = parser.createStatement(sqlString, createParsingOptions(session));
 
             Analyzer analyzer = new Analyzer(session, metadata, parser, groupProvider, accessControl, queryExplainer, parameters, parameterLookup, warningCollector, statsCalculator);
-            Analysis analysis = analyzer.analyze(statement, true);
+            Analysis analysis = analyzer.analyze(statement, DESCRIBE);
 
             Optional<Node> limit = Optional.empty();
             Row[] rows = analysis.getRootScope().getRelationType().getVisibleFields().stream().map(field -> createDescribeOutputRow(field, analysis)).toArray(Row[]::new);