zilliz 1.4.1

TUI and CLI tool for managing Zilliz Cloud clusters and Milvus operations
Documentation
use std::fs;
use std::path::{Path, PathBuf};

use anyhow::{Context, Result};
use ini::Ini;

const ORG_SECTION_PREFIX: &str = "org.";

/// Extract API key from an org JSON object.
/// The API returns `apiKeys: [{key: "sk_..."}]` (array), not `apiKey` (string).
fn extract_api_key(org: &serde_json::Value) -> String {
    // Try apiKeys array first (login API response format)
    if let Some(keys) = org.get("apiKeys").and_then(|v| v.as_array()) {
        if let Some(first) = keys.first() {
            if let Some(key) = first.get("key").and_then(|v| v.as_str()) {
                return key.to_string();
            }
        }
    }
    // Fallback: try apiKey as a plain string
    org.get("apiKey")
        .and_then(|v| v.as_str())
        .unwrap_or("")
        .to_string()
}

/// Manages ~/.zilliz/config and ~/.zilliz/credentials files.
/// Compatible with the Python zilliz-cli INI format.
pub struct ConfigManager {
    #[allow(dead_code)]
    pub config_dir: PathBuf,
    credentials_path: PathBuf,
    config_path: PathBuf,
}

impl ConfigManager {
    pub fn new(config_dir: Option<PathBuf>) -> Result<Self> {
        let config_dir = match config_dir {
            Some(d) => d,
            None => {
                if let Ok(env_dir) = std::env::var("ZILLIZ_CONFIG_DIR") {
                    PathBuf::from(env_dir)
                } else {
                    dirs::home_dir()
                        .context("Cannot determine home directory")?
                        .join(".zilliz")
                }
            }
        };

        fs::create_dir_all(&config_dir)
            .with_context(|| format!("Failed to create config dir: {}", config_dir.display()))?;

        let credentials_path = config_dir.join("credentials");
        let config_path = config_dir.join("config");

        Ok(Self {
            config_dir,
            credentials_path,
            config_path,
        })
    }

    // --- Credentials ---

    pub fn get_credential(&self, key: &str) -> Option<String> {
        let ini = self.read_ini(&self.credentials_path);
        ini.get_from(Some("default"), key).map(|s| {
            // Strip quotes and whitespace that INI parsers may preserve
            s.trim().trim_matches('"').trim_matches('\'').to_string()
        })
    }

    pub fn set_credential(&self, key: &str, value: &str) -> Result<()> {
        let mut ini = self.read_ini(&self.credentials_path);
        ini.with_section(Some("default")).set(key, value);
        self.write_ini(&self.credentials_path, &ini)?;
        Ok(())
    }

    /// Remove the credentials file entirely.
    pub fn clear_credentials(&self) -> Result<()> {
        if self.credentials_path.exists() {
            fs::remove_file(&self.credentials_path)
                .with_context(|| format!("Failed to remove {}", self.credentials_path.display()))?;
        }
        Ok(())
    }

    /// Return all credentials and config key-value pairs.
    #[allow(clippy::type_complexity)]
    pub fn list_all(&self) -> (Vec<(String, String)>, Vec<(String, String)>) {
        let mut credentials = Vec::new();
        let mut config = Vec::new();

        let cred_ini = self.read_ini(&self.credentials_path);
        if let Some(props) = cred_ini.section(Some("default")) {
            for (k, v) in props.iter() {
                credentials.push((k.to_string(), v.to_string()));
            }
        }

        let cfg_ini = self.read_ini(&self.config_path);
        if let Some(props) = cfg_ini.section(Some("default")) {
            for (k, v) in props.iter() {
                config.push((k.to_string(), v.to_string()));
            }
        }

        (credentials, config)
    }

    // --- Config ---

    pub fn get_config(&self, section: &str, key: &str) -> Option<String> {
        let ini = self.read_ini(&self.config_path);
        ini.get_from(Some(section), key).map(|s| s.to_string())
    }

    pub fn set_config(&self, section: &str, key: &str, value: &str) -> Result<()> {
        let mut ini = self.read_ini(&self.config_path);
        ini.with_section(Some(section)).set(key, value);
        self.write_ini(&self.config_path, &ini)?;
        Ok(())
    }

