harmony/harmony-k8s/src/client.rs

use std::sync::Arc;

use kube::config::{KubeConfigOptions, Kubeconfig};
use kube::{Client, Config, Discovery, Error};
use log::error;
use serde::Serialize;
use tokio::sync::{OnceCell, RwLock};

use crate::types::KubernetesDistribution;

// TODO not cool, should use a proper configuration mechanism
// cli arg, env var, config file
fn read_dry_run_from_env() -> bool {
    std::env::var("DRY_RUN")
        .map(|v| v == "true" || v == "1")
        .unwrap_or(false)
}

#[derive(Clone)]
pub struct K8sClient {
    pub(crate) client: Client,
    /// When `true` no mutation is sent to the API server; diffs are printed
    /// to stdout instead. Initialised from the `DRY_RUN` environment variable.
    pub(crate) dry_run: bool,
    pub(crate) k8s_distribution: Arc<OnceCell<KubernetesDistribution>>,
    /// API discovery cache. Wrapped in `RwLock` so it can be invalidated
    /// after installing CRDs or operators that register new API groups.
    pub(crate) discovery: Arc<RwLock<Option<Arc<Discovery>>>>,
}

impl Serialize for K8sClient {
    fn serialize<S>(&self, _serializer: S) -> Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
        todo!("K8sClient serialization is not meaningful; remove this impl if unused")
    }
}

impl std::fmt::Debug for K8sClient {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_fmt(format_args!(
            "K8sClient {{ namespace: {}, dry_run: {} }}",
            self.client.default_namespace(),
            self.dry_run,
        ))
    }
}

impl K8sClient {
    /// Create a client, reading `DRY_RUN` from the environment.
    pub fn new(client: Client) -> Self {
        Self {
            dry_run: read_dry_run_from_env(),
            client,
            k8s_distribution: Arc::new(OnceCell::new()),
            discovery: Arc::new(RwLock::new(None)),
        }
    }

    /// Create a client that always operates in dry-run mode, regardless of
    /// the environment variable.
    pub fn new_dry_run(client: Client) -> Self {
        Self {
            dry_run: true,
            ..Self::new(client)
        }
    }

    /// Returns `true` if this client is operating in dry-run mode.
    pub fn is_dry_run(&self) -> bool {
        self.dry_run
    }

    pub async fn try_default() -> Result<Self, Error> {
        Ok(Self::new(Client::try_default().await?))
    }

    pub async fn from_kubeconfig(path: &str) -> Option<Self> {
        Self::from_kubeconfig_with_opts(path, &KubeConfigOptions::default()).await
    }

    pub async fn from_kubeconfig_with_context(path: &str, context: Option<String>) -> Option<Self> {
        let opts = KubeConfigOptions {
            context,
            ..Default::default()
        };
        Self::from_kubeconfig_with_opts(path, &opts).await
    }

    pub async fn from_kubeconfig_with_opts(path: &str, opts: &KubeConfigOptions) -> Option<Self> {
        let k = match Kubeconfig::read_from(path) {
            Ok(k) => k,
            Err(e) => {
                error!("Failed to load kubeconfig from {path}: {e}");
                return None;
            }
        };
        Some(Self::new(
            Client::try_from(Config::from_custom_kubeconfig(k, opts).await.unwrap()).unwrap(),
        ))
    }
}