""" Execution Engine for Advanced Data Explorer. Validates an ExecutionPlan, fetches the required data, dispatches operations sequentially via OperationDispatcher, and packages results. """ import time import pandas as pd from typing import Any, Dict, List, Optional from dataclasses import dataclass, field from datetime import datetime from core import config from core.db_connector import get_db_connector from core.cached_queries import cached_get_pivot_data from analysis.operations import OperationDispatcher, OperationOutput # ── Result Dataclasses ─────────────────────────────────────────── @dataclass class ExecutionResult: """Complete result of executing a plan.""" success: bool query_type: str explanation: str source_data: Optional[pd.DataFrame] = None operation_results: List[OperationOutput] = field(default_factory=list) errors: List[str] = field(default_factory=list) warnings: List[str] = field(default_factory=list) execution_time_seconds: float = 0.0 table_used: str = "" row_count: int = 0 # ── Table label mapping (for display) ─────────────────────────── TABLE_LABELS = { config.TABLES['raw']: 'Raw (10Hz)', config.TABLES['agg_1sec']: '1-second aggregates', config.TABLES['agg_15sec']: '15-second aggregates', } # ── Execution Engine ───────────────────────────────────────────── class ExecutionEngine: """ Validate an ExecutionPlan, fetch data, run operations, package results. Usage: engine = ExecutionEngine() result = engine.execute(plan) """ def __init__(self): self.dispatcher = OperationDispatcher() def execute(self, plan: Dict) -> ExecutionResult: """Execute a validated plan end-to-end.""" t0 = time.time() errors = [] warnings = [] # ── 1. Validate plan structure ── qt = plan.get("query_type", "FETCH") explanation = plan.get("explanation", "") dr = plan.get("data_requirements", {}) operations = plan.get("operations", []) sensors = dr.get("sensors", []) start_time = dr.get("start_time") end_time = dr.get("end_time") resolution = dr.get("resolution", "auto") if not sensors: return self._fail("No sensors specified in plan", qt, explanation, t0) if not start_time or not end_time: return self._fail("Missing start_time or end_time in plan", qt, explanation, t0) # Ensure datetimes are datetime objects if isinstance(start_time, str): start_time = datetime.fromisoformat(start_time) if isinstance(end_time, str): end_time = datetime.fromisoformat(end_time) # ── 2. Resolve table from resolution ── table_override = None if resolution in config.RESOLUTION_MAP: table_override = config.RESOLUTION_MAP[resolution] elif resolution == "auto": # Let the DB connector decide via smart routing table_override = None else: warnings.append(f"Unknown resolution '{resolution}', using auto routing.") # ── 3. Safety: raw table time-range cap ── if table_override == config.TABLES['raw']: duration_h = (end_time - start_time).total_seconds() / 3600 if duration_h > 1.0: return self._fail( f"Raw resolution limited to 1 hour. Requested {duration_h:.1f} hours. " f"Use '1sec' or '15sec' for longer ranges.", qt, explanation, t0, ) # ── 4. Fetch data ── try: db = get_db_connector() df = cached_get_pivot_data( tag_names=tuple(sensors), start_time=start_time, end_time=end_time, table_override=table_override, ) except Exception as e: return self._fail(f"Data fetch failed: {e}", qt, explanation, t0) if df is None or df.empty: return self._fail( "No data returned for the specified sensors and time range. " "Check that the sensors exist and the time range is within March 14 – September 25, 2025.", qt, explanation, t0, ) # Determine which table was actually used (for display) if table_override: table_used = table_override else: table_used = db._select_table(start_time, end_time) table_label = TABLE_LABELS.get(table_used, table_used) # Check which sensors actually came back returned_sensors = [s for s in sensors if s in df.columns] missing = [s for s in sensors if s not in df.columns and s != "timestamp"] if missing: warnings.append(f"Sensors not found in data: {', '.join(missing)}") # ── 5. Run operations sequentially ── op_results = [] prev_result = None for i, op_spec in enumerate(operations): op_name = op_spec.get("op", "") op_params = op_spec.get("params", {}) op_label = op_spec.get("label", op_name) try: output = self.dispatcher.run( op_name=op_name, df=df, params=op_params, prev_result=prev_result, db_connector=get_db_connector(), ) op_results.append(output) prev_result = output # Check for operation-level errors if output.metadata.get("error"): warnings.append(f"Operation '{op_label}': {output.metadata['error']}") except Exception as e: error_msg = f"Operation '{op_label}' failed: {e}" errors.append(error_msg) # Create a placeholder output so downstream chaining can detect the failure op_results.append(OperationOutput( op_name=op_name, label=f"{op_label} (FAILED)", data=None, metadata={"error": str(e)}, )) prev_result = op_results[-1] # ── 6. Package result ── elapsed = time.time() - t0 success = len(errors) == 0 return ExecutionResult( success=success, query_type=qt, explanation=explanation, source_data=df, operation_results=op_results, errors=errors, warnings=warnings, execution_time_seconds=round(elapsed, 2), table_used=table_label, row_count=len(df), ) @staticmethod def _fail(msg: str, qt: str, explanation: str, t0: float) -> ExecutionResult: """Return a failed ExecutionResult.""" return ExecutionResult( success=False, query_type=qt, explanation=explanation, errors=[msg], execution_time_seconds=round(time.time() - t0, 2), )