c3c8e1b766
Update Druid activity to handle None values for last data timestamp, ensuring robust date parsing. Bump image tag in values.yaml to 0.2.3 for version consistency.
79 lines
2.5 KiB
Python
79 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 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"
|
|
|
|
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")
|