PYTHON-1419 Call endSessions on MongoClient.close.

Author: ShaneHarvey

pymongo/client_session.py

@@ -236,6 +236,12 @@ class _ServerSessionPool(collections.deque):
 
  This class is not thread-safe, access it while holding the Topology lock.
  """
+ def pop_all(self):
+ ids = []
+ while self:
+ ids.append(self.pop().session_id)
+ return ids
+
  def get_server_session(self, session_timeout_minutes):
  # Although the Driver Sessions Spec says we only clear stale sessions
  # in return_server_session, PyMongo can't take a lock when returning

pymongo/common.py

@@ -94,6 +94,10 @@
 # Error codes to ignore if GridFS calls createIndex on a secondary
 UNAUTHORIZED_CODES = (13, 16547, 16548)
 
+# Maximum number of sessions to send in a single endSessions command.
+# From the driver sessions spec.
+_MAX_END_SESSIONS = 10000
+
 
 def partition_node(node):
  """Split a host:port string into (host, int(port)) pair."""

pymongo/mongo_client.py

@@ -857,13 +857,42 @@ def _is_writable(self):
  except ConnectionFailure:
  return False
 
+ def _end_sessions(self, session_ids):
+ """Send endSessions command(s) with the given session ids."""
+ try:
+ # Use SocketInfo.command directly to avoid implicitly creating
+ # another session.
+ with self._socket_for_reads(
+ ReadPreference.PRIMARY_PREFERRED) as (sock_info, slave_ok):
+ if not sock_info.supports_sessions:
+ return
+
+ for i in range(0, len(session_ids), common._MAX_END_SESSIONS):
+ spec = SON([('endSessions',
+ session_ids[i:i + common._MAX_END_SESSIONS])])
+ sock_info.command(
+ 'admin', spec, slave_ok=slave_ok, client=self)
+ except PyMongoError:
+ # Drivers MUST ignore any errors returned by the endSessions
+ # command.
+ pass
+
  def close(self):
- """Disconnect from MongoDB.
+ """Cleanup client resources and disconnect from MongoDB.
+
+ On MongoDB >= 3.6, end all server sessions created by this client by
+ sending one or more endSessions commands.
 
  Close all sockets in the connection pools and stop the monitor threads.
  If this instance is used again it will be automatically re-opened and
  the threads restarted.
+
+ .. versionchanged:: 3.6
+ End all server sessions created by this client.
  """
+ session_ids = self._topology.pop_all_sessions()
+ if session_ids:
+ self._end_sessions(session_ids)
  # Run _process_periodic_tasks to send pending killCursor requests
  # before closing the topology.
  self._process_periodic_tasks()
@@ -1307,7 +1336,7 @@ def _process_periodic_tasks(self):
  except Exception:
  helpers._handle_exception()
 
- def start_session(self, **kwargs):
+ def start_session(self, causal_consistency=True):
  """Start a logical session.
 
  This method takes the same parameters as
@@ -1318,6 +1347,12 @@ def start_session(self, **kwargs):
  if this client has been authenticated to multiple databases using the
  deprecated method :meth:`~pymongo.database.Database.authenticate`.
 
+ A :class:`~pymongo.client_session.ClientSession` may only be used with
+ the MongoClient that started it.
+
+ :Returns:
+ An instance of :class:`~pymongo.client_session.ClientSession`.
+
  .. versionadded:: 3.6
  """
  # Driver Sessions Spec: "If startSession is called when multiple users
@@ -1330,8 +1365,10 @@ def start_session(self, **kwargs):
 
  # Raises ConfigurationError if sessions are not supported.
  server_session = self._get_server_session()
- opts = client_session.SessionOptions(**kwargs)
- return client_session.ClientSession(self, server_session, opts, authset)
+ opts = client_session.SessionOptions(
+ causal_consistency=causal_consistency)
+ return client_session.ClientSession(
+ self, server_session, opts, authset)
 
  def _get_server_session(self):
  """Internal: start or resume a _ServerSession."""

pymongo/topology.py

@@ -397,6 +397,11 @@ def close(self):
  def description(self):
  return self._description
 
+ def pop_all_sessions(self):
+ """Pop all session ids from the pool."""
+ with self._lock:
+ return self._session_pool.pop_all()
+
  def get_server_session(self):
  """Start or resume a server session, or raise ConfigurationError."""
  with self._lock:

test/test_session.py

@@ -21,6 +21,7 @@
 from bson.py3compat import StringIO
 from gridfs import GridFS, GridFSBucket
 from pymongo import ASCENDING, InsertOne, IndexModel, OFF, monitoring
+from pymongo.common import _MAX_END_SESSIONS
 from pymongo.errors import (ConfigurationError,
  InvalidOperation,
  OperationFailure)
@@ -943,6 +944,29 @@ def test_cluster_time_no_server_support(self):
  '$clusterTime')
  self.assertIsNone(after_cluster_time)
 
+ def test_end_sessions(self):
+ listener = SessionTestListener()
+ client = rs_or_single_client(event_listeners=[listener])
+ # Start many sessions.
+ sessions = [client.start_session()
+ for _ in range(_MAX_END_SESSIONS + 1)]
+ for s in sessions:
+ s.end_session()
+
+ # Closing the client should end all sessions and clear the pool.
+ self.assertEqual(len(client._topology._session_pool),
+ _MAX_END_SESSIONS + 1)
+ client.close()
+ self.assertEqual(len(client._topology._session_pool), 0)
+ end_sessions = [e for e in listener.results['started']
+ if e.command_name == 'endSessions']
+ self.assertEqual(len(end_sessions), 2)
+
+ # Closing again should not send any commands.
+ listener.results.clear()
+ client.close()
+ self.assertEqual(len(listener.results['started']), 0)
+
 
 class TestSessionsMultiAuth(IntegrationTest):
  @client_context.require_auth