Added support for user-defined primary keys (_ids)

docs/changelog.rst

@@ -2,7 +2,15 @@
 Changelog
 =========
 
-Changes is v0.1.3
+Changes in v0.2
+===============
+- Added Q class for building advanced queries
+- Added QuerySet methods for atomic updates to documents
+- Fields may now specify ``unique=True`` to enforce uniqueness across a collection
+- Added option for default document ordering
+- Fixed bug in index definitions
+
+Changes in v0.1.3
 =================
 - Added Django authentication backend
 - Added Document.meta support for indexes, which are ensured just before 

docs/userguide.rst

@@ -318,8 +318,25 @@ saved::
     >>> page.id
     ObjectId('123456789abcdef000000000')
 
-Alternatively, you may explicitly set the :attr:`id` before you save the
-document, but the id must be a valid PyMongo :class:`ObjectId`.
+Alternatively, you may define one of your own fields to be the document's
+"primary key" by providing ``primary_key=True`` as a keyword argument to a
+field's constructor. Under the hood, MongoEngine will use this field as the
+:attr:`id`; in fact :attr:`id` is actually aliased to your primary key field so
+you may still use :attr:`id` to access the primary key if you want::
+
+    >>> class User(Document):
+    ...     email = StringField(primary_key=True)
+    ...     name = StringField()
+    ...
+    >>> bob = User(email='bob@example.com', name='Bob')
+    >>> bob.save()
+    >>> bob.id == bob.email == 'bob@example.com'
+    True
+
+.. note::
+   If you define your own primary key field, the field implicitly becomes
+   required, so a :class:`ValidationError` will be thrown if you don't provide
+   it.
 
 Querying the database
 =====================

mongoengine/base.py

@@ -13,12 +13,13 @@ class BaseField(object):
     """
     
     def __init__(self, name=None, required=False, default=None, unique=False,
-                 unique_with=None):
-        self.name = name
-        self.required = required
+                 unique_with=None, primary_key=False):
+        self.name = name if not primary_key else '_id'
+        self.required = required or primary_key
         self.default = default
         self.unique = bool(unique or unique_with)
         self.unique_with = unique_with
+        self.primary_key = primary_key
 
     def __get__(self, instance, owner):
         """Descriptor for retrieving a value from a field in a document. Do 
@@ -72,7 +73,7 @@ class ObjectIdField(BaseField):
 
     def to_mongo(self, value):
         if not isinstance(value, pymongo.objectid.ObjectId):
-            return pymongo.objectid.ObjectId(value)
+            return pymongo.objectid.ObjectId(str(value))
         return value
 
     def prepare_query_value(self, value):
@@ -139,6 +140,7 @@ class TopLevelDocumentMetaclass(DocumentMetaclass):
         collection = name.lower()
         
         simple_class = True
+        id_field = None
 
         # Subclassed documents inherit collection from superclass
         for base in bases:
@@ -153,13 +155,16 @@ class TopLevelDocumentMetaclass(DocumentMetaclass):
                     simple_class = False
                 collection = base._meta['collection']
 
+                id_field = id_field or base._meta.get('id_field')
+
         meta = {
             'collection': collection,
             'allow_inheritance': True,
             'max_documents': None,
             'max_size': None,
             'ordering': [], # default ordering applied at runtime
-            'indexes': [] # indexes to be ensured at runtime
+            'indexes': [], # indexes to be ensured at runtime
+            'id_field': id_field,
         }
 
         # Apply document-defined meta options
@@ -172,16 +177,14 @@ class TopLevelDocumentMetaclass(DocumentMetaclass):
                              '"allow_inheritance" to False')
         attrs['_meta'] = meta
 
-        attrs['id'] = ObjectIdField(name='_id')
-
         # Set up collection manager, needs the class to have fields so use
         # DocumentMetaclass before instantiating CollectionManager object
         new_class = super_new(cls, name, bases, attrs)
         new_class.objects = QuerySetManager()
 
-        # Generate a list of indexes needed by uniqueness constraints
         unique_indexes = []
         for field_name, field in new_class._fields.items():
