Merge pull request #50 from atlarge-research/bug/bson-objectid

Use BSON ObjectId for model identifiers

1 files changed, 8 insertions, 3 deletions

diff --git a/api/opendc/models/model.py b/api/opendc/models/model.py
index bcb833ae..f9dfc9ad 100644
--- a/api/opendc/models/model.py
+++ b/api/opendc/models/model.py
@@ -1,4 +1,4 @@
-from uuid import uuid4
+from bson.objectid import ObjectId
 
 from opendc.util.database import DB
 from opendc.util.exceptions import ClientError
@@ -13,6 +13,11 @@ class Model:
     @classmethod
     def from_id(cls, _id):
         """Fetches the document with given ID from the collection."""
+        if isinstance(_id, str) and len(_id) == 24:
+            _id = ObjectId(_id)
+        elif not isinstance(_id, ObjectId):
+            return cls(None)
+
         return cls(DB.fetch_one({'_id': _id}, cls.collection_name))
 
     @classmethod
@@ -25,7 +30,7 @@ class Model:
 
     def get_id(self):
         """Returns the ID of the enclosed object."""
-        return str(self.obj['_id'])
+        return self.obj['_id']
 
     def check_exists(self):
         """Raises an error if the enclosed object does not exist."""
@@ -42,7 +47,7 @@ class Model:
 
     def insert(self):
         """Inserts the enclosed object and generates a UUID for it."""
-        self.obj['_id'] = str(uuid4())
+        self.obj['_id'] = ObjectId()
         DB.insert(self.obj, self.collection_name)
 
     def update(self):