e7b05ebf3a
Add logging for consumer creation in Kafka activity to enhance traceability of topic management.
106 lines
3.2 KiB
Python
106 lines
3.2 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:
|
|
self.info(
|
|
f"Creating consumer for topic: {topic}",
|
|
metadata=metadata
|
|
)
|
|
await self.create_consumer(topic)
|
|
|
|
consumer = self.consumers[topic]
|
|
|
|
consumer.subscribe(topics=[topic])
|
|
|
|
message_values = []
|
|
|
|
messages = await consumer.getmany(timeout_ms=self.polling_time)
|
|
|
|
for tp, msgs in messages.items():
|
|
msg_topic = tp.topic
|
|
if msg_topic == topic:
|
|
for msg in msgs:
|
|
message_values.append(msg.value)
|
|
|
|
self.info(
|
|
f"Loaded {len(message_values)} messages from topic: {topic}",
|
|
metadata=metadata
|
|
)
|
|
|
|
self.debug(
|
|
f"Loaded data: {message_values}",
|
|
metadata=metadata
|
|
)
|
|
|
|
consumer.unsubscribe()
|
|
|
|
if not message_values:
|
|
return {}
|
|
|
|
return DataFrame(message_values).to_dict()
|