feat(boto3): trace boto3 rds data cs 24 (#382)

Author: maorlx

epsagon/events/botocore.py

@@ -1837,6 +1837,95 @@ def get_secret_value_response(self):
         self._add_data_to_metadata('Response', response_data)
 
 
+class BotocoreAuroraServerlessEvent(BotocoreEvent):
+    """
+    Represents data API (aurora serverless) botocore event.
+    """
+
+    RESOURCE_TYPE = 'rdsdataservice'
+    RESOURCE_TYPE_UPDATE = 'database'
+    EXECUTE_STATEMENT_OPERATION = 'ExecuteStatement'
+    DATABASE_FIELD = 'database'
+    RESOURCE_ARN_FIELD = 'resourceArn'
+
+    def _get_resource_name(self):
+        """
+        Gets the relevant resource name - the managed RDS name
+        """
+        resource_arn = self.request_data.get(self.RESOURCE_ARN_FIELD, '')
+        return resource_arn.split(':')[-1]
+
+
+    def __init__(self, wrapped, instance, args, kwargs, start_time, response,
+                 exception):
+        """
+        Initialize.
+        :param wrapped: wrapt's wrapped
+        :param instance: wrapt's instance
+        :param args: wrapt's args
+        :param kwargs: wrapt's kwargs
+        :param start_time: Start timestamp (epoch)
+        :param response: response data
+        :param exception: Exception (if happened)
+        """
+        self.OPERATION_TO_FUNC.update({
+            self.EXECUTE_STATEMENT_OPERATION: self.execute_statement_op,
+        })
+
+        self.RESPONSE_TO_FUNC.update({
+            self.EXECUTE_STATEMENT_OPERATION: self.execute_statement_response,
+        })
+        _, request_data = args
+        self.request_data = request_data
+        self.response = response
+
+        super(BotocoreAuroraServerlessEvent, self).__init__(
+            wrapped,
+            instance,
+            args,
+            kwargs,
+            start_time,
+            response,
+            exception
+        )
+        self.resource['name'] = self._get_resource_name()
+        self.resource['type'] = self.RESOURCE_TYPE_UPDATE
+        self.OPERATION_TO_FUNC.get(self.resource['operation'], empty_func)()
+
+    def update_response(self, response):
+        """
+        Adds response data to event.
+        :param response: Response from botocore
+        """
+        super(BotocoreAuroraServerlessEvent, self).update_response(response)
+        self.RESPONSE_TO_FUNC.get(self.resource['operation'], empty_func)()
+
+    def execute_statement_op(self):
+        """
+        Handles execute statement operation.
+        """
+        self.resource['metadata']['database'] = self.request_data.get(
+            self.DATABASE_FIELD
+        )
+        self.resource['metadata']['Secret Arn'] = self.request_data.get(
+            'secretArn'
+        )
+        add_data_if_needed(
+            self.resource['metadata'], 'sql', self.request_data.get('sql')
+        )
+
+    def execute_statement_response(self):
+        """
+        Handles execute statement response.
+        """
+        add_data_if_needed(
+            self.resource['metadata'], 'records', self.response.get('records')
+        )
+        self.resource['metadata']['number of records updated'] = (
+            self.response.get('numberOfRecordsUpdated')
+        )
+
+
 class BotocoreEventFactory(object):
     """
     Factory class, generates botocore event.
@@ -1866,7 +1955,6 @@ def create_event(wrapped, instance, args, kwargs, start_time, response,
             instance_type,
             None
         )
-
         if event_class is not None:
             event = event_class(
                 wrapped,