332612ac84
Integrate Druid activity into the Activities class, adding support for Druid configuration and initialization. Update worker and connectors configuration to accommodate Druid, enhancing data processing capabilities.
75 lines
2.3 KiB
Python
75 lines
2.3 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 = datetime.strptime(
|
|
input_data['last_data_timestamp'], "%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")
|