add support for pk property in documents and filters

tests/document.py

@@ -355,12 +355,26 @@ class DocumentTest(unittest.TestCase):
 
         user_obj = User.objects.first()
         self.assertEqual(user_obj.id, 'test')
+        self.assertEqual(user_obj.pk, '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()
+        
+        user = User(pk='mongo', name='mongo user')
+        user.save()
+        
+        user_obj = User.objects.first()
+        self.assertEqual(user_obj.id, 'mongo')
+        self.assertEqual(user_obj.pk, 'mongo')
+        
+        user_son = User.objects._collection.find_one()
+        self.assertEqual(user_son['_id'], 'mongo')
+        self.assertTrue('username' not in user_son['_id'])
+        
+        User.drop_collection()
 
     def test_creation(self):
         """Ensure that document may be created using keyword arguments.
@@ -479,6 +493,18 @@ class DocumentTest(unittest.TestCase):
         collection = self.db[self.Person._meta['collection']]
         person_obj = collection.find_one({'name': 'Test User'})
         self.assertEqual(str(person_obj['_id']), '497ce96f395f2f052a494fd4')
+        
+    def test_save_custom_pk(self):
+        """Ensure that a document may be saved with a custom _id using pk alias.
+        """
+        # Create person object and save it to the database
+        person = self.Person(name='Test User', age=30, 
+                             pk='497ce96f395f2f052a494fd4')
+        person.save()
+        # Ensure that the object is in the database with the correct _id
+        collection = self.db[self.Person._meta['collection']]
+        person_obj = collection.find_one({'name': 'Test User'})
+        self.assertEqual(str(person_obj['_id']), '497ce96f395f2f052a494fd4')
 
     def test_save_list(self):
         """Ensure that a list field may be properly saved.

tests/queryset.py

@@ -1094,7 +1094,8 @@ class QuerySetTest(unittest.TestCase):
         BlogPost.drop_collection()
 
         data = {'title': 'Post 1', 'comments': [Comment(content='test')]}
-        BlogPost(**data).save()
+        post = BlogPost(**data)
+        post.save()
 
         self.assertTrue('postTitle' in
                         BlogPost.objects(title=data['title'])._query)
@@ -1102,12 +1103,33 @@ class QuerySetTest(unittest.TestCase):
                          BlogPost.objects(title=data['title'])._query)
         self.assertEqual(len(BlogPost.objects(title=data['title'])), 1)
 
+        self.assertTrue('_id' in BlogPost.objects(pk=post.id)._query)
+        self.assertEqual(len(BlogPost.objects(pk=post.id)), 1)
+
         self.assertTrue('postComments.commentContent' in
                         BlogPost.objects(comments__content='test')._query)
         self.assertEqual(len(BlogPost.objects(comments__content='test')), 1)
 
         BlogPost.drop_collection()
 
+    def test_query_pk_field_name(self):
+        """Ensure that the correct "primary key" field name is used when querying
+        """
+        class BlogPost(Document):
+            title = StringField(primary_key=True, db_field='postTitle')
+
+        BlogPost.drop_collection()
+
+        data = { 'title':'Post 1' }
+        post = BlogPost(**data)
+        post.save()
+
+        self.assertTrue('_id' in BlogPost.objects(pk=data['title'])._query)
+        self.assertTrue('_id' in BlogPost.objects(title=data['title'])._query)
+        self.assertEqual(len(BlogPost.objects(pk=data['title'])), 1)
+
+        BlogPost.drop_collection()
+
     def test_query_value_conversion(self):
         """Ensure that query values are properly converted when necessary.
         """