Add support for find functions, replace functions, count_documents, and other CRUD helpers (#6)

Author: juliusgeo

pymongoexplain/commands.py

@@ -20,32 +20,21 @@
 
 from bson.son import SON
 from pymongo.collection import Collection
+from .utils import convert_to_camelcase
+
 
 Document = Union[dict, SON]
 
+
 class BaseCommand():
-    def __init__(self, dictionary):
-        if type(self) == BaseCommand:
-            raise NotImplementedError
-        self.command_document = self.convert_to_camelcase(dictionary)
-
-    def convert_to_camelcase(self, d: dict) -> dict:
-        ret = dict()
-        for key in d.keys():
-            if d[key] is None:
-                continue
-            new_key = key
-            if "_" in key:
-                new_key = key.split("_")[0] + ''.join(
-                    [i.capitalize() for i in key.split("_")[1:]])
-            if type(d[key]) == list:
-                ret[new_key] = [self.convert_to_camelcase(i) for i in d[key]
-                                if type(i) == dict]
-            elif type(d[key]) == dict:
-                ret[new_key] = self.convert_to_camelcase(d[key])
-            else:
-                ret[new_key] = d[key]
-        return ret
+    def __init__(self, collection):
+        self.command_document = {}
+        self.collection = collection
+
+    @property
+    def command_name(self):
+        """New command classes will specify the command name here."""
+        return None
 
     def get_SON(self):
         cmd = SON([(self.command_name, self.collection)])
@@ -56,72 +45,99 @@ def get_SON(self):
 class UpdateCommand(BaseCommand):
     def __init__(self, collection: Collection, filter, update,
                  kwargs):
-        super().__init__(kwargs)
-        self.command_name = "update"
-        self.collection = collection.name
-        return_dictionary =  {"updates":[{"q": filter, "u": update}]}
-        for key, value in self.command_document.items():
-            if key == "bypassDocumentValidation":
-                return_dictionary[key] = value
+        super().__init__(collection.name)
+        return_document = {"updates":[{"q": filter, "u": update}]}
+        for key in kwargs:
+            value = kwargs[key]
+            if key == "bypass_document_validation":
+                return_document[key] = value
             else:
-                return_dictionary["updates"][0][key] = value
-        self.command_document = return_dictionary
+                return_document["updates"][0][key] = value
+        self.command_document = convert_to_camelcase(return_document)
+
+    @property
+    def command_name(self):
+        return "update"
 
 
 class DistinctCommand(BaseCommand):
     def __init__(self, collection: Collection, key, filter, session,
                  kwargs):
-        self.command_name = "distinct"
-        self.collection = collection.name
+        super().__init__(collection.name)
         self.command_document = {"key": key, "query": filter}
         for key, value in kwargs.items():
             self.command_document[key] = value
-        super().__init__(self.command_document)
+        self.command_document = convert_to_camelcase(self.command_document)
+
+    @property
+    def command_name(self):
+        return "distinct"
 
 
 class AggregateCommand(BaseCommand):
     def __init__(self, collection: Collection, pipeline, session,
                  cursor_options,
-                 kwargs):
-        self.command_name = "aggregate"
-        self.collection = collection.name
+                 kwargs, exclude_keys = []):
+        super().__init__(collection.name)
         self.command_document = {"pipeline": pipeline, "cursor": cursor_options}
         for key, value in kwargs.items():
             self.command_document[key] = value
-        super().__init__(self.command_document)
 
+        self.command_document = convert_to_camelcase(
+            self.command_document, exclude_keys=exclude_keys)
+
+    @property
+    def command_name(self):
+        return "aggregate"
 
 class CountCommand(BaseCommand):
     def __init__(self, collection: Collection, filter,
                  kwargs):
-        self.command_name = "count"
-        self.collection = collection.name
+        super().__init__(collection.name)
         self.command_document = {"query": filter}
         for key, value in kwargs.items():
-            self.dictionary[key] = value
-        super().__init__(self.command_document)
+            self.command_document[key] = value
+        self.command_document = convert_to_camelcase(self.command_document)
+
+    @property
+    def command_name(self):
+        return "count"
 
 
 class FindCommand(BaseCommand):
     def __init__(self, collection: Collection,
                  kwargs):
-        self.command_name = "find"
-        self.collection = collection.name
-        self.command_document={}
+        super().__init__(collection.name)
         for key, value in kwargs.items():
             self.command_document[key] = value
