Added support for user-defined primary keys (_ids)

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