From 946258bd1e47a11b997eca696af1bf55be6476ed Mon Sep 17 00:00:00 2001
From: Naoki Takezoe <takezoe@gmail.com>
Date: Tue, 11 Jul 2023 20:54:11 +0900
Subject: [PATCH] Fix resolving CTE columns in AliasedRelation

---
 .../airframe/sql/analyzer/TypeResolver.scala     |  3 +--
 .../airframe/sql/analyzer/TypeResolverTest.scala | 16 ++++++++++++++++
 2 files changed, 17 insertions(+), 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
index 2ea71b5184..06abe0d74d 100644
--- 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
index 203915d8e9..50548e384d 100644
--- 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"
+    }
+  }
 }