From a44e719dca39e1f10c648b58415e9c2586c2b8cd Mon Sep 17 00:00:00 2001
From: Shane Harvey <shane.harvey@mongodb.com>
Date: Wed, 28 Apr 2021 12:18:54 -0700
Subject: [PATCH] PYTHON-2533 Add support for sample_rate and filter in
 set_profiling_level (#605)

---
 pymongo/database.py   | 24 ++++++++++++++++++++----
 test/test_database.py | 29 +++++++++++++++++++++++++++++
 2 files changed, 49 insertions(+), 4 deletions(-)

diff --git a/pymongo/database.py b/pymongo/database.py
index a810c67a6..1291c2a14 100644
--- a/pymongo/database.py
+++ b/pymongo/database.py
@@ -914,7 +914,8 @@ class Database(common.BaseObject):
         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:
@@ -925,6 +926,10 @@ class Database(common.BaseObject):
             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:
 
@@ -942,6 +947,9 @@ class Database(common.BaseObject):
         (: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.
 
@@ -953,10 +961,18 @@ class Database(common.BaseObject):
         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.
diff --git a/test/test_database.py b/test/test_database.py
index e910c4b9c..0f64f2a46 100644
--- a/test/test_database.py
+++ b/test/test_database.py
@@ -397,6 +397,35 @@ class TestDatabase(IntegrationTest):
         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