photoncloud-monorepo/deployer/crates/deployer-server/src/storage.rs

//! ChainFire-backed node storage
//!
//! This module provides persistent storage for node configurations
//! using ChainFire as the backend.

use chainfire_client::Client as ChainFireClient;
use deployer_types::{EnrollmentRuleSpec, NodeClassSpec, NodeConfig, NodeInfo, NodePoolSpec};
use serde::de::DeserializeOwned;
use serde::Serialize;
use thiserror::Error;
use tracing::{debug, error, warn};

use crate::cluster::ClusterNodeRecord;

/// Storage errors
#[derive(Error, Debug)]
pub enum StorageError {
    #[error("ChainFire connection error: {0}")]
    Connection(String),
    #[error("Serialization error: {0}")]
    Serialization(#[from] serde_json::Error),
    #[error("ChainFire client error: {0}")]
    Client(String),
    #[error("Node mapping conflict: {0}")]
    Conflict(String),
}

impl From<chainfire_client::ClientError> for StorageError {
    fn from(e: chainfire_client::ClientError) -> Self {
        StorageError::Client(e.to_string())
    }
}

/// Node storage backed by ChainFire
pub struct NodeStorage {
    client: ChainFireClient,
    namespace: String,
}

impl NodeStorage {
    /// Connect to ChainFire and create a new storage instance
    pub async fn connect(endpoint: &str, namespace: &str) -> Result<Self, StorageError> {
        debug!(endpoint = %endpoint, namespace = %namespace, "Connecting to ChainFire");

        let client = ChainFireClient::connect(endpoint)
            .await
            .map_err(|e| StorageError::Connection(e.to_string()))?;

        Ok(Self {
            client,
            namespace: namespace.to_string(),
        })
    }

    /// Key for node config by machine_id
    fn config_key(&self, machine_id: &str) -> String {
        format!("{}/nodes/config/{}", self.namespace, machine_id)
    }

    /// Key for node info by node_id
    fn info_key(&self, node_id: &str) -> String {
        format!("{}/nodes/info/{}", self.namespace, node_id)
    }

    fn cluster_node_key(&self, cluster_namespace: &str, cluster_id: &str, node_id: &str) -> String {
        format!(
            "{}/clusters/{}/nodes/{}",
            cluster_namespace, cluster_id, node_id
        )
    }

    fn cluster_nodes_prefix(&self, cluster_namespace: &str, cluster_id: &str) -> String {
        format!("{}/clusters/{}/nodes/", cluster_namespace, cluster_id)
    }

    fn cluster_node_classes_prefix(&self, cluster_namespace: &str, cluster_id: &str) -> String {
        format!(
            "{}/clusters/{}/node-classes/",
            cluster_namespace, cluster_id
        )
    }

    fn cluster_pools_prefix(&self, cluster_namespace: &str, cluster_id: &str) -> String {
        format!("{}/clusters/{}/pools/", cluster_namespace, cluster_id)
    }

    fn cluster_enrollment_rules_prefix(&self, cluster_namespace: &str, cluster_id: &str) -> String {
        format!(
            "{}/clusters/{}/enrollment-rules/",
            cluster_namespace, cluster_id
        )
    }

    async fn list_cluster_objects<T: DeserializeOwned>(
        &mut self,
        prefix: String,
    ) -> Result<Vec<T>, StorageError> {
        let kvs = self.client.get_prefix(&prefix).await?;
        let mut values = Vec::with_capacity(kvs.len());
        for (_key, value) in kvs {
            match serde_json::from_slice::<T>(&value) {
                Ok(record) => values.push(record),
                Err(e) => {
                    warn!(error = %e, prefix = %prefix, "Failed to decode cluster object");
                }
            }
        }
        Ok(values)
    }

