SIENTIAPDE-1110
Refactor Activities class to remove Kafka and Druid dependencies, simplifying initialization. Update values.yaml to set replica count to 1 for reduced resource usage. Adjust Redis activity to set TTL to None for better data retention. Remove unused Kafka and Druid activity files and their associated tests, streamlining the codebase.
This commit is contained in:
vitor-aignosi
@@ -1,26 +1,22 @@
|
||||
from temporalio import activity, workflow
|
||||
from temporalio import workflow
|
||||
|
||||
with workflow.unsafe.imports_passed_through():
|
||||
from sientia_do.temporal.activities.postgres import Postgres
|
||||
from sientia_do.notifications.handlers import NotificationHandler
|
||||
from sientia_do.temporal.utils.logger import Logger
|
||||
from scouter.activities.redis import Redis
|
||||
from scouter.activities.kafka import Kafka
|
||||
from scouter.activities.gates import Gates
|
||||
from scouter.activities.mongodb import MongoDB
|
||||
from scouter.activities.pydruid import Druid
|
||||
from typing import Any
|
||||
|
||||
|
||||
class Activities(Postgres, Redis, Kafka, Gates, MongoDB, Druid):
|
||||
class Activities(Postgres, Redis, Gates, MongoDB,):
|
||||
"""Activities class that combines multiple services with proper initialization."""
|
||||
|
||||
def __init__(self,
|
||||
postgres_config: dict[str, Any],
|
||||
redis_config: dict[str, Any],
|
||||
kafka_config: dict[str, Any],
|
||||
mongodb_config: dict[str, Any],
|
||||
druid_config: dict[str, Any],
|
||||
logger: Logger,
|
||||
notification_handler: NotificationHandler):
|
||||
|
||||
@@ -49,16 +45,6 @@ class Activities(Postgres, Redis, Kafka, Gates, MongoDB, Druid):
|
||||
password=redis_config['password']
|
||||
)
|
||||
|
||||
# Initialize Kafka
|
||||
Kafka.__init__(
|
||||
self,
|
||||
bootstrap_servers=kafka_config['bootstrap_servers'],
|
||||
polling_time=kafka_config['polling_time'],
|
||||
group_id=kafka_config['group_id'],
|
||||
logger=logger,
|
||||
notification_handler=notification_handler
|
||||
)
|
||||
|
||||
# Initialize Gates
|
||||
Gates.__init__(
|
||||
self,
|
||||
@@ -75,16 +61,6 @@ class Activities(Postgres, Redis, Kafka, Gates, MongoDB, Druid):
|
||||
notification_handler=notification_handler
|
||||
)
|
||||
|
||||
# Initialize Druid
|
||||
Druid.__init__(
|
||||
self,
|
||||
host=druid_config['host'],
|
||||
port=druid_config['port'],
|
||||
logger=logger,
|
||||
notification_handler=notification_handler
|
||||
)
|
||||
|
||||
def shutdown(self):
|
||||
Postgres.close(self)
|
||||
Kafka.close(self)
|
||||
MongoDB.close(self)
|
||||
MongoDB.shutdown(self)
|
||||
|
||||
@@ -1,105 +0,0 @@
|
||||
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()
|
||||
@@ -1,71 +0,0 @@
|
||||
from temporalio import workflow, activity
|
||||
|
||||
with workflow.unsafe.imports_passed_through():
|
||||
import pandas as pd
|
||||
from typing import List, Optional, Any
|
||||
from datetime import datetime, timedelta
|
||||
from pydruid.client import PyDruid
|
||||
from pydruid.query import QueryBuilder
|
||||
from sqlalchemy.engine import create_engine
|
||||
from sqlalchemy import MetaData, Table, select, text
|
||||
from sientia_do.temporal.activities.base import BaseActivity
|
||||
from sientia_do.notifications.handlers import NotificationHandler
|
||||
from sientia_do.temporal.utils.logger import Logger
|
||||
|
||||
|
||||
class Druid(BaseActivity):
|
||||
def __init__(self, host: str, port: int,
|
||||
logger: Logger, notification_handler: NotificationHandler):
|
||||
|
||||
self.host = host
|
||||
self.port = port
|
||||
self.druid_engine = create_engine(
|
||||
f'druid://{self.host}:{self.port}/druid/v2/sql/')
|
||||
logger.info(
|
||||
f"Druid client initialized with host: {self.host}, port: {self.port}")
|
||||
|
||||
BaseActivity.__init__(self, logger=logger,
|
||||
notification_handler=notification_handler)
|
||||
|
||||
def shutdown(self):
|
||||
self.client.close()
|
||||
|
||||
def __del__(self):
|
||||
self.shutdown()
|
||||
|
||||
@activity.defn(name="load_latest_druid_data")
|
||||
async def load_latest_druid_data(self, input_data: dict[str, Any]) -> dict[str, Any]:
|
||||
"""
|
||||
Loads the latest data from Druid.
|
||||
"""
|
||||
metadata = input_data['metadata']
|
||||
datasource = f"raw_{input_data['schedule_name']}"
|
||||
last_data_timestamp = input_data['last_data_timestamp']
|
||||
last_data_timestamp = last_data_timestamp if last_data_timestamp is not None else '1970-01-01 00:00:00'
|
||||
self.debug(
|
||||
f"Loading data from Druid: {input_data}", metadata=metadata)
|
||||
|
||||
query = f'"__time" > TIMESTAMP \'{last_data_timestamp}\''
|
||||
|
||||
self.info(
|
||||
f"Loading data from Druid: {datasource} with query: {query}"
|
||||
)
|
||||
|
||||
places = Table(datasource, MetaData(), autoload_with=self.druid_engine)
|
||||
stmt = select(places).where(text(query))
|
||||
|
||||
result = pd.read_sql(stmt, self.druid_engine)
|
||||
|
||||
result["inserted_at"] = pd.to_datetime(result["__time"]).dt.strftime(
|
||||
"%Y-%m-%d %H:%M:%S.%f")
|
||||
|
||||
result.drop(columns=["__time"], inplace=True)
|
||||
|
||||
self.info(
|
||||
f"Loaded {len(result)} rows from Druid"
|
||||
)
|
||||
|
||||
self.debug(
|
||||
f"Druid query result: {result}", metadata=metadata)
|
||||
|
||||
return result.to_dict()
|
||||
@@ -55,7 +55,7 @@ class Redis(RedisBase):
|
||||
metadata=metadata
|
||||
)
|
||||
|
||||
self.set(key, last_data_timestamp)
|
||||
self.set(key, last_data_timestamp, ttl=None)
|
||||
|
||||
return last_data_timestamp
|
||||
|
||||
|
||||
Reference in New Issue
Block a user