PYTHON-2533 Add support for sample_rate and filter in set_profiling_level (#605)

(cherry picked from commit a44e719)

Author: ShaneHarvey

pymongo/database.py

@@ -1086,7 +1086,8 @@ def profiling_level(self, session=None):
         assert result["was"] >= 0 and result["was"] <= 2
         return result["was"]
 
-    def set_profiling_level(self, level, slow_ms=None, session=None):
+    def set_profiling_level(self, level, slow_ms=None, session=None,
+                            sample_rate=None, filter=None):
         """Set the database's profiling level.
 
         :Parameters:
@@ -1097,6 +1098,10 @@ def set_profiling_level(self, level, slow_ms=None, session=None):
             slower than the `slow_ms` level will get written to the logs.
           - `session` (optional): a
             :class:`~pymongo.client_session.ClientSession`.
+          - `sample_rate` (optional): The fraction of slow operations that
+            should be profiled or logged expressed as a float between 0 and 1.
+          - `filter` (optional): A filter expression that controls which
+            operations are profiled and logged.
 
         Possible `level` values:
 
@@ -1114,6 +1119,9 @@ def set_profiling_level(self, level, slow_ms=None, session=None):
         (:data:`~pymongo.OFF`, :data:`~pymongo.SLOW_ONLY`,
         :data:`~pymongo.ALL`).
 
+        .. versionchanged:: 3.12
+           Added the ``sample_rate`` and ``filter`` parameters.
+
         .. versionchanged:: 3.6
            Added ``session`` parameter.
 
@@ -1125,10 +1133,18 @@ def set_profiling_level(self, level, slow_ms=None, session=None):
         if slow_ms is not None and not isinstance(slow_ms, int):
             raise TypeError("slow_ms must be an integer")
 
+        if sample_rate is not None and not isinstance(sample_rate, float):
+            raise TypeError(
+                "sample_rate must be a float, not %r" % (sample_rate,))
+
+        cmd = SON(profile=level)
         if slow_ms is not None:
-            self.command("profile", level, slowms=slow_ms, session=session)
-        else:
-            self.command("profile", level, session=session)
+            cmd['slowms'] = slow_ms
+        if sample_rate is not None:
+            cmd['sampleRate'] = sample_rate
+        if filter is not None:
+            cmd['filter'] = filter
+        self.command(cmd, session=session)
 
     def profiling_info(self, session=None):
         """Returns a list containing current profiling information.

test/test_database.py

@@ -418,6 +418,35 @@ def test_profiling_levels(self):
         db.set_profiling_level(OFF, 100)  # back to default
         self.assertEqual(100, db.command("profile", -1)['slowms'])
 
+    @client_context.require_no_mongos
+    @client_context.require_version_min(3, 6)
+    def test_profiling_sample_rate(self):
+        db = self.client.pymongo_test
+        with self.assertRaises(TypeError):
+            db.set_profiling_level(SLOW_ONLY, 50, sample_rate='1')
+        with self.assertRaises(TypeError):
+            db.set_profiling_level(SLOW_ONLY, 50, sample_rate=1)
+
+        db.set_profiling_level(SLOW_ONLY, 50, sample_rate=0.0)
+        db.set_profiling_level(SLOW_ONLY, 50, sample_rate=1.0)
+        db.set_profiling_level(SLOW_ONLY, 50, sample_rate=0.5)
+        profile = db.command("profile", -1)
+        self.assertEqual(50, profile['slowms'])
+        self.assertEqual(0.5, profile['sampleRate'])
+        db.set_profiling_level(OFF, 100)  # back to default
+        self.assertEqual(100, db.command("profile", -1)['slowms'])
+
+    @client_context.require_no_mongos
+    @client_context.require_version_min(4, 4, 2)
+    def test_profiling_filter(self):
+        db = self.client.pymongo_test
+        db.set_profiling_level(ALL, filter={'ns': {'$eq': 'test.test'}})
+        profile = db.command("profile", -1)
+        self.assertEqual({'ns': {'$eq': 'test.test'}}, profile['filter'])
+        # filter='unset' resets the filter back to the default.
+        db.set_profiling_level(OFF, 100, filter='unset')
+        self.assertEqual(100, db.command("profile", -1)['slowms'])
+
     @client_context.require_no_mongos
     def test_profiling_info(self):
         db = self.client.pymongo_test