clawmates/crates/cm-db/src/repo/nodes.rs

//! Fleet node registry: user-connected local-hardware hosts that run the
//! clawmates-node daemon. The daemon authenticates its control channel with the
//! node `token`, then upserts its host-health snapshot on every heartbeat.

use cm_domain::{NodeId, WorkspaceId};
use sqlx::{PgPool, Row};
use time::OffsetDateTime;

use crate::DbError;

/// Latest host-health snapshot for a node.
#[derive(Debug, Clone)]
pub struct NodeHealth {
    pub cpu_pct: f64,
    pub mem_total: i64,
    pub mem_used: i64,
    pub mem_pressure: f64,
    pub swap_used: i64,
    pub disk_total: i64,
    pub disk_free: i64,
    pub load1: f64,
    pub load5: f64,
    pub load15: f64,
    pub container_count: i32,
}

/// A registered node plus its latest health (if it has reported one).
#[derive(Debug, Clone)]
pub struct NodeRow {
    pub id: NodeId,
    pub name: String,
    pub hostname: Option<String>,
    pub local_ip: Option<String>,
    pub status: String,
    pub agent_version: Option<String>,
    pub tailscale_ip: Option<String>,
    pub last_seen: Option<OffsetDateTime>,
    pub created_at: OffsetDateTime,
    pub health: Option<NodeHealth>,
    /// Latest Beszel metrics scalars (the rich bits beyond basic health).
    pub gpu_pct: Option<f64>,
    pub temp_max: Option<f64>,
}

/// Register a new (pending) node with its control-channel token.
pub async fn create(
    pool: &PgPool,
    workspace_id: WorkspaceId,
    name: &str,
    token: &str,
) -> Result<NodeId, DbError> {
    let id = NodeId::new();
    sqlx::query("INSERT INTO nodes (id, workspace_id, name, status, token) VALUES ($1, $2, $3, 'pending', $4)")
        .bind(id.as_uuid())
        .bind(workspace_id.as_uuid())
        .bind(name)
        .bind(token)
        .execute(pool)
        .await?;
    Ok(id)
}

/// Resolve a control-channel token to its node + workspace (daemon auth).
pub async fn auth(pool: &PgPool, token: &str) -> Result<Option<(NodeId, WorkspaceId)>, DbError> {
    let row = sqlx::query("SELECT id, workspace_id FROM nodes WHERE token = $1")
        .bind(token)
        .fetch_optional(pool)
        .await?;
    Ok(row.map(|r| {
        (
            NodeId::from(r.get::<uuid::Uuid, _>("id")),
            WorkspaceId::from(r.get::<uuid::Uuid, _>("workspace_id")),
        )
    }))
}

const SELECT_WITH_HEALTH: &str = "SELECT n.id, n.name, n.hostname, n.local_ip, n.status, n.agent_version, n.tailscale_ip, n.last_seen, n.created_at,
        h.node_id AS health_node, h.cpu_pct, h.mem_total, h.mem_used, h.mem_pressure, h.swap_used,
        h.disk_total, h.disk_free, h.load1, h.load5, h.load15, h.container_count,
        m.gpu_pct AS m_gpu_pct, m.temp_max AS m_temp_max
   FROM nodes n LEFT JOIN node_health h ON h.node_id = n.id
                LEFT JOIN node_metrics m ON m.node_id = n.id";

/// List a workspace's nodes (oldest first) with their latest health.
pub async fn list(pool: &PgPool, workspace_id: WorkspaceId) -> Result<Vec<NodeRow>, DbError> {
    let rows = sqlx::query(&format!(
        "{SELECT_WITH_HEALTH} WHERE n.workspace_id = $1 ORDER BY n.created_at"
    ))
    .bind(workspace_id.as_uuid())
    .fetch_all(pool)
    .await?;
    Ok(rows.into_iter().map(map_node).collect())
}

/// Fetch one node (workspace-scoped) with its latest health.
pub async fn get(
    pool: &PgPool,
    id: NodeId,
    workspace_id: WorkspaceId,
) -> Result<Option<NodeRow>, DbError> {
    let row = sqlx::query(&format!(
        "{SELECT_WITH_HEALTH} WHERE n.id = $1 AND n.workspace_id = $2"
    ))
    .bind(id.as_uuid())
    .bind(workspace_id.as_uuid())
    .fetch_optional(pool)
    .await?;
    Ok(row.map(map_node))
}

