Use slaves for group, distinct, count PYTHON-256

behackett · behackett · commit e5ddb13f213b · 2011-11-02T13:12:53.000-07:00
Or secondaries with ReplicaSetConnection.
diff --git a/pymongo/collection.py b/pymongo/collection.py
@@ -887,6 +887,13 @@ def group(self, key, condition, initial, reduce, finalize=None,
             containing a JavaScript function to be applied to each
             document, returning the key to group by.
 
+        With :class:`~pymongo.replica_set_connection.ReplicaSetConnection`
+        or :class:`~pymongo.master_slave_connection.MasterSlaveConnection`,
+        if the `read_preference` attribute of this instance is not set to
+        :attr:`pymongo.ReadPreference.PRIMARY` or the (deprecated)
+        `slave_okay` attribute of this instance is set to `True` the group
+        command will be sent to a secondary or slave.
+
         :Parameters:
           - `key`: fields to group by (see above description)
           - `condition`: specification of rows to be
@@ -921,7 +928,10 @@ def group(self, key, condition, initial, reduce, finalize=None,
         if finalize is not None:
             group["finalize"] = Code(finalize)
 
-        return self.__database.command("group", group)["retval"]
+        use_master = not self.slave_okay and not self.read_preference
+
+        return self.__database.command("group", group,
+                                       _use_master=use_master)["retval"]
 
     def rename(self, new_name, **kwargs):
         """Rename this collection.
diff --git a/pymongo/cursor.py b/pymongo/cursor.py
@@ -449,6 +449,12 @@ def count(self, with_limit_and_skip=False):
         `with_limit_and_skip` to ``True`` if that is the desired behavior.
         Raises :class:`~pymongo.errors.OperationFailure` on a database error.
 
+        With :class:`~pymongo.replica_set_connection.ReplicaSetConnection`
+        or :class:`~pymongo.master_slave_connection.MasterSlaveConnection`,
+        if `read_preference` is not :attr:`pymongo.ReadPreference.PRIMARY` or
+        (deprecated) `slave_okay` is `True` the count command will be sent to
+        a secondary or slave.
+
         :Parameters:
           - `with_limit_and_skip` (optional): take any :meth:`limit` or
             :meth:`skip` that has been applied to this cursor into account when
@@ -464,6 +470,9 @@ def count(self, with_limit_and_skip=False):
         """
         command = {"query": self.__spec, "fields": self.__fields}
 
+        use_master = not self.__slave_okay and not self.__read_preference
+        command['_use_master'] = use_master
+
         if with_limit_and_skip:
             if self.__limit:
                 command["limit"] = self.__limit
@@ -484,6 +493,12 @@ def distinct(self, key):
         Raises :class:`TypeError` if `key` is not an instance of
         :class:`basestring`.
 
+        With :class:`~pymongo.replica_set_connection.ReplicaSetConnection`
+        or :class:`~pymongo.master_slave_connection.MasterSlaveConnection`,
+        if `read_preference` is not :attr:`pymongo.ReadPreference.PRIMARY` or
+        (deprecated) `slave_okay` is `True` the distinct command will be sent
+        to a secondary or slave.
+
         :Parameters:
           - `key`: name of key for which we want to get the distinct values
 
@@ -500,6 +515,9 @@ def distinct(self, key):
         if self.__spec:
             options["query"] = self.__spec
 
+        use_master = not self.__slave_okay and not self.__read_preference
+        options['_use_master'] = use_master
+
         return self.__collection.database.command("distinct",
                                                   self.__collection.name,
                                                   **options)["values"]
diff --git a/pymongo/database.py b/pymongo/database.py
@@ -326,14 +326,16 @@ def command(self, command, value=1,
         if isinstance(command, basestring):
             command = SON([(command, value)])
 
+        use_master = kwargs.pop('_use_master', True)
+
         fields = kwargs.get('fields')
         if fields is not None and not isinstance(fields, dict):
                 kwargs['fields'] = helpers._fields_list_to_dict(fields)
 
         command.update(kwargs)
 
         result = self["$cmd"].find_one(command,
-                                       _must_use_master=True,
+                                       _must_use_master=use_master,
                                        _is_command=True)
 
         if check:
diff --git a/pymongo/master_slave_connection.py b/pymongo/master_slave_connection.py
@@ -66,8 +66,7 @@ def __init__(self, master, slaves=[], tz_aware=False):
                                 slave)
 
         super(MasterSlaveConnection,
-              self).__init__(slave_okay=True,
-                             read_preference=ReadPreference.SECONDARY,
+              self).__init__(read_preference=ReadPreference.SECONDARY,
                              safe=master.safe,
                              **(master.get_lasterror_options()))
 
diff --git a/test/test_master_slave_connection.py b/test/test_master_slave_connection.py
@@ -318,36 +318,44 @@ def cursor_count():
 
     def test_base_object(self):
         c = self.connection
-        self.assertTrue(c.slave_okay)
+        self.assertFalse(c.slave_okay)
+        self.assertTrue(bool(c.read_preference))
         self.assertFalse(c.safe)
         self.assertEqual({}, c.get_lasterror_options())
         db = c.test
-        self.assertTrue(db.slave_okay)
+        self.assertFalse(db.slave_okay)
+        self.assertTrue(bool(c.read_preference))
         self.assertFalse(db.safe)
         self.assertEqual({}, db.get_lasterror_options())
         coll = db.test
-        self.assertTrue(coll.slave_okay)
+        self.assertFalse(coll.slave_okay)
+        self.assertTrue(bool(c.read_preference))
         self.assertFalse(coll.safe)
         self.assertEqual({}, coll.get_lasterror_options())
         cursor = coll.find()
-        self.assertTrue(cursor._Cursor__slave_okay)
+        self.assertFalse(cursor._Cursor__slave_okay)
+        self.assertTrue(bool(cursor._Cursor__read_preference))
 
         c.safe = True
         c.set_lasterror_options(w=2, wtimeout=100)
-        self.assertTrue(c.slave_okay)
+        self.assertFalse(c.slave_okay)
+        self.assertTrue(bool(c.read_preference))
         self.assertTrue(c.safe)
         self.assertEqual({'w': 2, 'wtimeout': 100}, c.get_lasterror_options())
         db = c.test
-        self.assertTrue(db.slave_okay)
+        self.assertFalse(db.slave_okay)
+        self.assertTrue(bool(c.read_preference))
         self.assertTrue(db.safe)
         self.assertEqual({'w': 2, 'wtimeout': 100}, db.get_lasterror_options())
         coll = db.test
-        self.assertTrue(coll.slave_okay)
+        self.assertFalse(coll.slave_okay)
+        self.assertTrue(bool(c.read_preference))
         self.assertTrue(coll.safe)
         self.assertEqual({'w': 2, 'wtimeout': 100},
                          coll.get_lasterror_options())
         cursor = coll.find()
-        self.assertTrue(cursor._Cursor__slave_okay)
+        self.assertFalse(cursor._Cursor__slave_okay)
+        self.assertTrue(bool(cursor._Cursor__read_preference))
 
         coll.insert({'foo': 'bar'})
         self.assertEquals(1, coll.find({'foo': 'bar'}).count())
@@ -358,7 +366,8 @@ def test_base_object(self):
         # Set self.connection back to defaults
         c.safe = False
         c.unset_lasterror_options()
-        self.assertTrue(self.connection.slave_okay)
+        self.assertFalse(self.connection.slave_okay)
+        self.assertTrue(bool(self.connection.read_preference))
         self.assertFalse(self.connection.safe)
         self.assertEqual({}, self.connection.get_lasterror_options())
 
