fix doc.get_<field>_display + unit test inspired by #1279

Fix BaseQuerySet#sum and BaseQuerySet#average for fields that specify a db_field

README.rst

@@ -52,10 +52,14 @@ Some simple examples of what MongoEngine code looks like:
 
 .. code :: python
 
+    from mongoengine import *
+    connect('mydb')
+
     class BlogPost(Document):
         title = StringField(required=True, max_length=200)
         posted = DateTimeField(default=datetime.datetime.now)
         tags = ListField(StringField(max_length=50))
+        meta = {'allow_inheritance': True}
 
     class TextPost(BlogPost):
         content = StringField(required=True)

docs/changelog.rst

@@ -4,7 +4,9 @@ Changelog
 
 Changes in 0.10.8
 =================
-- Fill this in as PRs for v0.10.8 are merged
+- Added ability to specify an authentication mechanism (e.g. X.509) #1333
+- Added support for falsey primary keys (e.g. doc.pk = 0) #1354
+- Fixed BaseQuerySet#sum/average for fields w/ explicit db_field #1417
 
 Changes in 0.10.7
 =================

mongoengine/base/document.py

@@ -121,7 +121,7 @@ class BaseDocument(object):
                 else:
                     self._data[key] = value
 
-        # Set any get_fieldname_display methods
+        # Set any get_<field>_display methods
         self.__set_field_display()
 
         if self._dynamic:
@@ -1005,19 +1005,18 @@ class BaseDocument(object):
         return '.'.join(parts)
 
     def __set_field_display(self):
-        """Dynamically set the display value for a field with choices"""
-        for attr_name, field in self._fields.items():
-            if field.choices:
-                if self._dynamic:
-                    obj = self
-                else:
-                    obj = type(self)
-                setattr(obj,
-                        'get_%s_display' % attr_name,
-                        partial(self.__get_field_display, field=field))
+        """For each field that specifies choices, create a
+        get_<field>_display method.
+        """
+        fields_with_choices = [(n, f) for n, f in self._fields.items()
+                               if f.choices]
+        for attr_name, field in fields_with_choices:
+            setattr(self,
+                    'get_%s_display' % attr_name,
+                    partial(self.__get_field_display, field=field))
 
     def __get_field_display(self, field):
-        """Returns the display value for a choice field"""
+        """Return the display value for a choice field"""
         value = getattr(self, field.name)
         if field.choices and isinstance(field.choices[0], (list, tuple)):
             return dict(field.choices).get(value, value)

mongoengine/connection.py

