208491498f
Update Druid configuration in values.yaml to change port from 8081 to 8888. Refactor pydruid.py to use druid_engine for SQLAlchemy connections, enhancing clarity and consistency in data loading operations.
72 lines
2.5 KiB
Python
72 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.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()
|