Streaming execution events

"""
Stream Flyte node execution events and alert to Slack when:
 - A node enters FAILED status
 - A node has been QUEUED longer than QUEUED_THRESHOLD_SEC (default 300s)

Usage:
  python get_node_events.py

  export SLACK_WEBHOOK_URL="https://hooks.slack.com/services/..."
  export QUEUED_THRESHOLD_SEC=600  # Alert if queued > 10 min
  python get_node_events.py
"""

import asyncio
import json
import os
import sys
import time
import urllib.request

from union.remote import UnionRemote
from flytekit.configuration import Config
from flytekit.models.core.execution import WorkflowExecutionPhase

# Slack webhook (override via SLACK_WEBHOOK_URL or ALERT_WEBHOOK_URL env)
SLACK_WEBHOOK_URL = ""
QUEUED_THRESHOLD_SEC = int(os.getenv("QUEUED_THRESHOLD_SEC", "300"))
QUEUED_CHECK_INTERVAL_SEC = 60

# (execution_name, node_id) -> (first_seen_timestamp, project, domain)
queued_since = {}


def _post_webhook_sync(url: str, payload: dict) -> None:
    data = json.dumps(payload).encode("utf-8")
    req = urllib.request.Request(
        url, data=data,
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    with urllib.request.urlopen(req, timeout=10) as resp:
        resp.read()


def _get_webhook_url():
    return os.getenv("SLACK_WEBHOOK_URL") or os.getenv("ALERT_WEBHOOK_URL") or SLACK_WEBHOOK_URL


async def alert_on_failed(*, event_line: str, execution_name: str, node_id: str | None = None):
    if "FAILED" not in event_line:
        return
    node_info = f" Node: `{node_id}`" if node_id else ""
    print(f"[ALERT] Node FAILED: execution={execution_name}{node_info} - {event_line}", file=sys.stderr)
    webhook_url = _get_webhook_url()
    if not webhook_url:
        return
    payload = {
        "text": f"Flyte node execution *FAILED*{node_info}\nExecution: `{execution_name}`\n{event_line}",
    }
    await asyncio.to_thread(_post_webhook_sync, webhook_url, payload)


async def alert_on_queued_threshold(*, execution_name: str, node_id: str | None, project: str, domain: str, queued_sec: float):
    webhook_url = _get_webhook_url()
    if not webhook_url:
        return
    node_info = f" Node: `{node_id}`" if node_id else ""
    project_domain = f"{project}/{domain}" if project or domain else "-"
    payload = {
        "text": (
            f"Flyte node *queued too long*{node_info}\n"
            f"Execution: `{execution_name}`\nProject/Domain: `{project_domain}`\n"
            f"Queued for: {queued_sec:.0f}s (threshold: {QUEUED_THRESHOLD_SEC}s)"
        ),
    }
    await asyncio.to_thread(_post_webhook_sync, webhook_url, payload)


async def _check_queued_threshold():
    while True:
        await asyncio.sleep(QUEUED_CHECK_INTERVAL_SEC)
        now = time.time()
        to_alert = []
        for key, (first_seen, project, domain) in list(queued_since.items()):
            elapsed = now - first_seen
            if elapsed >= QUEUED_THRESHOLD_SEC:
                to_alert.append((key, project, domain, elapsed))
                del queued_since[key]
        for (execution_name, node_id), project, domain, queued_sec in to_alert:
            print(
                f"[ALERT] Node queued {queued_sec:.0f}s (threshold {QUEUED_THRESHOLD_SEC}s): "
                f"{project}/{domain} execution={execution_name} node={node_id}",
                file=sys.stderr,
            )
            await alert_on_queued_threshold(
                execution_name=execution_name, node_id=node_id,
                project=project, domain=domain, queued_sec=queued_sec,
            )


async def _process_event(raw_event):
    """Process a single streamed event (node execution)."""
    event_id = getattr(raw_event, "id", None)
    exec_id = getattr(event_id, "execution_id", None) if event_id else None
    if exec_id is None:
        exec_id = getattr(raw_event, "execution_id", None)

    execution_name = getattr(exec_id, "name", None) if exec_id else "unknown"
    project = getattr(exec_id, "project", "") if exec_id else ""
    domain = getattr(exec_id, "domain", "") if exec_id else ""

    phase_str = WorkflowExecutionPhase.enum_to_string(raw_event.phase)

    node_id_raw = getattr(event_id, "node_id", None) if event_id else getattr(raw_event, "node_id", None)
    node_id_str = None
    if node_id_raw is not None:
        node_id_str = (
            node_id_raw.node_id if hasattr(node_id_raw, "node_id") else
            node_id_raw if isinstance(node_id_raw, str) else str(node_id_raw)
        )

    key = (execution_name, node_id_str)

    if phase_str == "QUEUED":
        if key not in queued_since:
            queued_since[key] = (time.time(), project, domain)
    else:
        queued_since.pop(key, None)

    event_line = f"Node execution {execution_name} in {project}-{domain} is {phase_str}"
    if node_id_str:
        event_line = f"{event_line} (node={node_id_str})"
    print(event_line)
    await alert_on_failed(event_line=event_line, execution_name=execution_name, node_id=node_id_str)

async def main():
    union_remote = UnionRemote(
        config=Config.for_endpoint(endpoint="YOUR_UNION_ENDPOINT"),
        default_project="flytesnacks",
        default_domain="development",
    )

    asyncio.create_task(_check_queued_threshold())

    while True:
        try:
            async for response in union_remote.stream_execution_events(
                event_count=None,
                include_workflow_executions=False,
                include_task_executions=False,
                include_node_executions=True,
            ):
                # response is CloudEventNodeExecution (or workflow/task depending on filters)
                await _process_event(response)

if __name__ == "__main__":
    asyncio.run(main())