@@ -6,6 +6,7 @@ __all__ = ['ConnectionError', 'connect', 'register_connection',
 
 
 DEFAULT_CONNECTION_NAME = 'default'
+
 if IS_PYMONGO_3:
     READ_PREFERENCE = ReadPreference.PRIMARY
 else:
@@ -25,6 +26,7 @@ _dbs = {}
 def register_connection(alias, name=None, host=None, port=None,
                         read_preference=READ_PREFERENCE,
                         username=None, password=None, authentication_source=None,
+                        authentication_mechanism=None,
                         **kwargs):
     """Add a connection.
 
@@ -38,6 +40,9 @@ def register_connection(alias, name=None, host=None, port=None,
     :param username: username to authenticate with
     :param password: password to authenticate with
     :param authentication_source: database to authenticate against
+    :param authentication_mechanism: database authentication mechanisms.
+        By default, use SCRAM-SHA-1 with MongoDB 3.0 and later,
+        MONGODB-CR (MongoDB Challenge Response protocol) for older servers.
     :param is_mock: explicitly use mongomock for this connection
         (can also be done by using `mongomock://` as db host prefix)
     :param kwargs: allow ad-hoc parameters to be passed into the pymongo driver
@@ -53,9 +58,11 @@ def register_connection(alias, name=None, host=None, port=None,
         'read_preference': read_preference,
         'username': username,
         'password': password,
-        'authentication_source': authentication_source
+        'authentication_source': authentication_source,
+        'authentication_mechanism': authentication_mechanism
     }
 
+    # Handle uri style connections
     conn_host = conn_settings['host']
     # host can be a list or a string, so if string, force to a list
     if isinstance(conn_host, str_types):
@@ -82,6 +89,8 @@ def register_connection(alias, name=None, host=None, port=None,
                 conn_settings['replicaSet'] = True
             if 'authsource' in uri_options:
                 conn_settings['authentication_source'] = uri_options['authsource']
+            if 'authmechanism' in uri_options:
+                conn_settings['authentication_mechanism'] = uri_options['authmechanism']
         else:
             resolved_hosts.append(entity)
     conn_settings['host'] = resolved_hosts
@@ -123,6 +132,7 @@ def get_connection(alias=DEFAULT_CONNECTION_NAME, reconnect=False):
         conn_settings.pop('username', None)
         conn_settings.pop('password', None)
         conn_settings.pop('authentication_source', None)
+        conn_settings.pop('authentication_mechanism', None)
 
         is_mock = conn_settings.pop('is_mock', None)
         if is_mock:
@@ -157,6 +167,7 @@ def get_connection(alias=DEFAULT_CONNECTION_NAME, reconnect=False):
                 connection_settings.pop('username', None)
                 connection_settings.pop('password', None)
                 connection_settings.pop('authentication_source', None)
+                connection_settings.pop('authentication_mechanism', None)
                 if conn_settings == connection_settings and _connections.get(db_alias, None):
                     connection = _connections[db_alias]
                     break
@@ -176,11 +187,13 @@ def get_db(alias=DEFAULT_CONNECTION_NAME, reconnect=False):
         conn = get_connection(alias)
         conn_settings = _connection_settings[alias]
         db = conn[conn_settings['name']]
+        auth_kwargs = {'source': conn_settings['authentication_source']}
+        if conn_settings['authentication_mechanism'] is not None:
+            auth_kwargs['mechanism'] = conn_settings['authentication_mechanism']
         # Authenticate if necessary
-        if conn_settings['username'] and conn_settings['password']:
-            db.authenticate(conn_settings['username'],
-                            conn_settings['password'],
-                            source=conn_settings['authentication_source'])
+        if conn_settings['username'] and (conn_settings['password'] or
+                                          conn_settings['authentication_mechanism'] == 'MONGODB-X509'):
+            db.authenticate(conn_settings['username'], conn_settings['password'], **auth_kwargs)
         _dbs[alias] = db
     return _dbs[alias]
 

mongoengine/document.py

@@ -472,7 +472,7 @@ class Document(BaseDocument):
         Raises :class:`OperationError` if called on an object that has not yet
         been saved.
         """
-        if not self.pk:
+        if self.pk is None:
             if kwargs.get('upsert', False):
                 query = self.to_mongo()
                 if "_cls" in query:
@@ -604,7 +604,7 @@ class Document(BaseDocument):
         elif "max_depth" in kwargs:
             max_depth = kwargs["max_depth"]
 
-        if not self.pk:
+        if self.pk is None:
             raise self.DoesNotExist("Document does not exist")
         obj = self._qs.read_preference(ReadPreference.PRIMARY).filter(
             **self._object_key).only(*fields).limit(
@@ -655,7 +655,7 @@ class Document(BaseDocument):
     def to_dbref(self):
         """Returns an instance of :class:`~bson.dbref.DBRef` useful in
         `__raw__` queries."""
-        if not self.pk:
+        if self.pk is None:
             msg = "Only saved documents can have a valid dbref"
             raise OperationError(msg)
         return DBRef(self.__class__._get_collection_name(), self.pk)

mongoengine/queryset/base.py

@@ -1271,9 +1271,10 @@ class BaseQuerySet(object):
         :param field: the field to sum over; use dot notation to refer to
             embedded document fields
         """
+        db_field = self._fields_to_dbfields([field]).pop()
         pipeline = [
             {'$match': self._query},
-            {'$group': {'_id': 'sum', 'total': {'$sum': '$' + field}}}
+            {'$group': {'_id': 'sum', 'total': {'$sum': '$' + db_field}}}
         ]
 
         # if we're performing a sum over a list field, we sum up all the
@@ -1300,9 +1301,10 @@ class BaseQuerySet(object):
         :param field: the field to average over; use dot notation to refer to
             embedded document fields
         """
+        db_field = self._fields_to_dbfields([field]).pop()
         pipeline = [
             {'$match': self._query},
-            {'$group': {'_id': 'avg', 'total': {'$avg': '$' + field}}}
+            {'$group': {'_id': 'avg', 'total': {'$avg': '$' + db_field}}}
         ]
 
         # if we're performing an average over a list field, we average out

tests/document/instance.py

@@ -3202,5 +3202,20 @@ class InstanceTest(unittest.TestCase):
             self.assertEqual(b._instance, a)
         self.assertEqual(idx, 2)
 
+    def test_falsey_pk(self):
+        """Ensure that we can create and update a document with Falsey PK.
+        """
+        class Person(Document):
+            age = IntField(primary_key=True)
+            height = FloatField()
+
+        person = Person()
+        person.age = 0
+        person.height = 1.89
+        person.save()
+
+        person.update(set__height=2.0)
+
+
 if __name__ == '__main__':
     unittest.main()

tests/fields/fields.py

@@ -3001,28 +3001,32 @@ class FieldTest(unittest.TestCase):
                 ('S', 'Small'), ('M', 'Medium'), ('L', 'Large'),
                 ('XL', 'Extra Large'), ('XXL', 'Extra Extra Large')))
             style = StringField(max_length=3, choices=(
-                ('S', 'Small'), ('B', 'Baggy'), ('W', 'wide')), default='S')
+                ('S', 'Small'), ('B', 'Baggy'), ('W', 'Wide')), default='W')
 
         Shirt.drop_collection()
 
