MongoDB: replaceOne operation

mongodb/src/main/scala/akka/stream/alpakka/mongodb/DocumentReplace.scala

@@ -0,0 +1,9 @@
+/*
+ * Copyright (C) 2016-2020 Lightbend Inc. <https://www.lightbend.com>
+ */
+
+package akka.stream.alpakka.mongodb
+
+import org.bson.conversions.Bson
+
+case class DocumentReplace[T](filter: Bson, replacement: T)

mongodb/src/main/scala/akka/stream/alpakka/mongodb/scaladsl/MongoFlow.scala

@@ -7,8 +7,8 @@ package akka.stream.alpakka.mongodb.scaladsl
 import akka.stream.scaladsl.{Flow, Source}
 import akka.NotUsed
 import akka.annotation.InternalApi
-import akka.stream.alpakka.mongodb.DocumentUpdate
-import com.mongodb.client.model.{DeleteOptions, InsertManyOptions, InsertOneOptions, UpdateOptions}
+import akka.stream.alpakka.mongodb.{DocumentReplace, DocumentUpdate}
+import com.mongodb.client.model.{DeleteOptions, InsertManyOptions, InsertOneOptions, ReplaceOptions, UpdateOptions}
 import com.mongodb.client.result.{DeleteResult, UpdateResult}
 import com.mongodb.reactivestreams.client.MongoCollection
 import org.bson.conversions.Bson
@@ -29,6 +29,9 @@ object MongoFlow {
   /** Internal Api */
   @InternalApi private[mongodb] val DefaultDeleteOptions = new DeleteOptions()
 
+  /** Internal Api */
+  @InternalApi private[mongodb] val DefaultReplaceOptions = new ReplaceOptions()
+
   /**
    * A [[akka.stream.scaladsl.Flow Flow]] that will insert documents into a collection.
    *
@@ -103,4 +106,21 @@ object MongoFlow {
   def deleteMany[T](collection: MongoCollection[T],
                     options: DeleteOptions = DefaultDeleteOptions): Flow[Bson, (DeleteResult, Bson), NotUsed] =
     Flow[Bson].flatMapConcat(bson => Source.fromPublisher(collection.deleteMany(bson, options)).map(_ -> bson))
+
+  /**
+   * A [[akka.stream.scaladsl.Flow Flow]] that will replace document as defined by a [[DocumentReplace]].
+   *
+   * @param collection the mongo db collection to update.
+   * @param options options to apply to the operation
+   */
+  def replaceOne[T](
+      collection: MongoCollection[T],
+      options: ReplaceOptions = DefaultReplaceOptions
+  ): Flow[DocumentReplace[T], (UpdateResult, DocumentReplace[T]), NotUsed] =
+    Flow[DocumentReplace[T]].flatMapConcat(
+      documentReplace =>
+        Source
+          .fromPublisher(collection.replaceOne(documentReplace.filter, documentReplace.replacement, options))
+          .map(_ -> documentReplace)
+    )
 }

mongodb/src/main/scala/akka/stream/alpakka/mongodb/scaladsl/MongoSink.scala

@@ -6,14 +6,15 @@ package akka.stream.alpakka.mongodb.scaladsl
 
 import akka.stream.scaladsl.{Keep, Sink}
 import akka.Done
-import akka.stream.alpakka.mongodb.DocumentUpdate
+import akka.stream.alpakka.mongodb.{DocumentReplace, DocumentUpdate}
 import akka.stream.alpakka.mongodb.scaladsl.MongoFlow.{
   DefaultDeleteOptions,
   DefaultInsertManyOptions,
   DefaultInsertOneOptions,
+  DefaultReplaceOptions,
   DefaultUpdateOptions
 }
-import com.mongodb.client.model.{DeleteOptions, InsertManyOptions, InsertOneOptions, UpdateOptions}
+import com.mongodb.client.model.{DeleteOptions, InsertManyOptions, InsertOneOptions, ReplaceOptions, UpdateOptions}
 import com.mongodb.reactivestreams.client.MongoCollection
 import org.bson.conversions.Bson
 
@@ -63,6 +64,18 @@ object MongoSink {
   ): Sink[DocumentUpdate, Future[Done]] =
     MongoFlow.updateMany(collection, options).toMat(Sink.ignore)(Keep.right)
 
