Merge pull request #30 from havocp/support-immutable-checkID

Make SerializableBSONObject.checkID return a new doc, so it can work with immutable doc types
bwmcadams · Jul 5, 2011 · 35480ad · 35480ad
2 parents 49f930c + d2c527b
commit 35480ad
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 22 deletions.
diff --git a/bson-driver/src/main/scala/SerializableBSONObject.scala b/bson-driver/src/main/scala/SerializableBSONObject.scala
@@ -45,13 +45,15 @@ trait SerializableBSONObject[T] {
   def checkKeys(doc: T): Unit
 
   /**
-   * Checks for an ID and generates one.
+   * Checks for an ID and generates one, returning a new doc with the id.
+   * The new doc may be a mutation of the old doc, OR a new object
+   * if the old doc was immutable.
    * Not all implementers will need this, but it gets invoked nonetheless
    * as a signal to BSONDocument, etc implementations to verify an id is there 
    * and generate one if needed.
    *
    */
-  def checkID(doc: T): Unit
+  def checkID(doc: T): T
 
   def _id(doc: T): Option[AnyRef]
 

diff --git a/bson-driver/src/main/scala/collection/Implicits.scala b/bson-driver/src/main/scala/collection/Implicits.scala
@@ -82,22 +82,25 @@ object `package` {
      * as a signal to BSONDocument, etc implementations to verify an id is there 
      * and generate one if needed.
      */
-    def checkID(doc: T) = doc.get("_id") match {
-      case Some(oid: ObjectId) => {
-        log.debug("Found an existing OID")
-        oid.notNew()
-        //oid
-      }
-      case Some(other) => {
-        log.debug("Found a non-OID ID")
-        //other
-      }
-      case None => {
-        val oid = new ObjectId()
-        doc.put("_id", oid)
-        log.trace("no ObjectId. Generated: %s", doc.get("_id"))
-        //oid
+    def checkID(doc: T) : T = {
+      doc.get("_id") match {
+        case Some(oid: ObjectId) => {
+          log.debug("Found an existing OID")
+          oid.notNew()
+          //oid
+        }
+        case Some(other) => {
+          log.debug("Found a non-OID ID")
+          //other
+        }
+        case None => {
+          val oid = new ObjectId()
+          doc.put("_id", oid)
+          log.trace("no ObjectId. Generated: %s", doc.get("_id"))
+          //oid
+        }
       }
+      doc
     }
 
     def _id(doc: T): Option[AnyRef] = doc.getAs[AnyRef]("_id")

diff --git a/mongo-driver/src/main/scala/MongoConnection.scala b/mongo-driver/src/main/scala/MongoConnection.scala
@@ -177,11 +177,14 @@ abstract class MongoConnection extends Logging {
   // TODO - Immutable mode / support immutable objects
   def insert[T](db: String)(collection: String)(doc: T, validate: Boolean = true)(callback: WriteRequestFuture)(implicit concern: WriteConcern = this.writeConcern, m: SerializableBSONObject[T]) {
     log.trace("Inserting: %s to %s.%s with WriteConcern: %s", doc, db, collection, concern)
-    if (validate) {
+    val checked = if (validate) {
       m.checkObject(doc)
       m.checkID(doc)
-    } else log.info("Validation of objects disabled; no ID Gen.")
-    send(InsertMessage(db + "." + collection, doc), callback)
+    } else {
+      log.info("Validation of objects disabled; no ID Gen.")
+      doc
+    }
+    send(InsertMessage(db + "." + collection, checked), callback)
   }
 
   /**
@@ -195,11 +198,11 @@ abstract class MongoConnection extends Logging {
    */
   def batchInsert[T](db: String)(collection: String)(docs: T*)(callback: WriteRequestFuture)(implicit concern: WriteConcern = this.writeConcern, m: SerializableBSONObject[T]) {
     log.trace("Batch Inserting: %s to %s.%s with WriteConcern: %s", docs, db, collection, concern)
-    docs.foreach(x => {
+    val checked = docs.map(x => {
       m.checkObject(x)
       m.checkID(x)
     })
-    send(InsertMessage(db + "." + collection, docs: _*), callback)
+    send(InsertMessage(db + "." + collection, checked: _*), callback)
   }