-        shirt = Shirt()
+        shirt1 = Shirt()
+        shirt2 = Shirt()
 
-        self.assertEqual(shirt.get_size_display(), None)
-        self.assertEqual(shirt.get_style_display(), 'Small')
+        # Make sure get_<field>_display returns the default value (or None)
+        self.assertEqual(shirt1.get_size_display(), None)
+        self.assertEqual(shirt1.get_style_display(), 'Wide')
 
-        shirt.size = "XXL"
-        shirt.style = "B"
-        self.assertEqual(shirt.get_size_display(), 'Extra Extra Large')
-        self.assertEqual(shirt.get_style_display(), 'Baggy')
+        shirt1.size = 'XXL'
+        shirt1.style = 'B'
+        shirt2.size = 'M'
+        shirt2.style = 'S'
+        self.assertEqual(shirt1.get_size_display(), 'Extra Extra Large')
+        self.assertEqual(shirt1.get_style_display(), 'Baggy')
+        self.assertEqual(shirt2.get_size_display(), 'Medium')
+        self.assertEqual(shirt2.get_style_display(), 'Small')
 
         # Set as Z - an invalid choice
-        shirt.size = "Z"
-        shirt.style = "Z"
-        self.assertEqual(shirt.get_size_display(), 'Z')
-        self.assertEqual(shirt.get_style_display(), 'Z')
-        self.assertRaises(ValidationError, shirt.validate)
-
-        Shirt.drop_collection()
+        shirt1.size = 'Z'
+        shirt1.style = 'Z'
+        self.assertEqual(shirt1.get_size_display(), 'Z')
+        self.assertEqual(shirt1.get_style_display(), 'Z')
+        self.assertRaises(ValidationError, shirt1.validate)
 
     def test_simple_choices_validation(self):
         """Ensure that value is in a container of allowed values.

tests/queryset/queryset.py

@@ -2838,6 +2838,34 @@ class QuerySetTest(unittest.TestCase):
             sum([a for a in ages if a >= 50])
         )
 
+    def test_sum_over_db_field(self):
+        """Ensure that a field mapped to a db field with a different name
+        can be summed over correctly.
+        """
+        class UserVisit(Document):
+            num_visits = IntField(db_field='visits')
+
+        UserVisit.drop_collection()
+
+        UserVisit.objects.create(num_visits=10)
+        UserVisit.objects.create(num_visits=5)
+
+        self.assertEqual(UserVisit.objects.sum('num_visits'), 15)
+
+    def test_average_over_db_field(self):
+        """Ensure that a field mapped to a db field with a different name
+        can have its average computed correctly.
+        """
+        class UserVisit(Document):
+            num_visits = IntField(db_field='visits')
+
+        UserVisit.drop_collection()
+
+        UserVisit.objects.create(num_visits=20)
+        UserVisit.objects.create(num_visits=10)
+
+        self.assertEqual(UserVisit.objects.average('num_visits'), 15)
+
     def test_embedded_average(self):
         class Pay(EmbeddedDocument):
             value = DecimalField()