+            # Generate a list of indexes needed by uniqueness constraints
             if field.unique:
                 field.required = True
                 unique_fields = [field_name]
@@ -204,10 +207,25 @@ class TopLevelDocumentMetaclass(DocumentMetaclass):
                     unique_fields += unique_with
 
                 # Add the new index to the list
-                index = [(field, pymongo.ASCENDING) for field in unique_fields]
+                index = [(f, pymongo.ASCENDING) for f in unique_fields]
                 unique_indexes.append(index)
+
+            # Check for custom primary key
+            if field.primary_key:
+                if not new_class._meta['id_field']:
+                    new_class._meta['id_field'] = field_name
+                    # Make 'Document.id' an alias to the real primary key field
+                    new_class.id = field
+                    #new_class._fields['id'] = field
+                else:
+                    raise ValueError('Cannot override primary key field')
+
         new_class._meta['unique_indexes'] = unique_indexes
 
+        if not new_class._meta['id_field']:
+            new_class._meta['id_field'] = 'id'
+            new_class.id = new_class._fields['id'] = ObjectIdField(name='_id')
+
         return new_class
 
 

mongoengine/document.py

@@ -68,19 +68,22 @@ class Document(BaseDocument):
         except pymongo.errors.OperationFailure, err:
             raise OperationError('Tried to save duplicate unique keys (%s)'
                                   % str(err))
-        self.id = self._fields['id'].to_python(object_id)
+        id_field = self._meta['id_field']
+        self[id_field] = self._fields[id_field].to_python(object_id)
 
     def delete(self):
         """Delete the :class:`~mongoengine.Document` from the database. This
         will only take effect if the document has been previously saved.
         """
-        object_id = self._fields['id'].to_mongo(self.id)
-        self.__class__.objects(id=object_id).delete()
+        id_field = self._meta['id_field']
+        object_id = self._fields[id_field].to_mongo(self[id_field])
+        self.__class__.objects(**{id_field: object_id}).delete()
 
     def reload(self):
         """Reloads all attributes from the database.
         """
-        obj = self.__class__.objects(id=self.id).first()
+        id_field = self._meta['id_field']
+        obj = self.__class__.objects(**{id_field: self[id_field]}).first()
         for field in self._fields:
             setattr(self, field, obj[field])
 

mongoengine/queryset.py

@@ -270,10 +270,10 @@ class QuerySet(object):
     def with_id(self, object_id):
         """Retrieve the object matching the id provided.
         """
-        if not isinstance(object_id, pymongo.objectid.ObjectId):
-            object_id = pymongo.objectid.ObjectId(str(object_id))
+        id_field = self._document._meta['id_field']
+        object_id = self._document._fields[id_field].to_mongo(object_id)
 
-        result = self._collection.find_one(object_id)
+        result = self._collection.find_one({'_id': object_id})
         if result is not None:
             result = self._document._from_son(result)
         return result

tests/document.py

@@ -287,6 +287,39 @@ class DocumentTest(unittest.TestCase):
 
         BlogPost.drop_collection()
 
+    def test_custom_id_field(self):
+        """Ensure that documents may be created with custom primary keys.
+        """
+        class User(Document):
+            username = StringField(primary_key=True)
+            name = StringField()
+
+        User.drop_collection()
+
+        self.assertEqual(User._fields['username'].name, '_id')
+        self.assertEqual(User._meta['id_field'], 'username')
+
+        def create_invalid_user():
+            User(name='test').save() # no primary key field
+        self.assertRaises(ValidationError, create_invalid_user)
+
+        def define_invalid_user():
+            class EmailUser(User):
+                email = StringField(primary_key=True)
+        self.assertRaises(ValueError, define_invalid_user)
+        
+        user = User(username='test', name='test user')
+        user.save()
+
+        user_obj = User.objects.first()
+        self.assertEqual(user_obj.id, 'test')
+
+        user_son = User.objects._collection.find_one()
+        self.assertEqual(user_son['_id'], 'test')
+        self.assertTrue('username' not in user_son['_id'])
+        
+        User.drop_collection()
+
     def test_creation(self):
         """Ensure that document may be created using keyword arguments.
         """