704099c387
Update Kafka activity to use AIOKafkaConsumer for asynchronous message handling, enhancing consumer management and adding support for multiple topics.
105 lines
3.4 KiB
Python
105 lines
3.4 KiB
Python
from temporalio import workflow, activity
|
|
|
|
with workflow.unsafe.imports_passed_through():
|
|
from logging import Logger
|
|
from sientia_do.notifications.handlers import NotificationHandler
|
|
from sientia_do.temporal.activities.base import BaseActivity
|
|
from sientia_do.temporal.utils.logger import Logger
|
|
from typing import Any
|
|
from aiokafka import AIOKafkaConsumer
|
|
from pandas import DataFrame
|
|
import json
|
|
import asyncio
|
|
|
|
|
|
class Kafka(BaseActivity):
|
|
def __init__(self, bootstrap_servers: str, polling_time: int,
|
|
group_id: str, logger: Logger, notification_handler: NotificationHandler):
|
|
self.polling_time = polling_time
|
|
self.bootstrap_servers = bootstrap_servers
|
|
self.group_id = group_id
|
|
self.consumers = {}
|
|
self._consumer_tasks = {}
|
|
BaseActivity.__init__(self, logger, notification_handler)
|
|
|
|
async def close(self):
|
|
"""Closes all consumer connections."""
|
|
self.info("Closing Kafka connectors...")
|
|
for _topic, consumer in self.consumers.items():
|
|
await consumer.stop()
|
|
|
|
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]:
|
|
"""
|
|
Loads data from a kafka topic. Polls the topic for a given time and returns the data.
|
|
|
|
Args:
|
|
input_data (dict[str, Any]): The data to load. Contains:
|
|
topic (str): The topic to load data from.
|
|
Returns:
|
|
dict[str, Any]: The data loaded from the topic.
|
|
"""
|
|
|
|
metadata = input_data['metadata']
|
|
|
|
self.debug(
|
|
f"Loading data from topic: {input_data['topic']}",
|
|
metadata=metadata
|
|
)
|
|
|
|
topic = input_data["topic"]
|
|
|
|
if topic not in self.consumers:
|
|
await self.create_consumer(topic)
|
|
|
|
consumer = self.consumers[topic]
|
|
|
|
await consumer.subscribe([topic])
|
|
|
|
message_values = []
|
|
|
|
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)
|
|
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}",
|
|
metadata=metadata
|
|
)
|
|
|
|
self.debug(
|
|
f"Loaded data: {message_values}",
|
|
metadata=metadata
|
|
)
|
|
|
|
await consumer.unsubscribe()
|
|
|
|
if not message_values:
|
|
return {}
|
|
|
|
return DataFrame(message_values).to_dict()
|