    // --- Persisted control-plane endpoint ---
    //
    // Stored as `[default].endpoint` in `~/.zilliz/config`. Written by
    // `zilliz login --cn`, cleared by `zilliz logout` and by `zilliz login`
    // without `--cn`. Distinct from the per-cluster `[context].endpoint`.

    pub fn get_control_plane_endpoint(&self) -> Option<String> {
        let ini = self.read_ini(&self.config_path);
        ini.get_from(Some("default"), "endpoint")
            .map(|s| s.trim().trim_matches('"').trim_matches('\'').to_string())
    }

    pub fn set_control_plane_endpoint(&self, url: &str) -> Result<()> {
        self.set_config("default", "endpoint", url)
    }

    pub fn clear_control_plane_endpoint(&self) -> Result<()> {
        let mut ini = self.read_ini(&self.config_path);
        if ini.delete_from(Some("default"), "endpoint").is_none() {
            return Ok(());
        }
        self.write_ini(&self.config_path, &ini)?;
        Ok(())
    }

    // --- Context ---

    pub fn get_context(&self) -> super::context::ClusterContext {
        let ini = self.read_ini(&self.config_path);
        super::context::ClusterContext {
            cluster_id: ini
                .get_from(Some("context"), "cluster_id")
                .map(|s| s.to_string()),
            endpoint: ini
                .get_from(Some("context"), "endpoint")
                .map(|s| s.to_string()),
            database: ini
                .get_from(Some("context"), "database")
                .map(|s| s.to_string())
                .unwrap_or_else(|| "default".to_string()),
            plan: ini.get_from(Some("context"), "plan").map(|s| s.to_string()),
        }
    }

    pub fn set_context(&self, ctx: &super::context::ClusterContext) -> Result<()> {
        let mut ini = self.read_ini(&self.config_path);
        let mut section = ini.with_section(Some("context"));
        if let Some(ref id) = ctx.cluster_id {
            section.set("cluster_id", id);
        }
        if let Some(ref ep) = ctx.endpoint {
            section.set("endpoint", ep);
        }
        section.set("database", &ctx.database);
        if let Some(ref plan) = ctx.plan {
            section.set("plan", plan);
        }
        self.write_ini(&self.config_path, &ini)?;
        Ok(())
    }

    pub fn clear_context(&self) -> Result<()> {
        let mut ini = self.read_ini(&self.config_path);
        ini.delete(Some("context"));
        self.write_ini(&self.config_path, &ini)?;
        Ok(())
    }

    // --- Org management ---

    #[allow(dead_code)]
    pub fn get_all_orgs(&self) -> Vec<OrgInfo> {
        let ini = self.read_ini(&self.credentials_path);
        let mut orgs = Vec::new();
        for (section, _) in &ini {
            if let Some(name) = section {
                if let Some(org_id) = name.strip_prefix(ORG_SECTION_PREFIX) {
                    let org_name = ini
                        .get_from(Some(name), "org_name")
                        .unwrap_or("")
                        .to_string();
                    let api_key = ini.get_from(Some(name), "api_key").map(|s| s.to_string());
                    orgs.push(OrgInfo {
                        org_id: org_id.to_string(),
                        name: org_name,
                        api_key,
                    });
                }
            }
        }
        orgs
    }

    pub fn get_current_org(&self) -> Option<OrgInfo> {
        let ini = self.read_ini(&self.credentials_path);
        let org_id = ini.get_from(Some("default"), "org_id")?;
        Some(OrgInfo {
            org_id: org_id.to_string(),
            name: ini
                .get_from(Some("default"), "org_name")
                .unwrap_or("")
                .to_string(),
            api_key: ini
                .get_from(Some("default"), "api_key")
                .map(|s| s.to_string()),
        })
    }

    // --- User info (from OAuth login) ---

    pub fn get_user_info(&self) -> Option<UserInfo> {
        let ini = self.read_ini(&self.credentials_path);
        let email = ini.get_from(Some("user"), "email")?.to_string();
        Some(UserInfo {
            user_id: ini
                .get_from(Some("user"), "user_id")
                .unwrap_or("")
                .to_string(),
            email,
            name: ini.get_from(Some("user"), "name").unwrap_or("").to_string(),
        })
    }

