879c1fdc49
Add logging for data loading in Druid activity, enhancing traceability of data retrieval intervals.
83 lines
2.6 KiB
Python
83 lines
2.6 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 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,
|
|
endpoint: str = "druid/v2"):
|
|
|
|
self.host = host
|
|
self.port = port
|
|
self.endpoint = endpoint
|
|
self.client = PyDruid(
|
|
f"http://{self.host}:{self.port}", {self.endpoint}
|
|
)
|
|
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_str = input_data['last_data_timestamp']
|
|
if last_data_timestamp_str is None:
|
|
last_data_timestamp = datetime(1970, 1, 1, 0, 0, 0)
|
|
else:
|
|
last_data_timestamp = datetime.strptime(
|
|
last_data_timestamp_str, "%Y-%m-%d %H:%M:%S.%f")
|
|
|
|
self.debug(
|
|
f"Loading data from Druid: {input_data}", metadata=metadata)
|
|
|
|
end_time = datetime(9999, 12, 31, 23, 59, 59)
|
|
|
|
interval = f"{last_data_timestamp.isoformat()}Z/{end_time.isoformat()}Z"
|
|
|
|
self.info(
|
|
f"Loading data from Druid: {datasource} with interval: {interval}"
|
|
)
|
|
|
|
builder = QueryBuilder()
|
|
|
|
query = builder.scan(
|
|
{
|
|
"datasource": datasource,
|
|
"intervals": interval,
|
|
"columns": ["timestamp", "value", "tag"],
|
|
"limit": 10000,
|
|
}
|
|
)
|
|
|
|
result = query.export_pandas()
|
|
|
|
self.info(
|
|
f"Loaded {len(result)} rows from Druid"
|
|
)
|
|
|
|
self.debug(
|
|
f"Druid query result: {result}", metadata=metadata)
|
|
|
|
return result.to_dict(orient="records")
|