SIENTIAPDE-1110

Enhance Druid activity by adding configuration for Druid host and port in values.yaml. Refactor data loading logic to utilize SQLAlchemy for querying Druid, improving query efficiency and readability. Update timestamp handling to streamline data retrieval process.

scouter/activities/pydruid.py

@@ -6,6 +6,8 @@ with workflow.unsafe.imports_passed_through():
     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
@@ -13,15 +15,12 @@ with workflow.unsafe.imports_passed_through():
 
 class Druid(BaseActivity):
     def __init__(self, host: str, port: int,
-                 logger: Logger, notification_handler: NotificationHandler,
-                 endpoint: str = "druid/v2"):
+                 logger: Logger, notification_handler: NotificationHandler):
 
         self.host = host
         self.port = port
-        self.endpoint = endpoint
-        self.client = PyDruid(
-            f"http://{self.host}:{self.port}", {self.endpoint}
-        )
+        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}")
 
@@ -41,36 +40,25 @@ class Druid(BaseActivity):
         """
         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")
-
+        last_data_timestamp = input_data['last_data_timestamp']
         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"
+        query = f'"__time" > TIMESTAMP \'{last_data_timestamp}\''
 
         self.info(
-            f"Loading data from Druid: {datasource} with interval: {interval}"
+            f"Loading data from Druid: {datasource} with query: {query}"
         )
 
-        builder = QueryBuilder()
+        places = Table(datasource, MetaData(), autoload_with=self.engine)
+        stmt = select(places).where(text(query))
 
-        query = builder.scan(
-            {
-                "datasource": datasource,
-                "intervals": interval,
-                "columns": ["timestamp", "value", "tag"],
-                "limit": 10000,
-            }
-        )
+        result = pd.read_sql(stmt, self.engine)
 
-        result = query.export_pandas()
+        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"