3bbc6a6a79
Update Druid activity to handle None values for last data timestamp, ensuring default timestamp is set to epoch start. Enhance data loading debug logging for improved traceability.
73 lines
2.5 KiB
Python
73 lines
2.5 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']
|
|
last_data_timestamp = last_data_timestamp if last_data_timestamp is not None else datetime(
|
|
1970, 1, 1, 0, 0, 0)
|
|
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()
|