6caec63977
Refactor last data timestamp handling in Druid activity to use string format for default value, improving consistency in date representation.
72 lines
2.4 KiB
Python
72 lines
2.4 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 '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.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()
|