Replace SheathLeafClauses with simpler SqlQuery-oriented transforms

Author: deusaquilus

quill-engine/src/main/scala/io/getquill/sql/SqlQuery.scala

@@ -11,14 +11,27 @@ import io.getquill.sql.Common.ContainsImpurities
 
 final case class OrderByCriteria(ast: Ast, ordering: PropertyOrdering)
 
-sealed trait FromContext { def quat: Quat }
+sealed trait FromContext {
+  def quat: Quat
+  def mapAst(f: Ast => Ast): FromContext = this match {
+    case c: TableContext            => c
+    case QueryContext(query, alias) => QueryContext(query.mapAst(f), alias)
+    case c: InfixContext            => c.mapAsts(f)
+    case JoinContext(t, a, b, on)   => JoinContext(t, a.mapAst(f), b.mapAst(f), f(on))
+    case FlatJoinContext(t, a, on)  => FlatJoinContext(t, a.mapAst(f), f(on))
+  }
+}
 final case class TableContext(entity: Entity, alias: Ident) extends FromContext {
   override def quat: Quat = entity.quat
 }
 final case class QueryContext(query: SqlQuery, alias: Ident) extends FromContext {
   override def quat: Quat = query.quat
 }
-final case class InfixContext(infix: Infix, alias: Ident) extends FromContext { override def quat: Quat = infix.quat }
+final case class InfixContext(infix: Infix, alias: Ident) extends FromContext {
+  override def quat: Quat = infix.quat
+  def mapAsts(f: Ast => Ast): InfixContext =
+    copy(infix = infix.copy(params = infix.params.map(f)))
+}
 final case class JoinContext(t: JoinType, a: FromContext, b: FromContext, on: Ast) extends FromContext {
   override def quat: Quat = Quat.Tuple(a.quat, b.quat)
 }
@@ -29,6 +42,16 @@ final case class FlatJoinContext(t: JoinType, a: FromContext, on: Ast) extends F
 sealed trait SqlQuery {
   def quat: Quat
 
+  def mapAst(f: Ast => Ast): SqlQuery =
+    this match {
+      case flatten: FlattenSqlQuery =>
+        flatten.mapAsts(f)
+      case SetOperationSqlQuery(a, op, b) =>
+        SetOperationSqlQuery(a.mapAst(f), op, b.mapAst(f))(quat)
+      case UnaryOperationSqlQuery(op, q) =>
+        UnaryOperationSqlQuery(op, q.mapAst(f))(quat)
+    }
+
   override def toString: String = {
     import io.getquill.MirrorSqlDialect._
     import io.getquill.idiom.StatementInterpolator._
@@ -83,6 +106,17 @@ final case class FlattenSqlQuery(
 )(quatType: Quat)
     extends SqlQuery {
   override def quat: Quat = quatType
+
+  def mapAsts(f: Ast => Ast): FlattenSqlQuery =
+    copy(
+      from = from.map(_.mapAst(f)),
+      where = where.map(f),
+      groupBy = groupBy.map(f),
+      orderBy = orderBy.map(o => o.copy(ast = f(o.ast))),
+      limit = limit.map(f),
+      offset = offset.map(f),
+      select = select.map(s => s.copy(ast = f(s.ast)))
+    )(quatType)
 }
 
 object TakeDropFlatten {

quill-engine/src/main/scala/io/getquill/sql/norm/RemoveExtraAlias.scala

@@ -28,24 +28,26 @@ case class ValueizeSingleLeafSelects(strategy: NamingStrategy) extends Stateless
     def containsAlias(ast: Ast): Boolean =
       CollectAst.byType[Ident](ast).exists(id => fromContextAliases.contains(id))
 
-    // if it is a leaf add leaf.value
-    val select =
-      q.select.map {
-        // TODO need to do this kind of replacement in Join-by clauses the aggregations etc...
-        case sv: SelectValue if containsAlias(sv.ast) =>
-          val reductions     = CollectAst.byType[Ident](sv.ast).map(id => id -> valueize(id))
-          val newAst         = BetaReduction(sv.ast, TypeBehavior.ReplaceWithReduction, reductions: _*)
-          val newSelectValue = SelectValue(newAst, sv.alias, sv.concat)
-          newSelectValue match {
-            case sv @ SelectValue(LeafQuat(ast), _, _) => sv.copy(alias = Some("value"))
-            case _                                     => newSelectValue
-          }
-        case sv @ SelectValue(LeafQuat(ast), _, _) =>
-          sv.copy(alias = Some("value")) // TODO check if there is no alias already? Probably don't need to since aliasing only really happens in ExpandNestedQueries
-        case sv => sv
+    // If there is one single select clause that has a primitive (i.e. Leaf) quat then we can alias it to "value"
+    // This is the case of `SELECT primitive FROM (SELECT p.age from Person p) AS primitive`
+    // where we turn it into `SELECT p.name AS value FROM Person p`
+    def aliasSelects(selectValues: List[SelectValue]) =
+      selectValues match {
+        case List(sv @ SelectValue(LeafQuat(ast), _, _)) => List(sv.copy(alias = Some("value")))
+        case other                                       => other
+      }
 
+    val valuizedQuery =
+      q.copy(from = from)(q.quat).mapAsts { ast =>
+        if (containsAlias(ast)) {
+          val reductions = CollectAst.byType[Ident](ast).filter(id => fromContextAliases.contains(id)).map(id => id -> valueize(id))
+          BetaReduction(ast, TypeBehavior.ReplaceWithReduction, reductions: _*)
+        } else {
+          ast
+        }
       }
-    q.copy(select = select, from = from)(q.quat)
+
+    valuizedQuery.copy(select = aliasSelects(valuizedQuery.select))(q.quat)
   }
 
   // Turn every `FROM primitive-x` into a `FROM case-class(x.primitive)`
@@ -57,6 +59,17 @@ case class ValueizeSingleLeafSelects(strategy: NamingStrategy) extends Stateless
         other
     }
 
+  // protected def expandContext(s: FromContext): FromContext =
+  //   s match {
+  //     case QueryContext(q, alias) =>
+  //       QueryContext(apply(q, QueryLevel.Inner), alias)
+  //     case JoinContext(t, a, b, on) =>
+  //       JoinContext(t, expandContext(a), expandContext(b), on)
+  //     case FlatJoinContext(t, a, on) =>
+  //       FlatJoinContext(t, expandContext(a), on)
+  //     case _: TableContext | _: InfixContext => s
+  //   }
+
   private def collectAliases(contexts: List[FromContext]): List[Ident] =
     contexts.flatMap {
       case c: TableContext             => List(c.alias)