PYTHON-3333 Fix bug where non-cursor read operations fail in a transaction with directConnection=True on primary (#991)

Author: ShaneHarvey

pymongo/message.py

@@ -367,6 +367,8 @@ def as_command(self, sock_info, apply_timeout=False):
 
     def get_message(self, read_preference, sock_info, use_cmd=False):
         """Get a query message, possibly setting the secondaryOk bit."""
+        # Use the read_preference decided by _socket_from_server.
+        self.read_preference = read_preference
         if read_preference.mode:
             # Set the secondaryOk bit.
             flags = self.flags | 4

pymongo/mongo_client.py

@@ -1251,7 +1251,7 @@ def _socket_from_server(self, read_preference, server, session):
 
         with self._get_socket(server, session) as sock_info:
             if single:
-                if sock_info.is_repl:
+                if sock_info.is_repl and not (session and session.in_transaction):
                     # Use primary preferred to ensure any repl set member
                     # can handle the request.
                     read_preference = ReadPreference.PRIMARY_PREFERRED

test/test_transactions.py

@@ -35,6 +35,8 @@
 from gridfs import GridFS, GridFSBucket
 from pymongo import WriteConcern, client_session
 from pymongo.client_session import TransactionOptions
+from pymongo.command_cursor import CommandCursor
+from pymongo.cursor import Cursor
 from pymongo.errors import (
     CollectionInvalid,
     ConfigurationError,
@@ -351,6 +353,42 @@ def test_transaction_starts_with_batched_write(self):
             self.assertEqual(txn_number, event.command["txnNumber"])
         self.assertEqual(48, coll.count_documents({}))
 
+    @client_context.require_transactions
+    def test_transaction_direct_connection(self):
+        client = single_client()
+        self.addCleanup(client.close)
+        coll = client.pymongo_test.test
+
+        # Make sure the collection exists.
+        coll.insert_one({})
+        self.assertEqual(client.topology_description.topology_type_name, "Single")
+        ops = [
+            (coll.bulk_write, [[InsertOne({})]]),
+            (coll.insert_one, [{}]),
+            (coll.insert_many, [[{}, {}]]),
+            (coll.replace_one, [{}, {}]),
+            (coll.update_one, [{}, {"$set": {"a": 1}}]),
+            (coll.update_many, [{}, {"$set": {"a": 1}}]),
+            (coll.delete_one, [{}]),
+            (coll.delete_many, [{}]),
+            (coll.find_one_and_replace, [{}, {}]),
+            (coll.find_one_and_update, [{}, {"$set": {"a": 1}}]),
+            (coll.find_one_and_delete, [{}, {}]),
+            (coll.find_one, [{}]),
+            (coll.count_documents, [{}]),
+            (coll.distinct, ["foo"]),
+            (coll.aggregate, [[]]),
+            (coll.find, [{}]),
+            (coll.aggregate_raw_batches, [[]]),
+            (coll.find_raw_batches, [{}]),
+            (coll.database.command, ["find", coll.name]),
+        ]
+        for f, args in ops:
+            with client.start_session() as s, s.start_transaction():
+                res = f(*args, session=s)
+                if isinstance(res, (CommandCursor, Cursor)):
+                    list(res)
+
 
 class PatchSessionTimeout(object):
     """Patches the client_session's with_transaction timeout for testing."""