~ knotless Search/
↑↓ move↵ open⇧↵ all results

johnandersen777.bsky.social/knotless-runner-k8s

0
0
304 lines · 9.8 KB Raw History
//! The `kubectl` calls the proof of concept makes.
//!
//! The real executor would use a Kubernetes client crate. The proof of
//! concept shells out to `kubectl` for the same reason knotless's docker
//! executor shells out to `docker`: the transport is visible, and a fake
//! client in the tests can show every call.

use std::io::Write;
use std::path::{Path, PathBuf};
use std::process::{Child, Command, Stdio};
use std::time::Duration;

use crate::cli::Options;

/// A configured `kubectl` client.
#[derive(Clone, Debug)]
pub struct Kubectl {
    program: PathBuf,
    kubeconfig: Option<PathBuf>,
    context: Option<String>,
    namespace: String,
}

impl Kubectl {
    pub fn new(options: &Options) -> Self {
        Self {
            program: options.kubectl.clone(),
            kubeconfig: options.kubeconfig.clone(),
            context: options.context.clone(),
            namespace: options.namespace.clone(),
        }
    }

    /// The client with its kubeconfig, context and namespace.
    fn base(&self) -> Command {
        let mut command = Command::new(&self.program);
        if let Some(kubeconfig) = &self.kubeconfig {
            command.arg("--kubeconfig").arg(kubeconfig);
        }
        if let Some(context) = &self.context {
            command.arg("--context").arg(context);
        }
        command.args(["-n", &self.namespace]);
        command
    }

    /// A readable line for `--dry-run`, as a shell would take it.
    pub fn planned(&self, args: &[&str]) -> String {
        let mut parts = vec![self.program.display().to_string()];
        if let Some(kubeconfig) = &self.kubeconfig {
            parts.push("--kubeconfig".into());
            parts.push(kubeconfig.display().to_string());
        }
        if let Some(context) = &self.context {
            parts.push("--context".into());
            parts.push(context.clone());
        }
        parts.push("-n".into());
        parts.push(self.namespace.clone());
        parts.extend(args.iter().map(|a| shell_word(a)));
        parts.join(" ")
    }

    /// Applies the Job manifest from stdin.
    pub fn apply(&self, manifest: &str) -> Result<(), String> {
        let mut child = self
            .base()
            .args(["apply", "-f", "-"])
            .stdin(Stdio::piped())
            .spawn()
            .map_err(|err| format!("running {}: {err}", self.program.display()))?;
        {
            let mut stdin = child
                .stdin
                .take()
                .ok_or_else(|| "kubectl apply: no stdin".to_string())?;
            stdin
                .write_all(manifest.as_bytes())
                .map_err(|err| format!("writing the Job manifest: {err}"))?;
        }
        let status = child
            .wait()
            .map_err(|err| format!("waiting for kubectl apply: {err}"))?;
        if !status.success() {
            return Err("kubectl apply failed".to_string());
        }
        Ok(())
    }

    /// Waits for the Job's pod to be Ready.
    pub fn wait_pod(&self, job: &str, timeout: Duration) -> Result<(), String> {
        let selector = format!("job-name={job}");
        let timeout = format!("{}s", timeout.as_secs().max(1));
        let status = self
            .base()
            .args([
                "wait",
                "--for=condition=Ready",
                "pod",
                "-l",
                &selector,
                "--timeout",
                &timeout,
            ])
            .status()
            .map_err(|err| format!("running kubectl wait: {err}"))?;
        if !status.success() {
            return Err(format!("pod for Job {job} did not become Ready"));
        }
        Ok(())
    }

    /// The first pod of the Job.
    pub fn pod(&self, job: &str) -> Result<String, String> {
        let selector = format!("job-name={job}");
        let output = self
            .base()
            .args([
                "get",
                "pods",
                "-l",
                &selector,
                "-o",
                "jsonpath={.items[0].metadata.name}",
            ])
            .output()
            .map_err(|err| format!("running kubectl get pods: {err}"))?;
        if !output.status.success() {
            let stderr = String::from_utf8_lossy(&output.stderr);
            return Err(format!("kubectl get pods failed: {}", stderr.trim()));
        }
        let pod = String::from_utf8_lossy(&output.stdout).trim().to_string();
        if pod.is_empty() {
            return Err(format!("Job {job} has no pod"));
        }
        Ok(pod)
    }

    /// The command that runs `script` in `pod` as a step.
    pub fn exec(&self, pod: &str, script: &str) -> Command {
        let mut command = self.base();
        command.args(["exec", pod, "--", "sh", "-ec", script]);
        command
    }

