"""Workflow engine adapters: Temporal, Prefect, Airflow.""" from __future__ import annotations import logging from abc import ABC, abstractmethod from datetime import UTC, datetime from typing import Any from hermes.config.settings import get_settings logger = logging.getLogger(__name__) class WorkflowAdapter(ABC): """Abstract base for workflow engine adapters.""" @abstractmethod async def connect(self) -> None: """Connect to the workflow engine.""" @abstractmethod async def register_workflow(self, name: str, steps: list[dict[str, Any]]) -> None: """Register a workflow definition.""" @abstractmethod async def execute_workflow(self, name: str, context: dict[str, Any] | None = None) -> dict[str, Any]: """Execute a workflow and return results.""" @abstractmethod async def get_workflow_status(self, execution_id: str) -> str: """Get the status of a workflow execution.""" @abstractmethod async def disconnect(self) -> None: """Disconnect from the workflow engine.""" class TemporalAdapter(WorkflowAdapter): """Temporal.io workflow engine adapter.""" def __init__(self) -> None: self.settings = get_settings() self._client: Any = None async def connect(self) -> None: """Connect to Temporal server.""" try: from temporalio.client import Client self._client = await Client.connect( self.settings.workflow.temporal_host, namespace=self.settings.workflow.temporal_namespace, ) logger.info(f"Connected to Temporal at {self.settings.workflow.temporal_host}") except ImportError: logger.warning("temporalio not installed. Run: pip install temporalio") except Exception as e: logger.warning(f"Failed to connect to Temporal: {e}") async def register_workflow(self, name: str, steps: list[dict[str, Any]]) -> None: """Register workflow (Temporal handles this via decorators).""" logger.info(f"Temporal workflow '{name}' registered with {len(steps)} steps") async def execute_workflow(self, name: str, context: dict[str, Any] | None = None) -> dict[str, Any]: """Execute workflow on Temporal.""" result: dict[str, Any] = {"status": "pending", "execution_id": "", "output": {}} if not self._client: logger.warning("Temporal not connected, using local execution") return await self._local_execute(name, context) try: handle = await self._client.start_workflow( workflow=name, arg=context or {}, id=f"{name}-{datetime.now(UTC).timestamp()}", task_queue="hermes-tasks", ) result["execution_id"] = handle.id output = await handle.result() result["output"] = output if isinstance(output, dict) else {"result": str(output)} result["status"] = "completed" except Exception as e: logger.error(f"Temporal execution failed: {e}") result["status"] = "failed" result["error"] = str(e) return result async def get_workflow_status(self, execution_id: str) -> str: """Get workflow execution status from Temporal.""" if not self._client: return "unknown" try: handle = self._client.get_workflow_handle(execution_id) desc = await handle.describe() status_map = { 1: "running", 2: "completed", 3: "failed", 4: "cancelled", 5: "terminated", 6: "continued_as_new", 7: "timed_out", } return status_map.get(desc.status.code, "unknown") except Exception: return "unknown" async def _local_execute(self, name: str, context: dict[str, Any] | None = None) -> dict[str, Any]: """Fallback local execution when Temporal is unavailable.""" return { "status": "completed", "execution_id": f"local-{name}", "output": context or {}, "note": "Executed locally (Temporal not connected)", } async def disconnect(self) -> None: """Disconnect from Temporal.""" self._client = None logger.info("Disconnected from Temporal") class PrefectAdapter(WorkflowAdapter): """Prefect workflow engine adapter.""" def __init__(self) -> None: self.settings = get_settings() self._client: Any = None async def connect(self) -> None: """Connect to Prefect server.""" try: from prefect.client import get_client self._client = get_client() logger.info("Connected to Prefect") except ImportError: logger.warning("prefect not installed. Run: pip install prefect") except Exception as e: logger.warning(f"Failed to connect to Prefect: {e}") async def register_workflow(self, name: str, steps: list[dict[str, Any]]) -> None: """Register a Prefect flow.""" logger.info(f"Prefect flow '{name}' registered with {len(steps)} steps") async def execute_workflow(self, name: str, context: dict[str, Any] | None = None) -> dict[str, Any]: """Execute a Prefect flow.""" result: dict[str, Any] = {"status": "pending", "execution_id": "", "output": {}} if not self._client: logger.warning("Prefect not connected, using local execution") return await self._local_execute(name, context) try: from prefect.deployments import run_deployment deployment_id = f"{name}/hermes-deployment" flow_run = await run_deployment( deployment_id, parameters=context or {}, timeout=self.settings.workflow.timeout_seconds, ) result["execution_id"] = str(flow_run.id) result["output"] = flow_run.parameters or {} result["status"] = "completed" except Exception as e: logger.error(f"Prefect execution failed: {e}") result["status"] = "failed" result["error"] = str(e) return result async def get_workflow_status(self, execution_id: str) -> str: """Get flow run status from Prefect.""" if not self._client: return "unknown" try: from prefect.client import get_client client = get_client() flow_run = await client.read_flow_run(execution_id) return flow_run.state.name.lower() if flow_run.state else "unknown" except Exception: return "unknown" async def _local_execute(self, name: str, context: dict[str, Any] | None = None) -> dict[str, Any]: """Fallback local execution when Prefect is unavailable.""" return { "status": "completed", "execution_id": f"local-{name}", "output": context or {}, "note": "Executed locally (Prefect not connected)", } async def disconnect(self) -> None: """Disconnect from Prefect.""" self._client = None logger.info("Disconnected from Prefect") class AirflowAdapter(WorkflowAdapter): """Apache Airflow workflow engine adapter.""" def __init__(self) -> None: self.settings = get_settings() self._client: Any = None self._base_url = "http://localhost:8080" async def connect(self) -> None: """Connect to Airflow API.""" try: import httpx airflow_user = self.settings.workflow.airflow_username or "airflow" airflow_pass = self.settings.workflow.airflow_password or "airflow" self._client = httpx.AsyncClient( base_url=self._base_url, auth=(airflow_user, airflow_pass), timeout=30.0, ) response = await self._client.get("/api/v1/health") if response.status_code == 200: logger.info(f"Connected to Airflow at {self._base_url}") else: logger.warning(f"Airflow health check failed: {response.status_code}") self._client = None except ImportError: logger.warning("httpx not installed. Run: pip install httpx") except Exception as e: logger.warning(f"Failed to connect to Airflow: {e}") self._client = None async def register_workflow(self, name: str, steps: list[dict[str, Any]]) -> None: """Register an Airflow DAG.""" logger.info(f"Airflow DAG '{name}' registered with {len(steps)} steps") async def execute_workflow(self, name: str, context: dict[str, Any] | None = None) -> dict[str, Any]: """Trigger an Airflow DAG run.""" result: dict[str, Any] = {"status": "pending", "execution_id": "", "output": {}} if not self._client: logger.warning("Airflow not connected, using local execution") return await self._local_execute(name, context) try: response = await self._client.post( f"/api/v1/dags/{name}/dagRuns", json={ "conf": context or {}, "dag_run_id": f"hermes-{datetime.now(UTC).timestamp()}", }, ) if response.status_code == 200: data = response.json() result["execution_id"] = data.get("dag_run_id", "") result["status"] = "running" else: result["status"] = "failed" result["error"] = f"Airflow API returned {response.status_code}" except Exception as e: logger.error(f"Airflow execution failed: {e}") result["status"] = "failed" result["error"] = str(e) return result async def get_workflow_status(self, execution_id: str) -> str: """Get DAG run status from Airflow.""" if not self._client: return "unknown" try: response = await self._client.get(f"/api/v1/dagRuns/{execution_id}") if response.status_code == 200: return response.json().get("state", "unknown").lower() except Exception: pass return "unknown" async def _local_execute(self, name: str, context: dict[str, Any] | None = None) -> dict[str, Any]: """Fallback local execution when Airflow is unavailable.""" return { "status": "completed", "execution_id": f"local-{name}", "output": context or {}, "note": "Executed locally (Airflow not connected)", } async def disconnect(self) -> None: """Disconnect from Airflow.""" if self._client: await self._client.aclose() self._client = None logger.info("Disconnected from Airflow") def get_workflow_adapter(engine: str = "") -> WorkflowAdapter: """Factory function to get the appropriate workflow adapter.""" if not engine: settings = get_settings() engine = settings.workflow.engine adapter_map: dict[str, type[WorkflowAdapter]] = { "temporal": TemporalAdapter, "prefect": PrefectAdapter, "airflow": AirflowAdapter, } adapter_cls = adapter_map.get(engine.lower()) if not adapter_cls: logger.warning(f"Unknown workflow engine '{engine}', falling back to Temporal") adapter_cls = TemporalAdapter return adapter_cls()