Fix resolving CTE columns in AliasedRelation

wvlet · Jul 11, 2023 · 946258b · 946258b
1 parent 93034eb
commit 946258b
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 2 deletions.
diff --git a/airframe-sql/src/main/scala/wvlet/airframe/sql/analyzer/TypeResolver.scala b/airframe-sql/src/main/scala/wvlet/airframe/sql/analyzer/TypeResolver.scala
@@ -380,8 +380,7 @@ object TypeResolver extends LogSupport {
           a.copy(expr = resolved)
         }
       case SingleColumn(a: Attribute, qualifier, _) if a.resolved =>
-        // Optimizes the nested attributes, but preserves qualifier in the parent
-        a.setQualifierIfEmpty(qualifier)
+        a.withQualifier(qualifier)
       case m: MultiSourceColumn =>
         var changed = false
         val resolvedInputs = m.inputs.map {

diff --git a/airframe-sql/src/test/scala/wvlet/airframe/sql/analyzer/TypeResolverTest.scala b/airframe-sql/src/test/scala/wvlet/airframe/sql/analyzer/TypeResolverTest.scala
@@ -1070,4 +1070,20 @@ class TypeResolverTest extends AirSpec with ResolverTestHelper {
     analyze("select \"prénom\" from (select name as \"prénom\" from A)")
     // No error
   }
+
+  test("resolve CTE in AliasedRelation") {
+    val p1 = analyze("with t1 as (select id from A) select id from (select id from t1) t2")
+    p1.outputAttributes shouldMatch {
+      case List(col: ResolvedAttribute) =>
+        col.fullName shouldBe "id"
+        col.sourceColumn.head.fullName shouldBe "A.id"
+    }
+
+    val p2 = analyze("with t1 as (select id from A) select count(id) from (select id from t1) t2")
+    p2.outputAttributes shouldMatch {
+      case List(SingleColumn(FunctionCall("count", Seq(col: ResolvedAttribute), _, _, _, _), _, _)) =>
+        col.fullName shouldBe "t2.id"
+        col.sourceColumn.head.fullName shouldBe "A.id"
+    }
+  }
 }