X-Git-Url: http://nitlanguage.org

diff --git a/lib/mongodb/mongodb.nit b/lib/mongodb/mongodb.nit
index 1118741..1e9b9bd 100644
--- 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:
@@ -230,6 +256,7 @@ class MongoClient
 		assert is_alive
 		var res = new Array[String]
 		var nas = native.database_names
+		if nas == null then return res
 		var i = 0
 		var name = nas[i]
 		while not name.address_is_null do
@@ -277,6 +304,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.
@@ -318,6 +361,7 @@ class MongoDb
 		assert is_alive
 		var res = new Array[String]
 		var nas = native.collection_names
+		if nas == null then return res
 		var i = 0
 		var name = nas[i]
 		while not name.address_is_null do
@@ -400,6 +444,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 +468,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 +497,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`.
@@ -536,6 +601,7 @@ class MongoCollection
 	fun find(query: JsonObject): nullable JsonObject do
 		assert is_alive
 		var c = native.find(query.to_bson.native)
+		assert is_alive # FIXME used to avoid segfault (so `self` isn't garbage collected to soon)
 		if c == null then return null
 		var cursor = new MongoCursor(c)
 		if cursor.is_ok then