PYTHON-1332 - Send lsid with all commands

Author: ajdavis

pymongo/change_stream.py

@@ -48,12 +48,14 @@ class ChangeStream(object):
         per batch.
       - `collation` (optional): The :class:`~pymongo.collation.Collation`
         to use for the aggregation.
+      - `session` (optional): a
+        :class:`~pymongo.client_session.ClientSession`.
 
     .. versionadded: 3.6
     """
     def __init__(self, collection, pipeline, full_document,
                  resume_after=None, max_await_time_ms=None, batch_size=None,
-                 collation=None):
+                 collation=None, session=None):
         self._codec_options = collection.codec_options
         self._collection = collection.with_options(
             codec_options=DEFAULT_RAW_BSON_OPTIONS)
@@ -63,6 +65,7 @@ def __init__(self, collection, pipeline, full_document,
         self._max_await_time_ms = max_await_time_ms
         self._batch_size = batch_size
         self._collation = collation
+        self._session = session
         self._cursor = self._create_cursor()
 
     def _full_pipeline(self):
@@ -79,7 +82,7 @@ def _full_pipeline(self):
     def _create_cursor(self):
         """Initialize the cursor or raise a fatal error"""
         return self._collection.aggregate(
-            self._full_pipeline(), batchSize=self._batch_size,
+            self._full_pipeline(), self._session, batchSize=self._batch_size,
             collation=self._collation, maxAwaitTimeMS=self._max_await_time_ms)
 
     def close(self):

pymongo/client_session.py

@@ -169,6 +169,9 @@ def return_server_session(self, server_session, session_timeout_minutes):
         if not server_session.timed_out(session_timeout_minutes):
             self.appendleft(server_session)
 
+    def return_server_session_no_lock(self, server_session):
+        self.appendleft(server_session)
+
     def _clear_stale(self, session_timeout_minutes):
         # Clear stale sessions. The least recently used are on the right.
         while self:

pymongo/collection.py

@@ -1468,19 +1468,18 @@ def parallel_scan(self, num_cursors, session=None, **kwargs):
         cmd.update(kwargs)
 
         with self._socket_for_reads() as (sock_info, slave_ok):
-            # Avoid auto-injecting a session.
-            result = sock_info.command(
-                self.__database.name,
-                cmd,
-                slave_ok,
-                self.read_preference,
-                self.codec_options,
-                read_concern=self.read_concern,
-                session=session)
+            result = self._command(sock_info, cmd, slave_ok,
+                                   read_concern=self.read_concern,
+                                   session=session)
+
+        cursors = []
+        for cursor in result['cursors']:
+            s = self.__database.client._ensure_session(session)
+            cursors.append(CommandCursor(
+                self, cursor['cursor'], sock_info.address,
+                session=s, explicit_session=session is not None))
 
-        return [CommandCursor(self, cursor['cursor'], sock_info.address,
-                              session=session, explicit_session=True)
-                for cursor in result['cursors']]
+        return cursors
 
     def _count(self, cmd, collation=None, session=None):
         """Internal count helper."""
@@ -2041,38 +2040,25 @@ def _aggregate(self, pipeline, cursor_class, first_batch_size, session,
             cmd.update(kwargs)
             # Apply this Collection's read concern if $out is not in the
             # pipeline.
-            if sock_info.max_wire_version >= 4 and 'readConcern' not in cmd:
-                if dollar_out:
-                    # Avoid auto-injecting a session.
-                    result = sock_info.command(
-                        self.__database.name,
-                        cmd,
-                        slave_ok,
-                        self.read_preference,
-                        self.codec_options,
-                        parse_write_concern_error=True,
-                        collation=collation,
-                        session=session)
-                else:
-                    result = sock_info.command(
-                        self.__database.name,
-                        cmd,
-                        slave_ok,
-                        ReadPreference.PRIMARY,
-                        self.codec_options,
-                        read_concern=self.read_concern,
-                        collation=collation,
-                        session=session)
+            if (sock_info.max_wire_version >= 4
+                    and 'readConcern' not in cmd
+                    and not dollar_out):
+                read_concern = self.read_concern
             else:
-                result = sock_info.command(
-                    self.__database.name,
-                    cmd,
-                    slave_ok,
-                    self.read_preference,
-                    self.codec_options,
-                    parse_write_concern_error=dollar_out,
-                    collation=collation,
-                    session=session)
+                read_concern = DEFAULT_READ_CONCERN
+
+            # Avoid auto-injecting a session: aggregate() passes a session,
+            # aggregate_raw_batches() passes none.
+            result = sock_info.command(
+                self.__database.name,
+                cmd,
+                slave_ok,
+                self.read_preference,
+                self.codec_options,
+                parse_write_concern_error=dollar_out,
+                read_concern=read_concern,
+                collation=collation,
+                session=session)
 
             if "cursor" in result:
                 cursor = result["cursor"]
@@ -2202,7 +2188,8 @@ def aggregate_raw_batches(self, pipeline, **kwargs):
                                None, False, **kwargs)
 
     def watch(self, pipeline=None, full_document='default', resume_after=None,
-              max_await_time_ms=None, batch_size=None, collation=None):
+              max_await_time_ms=None, batch_size=None, collation=None,
+              session=None):
         """Watch changes on this collection.
 
         Performs an aggregation with an implicit initial
@@ -2266,6 +2253,8 @@ def watch(self, pipeline=None, full_document='default', resume_after=None,
             per batch.
           - `collation` (optional): The :class:`~pymongo.collation.Collation`
             to use for the aggregation.
+          - `session` (optional): a
+            :class:`~pymongo.client_session.ClientSession`.
 
         :Returns:
           A :class:`~pymongo.change_stream.ChangeStream` cursor.
@@ -2287,7 +2276,7 @@ def watch(self, pipeline=None, full_document='default', resume_after=None,
         common.validate_string_or_none('full_document', full_document)
 
         return ChangeStream(self, pipeline, full_document, resume_after,
-                            max_await_time_ms, batch_size, collation)
+                            max_await_time_ms, batch_size, collation, session)
 
     def group(self, key, condition, initial, reduce, finalize=None, **kwargs):
         """Perform a query similar to an SQL *group by* operation.

pymongo/topology.py

@@ -400,7 +400,7 @@ def return_server_session(self, server_session, lock):
                                                              session_timeout)
         else:
             # Called from a __del__ method, can't use a lock.
-            self._session_pool.append(server_session)
+            self._session_pool.return_server_session_no_lock(server_session)
 
     def _new_selection(self):
         """A Selection object, initially including all known servers.