Add ability to configure explain command options

pymongoexplain/explainable_collection.py

@@ -26,17 +26,27 @@
 
 
 class ExplainableCollection():
-    def __init__(self, collection):
+    def __init__(self, collection, verbosity=None, comment=None):
         self.collection = collection
         self.last_cmd_payload = None
+        self.verbosity = verbosity
+        self.comment = comment
 
     def _explain_command(self, command):
         command_son = command.get_SON()
         explain_command = SON([("explain", command_son)])
-        explain_command["verbosity"] = "queryPlanner"
+        explain_command["verbosity"] = self.verbosity or "queryPlanner"
+        if self.comment:
+            explain_command["comment"] = self.comment
         self.last_cmd_payload = command_son
         return self.collection.database.command(explain_command)
 
+    def update_settings(self, verbosity=None, comment=None):
+        if verbosity:
+            self.verbosity = verbosity
+        if comment:
+            self.comment = comment
+
     def update_one(self, filter, update, upsert=False,
                    bypass_document_validation=False,
                    collation=None, array_filters=None, hint=None,

test/test_collection.py

@@ -16,6 +16,7 @@
 import unittest
 import subprocess
 import os
+from collections import abc
 
 from pymongo import MongoClient
 from pymongo import monitoring
@@ -28,6 +29,7 @@
 class CommandLogger(monitoring.CommandListener):
     def __init__(self):
         self.cmd_payload = {}
+
     def started(self, event):
         self.cmd_payload = event.command
 
@@ -37,11 +39,13 @@ def succeeded(self, event):
     def failed(self, event):
         pass
 
+
 class TestExplainableCollection(unittest.TestCase):
     def setUp(self) -> None:
         self.logger = CommandLogger()
         self.client = MongoClient(serverSelectionTimeoutMS=1000,
-                                 event_listeners=[self.logger])
+                                  event_listeners=[self.logger])
+        self.server_version = self.client.server_info()["versionArray"]
         self.collection = self.client.db.products
         self.collection.insert_one({'x': 1})
         self.explain = ExplainCollection(self.collection)
@@ -50,6 +54,18 @@ def _compare_command_dicts(self, ours, theirs):
         for key in ours.keys():
             self.assertEqual(ours[key], theirs[key])
 
+    def _recursiveIn(self, member, container):
+        if isinstance(container, abc.Mapping):
+            if member in container:
+                return True
+            return any([self._recursiveIn(member, v) for v in
+                        container.values()])
+        if isinstance(container, list):
+            if member in container:
+                return True
+            return any([self._recursiveIn(member, v) for v in
+                        container])
+        return False
 
     def test_update_one(self):
         self.collection.update_one({"quantity": 1057, "category": "apparel"},
@@ -83,7 +99,7 @@ def test_count_documents(self):
         self.collection.count_documents({"ord_dt": {"$gt": 10}})
         last_logger_payload = self.logger.cmd_payload
         res = self.explain.count_documents({"ord_dt": {"$gt": 10}})
-        #self.assertIn("queryPlanner", res)
+        self.assertTrue(self._recursiveIn("queryPlanner", res))
         last_cmd_payload = self.explain.last_cmd_payload
         self._compare_command_dicts(last_cmd_payload, last_logger_payload)
 
@@ -94,7 +110,7 @@ def test_aggregate(self):
         last_logger_payload = self.logger.cmd_payload
         res = self.explain.aggregate([{"$project": {"tags": 1}}, {"$unwind":
                                                                  "$tags"}], None)
-        self.assertIn("queryPlanner", res["stages"][0]["$cursor"])
+        self.assertTrue(self._recursiveIn("queryPlanner", res))
         last_cmd_payload = self.explain.last_cmd_payload
         self._compare_command_dicts(last_cmd_payload, last_logger_payload)
 
@@ -116,13 +132,13 @@ def test_delete_many(self):
 
     def test_watch(self):
         res = self.explain.watch()
-        self.assertIn("queryPlanner", res["stages"][0]["$cursor"])
+        self.assertTrue(self._recursiveIn("queryPlanner", res))
         self.collection.watch(pipeline=[{"$project": {"tags": 1}}],
                               batch_size=10, full_document="updateLookup")
         last_logger_payload = self.logger.cmd_payload
         res = self.explain.watch(pipeline=[{"$project": {"tags": 1}}],
                                  batch_size=10, full_document="updateLookup")
-        self.assertIn("queryPlanner", res["stages"][0]["$cursor"])
+        self.assertTrue(self._recursiveIn("queryPlanner", res))
         last_cmd_payload = self.explain.last_cmd_payload
         self._compare_command_dicts(last_cmd_payload, last_logger_payload)
 
@@ -142,8 +158,6 @@ def test_find(self):
         last_cmd_payload = self.explain.last_cmd_payload
         self._compare_command_dicts(last_cmd_payload, last_logger_payload)
 
-
-
     def test_find_one(self):
         self.collection.find_one(projection=['a', 'b.c'])
         last_logger_payload = self.logger.cmd_payload
@@ -217,6 +231,29 @@ def test_imports(self):
         from pymongoexplain import ExplainableCollection
         self.assertEqual(ExplainableCollection, ExplainCollection)
 
+    def test_verbosity(self):
+        res = self.explain.find({})
+        self.assertFalse(self._recursiveIn("executionStats", res))
+        self.assertFalse(self._recursiveIn("allPlansExecution", res))
+        self.explain.update_settings(verbosity="executionStats")
+        res = self.explain.find({})
+        self.assertTrue(self._recursiveIn("executionStats", res))
+        self.assertFalse(self._recursiveIn("allPlansExecution", res))
+        self.explain.update_settings(verbosity="allPlansExecution")
+        res = self.explain.find({})
+        self.assertTrue(self._recursiveIn("executionStats", res))
+        self.assertTrue(self._recursiveIn("allPlansExecution", res))
+
+    def test_comment(self):
+        if tuple(self.server_version) < (5, 0, 0, 0):
+            self.skipTest("MongoDB 4.x does not embed the comment in the "
+                          "explain document")
+        res = self.explain.find({})
+        self.assertFalse(self._recursiveIn("comment", res))
+        self.explain.update_settings(comment="hey, I'm a comment")
+        res = self.explain.find({})
+        self.assertTrue(self._recursiveIn("comment", res))
+
 
 if __name__ == '__main__':
     unittest.main()