    /// Register or update node config for a machine_id
    pub async fn register_node(
        &mut self,
        machine_id: &str,
        config: &NodeConfig,
    ) -> Result<(), StorageError> {
        let config_key = self.config_key(machine_id);
        let config_json = serde_json::to_vec(config)?;

        if let Some(existing) = self.client.get(&config_key).await? {
            let existing_config: NodeConfig = serde_json::from_slice(&existing)?;
            if existing_config.assignment.node_id != config.assignment.node_id {
                return Err(StorageError::Conflict(format!(
                    "machine_id {} already mapped to {}",
                    machine_id, existing_config.assignment.node_id
                )));
            }
        }

        debug!(
            machine_id = %machine_id,
            node_id = %config.assignment.node_id,
            key = %config_key,
            "Registering node config in ChainFire"
        );

        self.client.put(&config_key, &config_json).await?;
        Ok(())
    }

    /// Lookup node config by machine_id
    pub async fn get_node_config(
        &mut self,
        machine_id: &str,
    ) -> Result<Option<NodeConfig>, StorageError> {
        let config_key = self.config_key(machine_id);

        debug!(machine_id = %machine_id, key = %config_key, "Looking up node config");

        // Get config
        match self.client.get(&config_key).await? {
            Some(bytes) => {
                let config: NodeConfig = serde_json::from_slice(&bytes)?;
                Ok(Some(config))
            }
            None => {
                debug!(machine_id = %machine_id, "No config found");
                Ok(None)
            }
        }
    }

    /// Store node info (runtime state)
    pub async fn store_node_info(&mut self, node_info: &NodeInfo) -> Result<(), StorageError> {
        let key = self.info_key(&node_info.id);
        let json = serde_json::to_vec(node_info)?;

        debug!(
            node_id = %node_info.id,
            key = %key,
            "Storing node info in ChainFire"
        );

        self.client.put(&key, &json).await?;
        Ok(())
    }

    /// Store cluster node state under ultracloud/clusters/{cluster_id}/nodes/{node_id}
    pub async fn store_cluster_node<T: Serialize>(
        &mut self,
        cluster_namespace: &str,
        cluster_id: &str,
        node_id: &str,
        node: &T,
    ) -> Result<(), StorageError> {
        let key = self.cluster_node_key(cluster_namespace, cluster_id, node_id);
        let json = serde_json::to_vec(node)?;
        debug!(
            cluster_namespace = %cluster_namespace,
            cluster_id = %cluster_id,
            node_id = %node_id,
            key = %key,
            "Storing cluster node in ChainFire"
        );
        self.client.put(&key, &json).await?;
        Ok(())
    }

    /// List cluster nodes under ultracloud/clusters/{cluster_id}/nodes/
    pub async fn list_cluster_nodes(
        &mut self,
        cluster_namespace: &str,
        cluster_id: &str,
    ) -> Result<Vec<ClusterNodeRecord>, StorageError> {
        let prefix = self.cluster_nodes_prefix(cluster_namespace, cluster_id);
        let kvs = self.client.get_prefix(&prefix).await?;

        let mut nodes = Vec::with_capacity(kvs.len());
        for (_key, value) in kvs {
            match serde_json::from_slice::<ClusterNodeRecord>(&value) {
                Ok(record) => nodes.push(record),
                Err(e) => {
                    warn!(error = %e, "Failed to decode cluster node record");
                }
            }
        }

        Ok(nodes)
    }

    pub async fn list_node_classes(
        &mut self,
        cluster_namespace: &str,
        cluster_id: &str,
    ) -> Result<Vec<NodeClassSpec>, StorageError> {
        self.list_cluster_objects(self.cluster_node_classes_prefix(cluster_namespace, cluster_id))
            .await
    }

    pub async fn list_pools(
        &mut self,
        cluster_namespace: &str,
        cluster_id: &str,
    ) -> Result<Vec<NodePoolSpec>, StorageError> {
        self.list_cluster_objects(self.cluster_pools_prefix(cluster_namespace, cluster_id))
            .await
    }

    pub async fn list_enrollment_rules(
        &mut self,
        cluster_namespace: &str,
        cluster_id: &str,
    ) -> Result<Vec<EnrollmentRuleSpec>, StorageError> {
        self.list_cluster_objects(
            self.cluster_enrollment_rules_prefix(cluster_namespace, cluster_id),
        )
        .await
    }

