feat(kafka.py): add support for kafka-python producer (#321)

Author: ranrib

README.md

@@ -461,8 +461,9 @@ Epsagon provides out-of-the-box instrumentation (tracing) for many popular frame
 |azure.cosmos        |`>=4.0.0`                  |
 |celery              |`>=4.0.0`                  |
 |grpc                |`>=0.3-10`                 |
-|greengrasssdk       |`>=1.4.0`                 |
-|SQLAlchemy                |`>=1.2.0`                 |
+|greengrasssdk       |`>=1.4.0`                  |
+|SQLAlchemy          |`>=1.2.0`                  |
+|kafka-python        |`>=1.4.0`                  |
 
 
 

epsagon/events/kafka.py

@@ -0,0 +1,117 @@
+"""
+kafka-python events.
+"""
+
+from __future__ import absolute_import
+import traceback
+from uuid import uuid4
+
+from epsagon.utils import add_data_if_needed
+from ..trace import trace_factory
+from ..event import BaseEvent
+from ..constants import EPSAGON_HEADER
+
+
+class KafkaEvent(BaseEvent):
+    """
+    Represents base Kafka event.
+    """
+
+    ORIGIN = 'kafka'
+    RESOURCE_TYPE = 'kafka'
+
+    # 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(KafkaEvent, self).__init__(start_time)
+        self.event_id = 'kafka-{}'.format(str(uuid4()))
+
+        topic = args[0]
+        headers = dict(kwargs['headers'])
+        servers = instance.config['bootstrap_servers']
+        if servers and isinstance(servers, list):
+            # Take the first server if it is a list
+            servers = servers[0]
+
+        self.resource['name'] = topic
+        self.resource['operation'] = 'send'
+        self.resource['metadata'] = {
+            'messaging.system': 'kafka',
+            'messaging.destination': topic,
+            'messaging.url': servers,
+            'messaging.message_payload_size_bytes': (
+                len(str(kwargs.get('value', '')))
+            ),
+        }
+        if instance.config.get('client_id'):
+            self.resource['metadata']['messaging.kafka.client_id'] = (
+                instance.config['client_id']
+            )
+        if headers.get(EPSAGON_HEADER):
+            self.resource['metadata'][EPSAGON_HEADER] = (
+                headers[EPSAGON_HEADER]
+            )
+        if kwargs['key']:
+            self.resource['metadata']['messaging.kafka.message_key'] = (
+                kwargs['key']
+            )
+
+        add_data_if_needed(
+            self.resource['metadata'],
+            'messaging.headers',
+            headers
+        )
+
+        add_data_if_needed(
+            self.resource['metadata'],
+            'messaging.message',
+            kwargs['value']
+        )
+
+        if getattr(response, 'value', None) is not None:
+            self.update_response(response.value)
+
+        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']['messaging.kafka.partition'] = (
+            response.partition
+        )
+
+
+class KafkaEventFactory(object):
+    """
+    Factory class, generates a kafka event.
+    """
+
+    @staticmethod
+    def create_event(wrapped, instance, args, kwargs, start_time, response,
+                     exception):
+        """Create an event"""
+        event = KafkaEvent(
+            wrapped,
+            instance,
+            args,
+            kwargs,
+            start_time,
+            response,
+            exception
+        )
+
+        trace_factory.add_event(event)

epsagon/modules/kafka.py

