Add support for $pop

leviysoft · Feb 4, 2024 · 9bc695a · 9bc695a
1 parent 0ce9ed8
commit 9bc695a
Show file tree

Hide file tree

Showing 8 changed files with 117 additions and 1 deletion.
diff --git a/README.md b/README.md
@@ -349,6 +349,17 @@ val q = update[Student](_.addToSetAll(_.courses, List(42, 44, 53)))
 // q us {"$addToSet": {"courses": {$each: [42, 44, 53] }}}
 ```
 
+2. $pop
+
+```scala
+case class Student(id: Int, courses: List[Int])
+
+val q = update[Student](_.popHead(_.courses)) // removes the first element
+// q us {"$pop": {"courses": 1}}
+
+val q1 = update[Student](_.popLast(_.courses)) // removes the last element
+// q1 us {"$pop": {"courses": -1}}
+```
 
 #### Projection
 

diff --git a/oolong-core/src/main/scala/oolong/AstParser.scala b/oolong-core/src/main/scala/oolong/AstParser.scala
@@ -313,6 +313,14 @@ private[oolong] class DefaultAstParser(using quotes: Quotes) extends AstParser {
           val prop  = parsePropSelector(selectProp)
           val value = getValueOrIterable(valueExpr)
           parseUpdater(updater, FieldUpdateExpr.AddToSet(UExpr.Prop(prop), value, multipleValues = false) :: acc)
+
+        case '{($updater: Updater[Doc]).popHead($selectProp)} =>
+          val prop  = parsePropSelector(selectProp)
+          parseUpdater(updater, FieldUpdateExpr.Pop(UExpr.Prop(prop), FieldUpdateExpr.Pop.Remove.First) :: acc)
+
+        case '{($updater: Updater[Doc]).popLast($selectProp)} =>
+          val prop  = parsePropSelector(selectProp)
+          parseUpdater(updater, FieldUpdateExpr.Pop(UExpr.Prop(prop), FieldUpdateExpr.Pop.Remove.Last) :: acc)
 
         case '{ $updater: Updater[Doc] } =>
           updater match {

diff --git a/oolong-core/src/main/scala/oolong/UExpr.scala b/oolong-core/src/main/scala/oolong/UExpr.scala
@@ -43,6 +43,14 @@ private[oolong] object UExpr {
     // array update operators
     case class AddToSet(prop: Prop, expr: UExpr, multipleValues: Boolean) extends FieldUpdateExpr(prop)
 
+    case class Pop(prop: Prop, remove: Pop.Remove) extends FieldUpdateExpr(prop)
+
+    object Pop {
+      enum Remove {
+        case First, Last
+      }
+    }
+
   }
 
 }
diff --git a/oolong-core/src/main/scala/oolong/dsl/Dsl.scala b/oolong-core/src/main/scala/oolong/dsl/Dsl.scala
@@ -55,4 +55,7 @@ sealed trait Updater[DocT] {
   def addToSetAll[PropT, ValueT](selectProp: DocT => Iterable[PropT], value: Iterable[ValueT])(using
       PropT =:= ValueT
   ): Updater[DocT] = useWithinMacro("addToSet")
+
+  def popHead(selectProp: DocT => Iterable[?]): Updater[DocT] = useWithinMacro("popHead")
+  def popLast(selectProp: DocT => Iterable[?]): Updater[DocT] = useWithinMacro("popLast")
 }
diff --git a/oolong-mongo-it/src/test/scala/oolong/mongo/OolongMongoUpdateSpec.scala b/oolong-mongo-it/src/test/scala/oolong/mongo/OolongMongoUpdateSpec.scala
@@ -32,6 +32,8 @@ class OolongMongoUpdateSpec extends AsyncFlatSpec with ForAllTestContainer with
       TestClass("3", 12, InnerClass("sdf", 1), Nil, None, List.empty),
       TestClass("4", 13, InnerClass("sdf", 1), List(1), None, List.empty),
       TestClass("12345", 12, InnerClass("sdf", 11), Nil, None, List.empty),
+      TestClass("popHead", 1, InnerClass("popHead", 1), List(1, 2, 3), None, List.empty),
+      TestClass("popTail", 1, InnerClass("popTail", 1), List(1, 2, 3), None, List.empty)
     )
 
     implicit val ec = ExecutionContext.global
@@ -168,4 +170,38 @@ class OolongMongoUpdateSpec extends AsyncFlatSpec with ForAllTestContainer with
     )
   }
 
+  it should "$pop the fist element" in {
+    for {
+      upd <- collection
+        .findOneAndUpdate(
+          query[TestClass](_.field1 == "popHead"),
+          update[TestClass](
+            _.popHead(_.field4)
+          ),
+          new FindOneAndUpdateOptions().returnDocument(ReturnDocument.AFTER)
+        )
+        .head()
+        .map(BsonDecoder[TestClass].fromBson(_).get)
+    } yield assert(
+      upd.field4 == List(2, 3)
+    )
+  }
+
+  it should "$pop the last element" in {
+    for {
+      upd <- collection
+        .findOneAndUpdate(
+          query[TestClass](_.field1 == "popTail"),
+          update[TestClass](
+            _.popLast(_.field4)
+          ),
+          new FindOneAndUpdateOptions().returnDocument(ReturnDocument.AFTER)
+        )
+        .head()
+        .map(BsonDecoder[TestClass].fromBson(_).get)
+    } yield assert(
+      upd.field4 == List(1, 2)
+    )
+  }
+
 }
diff --git a/oolong-mongo/src/main/scala/oolong/mongo/MongoUpdateCompiler.scala b/oolong-mongo/src/main/scala/oolong/mongo/MongoUpdateCompiler.scala
@@ -8,6 +8,7 @@ import oolong.*
 import oolong.UExpr.FieldUpdateExpr
 import oolong.bson.meta.QueryMeta
 import oolong.mongo.MongoUpdateNode.MongoUpdateOp
+import oolong.mongo.MongoUpdateNode.MongoUpdateOp.Pop.Remove
 import oolong.mongo.MongoUpdateNode as MU
 import org.mongodb.scala.bson.BsonArray
 import org.mongodb.scala.bson.BsonBoolean
@@ -48,6 +49,11 @@ object MongoUpdateCompiler extends Backend[UExpr, MU, BsonDocument] {
             case FieldUpdateExpr.Unset(prop) => MU.MongoUpdateOp.Unset(MU.Prop(renames.getOrElse(prop.path, prop.path)))
             case FieldUpdateExpr.AddToSet(prop, expr, each) =>
               MU.MongoUpdateOp.AddToSet(MU.Prop(renames.getOrElse(prop.path, prop.path)), rec(expr), each)
+            case FieldUpdateExpr.Pop(prop, remove) =>
+              val muRemove = remove match
+                case FieldUpdateExpr.Pop.Remove.First => Remove.First
+                case FieldUpdateExpr.Pop.Remove.Last  => Remove.Last
+              MU.MongoUpdateOp.Pop(MU.Prop(renames.getOrElse(prop.path, prop.path)), muRemove)
           })
         case UExpr.ScalaCode(code)      => MU.ScalaCode(code)
         case UExpr.Constant(t)          => MU.Constant(t)
@@ -91,7 +97,10 @@ object MongoUpdateCompiler extends Backend[UExpr, MU, BsonDocument] {
           )("$setOnInsert"),
           renderOps(
             ops.collect { case s: MU.MongoUpdateOp.AddToSet => s }.map(renderAddToSet)
-          )("$addToSet")
+          )("$addToSet"),
+          renderOps(
+            ops.collect { case s: MU.MongoUpdateOp.Pop => s }.map(op => render(op.prop) + ": " + render(op.value))
+          )("$pop")
         ).flatten
           .mkString("{\n", ",\n", "\n}")
 
@@ -153,6 +162,7 @@ object MongoUpdateCompiler extends Backend[UExpr, MU, BsonDocument] {
         val tRenames      = targetOps(ops.collect { case s: MU.MongoUpdateOp.Rename => s })
         val tSetOnInserts = targetOps(ops.collect { case s: MU.MongoUpdateOp.SetOnInsert => s })
         val tAddToSets    = targetOps(ops.collect { case s: MU.MongoUpdateOp.AddToSet => s })
+        val tPops         = targetOps(ops.collect { case s: MU.MongoUpdateOp.Pop => s })
 
         // format: off
         def updaterGroup(groupName: String, updaters: List[Expr[(String, BsonValue)]]): Option[Expr[(String, BsonDocument)]] =
@@ -173,6 +183,7 @@ object MongoUpdateCompiler extends Backend[UExpr, MU, BsonDocument] {
           updaterGroup("$rename", tRenames),
           updaterGroup("$setOnInsert", tSetOnInserts),
           updaterGroup("$addToSet", tAddToSets),
+          updaterGroup("$pop", tPops),
         ).flatten
 
         '{

diff --git a/oolong-mongo/src/main/scala/oolong/mongo/MongoUpdateNode.scala b/oolong-mongo/src/main/scala/oolong/mongo/MongoUpdateNode.scala
@@ -32,5 +32,16 @@ case object MongoUpdateNode {
 
     case class AddToSet(override val prop: Prop, override val value: MU, each: Boolean)
         extends MongoUpdateOp(prop, value)
+
+    case class Pop(override val prop: Prop, remove: Pop.Remove) extends MongoUpdateOp(prop, remove.toConstant)
+    object Pop {
+      enum Remove {
+        case First, Last
+
+        def toConstant: MU.Constant[Int] = this match
+          case Remove.First => MU.Constant(-1)
+          case Remove.Last  => MU.Constant(1)
+      }
+    }
   }
 }
diff --git a/oolong-mongo/src/test/scala/oolong/mongo/UpdateSpec.scala b/oolong-mongo/src/test/scala/oolong/mongo/UpdateSpec.scala
@@ -218,6 +218,34 @@ class UpdateSpec extends AnyFunSuite {
     )
   }
 
+  test("$pop first") {
+    val q    = update[TestClass](_.popHead(_.nestedListField))
+    val repr = renderUpdate[TestClass](_.popHead(_.nestedListField))
+    test(
+      q,
+      repr,
+      BsonDocument(
+        "$pop" -> BsonDocument(
+          "nestedListField" -> BsonInt32(-1)
+        )
+      ),
+    )
+  }
+
+  test("$pop last") {
+    val q    = update[TestClass](_.popLast(_.nestedListField))
+    val repr = renderUpdate[TestClass](_.popLast(_.nestedListField))
+    test(
+      q,
+      repr,
+      BsonDocument(
+        "$pop" -> BsonDocument(
+          "nestedListField" -> BsonInt32(1)
+        )
+      ),
+    )
+  }
+
   test("several update operators combined") {
     val q = update[TestClass](
       _.unset(_.dateField)