SIENTIAPDE-1110

Update Kafka activity to use AIOKafkaConsumer for asynchronous message handling, enhancing consumer management and adding support for multiple topics.

scouter/activities/kafka.py

@@ -6,9 +6,10 @@ with workflow.unsafe.imports_passed_through():
     from sientia_do.temporal.activities.base import BaseActivity
     from sientia_do.temporal.utils.logger import Logger
     from typing import Any
-    from kafka import KafkaConsumer
+    from aiokafka import AIOKafkaConsumer
     from pandas import DataFrame
     import json
+    import asyncio
 
 
 class Kafka(BaseActivity):
@@ -17,24 +18,32 @@ class Kafka(BaseActivity):
         self.polling_time = polling_time
         self.bootstrap_servers = bootstrap_servers
         self.group_id = group_id
-
-        self.kafka_connector = KafkaConsumer(
-            bootstrap_servers=bootstrap_servers,
-            auto_offset_reset="earliest",
-            enable_auto_commit=True,
-            group_id=group_id,
-            value_deserializer=lambda x: json.loads(x.decode("utf-8"))
-        )
-
+        self.consumers = {}
+        self._consumer_tasks = {}
         BaseActivity.__init__(self, logger, notification_handler)
 
-    def close(self):
-        """Closes the connector connection."""
-        self.info("Closing Kafka connector...")
-        self.kafka_connector.close()
+    async def close(self):
+        """Closes all consumer connections."""
+        self.info("Closing Kafka connectors...")
+        for _topic, consumer in self.consumers.items():
+            await consumer.stop()
 
-    def __del__(self):
-        self.close()
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc, tb):
+        await self.close()
+
+    async def create_consumer(self, topic: str):
+        consumer = AIOKafkaConsumer(
+            bootstrap_servers=self.bootstrap_servers,
+            auto_offset_reset="earliest",
+            enable_auto_commit=True,
+            group_id=f"{self.group_id}-{topic}",
+            value_deserializer=lambda x: json.loads(x.decode("utf-8"))
+        )
+        await consumer.start()
+        self.consumers[topic] = consumer
 
     @activity.defn(name="load_from_kafka")
     async def load_from_kafka(self, input_data: dict[str, Any]) -> dict[str, Any]:
@@ -57,28 +66,25 @@ class Kafka(BaseActivity):
 
         topic = input_data["topic"]
 
-        # Subscribe to the specified topic
-        self.kafka_connector.subscribe([topic])
+        if topic not in self.consumers:
+            await self.create_consumer(topic)
+
+        consumer = self.consumers[topic]
+
+        await consumer.subscribe([topic])
 
-        # List to store message values
         message_values = []
 
-        # Poll for messages
-        records = self.kafka_connector.poll(timeout_ms=self.polling_time)
-
-        self.debug(
-            f"Polled {len(records)} records from topic: {topic}",
-            metadata=metadata
-        )
-
-        # Process the polled records
-        for _topic_partition, msgs in records.items():
-            for msg in msgs:
+        end_time = asyncio.get_event_loop().time() + (self.polling_time / 1000)
+        while asyncio.get_event_loop().time() < end_time:
+            try:
+                msg = await asyncio.wait_for(consumer.getone(), timeout=(end_time - asyncio.get_event_loop().time()))
                 message_values.append(msg.value)
-
-        # Return empty dict if no messages were received
-        if not message_values:
-            return {}
+            except asyncio.TimeoutError:
+                break
+            except Exception as e:
+                self.error(f"Error while consuming: {e}", metadata=metadata)
+                break
 
         self.debug(
             f"Loaded {len(message_values)} messages from topic: {topic}",
@@ -90,6 +96,9 @@ class Kafka(BaseActivity):
             metadata=metadata
         )
 
-        self.kafka_connector.unsubscribe()
+        await consumer.unsubscribe()
+
+        if not message_values:
+            return {}
 
         return DataFrame(message_values).to_dict()