    /// Makes the directories the run streams into.
    pub fn mkdirs(&self, pod: &str) -> Result<(), String> {
        let status = self
            .base()
            .args(["exec", pod, "--", "mkdir", "-p", "/work", "/outputs", "/steps"])
            .status()
            .map_err(|err| format!("running kubectl exec mkdir: {err}"))?;
        if !status.success() {
            return Err("creating /work, /outputs and /steps in the pod failed".to_string());
        }
        Ok(())
    }

    /// Streams the contents of `source` on this host into `dest` in `pod`,
    /// over `kubectl exec -i` and tar, which Kubernetes has in place of a
    /// file copy API.
    pub fn tar_in(&self, pod: &str, dest: &str, source: &Path) -> Result<(), String> {
        let mut tar = Command::new("tar")
            .args(["-cf", "-", "-C"])
            .arg(source)
            .arg(".")
            .stdout(Stdio::piped())
            .spawn()
            .map_err(|err| format!("running tar for {}: {err}", source.display()))?;
        let stdout = tar
            .stdout
            .take()
            .ok_or_else(|| "tar: no stdout".to_string())?;
        let mut kubectl = self.base();
        kubectl
            .args(["exec", "-i", pod, "--", "tar", "-xf", "-", "-C", dest])
            .stdin(Stdio::from(stdout))
            .stdout(Stdio::null());
        let status = wait(
            kubectl
                .spawn()
                .map_err(|err| format!("running kubectl exec tar: {err}"))?,
            "kubectl exec tar",
        )?;
        let tar_status = wait(tar, "tar")?;
        if !status.success() {
            return Err(format!(
                "streaming {} into {dest} failed",
                source.display()
            ));
        }
        if !tar_status.success() {
            return Err(format!("tar reading {} failed", source.display()));
        }
        Ok(())
    }

    /// Streams the contents of `source` in `pod` into `dest` on this host.
    pub fn tar_out(&self, pod: &str, source: &str, dest: &Path) -> Result<(), String> {
        let mut kubectl = self.base();
        kubectl
            .args(["exec", pod, "--", "tar", "-cf", "-", "-C", source, "."])
            .stdout(Stdio::piped());
        let mut kubectl = kubectl
            .spawn()
            .map_err(|err| format!("running kubectl exec tar: {err}"))?;
        let stdout = kubectl
            .stdout
            .take()
            .ok_or_else(|| "kubectl exec tar: no stdout".to_string())?;
        let tar = Command::new("tar")
            .args(["-xf", "-", "-C"])
            .arg(dest)
            .stdin(Stdio::from(stdout))
            .spawn()
            .map_err(|err| format!("running tar for {}: {err}", dest.display()))?;
        let status = wait(kubectl, "kubectl exec tar")?;
        let tar_status = wait(tar, "tar")?;
        if !status.success() {
            return Err(format!("streaming {source} out of the pod failed"));
        }
        if !tar_status.success() {
            return Err(format!("tar writing {} failed", dest.display()));
        }
        Ok(())
    }

    /// Deletes the Job and its pod. Missing is fine.
    pub fn delete(&self, job: &str) -> Result<(), String> {
        let status = self
            .base()
            .args(["delete", "job", job, "--ignore-not-found", "--wait=false"])
            .status()
            .map_err(|err| format!("running kubectl delete: {err}"))?;
        if !status.success() {
            return Err(format!("kubectl delete job {job} failed"));
        }
        Ok(())
    }
}

fn wait(mut child: Child, what: &str) -> Result<std::process::ExitStatus, String> {
    child
        .wait()
        .map_err(|err| format!("waiting for {what}: {err}"))
}

/// A shell-ish word for a dry-run line, quoted only when it has spaces.
fn shell_word(word: &str) -> String {
    if word.is_empty()
        || word
            .chars()
            .any(|c| c.is_whitespace() || matches!(c, '\'' | '"' | '$' | '`' | '\\'))
    {
        format!("'{}'", word.replace('\'', "'\\''"))
    } else {
        word.to_string()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn client() -> Kubectl {
        Kubectl {
            program: PathBuf::from("kubectl"),
            kubeconfig: Some(PathBuf::from("/k")),
            context: Some("ctx".into()),
            namespace: "ci".into(),
        }
    }

    #[test]
    fn plans() {
        assert_eq!(
            client().planned(&["exec", "pod", "--", "sh", "-ec", "echo hi"]),
            "kubectl --kubeconfig /k --context ctx -n ci exec pod -- sh -ec 'echo hi'"
        );
    }

    #[test]
    fn exec_args() {
        let command = client().exec("pod-1", "make test");
        let args: Vec<_> = command.get_args().map(|a| a.to_string_lossy()).collect();
        assert_eq!(
            args,
            vec![
                "--kubeconfig",
                "/k",
                "--context",
                "ctx",
                "-n",
                "ci",
                "exec",
                "pod-1",
                "--",
                "sh",
                "-ec",
                "make test"
            ]
        );
    }
}