PYTHON-1332 - Session param for client methods

ajdavis · ajdavis · commit afb6a931d76f · 2017-09-11T17:46:44.000-04:00
diff --git a/pymongo/mongo_client.py b/pymongo/mongo_client.py
@@ -1234,19 +1234,37 @@ def _return_server_session(self, server_session):
         """Internal: return a _ServerSession to the pool."""
         return self._topology.return_server_session(server_session)
 
-    def server_info(self):
-        """Get information about the MongoDB server we're connected to."""
+    def server_info(self, session=None):
+        """Get information about the MongoDB server we're connected to.
+
+        :Parameters:
+          - `session` (optional): a
+            :class:`~pymongo.client_session.ClientSession`.
+
+        .. versionchanged:: 3.6
+           Added ``session`` parameter.
+        """
         return self.admin.command("buildinfo",
-                                  read_preference=ReadPreference.PRIMARY)
+                                  read_preference=ReadPreference.PRIMARY,
+                                  session=session)
+
+    def database_names(self, session=None):
+        """Get a list of the names of all databases on the connected server.
+
+        :Parameters:
+          - `session` (optional): a
+            :class:`~pymongo.client_session.ClientSession`.
 
-    def database_names(self):
-        """Get a list of the names of all databases on the connected server."""
+        .. versionchanged:: 3.6
+           Added ``session`` parameter.
+        """
         return [db["name"] for db in
                 self._database_default_options("admin").command(
                     SON([("listDatabases", 1),
-                         ("nameOnly", True)]))["databases"]]
+                         ("nameOnly", True)]),
+                    session=session)["databases"]]
 
-    def drop_database(self, name_or_database):
+    def drop_database(self, name_or_database, session=None):
         """Drop a database.
 
         Raises :class:`TypeError` if `name_or_database` is not an instance of
@@ -1257,6 +1275,11 @@ def drop_database(self, name_or_database):
           - `name_or_database`: the name of a database to drop, or a
             :class:`~pymongo.database.Database` instance representing the
             database to drop
+          - `session` (optional): a
+            :class:`~pymongo.client_session.ClientSession`.
+
+        .. versionchanged:: 3.6
+           Added ``session`` parameter.
 
         .. note:: The :attr:`~pymongo.mongo_client.MongoClient.write_concern` of
            this client is automatically applied to this operation when using
@@ -1284,7 +1307,8 @@ def drop_database(self, name_or_database):
                 slave_ok=slave_ok,
                 read_preference=ReadPreference.PRIMARY,
                 write_concern=self.write_concern,
-                parse_write_concern_error=True)
+                parse_write_concern_error=True,
+                session=session)
 
     def get_default_database(self):
         """DEPRECATED - Get the database named in the MongoDB connection URI.
@@ -1382,30 +1406,39 @@ def is_locked(self):
     def fsync(self, **kwargs):
         """Flush all pending writes to datafiles.
 
-        :Parameters:
-
-            Optional parameters can be passed as keyword arguments:
+        Optional parameters can be passed as keyword arguments:
+          - `lock`: If True lock the server to disallow writes.
+          - `async`: If True don't block while synchronizing.
+          - `session` (optional): a
+            :class:`~pymongo.client_session.ClientSession`.
 
-            - `lock`: If True lock the server to disallow writes.
-            - `async`: If True don't block while synchronizing.
+        .. versionchanged:: 3.6
+           Added ``session`` parameter.
 
-            .. warning:: `async` and `lock` can not be used together.
+        .. warning:: `async` and `lock` can not be used together.
 
-            .. warning:: MongoDB does not support the `async` option
-                         on Windows and will raise an exception on that
-                         platform.
+        .. warning:: MongoDB does not support the `async` option
+                     on Windows and will raise an exception on that
+                     platform.
         """
         self.admin.command("fsync",
                            read_preference=ReadPreference.PRIMARY, **kwargs)
 
-    def unlock(self):
+    def unlock(self, session=None):
         """Unlock a previously locked server.
+
+        :Parameters:
+          - `session` (optional): a
+            :class:`~pymongo.client_session.ClientSession`.
+
+        .. versionchanged:: 3.6
+           Added ``session`` parameter.
         """
         cmd = {"fsyncUnlock": 1}
         with self._socket_for_writes() as sock_info:
             if sock_info.max_wire_version >= 4:
                 try:
-                    sock_info.command("admin", cmd)
+                    sock_info.command("admin", cmd, session=session)
                 except OperationFailure as exc:
                     # Ignore "DB not locked" to replicate old behavior
                     if exc.code != 125:
diff --git a/test/test_session.py b/test/test_session.py
@@ -15,7 +15,9 @@
 """Test the client_session module."""
 from bson import DBRef
 from pymongo import OFF, InsertOne, IndexModel
-from pymongo.errors import InvalidOperation, ConfigurationError
+from pymongo.errors import (ConfigurationError,
+                            InvalidOperation,
+                            OperationFailure)
 from pymongo.monitoring import _SENSITIVE_COMMANDS
 from test import IntegrationTest, client_context
 from test.utils import ignore_deprecations, rs_or_single_client, EventListener
@@ -74,6 +76,46 @@ def test_sessions_not_supported(self):
                 ConfigurationError, "Sessions are not supported"):
             self.client.start_session()
 
+    @client_context.require_sessions
+    def test_client(self):
+        listener = SessionTestListener()
+        client = rs_or_single_client(event_listeners=[listener])
+
+        # Make sure if the test fails we unlock the server.
+        def unlock():
+            try:
+                client.unlock()
+            except OperationFailure:
+                pass
+
+        self.addCleanup(unlock)
+
+        ops = [
+            (client.server_info, [], {}),
+            (client.database_names, [], {}),
+            (client.drop_database, ['pymongo_test'], {}),
+            (client.fsync, [], {'lock': True}),
+            (client.unlock, [], {}),
+        ]
+
+        with client.start_session() as s:
+            for f, args, kwargs in ops:
+                listener.results.clear()
+                kwargs['session'] = s
+                f(*args, **kwargs)
+                self.assertGreaterEqual(len(listener.results['started']), 1)
+                for event in listener.results['started']:
+                    self.assertTrue(
+                        'lsid' in event.command,
+                        "%s sent no lsid with %s" % (
+                            f.__name__, event.command_name))
+
+                    self.assertEqual(
+                        s.session_id,
+                        event.command['lsid'],
+                        "%s sent wrong lsid with %s" % (
+                            f.__name__, event.command_name))
+
     @client_context.require_sessions
     def test_database(self):
         listener = SessionTestListener()
