Test changing transaction readPreference

Add original readPreference to getMores to simplify code check.

Author: ShaneHarvey

pymongo/bulk.py

@@ -37,6 +37,7 @@
                              _do_batched_write_command,
                              _randint,
                              _BulkWriteContext)
+from pymongo.read_preferences import ReadPreference
 from pymongo.write_concern import WriteConcern
 
 
@@ -266,7 +267,7 @@ def _execute_command(self, generator, write_concern, session,
             results = []
             while run.idx_offset < len(run.ops):
                 if session:
-                    session._apply_to(cmd, retryable)
+                    session._apply_to(cmd, retryable, ReadPreference.PRIMARY)
                 sock_info.send_cluster_time(cmd, session, client)
                 check_keys = run.op_type == _INSERT
                 ops = islice(run.ops, run.idx_offset, None)

pymongo/client_session.py

@@ -117,6 +117,7 @@ def __init__(self, client, server_session, options, authset):
         self._authset = authset
         self._cluster_time = None
         self._operation_time = None
+        self._current_txn_read_pref = None
         if self.options.auto_start_transaction:
             # TODO: Get transaction options from self.options.
             self._current_transaction_opts = TransactionOptions()
@@ -234,6 +235,7 @@ def _finish_transaction(self, command_name):
                 stmtId=self._server_session.statement_id,
                 session=self,
                 write_concern=self._current_transaction_opts.write_concern,
+                read_preference=self._current_txn_read_pref,
                 parse_write_concern_error=True)
         finally:
             self._server_session.reset_transaction()
@@ -293,7 +295,7 @@ def in_transaction(self):
         """True if this session has an active multi-statement transaction."""
         return self._current_transaction_opts is not None
 
-    def _apply_to(self, command, is_retryable):
+    def _apply_to(self, command, is_retryable, read_preference):
         self._check_ended()
 
         if self.options.auto_start_transaction and not self.in_transaction:
@@ -310,9 +312,12 @@ def _apply_to(self, command, is_retryable):
         if self._current_transaction_opts:
             if self._server_session.statement_id == 0:
                 # First statement begins a new transaction.
+                self._current_txn_read_pref = read_preference
                 self._server_session._transaction_id += 1
                 command['readConcern'] = {'level': 'snapshot'}
                 command['autocommit'] = False
+            elif read_preference != self._current_txn_read_pref:
+                raise InvalidOperation('Transaction readPreference changed')
 
             command['txnNumber'] = self._server_session.transaction_id
 

pymongo/collection.py

@@ -1930,7 +1930,8 @@ def list_indexes(self, session=None):
         .. versionadded:: 3.0
         """
         codec_options = CodecOptions(SON)
-        coll = self.with_options(codec_options)
+        coll = self.with_options(codec_options=codec_options,
+                                 read_preference=ReadPreference.PRIMARY)
         with self._socket_for_primary_reads() as (sock_info, slave_ok):
             cmd = SON([("listIndexes", self.__name), ("cursor", {})])
             if sock_info.max_wire_version > 2:

pymongo/command_cursor.py

@@ -229,6 +229,7 @@ def _refresh(self):
                                     self.__batch_size,
                                     self.__id,
                                     self.__collection.codec_options,
+                                    self.__collection.read_preference,
                                     self.__session,
                                     self.__collection.database.client,
                                     self.__max_await_time_ms))

pymongo/cursor.py

@@ -1100,6 +1100,7 @@ def _refresh(self):
                                         limit,
                                         self.__id,
                                         self.__codec_options,
+                                        self.__read_preference,
                                         self.__session,
                                         self.__collection.database.client,
                                         self.__max_await_time_ms)

pymongo/database.py

@@ -534,7 +534,8 @@ def command(self, command, value=1, check=True,
     def _list_collections(self, sock_info, slave_okay, session=None, **kwargs):
         """Internal listCollections helper."""
 
-        coll = self["$cmd"]
+        coll = self.get_collection(
+            "$cmd", read_preference=ReadPreference.PRIMARY)
         if sock_info.max_wire_version > 2:
             cmd = SON([("listCollections", 1),
                        ("cursor", {})])

pymongo/message.py

@@ -283,7 +283,7 @@ def as_command(self, sock_info):
             cmd = SON([('explain', cmd)])
         session = self.session
         if session:
-            session._apply_to(cmd, False)
+            session._apply_to(cmd, False, self.read_preference)
             # Explain does not support readConcern.
             if (not explain and session.options.causal_consistency
                     and session.operation_time is not None
@@ -333,17 +333,19 @@ class _GetMore(object):
     """A getmore operation."""
 
     __slots__ = ('db', 'coll', 'ntoreturn', 'cursor_id', 'max_await_time_ms',
-                 'codec_options', 'session', 'client', '__as_command')
+                 'codec_options', 'read_preference', 'session', 'client',
+                 '__as_command')
 
     name = 'getMore'
 
-    def __init__(self, db, coll, ntoreturn, cursor_id, codec_options, session,
-                 client, max_await_time_ms=None):
+    def __init__(self, db, coll, ntoreturn, cursor_id, codec_options,
+                 read_preference, session, client, max_await_time_ms=None):
         self.db = db
         self.coll = coll
         self.ntoreturn = ntoreturn
         self.cursor_id = cursor_id
         self.codec_options = codec_options
+        self.read_preference = read_preference
         self.session = session
         self.client = client
         self.max_await_time_ms = max_await_time_ms
@@ -364,7 +366,7 @@ def as_command(self, sock_info):
                                     self.max_await_time_ms)
 
         if self.session:
-            self.session._apply_to(cmd, False)
+            self.session._apply_to(cmd, False, self.read_preference)
         sock_info.send_cluster_time(cmd, self.session, self.client)
         self.__as_command = cmd, self.db
         return self.__as_command

pymongo/mongo_client.py

@@ -1475,15 +1475,15 @@ def list_databases(self, session=None, **kwargs):
         """
         cmd = SON([("listDatabases", 1)])
         cmd.update(kwargs)
-        res = self._database_default_options(
-            "admin").command(cmd, session=session)
+        admin = self._database_default_options("admin")
+        res = admin.command(cmd, session=session)
         # listDatabases doesn't return a cursor (yet). Fake one.
         cursor = {
             "id": 0,
             "firstBatch": res["databases"],
             "ns": "admin.$cmd",
         }
-        return CommandCursor(self.admin["$cmd"], cursor, None)
+        return CommandCursor(admin["$cmd"], cursor, None)
 
     def list_database_names(self, session=None):
         """Get a list of the names of all databases on the connected server.

pymongo/pool.py

@@ -502,7 +502,7 @@ def command(self, dbname, spec, slave_ok=False,
             # Ensure command name remains in first place.
             spec = SON(spec)
         if session:
-            session._apply_to(spec, retryable_write)
+            session._apply_to(spec, retryable_write, read_preference)
         self.send_cluster_time(spec, session, client)
         listeners = self.listeners if publish_events else None
         try:

test/test_client.py

@@ -1130,6 +1130,7 @@ def test_stale_getmore(self):
             client._send_message_with_response(
                 operation=message._GetMore('pymongo_test', 'collection',
                                            101, 1234, client.codec_options,
+                                           ReadPreference.PRIMARY,
                                            None, client),
                 address=('not-a-member', 27017))