zilliz 1.4.0

TUI and CLI tool for managing Zilliz Cloud clusters and Milvus operations
Documentation
use std::io::Write;
use std::path::PathBuf;

use anyhow::{bail, Context, Result};
use clap::CommandFactory;
use clap_complete::{generate, Shell};

use crate::cli::args::{Cli, CompletionCommands};

const COMPLETION_MARKER: &str = "# zilliz-cli completion";

pub fn run(cmd: CompletionCommands) -> Result<()> {
    match cmd {
        CompletionCommands::Install { shell, apply } => install(shell, apply),
        CompletionCommands::Uninstall { shell } => uninstall(shell),
        CompletionCommands::Status { shell } => status(shell),
        CompletionCommands::Show { shell } => show(shell),
    }
}

fn detect_shell() -> Result<Shell> {
    let shell_env = std::env::var("SHELL").unwrap_or_default();
    if shell_env.contains("zsh") {
        Ok(Shell::Zsh)
    } else if shell_env.contains("bash") {
        Ok(Shell::Bash)
    } else if shell_env.contains("fish") {
        Ok(Shell::Fish)
    } else {
        bail!("Cannot detect shell. Specify one: bash, zsh, or fish")
    }
}

fn resolve_shell(shell: Option<Shell>) -> Result<Shell> {
    match shell {
        Some(s) => Ok(s),
        None => detect_shell(),
    }
}

fn rc_path(shell: Shell) -> PathBuf {
    let home = dirs::home_dir().unwrap_or_else(|| PathBuf::from("~"));
    match shell {
        Shell::Bash => home.join(".bashrc"),
        Shell::Zsh => home.join(".zshrc"),
        Shell::Fish => home.join(".config/fish/config.fish"),
        _ => home.join(".bashrc"),
    }
}

fn completion_script_line(shell: Shell) -> String {
    match shell {
        Shell::Bash => r#"eval "$(zilliz completion show bash)""#.to_string(),
        Shell::Zsh => r#"eval "$(zilliz completion show zsh)""#.to_string(),
        Shell::Fish => "zilliz completion show fish | source".to_string(),
        _ => r#"eval "$(zilliz completion show bash)""#.to_string(),
    }
}

fn is_installed(shell: Shell) -> bool {
    let path = rc_path(shell);
    if !path.exists() {
        return false;
    }
    let content = std::fs::read_to_string(&path).unwrap_or_default();
    content.contains(COMPLETION_MARKER) || content.contains("zilliz completion show")
}

fn install(shell: Option<Shell>, apply: bool) -> Result<()> {
    let shell = resolve_shell(shell)?;

    if is_installed(shell) {
        println!("Completion already installed for {:?}.", shell);
        return Ok(());
    }

    let path = rc_path(shell);
    let script = completion_script_line(shell);

    if apply {
        // Create parent dirs for fish
        if let Some(parent) = path.parent() {
            std::fs::create_dir_all(parent)?;
        }
        let mut f = std::fs::OpenOptions::new()
            .append(true)
            .create(true)
            .open(&path)
            .with_context(|| format!("Cannot write to {}", path.display()))?;
        writeln!(f, "\n{}", COMPLETION_MARKER)?;
        writeln!(f, "{}", script)?;
        println!("Completion installed to {}", path.display());
        println!("Restart your shell or run: source {}", path.display());
    } else {
        println!("Add this to {}:", path.display());
        println!();
        println!("  {}", script);
        println!();
        println!("Or run: zilliz completion install {:?} --apply", shell);
    }

    Ok(())
}

fn uninstall(shell: Option<Shell>) -> Result<()> {
    let shell = resolve_shell(shell)?;
    let path = rc_path(shell);

    if !path.exists() {
        println!("RC file not found: {}", path.display());
        return Ok(());
    }

    let content = std::fs::read_to_string(&path)?;
    let lines: Vec<&str> = content.lines().collect();
    let mut new_lines = Vec::new();
    let mut skip_next = false;

    for line in &lines {
        if skip_next {
            skip_next = false;
            continue;
        }
        if line.contains(COMPLETION_MARKER) {
            skip_next = true; // skip the marker and the next line (the script)
            continue;
        }
        if line.contains("zilliz completion show") {
            continue;
        }
        new_lines.push(*line);
    }

    std::fs::write(&path, new_lines.join("\n") + "\n")?;
    println!("Completion removed from {}", path.display());

    Ok(())
}

fn status(shell: Option<Shell>) -> Result<()> {
    let shell = resolve_shell(shell)?;
    if is_installed(shell) {
        println!("Completion is installed for {:?}.", shell);
    } else {
        println!("Completion is NOT installed for {:?}.", shell);
        println!("Run: zilliz completion install {:?} --apply", shell);
    }
    Ok(())
}

fn show(shell: Shell) -> Result<()> {
    let mut cli = Cli::command();
    generate(shell, &mut cli, "zilliz", &mut std::io::stdout());
    Ok(())
}