-        super().__init__(self.command_document)
+        self.command_document = convert_to_camelcase(self.command_document)
+
+    @property
+    def command_name(self):
+        return "find"
+
+class FindAndModifyCommand(BaseCommand):
+    def __init__(self, collection: Collection,
+                 kwargs):
+        super().__init__(collection.name)
+        for key, value in kwargs.items():
+            self.command_document[key] = value
+        self.command_document = convert_to_camelcase(self.command_document)
+
+    @property
+    def command_name(self):
+        return "findAndModify"
 
 
 class DeleteCommand(BaseCommand):
     def __init__(self, collection: Collection, filter,
                  limit, collation, kwargs):
-        super().__init__(kwargs)
-        self.command_name = "delete"
-        self.collection = collection.name
-        return_dictionary = {"deletes": [{"q": filter, "limit": limit,
-                                         "collation": collation}]}
-
-        for key, value in self.command_document.items():
-            return_dictionary[key] = value
-        self.command_document = return_dictionary
+        super().__init__(collection.name)
+        self.command_document = {"deletes": [SON({"q": filter, "limit": limit})]}
+        for key, value in kwargs.items():
+            self.command_document[key] = value
+        self.command_document = convert_to_camelcase(self.command_document)
+
+    @property
+    def command_name(self):
+        return "delete"

pymongoexplain/explainable_collection.py

@@ -13,14 +13,13 @@
 # limitations under the License.
 
 
-from typing import Union, List
+from typing import Union, List, Dict
 
 import pymongo
 from bson.son import SON
 
 from .commands import AggregateCommand, FindCommand, CountCommand, \
-    UpdateCommand, DistinctCommand, DeleteCommand
-
+    UpdateCommand, DistinctCommand, DeleteCommand, FindAndModifyCommand
 
 Document = Union[dict, SON]
 
@@ -29,10 +28,11 @@ def __init__(self, collection):
         self.collection = collection
         self.last_cmd_payload = None
 
-    def _explain_command(self, command: Document):
-        explain_command = SON([("explain", command.get_SON())])
+    def _explain_command(self, command):
+        command_son = command.get_SON()
+        explain_command = SON([("explain", command_son)])
         explain_command["verbosity"] = "queryPlanner"
