lib/mongodb: set automatic last id when inserting or saving a JsonObject

Signed-off-by: Alexandre Terrasa <alexandre@moz-code.org>

diff --git a/lib/mongodb/mongodb.nit b/lib/mongodb/mongodb.nit
index 1118741..db83acd 100644 (file)
--- a/lib/mongodb/mongodb.nit
+++ b/lib/mongodb/mongodb.nit
@@ -40,8 +40,6 @@
 # var res = col.find(query)
 # assert res["bar"] == "bar"
 # ~~~
-#
-# TODO Get last Object_ID
 module mongodb
 
 import json
@@ -175,6 +173,34 @@ class MongoError
        redef fun to_s do return "{message} (code: {code})"
 end
 
+# MongoDB Object ID representation.
+#
+# For ObjectIDs, MongoDB uses the `ObjectId("hash")` notation.
+# This notation is replicated by the `to_s` service.
+#
+# Since the MongoDB notation is not JSON complient, the mongoc wrapper uses
+# a JSON based notation like `{"$oid": "hash"}`.
+# This is the notation returned by the `to_json` service.
+private class MongoObjectId
+
+       var native: BSONObjectId
+
+       # The unique ID as an MongoDB Object ID string.
+       fun id: String do return native.id
+
+       # Internal JSON representation of this Object ID.
+       #
+       # Something like `{"$oid": "5578e5dcf344225cc2378051"}`.
+       fun to_json: JsonObject do
+               var obj = new JsonObject
+               obj["$oid"] = id
+               return obj
+       end
+
+       # Formatted as `ObjectId("5578e5dcf344225cc2378051")`
+       redef fun to_s do return "ObjectId({id})"
+end
+
 # The MongoClient is used to connect to the mongo server and send queries.
 #
 # Usage:
@@ -277,6 +303,22 @@ class MongoClient
                if last_error == null then return null
                return new MongoError(last_error)
        end
+
+       # Last auto generated id.
+       private fun last_id: nullable MongoObjectId do
+               var last_id = sys.last_mongoc_id
+               if last_id == null then return null
+               return new MongoObjectId(last_id)
+       end
+
+       # Set the last generated id or `null` to unset once used.
+       private fun last_id=(id: nullable MongoObjectId) do
+               if id == null then
+                       sys.last_mongoc_id = null
+               else
+                       sys.last_mongoc_id = id.native
+               end
+       end
 end
 
 # A MongoDb database.
@@ -400,6 +442,15 @@ class MongoCollection
                        name.to_cstring)
        end
 
+       # Set the autogenerated last id if the `doc` does not contain one already.
+       private fun set_id(doc: JsonObject) do
+               var last_id = database.client.last_id
+               if last_id != null then
+                       doc["_id"] = last_id.to_json
+                       database.client.last_id = null
+               end
+       end
+
        # Inserts a new document in the collection.
        #
        # If no _id element is found in document, then a new one be generated locally
@@ -415,10 +466,13 @@ class MongoCollection
        # doc["bar"] = "bar"
        # doc["baz"] = new JsonArray
        # assert col.insert(doc)
+       # assert doc.has_key("_id")
        # ~~~
        fun insert(doc: JsonObject): Bool do
                assert is_alive
-               return native.insert(doc.to_bson.native)
+               var res = native.insert(doc.to_bson.native)
+               if res then set_id(doc)
+               return res
        end
 
        # Inserts multiple documents in the collection.
@@ -441,15 +495,24 @@ class MongoCollection
        # ~~~
        # var client = new MongoClient("mongodb://localhost:27017/")
        # var col = client.database("test").collection("test")
+       #
        # var doc = new JsonObject
        # doc["foo"] = 10
        # doc["bar"] = "bar"
        # doc["baz"] = new JsonArray
+       #
        # assert col.save(doc) # will be inserted
+       # assert doc.has_key("_id")
+       #
+       # var id = doc["_id"]
+       # assert col.save(doc) # will be updated
+       # assert doc["_id"] == id
        # ~~~
        fun save(doc: JsonObject): Bool do
                assert is_alive
-               return native.save(doc.to_bson.native)
+               var res = native.save(doc.to_bson.native)
+               if res then set_id(doc)
+               return res
        end
 
        # Removes the first document that matches `selector`.