430fe04202
Enhance Druid activity by adding configuration for Druid host and port in values.yaml. Refactor data loading logic to utilize SQLAlchemy for querying Druid, improving query efficiency and readability. Update timestamp handling to streamline data retrieval process.
71 lines
2.3 KiB
Python
71 lines
2.3 KiB
Python
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.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']
|
|
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.engine)
|
|
stmt = select(places).where(text(query))
|
|
|
|
result = pd.read_sql(stmt, self.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()
|