/// Record a heartbeat: mark the node online + refresh its version/tailscale IP,
/// and upsert its latest host-health snapshot.
#[allow(clippy::too_many_arguments)]
pub async fn heartbeat(
    pool: &PgPool,
    id: NodeId,
    agent_version: Option<&str>,
    tailscale_ip: Option<&str>,
    hostname: Option<&str>,
    local_ip: Option<&str>,
    h: &NodeHealth,
) -> Result<(), DbError> {
    sqlx::query(
        // Preserve a rule-/operator-set `draining` state across heartbeats; a node
        // is only un-drained by an explicit set_status.
        "UPDATE nodes SET status = CASE WHEN status = 'draining' THEN 'draining' ELSE 'online' END,
            last_seen = now(),
            agent_version = COALESCE($2, agent_version),
            tailscale_ip = COALESCE($3, tailscale_ip),
            hostname = COALESCE($4, hostname),
            local_ip = COALESCE($5, local_ip)
         WHERE id = $1",
    )
    .bind(id.as_uuid())
    .bind(agent_version)
    .bind(tailscale_ip)
    .bind(hostname)
    .bind(local_ip)
    .execute(pool)
    .await?;
    sqlx::query(
        "INSERT INTO node_health
            (node_id, captured_at, cpu_pct, mem_total, mem_used, mem_pressure, swap_used,
             disk_total, disk_free, load1, load5, load15, container_count)
         VALUES ($1, now(), $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12)
         ON CONFLICT (node_id) DO UPDATE SET
            captured_at = now(), cpu_pct = excluded.cpu_pct, mem_total = excluded.mem_total,
            mem_used = excluded.mem_used, mem_pressure = excluded.mem_pressure,
            swap_used = excluded.swap_used, disk_total = excluded.disk_total,
            disk_free = excluded.disk_free, load1 = excluded.load1, load5 = excluded.load5,
            load15 = excluded.load15, container_count = excluded.container_count",
    )
    .bind(id.as_uuid())
    .bind(h.cpu_pct)
    .bind(h.mem_total)
    .bind(h.mem_used)
    .bind(h.mem_pressure)
    .bind(h.swap_used)
    .bind(h.disk_total)
    .bind(h.disk_free)
    .bind(h.load1)
    .bind(h.load5)
    .bind(h.load15)
    .bind(h.container_count)
    .execute(pool)
    .await?;
    Ok(())
}

/// Set a node's status (e.g. 'offline' when its channel drops, 'draining' on
/// deregister).
pub async fn set_status(pool: &PgPool, id: NodeId, status: &str) -> Result<(), DbError> {
    sqlx::query("UPDATE nodes SET status = $2 WHERE id = $1")
        .bind(id.as_uuid())
        .bind(status)
        .execute(pool)
        .await?;
    Ok(())
}

/// Mark online nodes whose last heartbeat is older than `secs` as offline.
pub async fn mark_stale_offline(pool: &PgPool, secs: i64) -> Result<(), DbError> {
    sqlx::query(
        "UPDATE nodes SET status = 'offline'
         WHERE status = 'online'
           AND (last_seen IS NULL OR last_seen < now() - ($1 * interval '1 second'))",
    )
    .bind(secs)
    .execute(pool)
    .await?;
    Ok(())
}

/// Remove a node from the registry (workspace-scoped).
pub async fn delete(pool: &PgPool, id: NodeId, workspace_id: WorkspaceId) -> Result<(), DbError> {
    sqlx::query("DELETE FROM nodes WHERE id = $1 AND workspace_id = $2")
        .bind(id.as_uuid())
        .bind(workspace_id.as_uuid())
        .execute(pool)
        .await?;
    Ok(())
}

/// A node's current status (for metrics-aware placement: skip 'draining').
pub async fn status_of(pool: &PgPool, id: NodeId) -> Result<Option<String>, DbError> {
    Ok(
        sqlx::query_scalar::<_, String>("SELECT status FROM nodes WHERE id = $1")
            .bind(id.as_uuid())
            .fetch_optional(pool)
            .await?,
    )
}

fn map_node(r: sqlx::postgres::PgRow) -> NodeRow {
    let health = r
        .get::<Option<uuid::Uuid>, _>("health_node")
        .map(|_| NodeHealth {
            cpu_pct: r.get("cpu_pct"),
            mem_total: r.get("mem_total"),
            mem_used: r.get("mem_used"),
            mem_pressure: r.get("mem_pressure"),
            swap_used: r.get("swap_used"),
            disk_total: r.get("disk_total"),
            disk_free: r.get("disk_free"),
            load1: r.get("load1"),
            load5: r.get("load5"),
            load15: r.get("load15"),
            container_count: r.get("container_count"),
        });
    NodeRow {
        id: NodeId::from(r.get::<uuid::Uuid, _>("id")),
        name: r.get("name"),
        hostname: r.get("hostname"),
        local_ip: r.get("local_ip"),
        status: r.get("status"),
        agent_version: r.get("agent_version"),
        tailscale_ip: r.get("tailscale_ip"),
        last_seen: r.get("last_seen"),
        created_at: r.get("created_at"),
        health,
        gpu_pct: r.get("m_gpu_pct"),
        temp_max: r.get("m_temp_max"),
    }
}