@@ -0,0 +1,54 @@
+"""
+kafka-python patcher module.
+"""
+
+from __future__ import absolute_import
+import wrapt
+from epsagon.modules.general_wrapper import wrapper
+from ..events.kafka import KafkaEventFactory
+from ..constants import EPSAGON_HEADER
+from ..utils import get_epsagon_http_trace_id
+
+
+def _parse_args(
+    topic,
+    value=None,
+    key=None,
+    headers=None,
+    partition=None,
+    timestamp_ms=None
+):
+    """Sort and return args and kwargs according to the original signature"""
+    return (topic, ), {
+        'value': value,
+        'key': key,
+        'headers': headers,
+        'partition': partition,
+        'timestamp_ms': timestamp_ms
+    }
+
+
+def _wrapper(wrapped, instance, args, kwargs):
+    """KafkaProducer.send wrapper"""
+    new_args, new_kwargs = _parse_args(*args, **kwargs)
+
+    # Adds epsagon header
+    if not new_kwargs.get('headers'):
+        new_kwargs['headers'] = []
+    new_kwargs['headers'].append(
+        (EPSAGON_HEADER, get_epsagon_http_trace_id().encode())
+    )
+
+    return wrapper(KafkaEventFactory, wrapped, instance, new_args, new_kwargs)
+
+
+def patch():
+    """
+    patch module.
+    :return: None
+    """
+    wrapt.wrap_function_wrapper(
+        'kafka.producer.kafka',
+        'KafkaProducer.send',
+        _wrapper
+    )

epsagon/trace.py

@@ -53,6 +53,8 @@ def _decimal_serializer(o):
     """ Taken from `bootstrap.py` of AWS Lambda Python runtime """
     if isinstance(o, decimal.Decimal):
         return _number_str(o)
+    if isinstance(o, bytes):
+        return o.decode('utf-8')
     raise TypeError(repr(o) + ' is not JSON serializable')
 
 
@@ -552,7 +554,6 @@ def send_traces(self, trace=None):
                 if not trace_sent:
                     self.pop_trace(trace=trace)
 
-
     def prepare(self):
         """
         Prepare the relevant trace.

epsagon/trace_encoder.py

@@ -14,6 +14,8 @@ def default(self, o):  # pylint: disable=method-hidden
             return list(o)
         if isinstance(o, (datetime, date)):
             return o.isoformat()
+        if isinstance(o, bytes):
+            return o.decode('utf-8')
 
         output = repr(o)
         try:

requirements-dev.txt

@@ -20,3 +20,4 @@ httpx; python_version >= '3.5'
 asynctest; python_version >= '3.5'
 moto
 tornado
+kafka-python

tests/events/test_kafka.py

@@ -0,0 +1,51 @@
+import json
+import epsagon.wrappers.python_function
+import epsagon.runners.python_function
+import epsagon.constants
+import mock
+from kafka import KafkaProducer
+
+TEST_URL = 'https://example.test/'
+
+def test(*args, **kwargs):
+    return [{}, False, False]
+
+
+@mock.patch('epsagon.trace.TraceFactory.add_event')
+@mock.patch('kafka.producer.kafka.KafkaProducer._wait_on_metadata')
+@mock.patch('kafka.producer.kafka.KafkaProducer._partition')
+@mock.patch('kafka.producer.record_accumulator.RecordAccumulator.append', side_effect=test)
+def test_sanity(append_mock, partition_mock, wait_on_metadata_mock, add_event_mock):
+    retval = 'success'
+    body = {'test': 1}
+
+    @epsagon.wrappers.python_function.python_wrapper
+    def wrapped_function():
+        producer = KafkaProducer(
+            bootstrap_servers=['host:10'],
+            client_id='test_client_id',
+            api_version=(0, 11, 5),
+            value_serializer=lambda x: json.dumps(x).encode('ascii'),
+        )
+        response = producer.send('topic', body, headers=[('content-encoding', b'base64')])
+        return retval
+    assert wrapped_function() == retval
+    wait_on_metadata_mock.assert_called()
+    partition_mock.assert_called()
+    append_mock.assert_called()
+    add_event_mock.assert_called()
+    event = add_event_mock.call_args_list[0].args[0]
+    assert event.resource['name'] == 'topic'
+    assert event.resource['operation'] == 'send'
+    assert event.resource['type'] == 'kafka'
+    assert event.resource['metadata']['messaging.kafka.client_id'] == (
+        'test_client_id'
+    )
+    assert event.resource['metadata'][
+       'messaging.message_payload_size_bytes'
+   ] == len(str(body))
+    assert event.resource['metadata']['messaging.message'] == body
+    assert (
+        epsagon.constants.EPSAGON_HEADER in
+        event.resource['metadata']['messaging.headers']
+    )