    /// Get node info by node_id
    pub async fn get_node_info(&mut self, node_id: &str) -> Result<Option<NodeInfo>, StorageError> {
        let key = self.info_key(node_id);

        match self.client.get(&key).await? {
            Some(bytes) => {
                let info: NodeInfo = serde_json::from_slice(&bytes)?;
                Ok(Some(info))
            }
            None => Ok(None),
        }
    }

    /// List all registered nodes
    pub async fn list_nodes(&mut self) -> Result<Vec<NodeInfo>, StorageError> {
        let prefix = format!("{}/nodes/info/", self.namespace);

        let kvs = self.client.get_prefix(&prefix).await?;

        let mut nodes = Vec::with_capacity(kvs.len());
        for (_, value) in kvs {
            match serde_json::from_slice::<NodeInfo>(&value) {
                Ok(info) => nodes.push(info),
                Err(e) => {
                    error!(error = %e, "Failed to deserialize node info");
                }
            }
        }

        Ok(nodes)
    }

    /// List all pre-registered machine configs (machine_id -> config)
    pub async fn list_machine_configs(
        &mut self,
    ) -> Result<Vec<(String, NodeConfig)>, StorageError> {
        let config_prefix = format!("{}/nodes/config/", self.namespace);
        let configs = self.client.get_prefix(&config_prefix).await?;

        let mut results = Vec::new();
        for (key, value) in configs {
            let key_str = String::from_utf8_lossy(&key);
            if let Some(machine_id) = key_str.strip_prefix(&config_prefix) {
                if let Ok(config) = serde_json::from_slice::<NodeConfig>(&value) {
                    results.push((machine_id.to_string(), config));
                } else {
                    warn!(key = %key_str, "Failed to deserialize node config");
                }
            }
        }
        Ok(results)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use deployer_types::{BootstrapPlan, BootstrapSecrets, NodeAssignment};

    // Note: Integration tests require a running ChainFire instance.
    // These unit tests verify serialization and key generation.

    #[test]
    fn test_key_generation() {
        // Can't test connect without ChainFire, but we can verify key format
        let namespace = "deployer";
        let machine_id = "abc123";
        let node_id = "node01";

        let config_key = format!("{}/nodes/config/{}", namespace, machine_id);
        let info_key = format!("{}/nodes/info/{}", namespace, node_id);

        assert_eq!(config_key, "deployer/nodes/config/abc123");
        assert_eq!(info_key, "deployer/nodes/info/node01");

        let cluster_namespace = "ultracloud";
        let cluster_id = "cluster-a";
        let cluster_key = format!(
            "{}/clusters/{}/nodes/{}",
            cluster_namespace, cluster_id, node_id
        );
        assert_eq!(cluster_key, "ultracloud/clusters/cluster-a/nodes/node01");
    }

    #[test]
    fn test_node_config_serialization() {
        let config = NodeConfig::from_parts(
            NodeAssignment {
                node_id: "node01".to_string(),
                hostname: "node01".to_string(),
                role: "control-plane".to_string(),
                ip: "10.0.1.10".to_string(),
                labels: std::collections::HashMap::new(),
                pool: None,
                node_class: None,
                failure_domain: None,
            },
            BootstrapPlan {
                services: vec!["chainfire".to_string(), "flaredb".to_string()],
                nix_profile: None,
                install_plan: None,
            },
            BootstrapSecrets::default(),
        );

        let json = serde_json::to_vec(&config).unwrap();
        let deserialized: NodeConfig = serde_json::from_slice(&json).unwrap();

        assert_eq!(deserialized.assignment.hostname, "node01");
        assert_eq!(deserialized.assignment.role, "control-plane");
        assert_eq!(deserialized.bootstrap_plan.services.len(), 2);
        assert!(deserialized
            .bootstrap_secrets
            .ssh_authorized_keys
            .is_empty());
    }
}