crete/src/cluster.rs

use std::collections::HashMap;
use std::net::Ipv4Addr;
use std::path::PathBuf;
use std::str::FromStr;

use async_walkdir::WalkDir;
use futures::StreamExt;
use futures::stream::FuturesUnordered;
use minijinja::Environment;
use optional_struct::{Applicable, optional_struct};
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
use tokio::process::Command;

use crate::environment::PathEnvironment;
use crate::node::{Node, OptionalNodeDeserialize};
use crate::patch::Patches;
use crate::talosconfig::TalosConfig;
use crate::{get_configs_path, get_talos_path, run_command};

#[optional_struct]
#[derive(Debug, Deserialize, JsonSchema, Clone)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub(crate) struct Base {
    #[serde(default)]
    pub(crate) kernel_args: Vec<String>,
    #[serde(default)]
    pub(crate) patches: Patches<String>,
}

#[optional_struct]
#[derive(Debug, Serialize, Deserialize, JsonSchema, Clone)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct Version {
    kubernetes: semver::Version,
    talos: semver::Version,
}

impl Version {
    pub(crate) fn kubernetes(&self) -> String {
        format!("v{}", self.kubernetes)
    }

    pub(crate) fn talos(&self) -> String {
        format!("v{}", self.talos)
    }
}

#[derive(Debug, Serialize, Deserialize, JsonSchema, Clone, Copy, PartialEq, Eq)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub enum ClusterEnv {
    Production,
    Staging,
}

#[optional_struct]
#[derive(Debug, Deserialize, JsonSchema, Clone)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
struct ClusterDeserialize {
    #[optional_rename(OptionalVersion)]
    #[optional_wrap]
    version: Version,
    nodes: Vec<String>,
    cluster_env: ClusterEnv,
    control_plane_ip: Ipv4Addr,
    secrets_file: PathBuf,
    #[serde(default)]
    #[optional_skip_wrap]
    default: OptionalNodeDeserialize,
    #[optional_rename(OptionalBase)]
    #[optional_wrap]
    base: Base,
}

#[derive(Debug, Serialize, Clone)]
#[serde(rename_all = "camelCase")]
pub struct Cluster {
    pub(crate) name: String,
    pub(crate) version: Version,
    nodes: HashMap<String, Node>,
    cluster_env: ClusterEnv,
    pub(crate) control_plane_ip: Ipv4Addr,
    pub(crate) secrets_file: PathBuf,
}

impl Cluster {
    pub(crate) async fn get(cluster_name: &str, env: &Environment<'static>) -> Self {
        let path = get_talos_path().join("clusters").join("default.yaml");

        let default: OptionalClusterDeserialize = if path.exists() {
            let content = tokio::fs::read_to_string(path).await.unwrap();

            // The deserialization makes a blocking web request
            tokio::task::spawn_blocking(move || serde_yaml::from_str(&content).unwrap())
                .await
                .unwrap()
        } else {
            Default::default()
        };

        let mut path = get_talos_path().join("clusters").join(cluster_name);
        path.add_extension("yaml");
        let content = tokio::fs::read_to_string(path).await.unwrap();

        let mut cluster: OptionalClusterDeserialize = serde_yaml::from_str(&content).unwrap();

        // For some reason apply on the cluster does not properly apply to the default settings...
        // So we manually apply it here first
        cluster.default = default.default.clone().apply(cluster.default);

        let cluster: ClusterDeserialize = default.apply(cluster).try_into().unwrap();
        let nodes = cluster.nodes;
        let default = cluster.default;
        let base = cluster.base;

        let secrets_file = get_talos_path()
            .join("secrets")
            .join(cluster.secrets_file)
            .absolute()
            .unwrap();

        let cluster = Self {
            name: cluster_name.into(),
            version: cluster.version,
            nodes: Default::default(),
            cluster_env: cluster.cluster_env,
            control_plane_ip: cluster.control_plane_ip,
            secrets_file,
        };

        let nodes = nodes.iter().map(|name| async {
            let node = Node::get(name, env, &cluster, &default, &base).await;
            (node.hostname().to_owned(), node)
        });

        let nodes = FuturesUnordered::from_iter(nodes).collect().await;

        Self { nodes, ..cluster }
    }

    pub async fn get_all() -> HashMap<String, Cluster> {
        let env = PathEnvironment::new_patches();

        let names: Vec<_> = WalkDir::new(get_talos_path().join("clusters"))
            .filter_map(|e| async { e.ok() })
            .filter(|e| {
                let e = e.clone();
                async move { e.metadata().await.unwrap().is_file() }
            })
            .filter_map(|e| async move {
                let mut path = PathBuf::from_str(&e.file_name().to_string_lossy()).unwrap();
                path.set_extension("");

                let name = path.to_string_lossy().to_string();

                if name == "default" {
                    return None;
                }

                Some(name)
            })
            .collect()
            .await;

        // Get the clusters concurrently
        FuturesUnordered::from_iter(names.iter().map(|name| {
            let env = env.clone();
            let name = name.clone();

            async move {
                let cluster = Cluster::get(&name, &env).await;

                (name, cluster)
            }
        }))
        .collect()
        .await
    }

    pub fn nodes(&self) -> impl Iterator<Item = &Node> {
        self.nodes.values()
    }

    pub fn node(&self, node: &str) -> Option<&Node> {
        self.nodes.get(node)
    }

    pub async fn talosconfig(&self) -> String {
        let mut command = Command::new("talosctl");
        command.args([
            "gen",
            "config",
            &self.name,
            &format!("https://{}:6443", self.control_plane_ip),
            "--with-secrets",
            self.secrets_file.to_str().expect("Path should be valid"),
            "--output-types",
            "talosconfig",
            "-o",
            "-",
        ]);

        run_command(command).await
    }

    pub fn talosconfig_path(&self) -> PathBuf {
        get_configs_path().join(&self.name).join("talosconfig")
    }

    async fn talosconfig_add_endpoint(&self) {
        let mut command = Command::new("talosctl");
        command.args([
            "config",
            "--talosconfig",
            self.talosconfig_path()
                .to_str()
                .expect("Path should be valid utf-8"),
            "endpoint",
            &self.control_plane_ip.to_string(),
        ]);

        run_command(command).await;
    }

    pub async fn talosconfig_write(&self) {
        let config = self.talosconfig().await;

        let path = self.talosconfig_path();
        tokio::fs::create_dir_all(path.parent().expect("talosconfig should have parent"))
            .await
            .unwrap();

        tokio::fs::write(path, config).await.unwrap();

        self.talosconfig_add_endpoint().await
    }

    pub async fn talosconfig_merge(&self) {
        let mut command = Command::new("talosctl");
        command.args([
            "config",
            "--talosconfig",
            TalosConfig::merged_path()
                .to_str()
                .expect("Path should be valid utf-8"),
            "merge",
            self.talosconfig_path()
                .to_str()
                .expect("Path should be valid utf-8"),
        ]);

        run_command(command).await;
    }
}

impl JsonSchema for Cluster {
    fn schema_name() -> std::borrow::Cow<'static, str> {
        "Cluster".into()
    }

    fn json_schema(generator: &mut schemars::SchemaGenerator) -> schemars::Schema {
        OptionalClusterDeserialize::json_schema(generator)
    }
}