fix(pymongo): extend support for more operations (#394)

Author: sagivr2020

epsagon/events/pymongo.py

@@ -5,7 +5,6 @@
 from __future__ import absolute_import
 from uuid import uuid4
 import traceback
-from past.builtins import range
 
 from epsagon.utils import add_data_if_needed
 from ..event import BaseEvent
@@ -19,6 +18,10 @@ class PyMongoEvent(BaseEvent):
 
     ORIGIN = 'pymongo'
     RESOURCE_TYPE = 'pymongo'
+    INSERT_ONE = 'insert_one'
+    INSERT_MANY = 'insert_many'
+    INSERT_OPERATIONS = (INSERT_ONE, INSERT_MANY)
+    FILTER_OPERATIONS = ('find', 'update_one', 'delete_many')
 
     #pylint: disable=W0613
     def __init__(self, wrapped, instance, args, kwargs, start_time, response,
@@ -33,30 +36,25 @@ def __init__(self, wrapped, instance, args, kwargs, start_time, response,
         :param response: response data
         :param exception: Exception (if happened)
         """
-
         super(PyMongoEvent, self).__init__(start_time)
 
-        documents = args[0]
+        self.resource['operation'] = getattr(wrapped, '__name__')
 
         self.event_id = 'mongo-{}'.format(str(uuid4()))
-        self.resource['name'] = instance.full_name
-
-        self.resource['operation'] = \
-            'insert_many' if isinstance(documents, list) else 'insert_one'
+        self.resource['name'] = instance.name
         address = list(getattr(
             instance.database.client,
             '_topology_settings'
-        ).seeds)[0]
-
-        if self.resource['operation'] == 'insert_one':
-            documents = [documents]
+            ).seeds)[0]
 
         self.resource['metadata'] = {
-            'DB URL': ':'.join([str(x) for x in address]),
-            'DB Name': str(instance.database.name),
-            'Collection Name': str(instance.name),
-        }
-        add_data_if_needed(self.resource['metadata'], 'Items', documents)
+                'DB URL': ':'.join([str(x) for x in address]),
+                'DB Name': str(instance.database.name),
+                'Collection Name': str(instance.collection.name),
+            }
+
+        if args:
+            self.handle_request_payload(args)
 
         if response is not None:
             self.update_response(response)
@@ -70,17 +68,74 @@ def update_response(self, response):
         :param response: Response from botocore
         :return: None
         """
-        for i in range(len(self.resource['metadata'].get('items', []))):
-            self.resource['metadata']['items'][i]['_id'] = str(
-                self.resource['metadata']['items'][i]['_id']
-            )
 
-        if self.resource['operation'] == 'insert_many':
-            self.resource['metadata']['inserted_ids'] = \
-                [str(x) for x in response.inserted_ids]
-        elif self.resource['operation'] == 'insert_one':
+        if self.resource['operation'] in PyMongoEvent.INSERT_OPERATIONS:
+            self.handle_insert_operations_response(response)
+
+        elif self.resource['operation'] in PyMongoEvent.FILTER_OPERATIONS:
+            self.handle_filter_operations_response(response)
+
+
+    def handle_request_payload(self, input_args):
+        """
+        Handle input data before add to event.
+        :param input_args: input from botocore
+        :return: None
+        """
+
+        if self.resource['operation'] == PyMongoEvent.INSERT_MANY:
+            add_data_if_needed(self.resource['metadata'], 'Items',
+                input_args[0])
+
+        elif self.resource['operation'] == PyMongoEvent.INSERT_ONE:
+            add_data_if_needed(self.resource['metadata'], 'Item',
+                input_args[0])
+
+        elif self.resource['operation'] in PyMongoEvent.FILTER_OPERATIONS:
+            add_data_if_needed(self.resource['metadata'], 'Filter',
+                input_args[0])
+
+            if self.resource['operation'] == 'update_one':
+                add_data_if_needed(self.resource['metadata'], 'New Values',
+                    input_args[1])
+
+
+    def handle_insert_operations_response(self, response):
+        """
+        Handle response data before add to event.
+        :param response: response from botocore
+        :return: None
+        """
+
+        if self.resource['operation'] == PyMongoEvent.INSERT_MANY:
             self.resource['metadata']['inserted_ids'] = \
-                [str(response.inserted_id)]
+                    [str(x) for x in response.inserted_ids]
+
+        elif self.resource['operation'] == PyMongoEvent.INSERT_ONE:
+            self.resource['metadata']['inserted_id'] = \
+                    [str(response.inserted_id)]
+
+
+    def handle_filter_operations_response(self, response):
+        """
+        Handle response data before add to event.
+        :param response: response from botocore
+        :return: None
+        """
+
+        if self.resource['operation'] == 'find':
+            self.resource['metadata']['Results'] = \
+                        list(response)
+
+        elif self.resource['operation'] == 'update_one':
+            self.resource['metadata']['matched_count'] = \
+                        response.matched_count
+            self.resource['metadata']['modified_count'] = \
+                        response.modified_count
+
+        elif self.resource['operation'] == 'delete_many':
+            self.resource['metadata']['deleted_count'] = \
+                        response.deleted_count
 
 
 class PyMongoEventFactory(object):

epsagon/modules/pymongo.py

@@ -37,3 +37,18 @@ def patch():
         'Collection.insert_many',
         _wrapper
     )
+    wrapt.wrap_function_wrapper(
+        'pymongo.collection',
+        'Collection.find',
+        _wrapper
+    )
+    wrapt.wrap_function_wrapper(
+        'pymongo.collection',
+        'Collection.update_one',
+        _wrapper
+    )
+    wrapt.wrap_function_wrapper(
+        'pymongo.collection',
+        'Collection.delete_many',
+        _wrapper
+    )