From f103f3a070a497334c4f10495cd33ec902173b74 Mon Sep 17 00:00:00 2001
From: sagivr2020 <72692043+sagivr2020@users.noreply.github.com>
Date: Tue, 23 Aug 2022 16:04:43 +0300
Subject: [PATCH] feat(pyqldb): add support for qldb driver cs 273 (#417)

---
 epsagon/events/pyqldb.py        | 100 ++++++++++++++++++++++++++++++++
 epsagon/modules/pyqldb.py       |  31 ++++++++++
 scripts/run_acceptance_tests.sh |   2 +-
 3 files changed, 132 insertions(+), 1 deletion(-)
 create mode 100644 epsagon/events/pyqldb.py
 create mode 100644 epsagon/modules/pyqldb.py

diff --git a/epsagon/events/pyqldb.py b/epsagon/events/pyqldb.py
new file mode 100644
index 00000000..9cbf28e1
--- /dev/null
+++ b/epsagon/events/pyqldb.py
@@ -0,0 +1,100 @@
+"""
+pyqldb events module.
+"""
+
+from __future__ import absolute_import
+from uuid import uuid4
+import traceback
+
+from epsagon.utils import add_data_if_needed
+from ..event import BaseEvent
+from ..trace import trace_factory
+
+
+class QldbEvent(BaseEvent):
+    """
+    Represents base pyqldb event.
+    """
+
+    ORIGIN = 'qldb'
+    RESOURCE_TYPE = 'qldb'
+
+    #pylint: disable=W0613
+    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)
+        """
+        super(QldbEvent, self).__init__(start_time)
+
+        self.event_id = 'qldb-{}'.format(str(uuid4()))
+        self.resource['name'] = \
+            getattr(instance.__getattribute__('_transaction')._session,# pylint: disable=W0212
+                                                '_ledger_name')
+        self.resource['operation'] = wrapped.__func__.__name__
+
+        self.resource['metadata']['query'] = args[0]
+        add_data_if_needed(self.resource['metadata'], 'parameters',
+                            [args[i] for i in range(1, len(args))])
+
+        add_data_if_needed(self.resource['metadata'], 'transaction_id',
+                            getattr(instance, 'transaction_id'))
+
+        if response is not None:
+            self.update_response(response)
+
+        if exception is not None:
+            self.set_exception(exception, traceback.format_exc())
+
+
+    def update_response(self, response):
+        """
+        Adds response data to event.
+        :param response: Response from botocore
+        :return: None
+        """
+
+        self.resource['metadata']['Results'] = [str(x) for x in response]
+        self.resource['metadata']['response.consumed_information'] = \
+            response.get_consumed_ios()
+        self.resource['metadata']['response.timing_information'] = \
+            response.get_timing_information()
+
+
+
+class QldbEventFactory(object):
+    """
+    Factory class, generates Qldb event.
+    """
+
+    @staticmethod
+    def create_event(wrapped, instance, args, kwargs, start_time, response,
+                     exception):
+        """
+        Create a Qldb event.
+        :param wrapped:
+        :param instance:
+        :param args:
+        :param kwargs:
+        :param start_time:
+        :param response:
+        :param exception:
+        :return:
+        """
+        event = QldbEvent(
+            wrapped,
+            instance,
+            args,
+            kwargs,
+            start_time,
+            response,
+            exception
+        )
+        trace_factory.add_event(event)
diff --git a/epsagon/modules/pyqldb.py b/epsagon/modules/pyqldb.py
new file mode 100644
index 00000000..24dba04c
--- /dev/null
+++ b/epsagon/modules/pyqldb.py
@@ -0,0 +1,31 @@
+"""
+pyqldb patcher module
+"""
+from __future__ import absolute_import
+import wrapt
+from epsagon.modules.general_wrapper import wrapper
+from ..events.pyqldb import QldbEventFactory
+
+
+def _wrapper(wrapped, instance, args, kwargs):
+    """
+    General wrapper for Pyqldb instrumentation.
+    :param wrapped: wrapt's wrapped
+    :param instance: wrapt's instance
+    :param args: wrapt's args
+    :param kwargs: wrapt's kwargs
+    :return: None
+    """
+    return wrapper(QldbEventFactory, wrapped, instance, args, kwargs)
+
+
+def patch():
+    """
+    patch module.
+    :return: None
+    """
+    wrapt.wrap_function_wrapper(
+        'pyqldb.execution.executor',
+        'Executor.execute_statement',
+        _wrapper
+    )
diff --git a/scripts/run_acceptance_tests.sh b/scripts/run_acceptance_tests.sh
index 9c11ea45..d85a2db9 100755
--- a/scripts/run_acceptance_tests.sh
+++ b/scripts/run_acceptance_tests.sh
@@ -2,7 +2,7 @@
 if [ -z $AWS_ACCESS_KEY_ID ] || [ -z $AWS_SECRET_ACCESS_KEY ]; then
     echo "AWS credentials must be set in order to run acceptance tests"
     exit 1
-elif [ $TRAVIS_PYTHON_VERSION != "2.7" ]; then
+elif [ $TRAVIS_PYTHON_VERSION != "2.7" ] && [ $TRAVIS_PYTHON_VERSION != "3.6" ]; then
     npm install && export PATH=$(pwd)/node_modules/.bin:$PATH
     ./acceptance/run.sh $TRAVIS_BUILD_NUMBER $TRAVIS_PYTHON_VERSION
 fi