Top-Level quat passing to query expansions (#2420)

zio · Feb 3, 2022 · f8bdc4d · f8bdc4d
1 parent 629ba73
commit f8bdc4d
Show file tree

Hide file tree

Showing 54 changed files with 422 additions and 262 deletions.
diff --git a/quill-cassandra/src/main/scala/io/getquill/context/cassandra/Ops.scala b/quill-cassandra/src/main/scala/io/getquill/context/cassandra/Ops.scala
@@ -12,7 +12,7 @@ trait Ops {
   }
 
   implicit class QueryOps[Q <: Query[_]](q: Q) {
-    def allowFiltering = quote(infix"$q ALLOW FILTERING".generic.pure.as[Q])
+    def allowFiltering = quote(infix"$q ALLOW FILTERING".transparent.pure.as[Q])
   }
 
   implicit class EntityOps[A <: EntityQuery[_]](q: A)

diff --git a/quill-cassandra/src/test/scala/io/getquill/context/cassandra/CqlIdiomSpec.scala b/quill-cassandra/src/test/scala/io/getquill/context/cassandra/CqlIdiomSpec.scala
@@ -5,6 +5,7 @@ import io.getquill.idiom.StatementInterpolator._
 import io.getquill.ast.{ Action => AstAction, Query => _, _ }
 import io.getquill.idiom.StringToken
 import io.getquill.Query
+import io.getquill.context.ExecutionType
 import io.getquill.quat.Quat
 
 class CqlIdiomSpec extends Spec {
@@ -379,15 +380,15 @@ class CqlIdiomSpec extends Spec {
 
     "ident" in {
       val a: Ast = Ident("a")
-      translate(a) mustBe (a -> stmt"a")
+      translate(a, Quat.Unknown, ExecutionType.Unknown) mustBe ((a, stmt"a", ExecutionType.Unknown))
     }
     "assignment" in {
       val a: Ast = Assignment(Ident("a"), Ident("b"), Ident("c"))
-      translate(a: Ast) mustBe (a -> stmt"b = c")
+      translate(a: Ast, Quat.Unknown, ExecutionType.Unknown) mustBe ((a, stmt"b = c", ExecutionType.Unknown))
     }
     "assignmentDual" in {
       val a: Ast = AssignmentDual(Ident("a1"), Ident("a2"), Ident("b"), Ident("c"))
-      translate(a: Ast) mustBe (a -> stmt"b = c")
+      translate(a: Ast, Quat.Unknown, ExecutionType.Unknown) mustBe ((a, stmt"b = c", ExecutionType.Unknown))
     }
     "aggregation" in {
       val t = implicitly[Tokenizer[AggregationOperator]]

diff --git a/quill-core/src/main/scala/io/getquill/context/ActionMacro.scala b/quill-core/src/main/scala/io/getquill/context/ActionMacro.scala
@@ -1,7 +1,8 @@
 package io.getquill.context
 
-import io.getquill.ast._ // Only .returning(r => r.prop) or .returning(r => OneElementCaseClass(r.prop)) is allowed.
+import io.getquill.ast._
 import io.getquill.norm.BetaReduction
+import io.getquill.quat.Quat
 import io.getquill.quotation.ReifyLiftings
 import io.getquill.util.MacroContextExt._
 
@@ -21,7 +22,7 @@ class ActionMacro(val c: MacroContext)
     c.untypecheck {
       q"""
         ..${EnableReflectiveCalls(c)}
-        val expanded = ${expand(extractAst(quoted))}
+        val expanded = ${expand(extractAst(quoted), inferQuat(quoted.tpe))}
         ${c.prefix}.translateQuery(
           expanded.string,
           expanded.prepare,
@@ -55,7 +56,7 @@ class ActionMacro(val c: MacroContext)
     c.untypecheck {
       q"""
         ..${EnableReflectiveCalls(c)}
-        val expanded = ${expand(extractAst(quoted))}
+        val expanded = ${expand(extractAst(quoted), Quat.Value)}
         ${c.prefix}.executeAction(
           expanded.string,
           expanded.prepare
@@ -67,7 +68,7 @@ class ActionMacro(val c: MacroContext)
     c.untypecheck {
       q"""
         ..${EnableReflectiveCalls(c)}
-        val expanded = ${expand(extractAst(quoted))}
+        val expanded = ${expand(extractAst(quoted), inferQuat(t.tpe))}
         ${c.prefix}.executeActionReturning(
           expanded.string,
           expanded.prepare,
@@ -134,7 +135,7 @@ class ActionMacro(val c: MacroContext)
               }
             val (ast, _) = reifyLiftings(BetaReduction(body, alias -> nestedLift))
             c.untypecheck {
-              call(batch, param, expand(ast))
+              call(batch, param, expand(ast, Quat.Unknown))
             }
         }
       case other =>
@@ -146,7 +147,7 @@ class ActionMacro(val c: MacroContext)
       (expanded.ast match {
         case ret: io.getquill.ast.ReturningAction =>
             io.getquill.norm.ExpandReturning.applyMap(ret)(
-              (ast, statement) => io.getquill.context.Expand(${c.prefix}, ast, statement, idiom, naming).string
+              (ast, statement) => io.getquill.context.Expand(${c.prefix}, ast, statement, idiom, naming, io.getquill.context.ExecutionType.Unknown).string
             )(idiom, naming)
         case ast =>
           io.getquill.util.Messages.fail(s"Can't find returning column. Ast: '$$ast'")
@@ -157,7 +158,7 @@ class ActionMacro(val c: MacroContext)
     c.untypecheck {
       q"""
         ..${EnableReflectiveCalls(c)}
-        val expanded = ${expand(extractAst(quoted))}
+        val expanded = ${expand(extractAst(quoted), Quat.Value)}
         ${c.prefix}.prepareAction(
           expanded.string,
           expanded.prepare

diff --git a/quill-core/src/main/scala/io/getquill/context/ContextMacro.scala b/quill-core/src/main/scala/io/getquill/context/ContextMacro.scala
@@ -7,6 +7,7 @@ import io.getquill.util.LoadObject
 import io.getquill.util.MacroContextExt._
 import io.getquill.NamingStrategy
 import io.getquill.idiom._
+import io.getquill.quat.Quat
 
 import scala.util.Success
 import scala.util.Failure
@@ -15,11 +16,11 @@ trait ContextMacro extends Quotation {
   val c: MacroContext
   import c.universe.{ Function => _, Ident => _, _ }
 
-  protected def expand(ast: Ast): Tree =
+  protected def expand(ast: Ast, topLevelQuat: Quat): Tree =
     q"""
       val (idiom, naming) = ${idiomAndNamingDynamic}
-      val (ast, statement) = ${translate(ast)}
-      io.getquill.context.Expand(${c.prefix}, ast, statement, idiom, naming)
+      val (ast, statement, executionType) = ${translate(ast, topLevelQuat)}
+      io.getquill.context.Expand(${c.prefix}, ast, statement, idiom, naming, executionType)
     """
 
   protected def extractAst[T](quoted: Tree): Ast =
@@ -29,10 +30,10 @@ trait ContextMacro extends Quotation {
         Dynamic(quoted)
       }
 
-  private def translate(ast: Ast): Tree =
+  private def translate(ast: Ast, topLevelQuat: Quat): Tree =
     IsDynamic(ast) match {
-      case false => translateStatic(ast)
-      case true  => translateDynamic(ast)
+      case false => translateStatic(ast, topLevelQuat)
+      case true  => translateDynamic(ast, topLevelQuat)
     }
 
   abstract class TokenLift(numQuatFields: Int) extends LiftUnlift(numQuatFields) {
@@ -48,13 +49,13 @@ trait ContextMacro extends Quotation {
     }
   }
 
-  private def translateStatic(ast: Ast): Tree = {
+  private def translateStatic(ast: Ast, topLevelQuat: Quat): Tree = {
     val liftUnlift = new { override val mctx: c.type = c } with TokenLift(ast.countQuatFields)
     import liftUnlift._
 
     idiomAndNamingStatic match {
       case Success((idiom, naming)) =>
-        val (normalizedAst, statement) = idiom.translate(ast)(naming)
+        val (normalizedAst, statement, _) = idiom.translate(ast, topLevelQuat, ExecutionType.Static)(naming)
 
         val (string, _) =
           ReifyStatement(
@@ -68,23 +69,24 @@ trait ContextMacro extends Quotation {
 
         c.query(string, idiom)
 
-        q"($normalizedAst, ${statement: Token})"
+        q"($normalizedAst, ${statement: Token}, io.getquill.context.ExecutionType.Static)"
       case Failure(ex) =>
         c.info(s"Can't translate query at compile time because the idiom and/or the naming strategy aren't known at this point.")
-        translateDynamic(ast)
+        translateDynamic(ast, topLevelQuat)
     }
   }
 
-  private def translateDynamic(ast: Ast): Tree = {
+  private def translateDynamic(ast: Ast, topLevelQuat: Quat): Tree = {
     val liftUnlift = new { override val mctx: c.type = c } with TokenLift(ast.countQuatFields)
     import liftUnlift._
+    val liftQuat: Liftable[Quat] = liftUnlift.quatLiftable
     c.info("Dynamic query")
     val translateMethod = if (io.getquill.util.Messages.cacheDynamicQueries) {
       q"idiom.translateCached"
     } else q"idiom.translate"
     q"""
       val (idiom, naming) = ${idiomAndNamingDynamic}
-      $translateMethod(_root_.io.getquill.norm.RepropagateQuats($ast))(naming)
+      $translateMethod(_root_.io.getquill.norm.RepropagateQuats($ast), ${liftQuat(topLevelQuat)}, io.getquill.context.ExecutionType.Dynamic)(naming)
     """
   }
 

diff --git a/quill-core/src/main/scala/io/getquill/context/Expand.scala b/quill-core/src/main/scala/io/getquill/context/Expand.scala
@@ -7,11 +7,12 @@ import io.getquill.NamingStrategy
 import io.getquill.idiom.Idiom
 
 case class Expand[C <: Context[_, _]](
-  val context: C,
-  val ast:     Ast,
-  statement:   Statement,
-  idiom:       Idiom,
-  naming:      NamingStrategy
+  val context:   C,
+  val ast:       Ast,
+  statement:     Statement,
+  idiom:         Idiom,
+  naming:        NamingStrategy,
+  executionType: ExecutionType
 ) {
 
   val (string, externals) =

diff --git a/quill-core/src/main/scala/io/getquill/context/QueryMacro.scala b/quill-core/src/main/scala/io/getquill/context/QueryMacro.scala
@@ -1,9 +1,10 @@
 package io.getquill.context
 
 import io.getquill.ast._
+import io.getquill.quat.Quat
+
 import scala.reflect.macros.whitebox.{ Context => MacroContext }
-import io.getquill.util.OptionalTypecheck
-import io.getquill.util.EnableReflectiveCalls
+import io.getquill.util.{ EnableReflectiveCalls, Messages, OptionalTypecheck }
 
 class QueryMacro(val c: MacroContext) extends ContextMacro {
   import c.universe.{ Ident => _, _ }
@@ -45,13 +46,15 @@ class QueryMacro(val c: MacroContext) extends ContextMacro {
   def prepareQuery[T](quoted: Tree)(implicit t: WeakTypeTag[T]): Tree =
     expandQuery[T](quoted, PrepareQuery)
 
-  private def expandQuery[T](quoted: Tree, method: ContextMethod)(implicit t: WeakTypeTag[T]) =
+  private def expandQuery[T](quoted: Tree, method: ContextMethod)(implicit t: WeakTypeTag[T]) = {
+    val topLevelQuat = inferQuat(t.tpe)
     OptionalTypecheck(c)(q"implicitly[${c.prefix}.Decoder[$t]]") match {
-      case Some(decoder) => expandQueryWithDecoder(quoted, method, decoder)
-      case None          => expandQueryWithMeta[T](quoted, method)
+      case Some(decoder) => expandQueryWithDecoder(quoted, method, decoder, topLevelQuat)
+      case None          => expandQueryWithMeta[T](quoted, method, topLevelQuat)
     }
+  }
 
-  private def expandQueryWithDecoder(quoted: Tree, method: ContextMethod, decoder: Tree) = {
+  private def expandQueryWithDecoder(quoted: Tree, method: ContextMethod, decoder: Tree, topLevelQuat: Quat) = {
     val extractedAst = extractAst(quoted)
     val ast = Map(extractedAst, Ident("x", extractedAst.quat), Ident("x", extractedAst.quat))
     val invocation =
@@ -63,7 +66,7 @@ class QueryMacro(val c: MacroContext) extends ContextMacro {
               expanded.string,
               expanded.prepare,
               (row, session) => $decoder(0, row, session)
-            )(io.getquill.context.ExecutionInfo.unknown, ())
+            )(io.getquill.context.ExecutionInfo(expanded.executionType, expanded.ast, staticTopLevelQuat), ())
            """
         case StreamQuery(UsesDefaultFetch) =>
           q"""
@@ -72,15 +75,15 @@ class QueryMacro(val c: MacroContext) extends ContextMacro {
               expanded.string,
               expanded.prepare,
               (row, session) => $decoder(0, row, session)
-            )(io.getquill.context.ExecutionInfo.unknown, ())
+            )(io.getquill.context.ExecutionInfo(expanded.executionType, expanded.ast, staticTopLevelQuat), ())
            """
         case StreamQuery(DoesNotUseFetch) =>
           q"""
             ${c.prefix}.${TermName(method.name)}(
               expanded.string,
               expanded.prepare,
               (row, session) => $decoder(0, row, session)
-            )(io.getquill.context.ExecutionInfo.unknown, ())
+            )(io.getquill.context.ExecutionInfo(expanded.executionType, expanded.ast, staticTopLevelQuat), ())
            """
         case TranslateQuery(ExplicitPrettyPrint(argValue)) =>
           q"""
@@ -89,7 +92,7 @@ class QueryMacro(val c: MacroContext) extends ContextMacro {
               expanded.prepare,
               (row, session) => $decoder(0, row, session),
               prettyPrint = ${argValue}
-            )(io.getquill.context.ExecutionInfo.unknown, ())
+            )(io.getquill.context.ExecutionInfo(expanded.executionType, expanded.ast, staticTopLevelQuat), ())
            """
         case TranslateQuery(DefaultPrint) =>
           q"""
@@ -98,35 +101,38 @@ class QueryMacro(val c: MacroContext) extends ContextMacro {
               expanded.prepare,
               (row, session) => $decoder(0, row, session),
               prettyPrint = false
-            )(io.getquill.context.ExecutionInfo.unknown, ())
+            )(io.getquill.context.ExecutionInfo(expanded.executionType, expanded.ast, staticTopLevelQuat), ())
            """
         case PrepareQuery =>
           q"""
             ${c.prefix}.${TermName(method.name)}(
               expanded.string,
               expanded.prepare
-            )(io.getquill.context.ExecutionInfo.unknown, ())
+            )(io.getquill.context.ExecutionInfo(expanded.executionType, expanded.ast, staticTopLevelQuat), ())
            """
         case _ =>
           q"""
             ${c.prefix}.${TermName(method.name)}(
               expanded.string,
               expanded.prepare,
               (row, session) => $decoder(0, row, session)
-            )(io.getquill.context.ExecutionInfo.unknown, ())
+            )(io.getquill.context.ExecutionInfo(expanded.executionType, expanded.ast, staticTopLevelQuat), ())
            """
       }
 
+    val liftUnlift = new { override val mctx: c.type = c } with TokenLift(0) // Don't serialize quats for the top-level
+    val liftQuat: Liftable[Quat] = liftUnlift.quatLiftable
     c.untypecheck {
       q"""
         ..${EnableReflectiveCalls(c)}
-        val expanded = ${expand(ast)}
+        val staticTopLevelQuat = ${if (Messages.attachTopLevelQuats) liftQuat(topLevelQuat) else q"io.getquill.quat.Quat.Unknown"}
+        val expanded = ${expand(ast, topLevelQuat)}
         ${invocation}
       """
     }
   }
 
-  private def expandQueryWithMeta[T](quoted: Tree, method: ContextMethod)(implicit t: WeakTypeTag[T]) = {
+  private def expandQueryWithMeta[T](quoted: Tree, method: ContextMethod, topLevelQuat: Quat)(implicit t: WeakTypeTag[T]) = {
     val metaTpe = c.typecheck(tq"${c.prefix}.QueryMeta[$t]", c.TYPEmode).tpe
     val meta = c.inferImplicitValue(metaTpe).orElse(q"${c.prefix}.materializeQueryMeta[$t]")
     val ast = extractAst(c.typecheck(q"${c.prefix}.quote($meta.expand($quoted))"))
@@ -139,7 +145,7 @@ class QueryMacro(val c: MacroContext) extends ContextMacro {
               expanded.string,
               expanded.prepare,
               $meta.extract
-            )(io.getquill.context.ExecutionInfo.unknown, ())
+            )(io.getquill.context.ExecutionInfo(expanded.executionType, expanded.ast, staticTopLevelQuat), ())
            """
         case StreamQuery(UsesDefaultFetch) =>
           q"""
@@ -148,15 +154,15 @@ class QueryMacro(val c: MacroContext) extends ContextMacro {
               expanded.string,
               expanded.prepare,
               $meta.extract
-            )(io.getquill.context.ExecutionInfo.unknown, ())
+            )(io.getquill.context.ExecutionInfo(expanded.executionType, expanded.ast, staticTopLevelQuat), ())
            """
         case StreamQuery(DoesNotUseFetch) =>
           q"""
             ${c.prefix}.${TermName(method.name)}(
               expanded.string,
               expanded.prepare,
               $meta.extract
-            )(io.getquill.context.ExecutionInfo.unknown, ())
+            )(io.getquill.context.ExecutionInfo(expanded.executionType, expanded.ast, staticTopLevelQuat), ())
            """
         case TranslateQuery(ExplicitPrettyPrint(argValue)) =>
           q"""
@@ -165,7 +171,7 @@ class QueryMacro(val c: MacroContext) extends ContextMacro {
               expanded.prepare,
               $meta.extract,
               prettyPrint = ${argValue}
-            )(io.getquill.context.ExecutionInfo.unknown, ())
+            )(io.getquill.context.ExecutionInfo(expanded.executionType, expanded.ast, staticTopLevelQuat), ())
            """
         case TranslateQuery(DefaultPrint) =>
           q"""
@@ -174,29 +180,32 @@ class QueryMacro(val c: MacroContext) extends ContextMacro {
               expanded.prepare,
               $meta.extract,
               prettyPrint = false
-            )(io.getquill.context.ExecutionInfo.unknown, ())
+            )(io.getquill.context.ExecutionInfo(expanded.executionType, expanded.ast, staticTopLevelQuat), ())
            """
         case PrepareQuery =>
           q"""
             ${c.prefix}.${TermName(method.name)}(
               expanded.string,
               expanded.prepare
-            )(io.getquill.context.ExecutionInfo.unknown, ())
+            )(io.getquill.context.ExecutionInfo(expanded.executionType, expanded.ast, staticTopLevelQuat), ())
            """
         case _ =>
           q"""
             ${c.prefix}.${TermName(method.name)}(
               expanded.string,
               expanded.prepare,
               $meta.extract
-            )(io.getquill.context.ExecutionInfo.unknown, ())
+            )(io.getquill.context.ExecutionInfo(expanded.executionType, expanded.ast, staticTopLevelQuat), ())
            """
       }
 
+    val liftUnlift = new { override val mctx: c.type = c } with TokenLift(0) // Don't serialize quats for the top level
+    val liftQuat: Liftable[Quat] = liftUnlift.quatLiftable
     c.untypecheck {
       q"""
         ..${EnableReflectiveCalls(c)}
-        val expanded = ${expand(ast)}
+        val staticTopLevelQuat = ${if (Messages.attachTopLevelQuats) liftQuat(topLevelQuat) else q"io.getquill.quat.Quat.Unknown"}
+        val expanded = ${expand(ast, topLevelQuat)}
         ${invocation}
       """
     }