Fix the scope analysis of the CTEs (#729)

Author: goldmedal

wren-base/src/main/java/io/wren/base/sqlrewrite/analyzer/StatementAnalyzer.java

@@ -160,7 +160,8 @@ protected Scope visitTable(Table node, Optional<Scope> scope)
                 if (withQuery.isPresent()) {
                     // currently we only care about the table that is actually a model instead of a alias table that use cte table
                     // return empty scope here.
-                    Scope outputScope = createScopeForCommonTableExpression(node, withQuery.get(), scope);
+                    Optional<Scope> withScope = analysis.tryGetScope(withQuery.get().getQuery());
+                    Scope outputScope = createScopeForCommonTableExpression(node, withQuery.get(), withScope);
                     analysis.setScope(node, outputScope);
                     return outputScope;
                 }
@@ -202,8 +203,7 @@ protected Scope visitTable(Table node, Optional<Scope> scope)
         private Scope createScopeForCommonTableExpression(Table table, WithQuery withQuery, Optional<Scope> scope)
         {
             Query query = withQuery.getQuery();
-            Analysis analyzed = new Analysis(query);
-            Optional<Scope> queryScope = Optional.ofNullable(analyze(analyzed, query, sessionContext, wrenMDL));
+            Optional<Scope> queryScope = analysis.tryGetScope(query);
             List<Field> fields;
             Optional<List<Identifier>> columnNames = withQuery.getColumnNames();
             if (columnNames.isPresent()) {
@@ -243,7 +243,7 @@ private List<Field> createScopeForQuery(Query query, QualifiedName scopeName, Op
                     else {
                         SingleColumn singleColumn = (SingleColumn) selectItem;
                         String name = singleColumn.getAlias().map(Identifier::getValue)
-                                .or(() -> Optional.ofNullable(QueryUtil.getQualifiedName(singleColumn.getExpression())).map(QualifiedName::toString))
+                                .or(() -> Optional.ofNullable(QueryUtil.getQualifiedName(singleColumn.getExpression()).getSuffix()))
                                 .orElse(singleColumn.getExpression().toString());
                         if (scope.isPresent()) {
                             Optional<Field> fieldOptional = scope.get().getRelationType().resolveAnyField(QueryUtil.getQualifiedName(singleColumn.getExpression()));
@@ -402,7 +402,7 @@ private Scope analyzeFrom(QuerySpecification node, Optional<Scope> scope)
 
         private void analyzeWhere(Expression node, Scope scope)
         {
-            ExpressionAnalysis expressionAnalysis = analyzeExpression(node, scope);
+            analyzeExpression(node, scope);
         }
 
         private void analyzeWindowSpecification(WindowSpecification windowSpecification, Scope scope)
@@ -561,7 +561,7 @@ protected Scope visitTableSubquery(TableSubquery node, Optional<Scope> scope)
         private Optional<Scope> analyzeWith(Query node, Optional<Scope> scope)
         {
             if (node.getWith().isEmpty()) {
-                return Optional.empty();
+                return scope.map(s -> Scope.builder().parent(Optional.of(s)).build());
             }
 
             With with = node.getWith().get();
@@ -572,7 +572,8 @@ private Optional<Scope> analyzeWith(Query node, Optional<Scope> scope)
                 if (withScopeBuilder.containsNamedQuery(name)) {
                     throw new IllegalArgumentException(format("WITH query name '%s' specified more than once", name));
                 }
-                process(withQuery.getQuery(), withScopeBuilder.build());
+                Scope withQueryScope = process(withQuery.getQuery(), withScopeBuilder.build());
+                analysis.setScope(withQuery.getQuery(), withQueryScope);
                 withScopeBuilder.namedQuery(name, withQuery);
             }
 

wren-base/src/test/java/io/wren/base/sqlrewrite/analyzer/TestDecisionPointAnalyzer.java

@@ -676,6 +676,58 @@ WITH t1 as (SELECT * FROM customer), t2 as (SELECT * FROM orders)
                     new ExprSource("t1.custkey", "customer", "custkey", new NodeLocation(2, 33)),
                     new ExprSource("t2.custkey", "orders", "custkey", new NodeLocation(2, 50))));
         }
+
+        statement = parseSql("""
+                WITH t1 as (SELECT customer.custkey FROM customer),
+                t2 as (SELECT t1.custkey FROM t1)
+                SELECT t2.custkey FROM t2
+                """);
+        result = DecisionPointAnalyzer.analyze(statement, DEFAULT_SESSION_CONTEXT, mdl);
+        assertThat(result.size()).isEqualTo(3);
+        QueryAnalysis queryAnalysis = result.get(2);
+        assertThat(queryAnalysis.getSelectItems().get(0).getExprSources().size()).isEqualTo(1);
+
+        statement = parseSql("""
+                WITH t1 as (SELECT customer.custkey FROM customer),
+                t2 as (SELECT t1.custkey as custkey_alias FROM t1)
+                SELECT custkey_alias FROM t2
+                """);
+        result = DecisionPointAnalyzer.analyze(statement, DEFAULT_SESSION_CONTEXT, mdl);
+        assertThat(result.size()).isEqualTo(3);
+        queryAnalysis = result.get(2);
+        assertThat(queryAnalysis.getSelectItems().get(0).getExprSources().size()).isEqualTo(1);
+
+        statement = parseSql("""
+                WITH t1 as (SELECT customer.custkey FROM customer),
+                t2 as (SELECT * FROM t1)
+                SELECT custkey FROM t2
+                """);
+        result = DecisionPointAnalyzer.analyze(statement, DEFAULT_SESSION_CONTEXT, mdl);
+        assertThat(result.size()).isEqualTo(3);
+        queryAnalysis = result.get(2);
+        assertThat(queryAnalysis.getSelectItems().get(0).getExprSources().size()).isEqualTo(1);
+
+        statement = parseSql("""
+                WITH t1 as (SELECT customer.custkey FROM customer),
+                t2 as (SELECT * FROM t1)
+                SELECT custkey FROM t2
+                """);
+        result = DecisionPointAnalyzer.analyze(statement, DEFAULT_SESSION_CONTEXT, mdl);
+        assertThat(result.size()).isEqualTo(3);
+        queryAnalysis = result.get(2);
+        assertThat(queryAnalysis.getSelectItems().get(0).getExprSources().size()).isEqualTo(1);
+
+
+        // we only analyze the top-level expression
+        statement = parseSql("""
+                WITH t1 as (SELECT customer.custkey FROM customer),
+                t2 as (SELECT (t1.custkey + 1) as custkey_plus FROM t1)
+                SELECT custkey_plus FROM t2
+                """);
+        result = DecisionPointAnalyzer.analyze(statement, DEFAULT_SESSION_CONTEXT, mdl);
+        assertThat(result.size()).isEqualTo(3);
+        queryAnalysis = result.get(2);
+        assertThat(queryAnalysis.getSelectItems().get(0).getExprSources().size()).isEqualTo(0);
     }
 
     @Test