+  /**
+   * A [[akka.stream.scaladsl.Sink Sink]] that will replace document as defined by a [[akka.stream.alpakka.mongodb.DocumentReplace]].
+   *
+   * @param collection the mongo db collection to update.
+   * @param options options to apply to the operation
+   */
+  def replaceOne[T](
+      collection: MongoCollection[T],
+      options: ReplaceOptions = DefaultReplaceOptions
+  ): Sink[DocumentReplace[T], Future[Done]] =
+    MongoFlow.replaceOne(collection, options).toMat(Sink.ignore)(Keep.right)
+
   /**
    * A [[akka.stream.scaladsl.Sink Sink]] that will delete individual documents as defined by a [[org.bson.conversions.Bson Bson]] filter query.
    *

mongodb/src/test/scala/docs/scaladsl/MongoSinkSpec.scala

@@ -6,7 +6,7 @@ package docs.scaladsl
 
 import akka.actor.ActorSystem
 import akka.stream.ActorMaterializer
-import akka.stream.alpakka.mongodb.DocumentUpdate
+import akka.stream.alpakka.mongodb.{DocumentReplace, DocumentUpdate}
 import akka.stream.alpakka.mongodb.scaladsl.MongoSink
 import akka.stream.alpakka.testkit.scaladsl.LogCapturing
 import akka.stream.scaladsl.{Sink, Source}
@@ -35,8 +35,9 @@ class MongoSinkSpec
 
   // case class and codec for mongodb macros
   case class Number(_id: Int)
+  case class DomainObject(_id: Int, firstProperty: String, secondProperty: String)
 
-  val codecRegistry = fromRegistries(fromProviders(classOf[Number]), DEFAULT_CODEC_REGISTRY)
+  val codecRegistry = fromRegistries(fromProviders(classOf[Number], classOf[DomainObject]), DEFAULT_CODEC_REGISTRY)
 
   implicit val system = ActorSystem()
   implicit val mat = ActorMaterializer()
@@ -48,6 +49,8 @@ class MongoSinkSpec
   private val db = client.getDatabase("MongoSinkSpec").withCodecRegistry(codecRegistry)
   private val numbersColl: MongoCollection[Number] =
     db.getCollection("numbersSink", classOf[Number]).withCodecRegistry(codecRegistry)
+  private val domainObjectsColl: MongoCollection[DomainObject] =
+    db.getCollection("domainObjectsSink", classOf[DomainObject]).withCodecRegistry(codecRegistry)
   private val numbersDocumentColl = db.getCollection("numbersSink")
 
   implicit val defaultPatience =
@@ -67,6 +70,16 @@ class MongoSinkSpec
       .runWith(Sink.head)
       .futureValue
 
+  def insertDomainObjectsRange(): Unit =
+    Source
+      .fromPublisher(
+        domainObjectsColl.insertMany(
+          testRange.map(i => DomainObject(i, s"first-property-$i", s"second-property-$i")).asJava
+        )
+      )
+      .runWith(Sink.head)
+      .futureValue
+
   "MongoSinkSpec" must {
 
     "save with insertOne" in assertAllStagesStopped {
@@ -204,6 +217,29 @@ class MongoSinkSpec
 
       found mustBe empty
     }
+
+    "replace with replaceOne and codec support" in assertAllStagesStopped {
+      insertDomainObjectsRange()
+      val updatedObjects =
+        testRange.map(i => DomainObject(i, s"updated-first-property-$i", s"updated-second-property-$i"))
+
+      // #replace-one
+      val source = Source(testRange).map(
+        i =>
+          DocumentReplace[DomainObject](
+            filter = Filters.eq("_id", i),
+            replacement = DomainObject(i, s"updated-first-property-$i", s"updated-second-property-$i")
+          )
+      )
+      val completion = source.runWith(MongoSink.replaceOne[DomainObject](domainObjectsColl))
+      // #replace-one
+
+      completion.futureValue
+
+      val found = Source.fromPublisher(domainObjectsColl.find()).runWith(Sink.seq).futureValue
+
+      found must contain theSameElementsAs updatedObjects
+    }
   }
 
 }