Support for 'inline def' inside 'query()' blocks

Tinkoff · Jun 29, 2022 · dc57979 · dc57979
1 parent 26f3c2f
commit dc57979
Show file tree

Hide file tree

Showing 4 changed files with 131 additions and 6 deletions.
diff --git a/README.md b/README.md
@@ -74,18 +74,17 @@ val q = query[Person](_.name == "Joe" && unchecked(
 ))
 ```
 
-### Composing queries
+### Reusing queries
 
-At the moment query composition is only supported via `unchecked`:
+It's possible to reuse a query by defining an 'inline def':
 ```scala
-val cityFilter: BsonDocument = query[Person](_.address.!!.city == "Amsterdam")
+inline def cityFilter(doc: Person) = doc.address.!!.city == "Amsterdam"
 
-val q = query[Person](_.name == "Joe" && unchecked(cityFilter))
+val q = query[Person](p => p.name == "Joe" && cityFilter(p))
 ```
 
 ## Coming soon
 
-- better query composition
 - elasticsearch support
 - field renaming 
 - aggregation pipelines for Mongo

diff --git a/oolong-core/src/main/scala/ru/tinkoff/oolong/AstParser.scala b/oolong-core/src/main/scala/ru/tinkoff/oolong/AstParser.scala
@@ -146,6 +146,9 @@ private[oolong] class DefaultAstParser(using quotes: Quotes) extends AstParser {
       case AsTerm(Literal(BooleanConstant(c))) =>
         makeConst(c)
 
+      case InlinedSubquery(term) =>
+        parse(term.asExpr)
+
       case _ =>
         report.errorAndAbort("Unexpected expr while parsing AST: " + input.show + s"; term: ${showTerm(input.asTerm)}")
     }

diff --git a/oolong-core/src/main/scala/ru/tinkoff/oolong/Utils.scala b/oolong-core/src/main/scala/ru/tinkoff/oolong/Utils.scala
@@ -37,6 +37,24 @@ private[oolong] object Utils {
     }
   }
 
+  object InlinedSubquery {
+    def unapply(using quotes: Quotes)(
+        term: quotes.reflect.Term
+    ): Option[quotes.reflect.Term] = {
+      import quotes.reflect.*
+      term match {
+        case Inlined(_, _, expansion)          => unapply(expansion)
+        case Typed(term, _) => Some(term)
+        case _                                 => None
+      }
+    }
+
+    def unapply(expr: Expr[Any])(using quotes: Quotes): Option[quotes.reflect.Term] = {
+      import quotes.reflect.*
+      unapply(expr.asTerm)
+    }
+  }
+
   object AsTerm {
     def unapply(using quotes: Quotes)(
         expr: Expr[Any]
@@ -60,6 +78,11 @@ private[oolong] object Utils {
             loop(next, acc)
           case Ident(name) =>
             Some((name, acc))
+
+          // Ident() can be inlined if queries are composed via "inline def"
+          case Inlined(_, _, expansion) =>
+            loop(expansion, acc)
+
           case _ =>
             None
         }

diff --git a/oolong-mongo/src/test/scala/ru/tinkoff/oolong/mongo/QuerySpec.scala b/oolong-mongo/src/test/scala/ru/tinkoff/oolong/mongo/QuerySpec.scala
@@ -20,6 +20,10 @@ import ru.tinkoff.oolong.dsl.*
 
 class QuerySpec extends AnyFunSuite {
 
+  trait TestClassAncestor {
+    def intField: Int
+  }
+
   case class TestClass(
       intField: Int,
       stringField: String,
@@ -28,7 +32,7 @@ class QuerySpec extends AnyFunSuite {
       optionField: Option[Long],
       optionInnerClassField: Option[InnerClass],
       listField: List[Double]
-  )
+  ) extends TestClassAncestor
 
   case class InnerClass(
       fieldOne: String,
@@ -303,4 +307,100 @@ class QuerySpec extends AnyFunSuite {
     )
   }
 
+  inline def mySubquery1(doc: TestClass): Boolean = doc.intField == 123
+
+  test("calling an 'inline def' with the '(_)' syntax") {
+
+    val q = query[TestClass](mySubquery1(_))
+
+    assert(
+      q == BsonDocument(
+        "intField" -> BsonDocument("$eq" -> BsonInt32(123))
+      )
+    )
+  }
+
+  test("calling an 'inline def' with the '(x => f(x))' syntax") {
+
+    val q = query[TestClass](x => mySubquery1(x))
+
+    assert(
+      q == BsonDocument(
+        "intField" -> BsonDocument("$eq" -> BsonInt32(123))
+      )
+    )
+  }
+
+  test("'inline def' with '!!'") {
+
+    inline def myFilter(doc: TestClass): Boolean = doc.optionField.!! == 123L
+
+    val q = query[TestClass](myFilter(_))
+
+    assert(
+      q == BsonDocument(
+        "optionField" -> BsonDocument("$eq" -> BsonInt64(123))
+      )
+    )
+  }
+
+  test("generic 'inline def' with '<:' constraint") {
+
+    inline def genericSubquery[A <: TestClassAncestor](doc: A): Boolean = doc.intField == 123
+
+    val q = query[TestClass](genericSubquery(_))
+
+    assert(
+      q == BsonDocument(
+        "intField" -> BsonDocument("$eq" -> BsonInt32(123))
+      )
+    )
+  }
+
+  test("'inline def' without explicit return type") {
+
+    inline def myFilter(doc: TestClass) = doc.intField == 123
+
+    val q = query[TestClass](myFilter(_))
+
+    assert(
+      q == BsonDocument(
+        "intField" -> BsonDocument("$eq" -> BsonInt32(123))
+      )
+    )
+  }
+
+  test("composing queries via 'inline def' #1") {
+
+    val q = query[TestClass](x => x.stringField == "qqq" && mySubquery1(x))
+
+    assert(
+      q == BsonDocument(
+        "$and" -> BsonArray.fromIterable(
+          List(
+            BsonDocument("stringField" -> BsonDocument("$eq" -> BsonString("qqq"))),
+            BsonDocument("intField"    -> BsonDocument("$eq" -> BsonInt32(123)))
+          )
+        )
+      )
+    )
+  }
+
+  test("composing queries via 'inline def' #2") {
+
+    inline def mySubquery2(tc: TestClass): Boolean = mySubquery1(tc) || tc.intField == 456
+
+    val q = query[TestClass](mySubquery2(_))
+
+    assert(
+      q == BsonDocument(
+        "$or" -> BsonArray.fromIterable(
+          List(
+            BsonDocument("intField" -> BsonDocument("$eq" -> BsonInt32(123))),
+            BsonDocument("intField" -> BsonDocument("$eq" -> BsonInt32(456)))
+          )
+        )
+      )
+    )
+  }
 }