feat: support batch (getmany) in aiokafka instrumentation (#3257)

* feat: support batch (getmany) in aiokafka instrumentation

* test: fix unclosed resources and typing

* test: add test_wrap_getmany

* fix: get unique topic list in batch

* fix: update typing, run pyupgrade

* fix: remove json.dumps from SERVER_ADDRESS attribute

* fix pylint

* fix: sync span_kind with spec

* add CHANGELOG entry

* remove changes not from this issue

* move types under TYPE_CHECKING

* move CHANGELOG entry to unreleased

* enable pyright for aiokafka, fix key type

* Update CHANGELOG.md

---------

Co-authored-by: Riccardo Magliocchetti <[email protected]>
Co-authored-by: Emídio Neto <[email protected]>

Author: 3 people

CHANGELOG.md

@@ -11,6 +11,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## Unreleased
 
+### Added
+
+- `opentelemetry-instrumentation-aiokafka` Add instrumentation of `consumer.getmany` (batch)
+  ([#3257](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/3257))
+
 ### Fixed
 
 - `opentelemetry-instrumentation-system-metrics`: fix loading on Google Cloud Run

instrumentation/opentelemetry-instrumentation-aiokafka/pyproject.toml

@@ -28,6 +28,7 @@ dependencies = [
   "opentelemetry-api ~= 1.27",
   "opentelemetry-instrumentation == 0.55b0.dev",
   "opentelemetry-semantic-conventions == 0.55b0.dev",
+  "typing_extensions ~= 4.1",
 ]
 
 [project.optional-dependencies]

instrumentation/opentelemetry-instrumentation-aiokafka/src/opentelemetry/instrumentation/aiokafka/__init__.py

@@ -93,15 +93,20 @@ async def produce():
 ___
 """
 
+from __future__ import annotations
+
 from asyncio import iscoroutinefunction
-from typing import Collection
+from typing import TYPE_CHECKING, Collection
 
 import aiokafka
-from wrapt import wrap_function_wrapper
+from wrapt import (
+    wrap_function_wrapper,  # type: ignore[reportUnknownVariableType]
+)
 
 from opentelemetry import trace
 from opentelemetry.instrumentation.aiokafka.package import _instruments
 from opentelemetry.instrumentation.aiokafka.utils import (
+    _wrap_getmany,
     _wrap_getone,
     _wrap_send,
 )
@@ -110,6 +115,21 @@ async def produce():
 from opentelemetry.instrumentation.utils import unwrap
 from opentelemetry.semconv.schemas import Schemas
 
+if TYPE_CHECKING:
+    from typing import TypedDict
+
+    from typing_extensions import Unpack
+
+    from .utils import ConsumeHookT, ProduceHookT
+
+    class InstrumentKwargs(TypedDict, total=False):
+        tracer_provider: trace.TracerProvider
+        async_produce_hook: ProduceHookT
+        async_consume_hook: ConsumeHookT
+
+    class UninstrumentKwargs(TypedDict, total=False):
+        pass
+
 
 class AIOKafkaInstrumentor(BaseInstrumentor):
     """An instrumentor for kafka module
@@ -119,7 +139,7 @@ class AIOKafkaInstrumentor(BaseInstrumentor):
     def instrumentation_dependencies(self) -> Collection[str]:
         return _instruments
 
-    def _instrument(self, **kwargs):
+    def _instrument(self, **kwargs: Unpack[InstrumentKwargs]):
         """Instruments the kafka module
 
         Args:
@@ -155,7 +175,13 @@ def _instrument(self, **kwargs):
             "getone",
             _wrap_getone(tracer, async_consume_hook),
         )
+        wrap_function_wrapper(
+            aiokafka.AIOKafkaConsumer,
+            "getmany",
+            _wrap_getmany(tracer, async_consume_hook),
+        )
 
-    def _uninstrument(self, **kwargs):
+    def _uninstrument(self, **kwargs: Unpack[UninstrumentKwargs]):
         unwrap(aiokafka.AIOKafkaProducer, "send")
         unwrap(aiokafka.AIOKafkaConsumer, "getone")
+        unwrap(aiokafka.AIOKafkaConsumer, "getmany")