    /// Save login data from CLI login API response.
    /// Compatible with zilliz-cli INI format.
    pub fn save_login_data(
        &self,
        user_id: &str,
        email: &str,
        name: &str,
        orgs: &[serde_json::Value],
    ) -> Result<()> {
        let mut ini = self.read_ini(&self.credentials_path);

        // Clear old org sections
        let old_sections: Vec<String> = ini
            .iter()
            .filter_map(|(s, _)| s.map(|s| s.to_string()))
            .filter(|s| s.starts_with(ORG_SECTION_PREFIX) || s == "orgs")
            .collect();
        for section in &old_sections {
            ini.delete(Some(section.as_str()));
        }

        // Save user info
        ini.with_section(Some("user"))
            .set("user_id", user_id)
            .set("email", email)
            .set("name", name);

        // Save orgs and set first org as default
        if let Some(first_org) = orgs.first() {
            let org_id = first_org
                .get("orgId")
                .and_then(|v| v.as_str())
                .unwrap_or("");
            let org_name = first_org.get("name").and_then(|v| v.as_str()).unwrap_or("");
            let api_key = extract_api_key(first_org);

            let mut section = ini.with_section(Some("default"));
            section.set("org_id", org_id).set("org_name", org_name);
            if !api_key.is_empty() {
                section.set("api_key", &api_key);
            }
        }

        for org in orgs {
            let org_id = org
                .get("orgId")
                .and_then(|v| v.as_str())
                .unwrap_or_default();
            let org_name = org.get("name").and_then(|v| v.as_str()).unwrap_or_default();
            let api_key = extract_api_key(org);

            if !org_id.is_empty() {
                let section_name = format!("{}{}", ORG_SECTION_PREFIX, org_id);
                let mut section = ini.with_section(Some(section_name));
                section.set("org_name", org_name);
                if !api_key.is_empty() {
                    section.set("api_key", &api_key);
                }
            }
        }

        self.write_ini(&self.credentials_path, &ini)?;
        Ok(())
    }

    /// Save only an API key (manual login without OAuth).
    pub fn save_api_key_only(&self, api_key: &str) -> Result<()> {
        let mut ini = self.read_ini(&self.credentials_path);

        // Clear user and org sections
        ini.delete(Some("user"));
        let old_sections: Vec<String> = ini
            .iter()
            .filter_map(|(s, _)| s.map(|s| s.to_string()))
            .filter(|s| s.starts_with(ORG_SECTION_PREFIX) || s == "orgs")
            .collect();
        for section in &old_sections {
            ini.delete(Some(section.as_str()));
        }

        ini.with_section(Some("default")).set("api_key", api_key);
        self.write_ini(&self.credentials_path, &ini)?;
        Ok(())
    }

    /// Clear all login data (user info, orgs, api key).
    pub fn clear_login_data(&self) -> Result<()> {
        let mut ini = self.read_ini(&self.credentials_path);

        ini.delete(Some("user"));
        ini.delete(Some("default"));

        let old_sections: Vec<String> = ini
            .iter()
            .filter_map(|(s, _)| s.map(|s| s.to_string()))
            .filter(|s| s.starts_with(ORG_SECTION_PREFIX) || s == "orgs")
            .collect();
        for section in &old_sections {
            ini.delete(Some(section.as_str()));
        }

        self.write_ini(&self.credentials_path, &ini)?;
        Ok(())
    }

    // --- Internal ---

    fn read_ini(&self, path: &Path) -> Ini {
        if path.exists() {
            Ini::load_from_file(path).unwrap_or_default()
        } else {
            Ini::new()
        }
    }

    fn write_ini(&self, path: &Path, ini: &Ini) -> Result<()> {
        ini.write_to_file(path)
            .with_context(|| format!("Failed to write {}", path.display()))?;

        // Restrict credentials file permissions (owner read/write only)
        #[cfg(unix)]
        if path == self.credentials_path {
            use std::os::unix::fs::PermissionsExt;
            fs::set_permissions(path, fs::Permissions::from_mode(0o600))?;
        }

        Ok(())
    }
}

#[derive(Debug, Clone)]
#[allow(dead_code)]
pub struct OrgInfo {
    pub org_id: String,
    pub name: String,
    pub api_key: Option<String>,
}

#[derive(Debug, Clone)]
#[allow(dead_code)]
pub struct UserInfo {
    pub user_id: String,
    pub email: String,
    pub name: String,
}