-        self.last_cmd_payload = explain_command
+        self.last_cmd_payload = command_son
         return self.collection.database.command(explain_command)
 
     def update_one(self, filter, update, upsert=False,
@@ -41,6 +41,9 @@ def update_one(self, filter, update, upsert=False,
                    session=None, **kwargs):
         kwargs.update(locals())
         del kwargs["self"], kwargs["kwargs"], kwargs["filter"], kwargs["update"]
+        kwargs["multi"] = False
+        if bypass_document_validation == False:
+            del kwargs["bypass_document_validation"]
         command = UpdateCommand(self.collection, filter, update, kwargs)
         return self._explain_command(command)
 
@@ -49,6 +52,8 @@ def update_many(self, filter: Document, update: Document, upsert=False,
         kwargs.update(locals())
         del kwargs["self"], kwargs["kwargs"], kwargs["filter"], kwargs["update"]
         kwargs["multi"] = True
+        if bypass_document_validation == False:
+            del kwargs["bypass_document_validation"]
         command = UpdateCommand(self.collection, filter, update, kwargs)
         return self._explain_command(command)
 
@@ -61,8 +66,19 @@ def aggregate(self, pipeline: List[Document], session=None, **kwargs):
                                    {},kwargs)
         return self._explain_command(command)
 
-    def count_documents(self, filter: Document, session=None, **kwargs):
-        command = CountCommand(self.collection, filter,kwargs)
+    def estimated_document_count(self,
+                                 **kwargs):
+
+        command = CountCommand(self.collection, None, kwargs)
+        return self._explain_command(command)
+
+    def count_documents(self, filter: Document, session=None,
+                                 **kwargs):
+
+        command = AggregateCommand(self.collection, [{'$match': filter},
+                                                      {'$group': {'n': {'$sum': 1}, '_id': 1}}],
+                                   session, {}, kwargs,
+                                   exclude_keys=filter.keys())
         return self._explain_command(command)
 
     def delete_one(self, filter: Document, collation=None, session=None,
@@ -73,8 +89,11 @@ def delete_one(self, filter: Document, collation=None, session=None,
         return self._explain_command(command)
 
     def delete_many(self, filter: Document, collation=None,
-                    session=None, **kwargs):
+                    session=None, **kwargs: Dict[str, Union[int, str,
+                                                            Document,
+                                                      bool]]):
         limit = 0
+        kwargs["session"] = session
         command = DeleteCommand(self.collection, filter, limit, collation,
         kwargs)
         return self._explain_command(command)
@@ -100,20 +119,79 @@ def watch(self, pipeline: Document = None, full_document: Document = None,
                                        max_await_time_ms})
         return self._explain_command(command)
 
-    def find(self, filter: Document = None, projection: list = None,
-             skip: int = 0, limit: int = 0, no_cursor_timeout: bool = False,
-             sort: Document = None, allow_partial_results: bool = False,
-             oplog_replay: bool = False, batch_size: int=0,
-             collation: Document = None, hint: Union[Document, str] = None,
-             max_time_ms: int = None, max: Document = None, min: Document =
-             None, return_key: bool = False,
-             show_record_id: bool = False, comment: str = None,
-             session:Document = None, **kwargs: Union[int, str, Document,
-                                                      bool]):
+    def find(self, filter: Document = None,
+             **kwargs: Dict[str, Union[int, str,Document, bool]]):
         kwargs.update(locals())
         del kwargs["self"], kwargs["kwargs"]
         command = FindCommand(self.collection,
                                 kwargs)
         return self._explain_command(command)
 
+    def find_one(self, filter: Document = None, **kwargs: Dict[str,
+                                                               Union[int, str,
+                                                                Document, bool]]):
+        kwargs.update(locals())
+        del kwargs["self"], kwargs["kwargs"]
+        kwargs["limit"] = 1
+        command = FindCommand(self.collection, kwargs)
+        return self._explain_command(command)
+
+    def find_one_and_delete(self, filter: Document, projection: list = None,
+                            sort: Document=None, session=None,
+                        **kwargs):
+        kwargs["query"] = filter
+        kwargs["fields"] = projection
+        kwargs["sort"] = sort
+        kwargs["remove"] = True
+        kwargs["session"] = session
+
+        command = FindAndModifyCommand(self.collection,
+                                       kwargs)
+        return self._explain_command(command)
+
+    def find_one_and_replace(self, filter: Document, replacement: Document,
+                            projection: list = None, sort=None,
+                             return_document=pymongo.ReturnDocument.BEFORE,
+                             session=None, **kwargs):
+        kwargs["query"] = filter
+        kwargs["fields"] = projection
+        kwargs["sort"] = sort
+        kwargs["new"] = False
+        kwargs["update"] = replacement
+        kwargs["session"] = session
+        command = FindAndModifyCommand(self.collection,
+                                       kwargs)
+        return self._explain_command(command)
+
+    def find_one_and_update(self, filter: Document, replacement: Document,
+                            projection: list = None, sort=None,
+                             return_document=pymongo.ReturnDocument.BEFORE,
+                             session=None, **kwargs):
+        kwargs["query"] = filter
+        kwargs["fields"] = projection
+        kwargs["sort"] = sort
+        kwargs["upsert"] = False
+        kwargs["update"] = replacement
+        kwargs["session"] = session
+
+        command = FindAndModifyCommand(self.collection,
+                                       kwargs)
+        return self._explain_command(command)
+
+    def replace_one(self, filter: Document, replacement: Document,
+                    upsert=False, bypass_document_validation=False,
+                    collation=None, session=None, **kwargs):
+        kwargs.update(locals())
+        del kwargs["self"], kwargs["kwargs"], kwargs["filter"], kwargs[
+            "replacement"]
+        kwargs["multi"] = False
+        if not bypass_document_validation:
+            del kwargs["bypass_document_validation"]
+        update = replacement
+        command = UpdateCommand(self.collection, filter, update, kwargs)
+
+        return self._explain_command(command)
+
+
+
 

pymongoexplain/utils.py

@@ -0,0 +1,41 @@
+# Copyright 2020-present MongoDB, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+"""Utility functions"""
+
+
+def convert_to_camelcase(d, exclude_keys=[]):
+    if not isinstance(d, dict):
+        return d
+    ret = dict()
+    for key in d.keys():
+        if d[key] is None:
+            continue
+        if key in exclude_keys:
+            ret[key] = d[key]
+            continue
+        new_key = key
+        if "_" in key and key[0] != "_":
+            new_key = key.split("_")[0] + ''.join(
+                [i.capitalize() for i in key.split("_")[1:]])
+        if isinstance(d[key], list):
+            ret[new_key] = [convert_to_camelcase(
+                i, exclude_keys=exclude_keys) for i in d[key]]
+        elif isinstance(d[key], dict):
+            ret[new_key] = convert_to_camelcase(d[key],
+                                                     exclude_keys=exclude_keys)
+        else:
+            ret[new_key] = d[key]
+    return ret