replace dummy hook with callable check, add hook documentation

Author: ItayGibel-helios

instrumentation/opentelemetry-instrumentation-kafka-python/src/opentelemetry/instrumentation/kafka/__init__.py

@@ -38,6 +38,35 @@
     for message in consumer:
         # process message
 
+The `_instrument` method accepts the following keyword args:
+tracer_provider (TracerProvider) - an optional tracer provider
+produce_hook (Callable) - a function with extra user-defined logic to be performed before sending the message
+                          this function signature is:
+                          def produce_hook(span: Span, args, kwargs)
+consume_hook (Callable) - a function with extra user-defined logic to be performed after consuming a message
+                          this function signature is:
+                          def consume
+                          _hook(span: Span, record: , kafka.record.ABCRecord, args, kwargs)
+for example:
+.. code: python
+    from opentelemetry.instrumentation.kafka import KafkaInstrumentor
+    from kafka import KafkaProducer, KafkaConsumer
+
+    def produce_hook(span, args, kwargs):
+        if span and span.is_recording():
+            span.set_attribute("custom_user_attribute_from_produce_hook", "some-value")
+    def consume_hook(span, record, args, kwargs):
+        if span and span.is_recording():
+            span.set_attribute("custom_user_attribute_from_consume_hook", "some-value")
+
+    # instrument kafka with produce and consume hooks
+    KafkaInstrumentor().instrument(produce_hook=produce_hook, consume_hook=consume_hook)
+
+    # Using kafka as normal now will automatically generate spans,
+    # including user custom attributes added from the hooks
+    producer = KafkaProducer(bootstrap_servers=['localhost:9092'])
+    producer.send('my-topic', b'raw_bytes')
+
 API
 ___
 """
@@ -52,7 +81,6 @@
 from opentelemetry.instrumentation.kafka.utils import (
     _wrap_next,
     _wrap_send,
-    dummy_callback,
 )
 from opentelemetry.instrumentation.kafka.version import __version__
 from opentelemetry.instrumentation.utils import unwrap
@@ -76,8 +104,8 @@ def _instrument(self, **kwargs):
                 ``consume_hook``: a callable to be executed just after consuming a message
         """
         tracer_provider = kwargs.get("tracer_provider")
-        produce_hook = kwargs.get("produce_hook", dummy_callback)
-        consume_hook = kwargs.get("consume_hook", dummy_callback)
+        produce_hook = kwargs.get("produce_hook")
+        consume_hook = kwargs.get("consume_hook")
 
         tracer = trace.get_tracer(
             __name__, __version__, tracer_provider=tracer_provider

instrumentation/opentelemetry-instrumentation-kafka-python/src/opentelemetry/instrumentation/kafka/utils.py

@@ -1,6 +1,7 @@
 import json
 from logging import getLogger
 from typing import Callable, Dict, List, Optional
+from kafka.record.abc import ABCRecord
 
 from opentelemetry import context, propagate, trace
 from opentelemetry.propagators import textmap
@@ -76,11 +77,8 @@ def extract_send_partition(instance, args, kwargs):
         )
 
 
-HookT = Callable[[Span, List, Dict], None]
-
-
-def dummy_callback(span, args, kwargs):
-    ...
+ProduceHookT = Optional[Callable[[Span, List, Dict], None]]
+ConsumeHookT = Optional[Callable[[Span, ABCRecord, List, Dict], None]]
 
 
 class KafkaContextGetter(textmap.Getter):
@@ -130,7 +128,7 @@ def _get_span_name(operation: str, topic: str):
     return f"{topic} {operation}"
 
 
-def _wrap_send(tracer: Tracer, produce_hook: HookT) -> Callable:
+def _wrap_send(tracer: Tracer, produce_hook: ProduceHookT) -> Callable:
     def _traced_send(func, instance, args, kwargs):
         headers = KafkaPropertiesExtractor.extract_send_headers(args, kwargs)
         if headers is None:
@@ -155,7 +153,8 @@ def _traced_send(func, instance, args, kwargs):
                 setter=_kafka_setter,
             )
             try:
-                produce_hook(span, args, kwargs)
+                if callable(produce_hook):
+                    produce_hook(span, args, kwargs)
             except Exception as hook_exception:  # pylint: disable=W0703
                 _LOG.exception(hook_exception)
 
@@ -166,7 +165,7 @@ def _traced_send(func, instance, args, kwargs):
 
 def _wrap_next(
     tracer: Tracer,
-    consume_hook: HookT,
+    consume_hook: ConsumeHookT,
 ) -> Callable:
     def _traced_next(func, instance, args, kwargs):
 
@@ -194,7 +193,8 @@ def _traced_next(func, instance, args, kwargs):
                     span, bootstrap_servers, record.topic, record.partition
                 )
                 try:
-                    consume_hook(span, args, kwargs)
+                    if callable(consume_hook):
+                        consume_hook(span, record, args, kwargs)
                 except Exception as hook_exception:  # pylint: disable=W0703
                     _LOG.exception(hook_exception)
                 context.detach(token)