mz_deploy/cli/commands/

wait.rs

// Copyright Materialize, Inc. and contributors. All rights reserved.
//
// Use of this software is governed by the Business Source License
// included in the LICENSE file.
//
// As of the Change Date specified in that file, in accordance with
// the Business Source License, use of this software will be governed
// by the Apache License, Version 2.0.

//! Ready command - wait for staging deployment cluster hydration.

use crate::cli::CliError;
use crate::client::{
    Client, ClusterDeploymentStatus, ClusterStatusContext, FailureReason, HydrationStatusUpdate,
};
use crate::config::Settings;
use crate::log;
use crate::{info, info_nonl};
use crossterm::{
    cursor::{Hide, MoveToColumn, MoveUp, Show},
    execute,
    terminal::{Clear, ClearType},
};
use futures::StreamExt;
use owo_colors::{OwoColorize, Stream, Style};
use std::collections::BTreeMap;
use std::io::{self, Write};
use std::pin::pin;
use std::time::{Duration, Instant};

/// Wait for a staging deployment to become ready by monitoring hydration status.
///
/// This command:
/// - Validates the staging deployment exists and hasn't been promoted
/// - Subscribes to cluster hydration status
/// - Shows a live dashboard tracking hydration, lag, and health for each cluster
/// - Exits when all clusters are ready or timeout is reached
///
/// # Arguments
/// * `settings` - Resolved CLI settings (profile, project directory, etc.)
/// * `deploy_id` - Staging deployment ID
/// * `once` - If true, check once and exit; if false, track continuously
/// * `timeout` - Optional timeout in seconds
/// * `allowed_lag_secs` - Maximum allowed lag in seconds before marking as "lagging"
///
/// # Returns
/// `Ok(())` once the deployment is ready.
///
/// # Errors
/// Surfaces connection errors (including unknown or already-promoted
/// deployments) and `CliError::ReadyTimeout` when the timeout is reached.
pub async fn run(
    settings: &Settings,
    deploy_id: &str,
    once: bool,
    timeout: Option<u64>,
    allowed_lag_secs: i64,
) -> Result<(), CliError> {
    let profile = settings.connection();
    // Connect to database
    let mut client = Client::connect_with_profile(profile.clone())
        .await
        .map_err(CliError::Connection)?;
    // Validate staging deployment exists and is not promoted
    client.deployments().validate_staging(deploy_id).await?;

    if log::json_output_enabled() {
        if once {
            return run_snapshot_json(deploy_id, &client, allowed_lag_secs).await;
        } else {
            return run_continuous_json(deploy_id, &mut client, timeout, allowed_lag_secs).await;
        }
    }

    if once {
        // Snapshot mode: query once and display status
        run_snapshot(deploy_id, &client, allowed_lag_secs).await
    } else {
        // Continuous mode: subscribe and track with live dashboard
        run_continuous(deploy_id, &mut client, timeout, allowed_lag_secs).await
    }
}

/// Run in snapshot mode: query hydration status once and display.
async fn run_snapshot(
    deploy_id: &str,
    client: &Client,
    allowed_lag_secs: i64,
) -> Result<(), CliError> {
    let statuses = client
        .deployments()
        .get_deployment_hydration_status_with_lag(deploy_id, allowed_lag_secs)
        .await?;

    if statuses.is_empty() {
        info!("No clusters found in staging deployment '{}'", deploy_id);
        return Ok(());
    }

    info!();
    let style = Style::new().cyan().bold();
    info!(
        "{}",
        format!("  Deployment: {}", deploy_id)
            .if_supports_color(Stream::Stderr, |t| style.style(t))
    );
    info!();

    let mut has_failures = false;
    let mut has_non_ready = false;

    for ctx in &statuses {
        print_cluster_status(ctx, allowed_lag_secs);

        match &ctx.status {
            ClusterDeploymentStatus::Failing { .. } => has_failures = true,
            ClusterDeploymentStatus::Ready => {}
            _ => has_non_ready = true,
        }
    }

    info!();
    print_summary(&statuses);
    info!();

    if has_failures {
        Err(CliError::DeploymentFailing {
            name: deploy_id.to_string(),
        })
    } else if has_non_ready {
        Err(CliError::ClustersHydrating)
    } else {
        let style = Style::new().green().bold();
        info!(
            "{}",
            "  All clusters are ready!".if_supports_color(Stream::Stderr, |t| style.style(t))
        );
        Ok(())
    }
}

/// Run in snapshot mode with JSON output.
async fn run_snapshot_json(
    deploy_id: &str,
    client: &Client,
    allowed_lag_secs: i64,
) -> Result<(), CliError> {
    let statuses = client
        .deployments()
        .get_deployment_hydration_status_with_lag(deploy_id, allowed_lag_secs)
        .await?;

    let all_ready = statuses
        .iter()
        .all(|ctx| matches!(ctx.status, ClusterDeploymentStatus::Ready));

    let json = serde_json::json!({
        "deploy_id": deploy_id,
        "ready": all_ready,
        "clusters": statuses,
    });
    log::output_json(&json);

    if statuses
        .iter()
        .any(|ctx| matches!(ctx.status, ClusterDeploymentStatus::Failing { .. }))
    {
        Err(CliError::DeploymentFailing {
            name: deploy_id.to_string(),
        })
    } else if !all_ready {
        Err(CliError::ClustersHydrating)
    } else {
        Ok(())
    }
}

/// Run in continuous mode with NDJSON output.
async fn run_continuous_json(
    deploy_id: &str,
    client: &mut Client,
    timeout: Option<u64>,
    allowed_lag_secs: i64,
) -> Result<(), CliError> {
    let monitor_future = monitor_hydration_ndjson(deploy_id, client, allowed_lag_secs);

    if let Some(secs) = timeout {
        match tokio::time::timeout(Duration::from_secs(secs), monitor_future).await {
            Ok(result) => result,
            Err(_) => Err(CliError::ReadyTimeout {
                name: deploy_id.to_string(),
                seconds: secs,
            }),
        }
    } else {
        monitor_future.await
    }
}

/// Monitor hydration via SUBSCRIBE, emitting one NDJSON line per update.
async fn monitor_hydration_ndjson(
    deploy_id: &str,
    client: &mut Client,
    allowed_lag_secs: i64,
) -> Result<(), CliError> {
    let initial_statuses = client
        .deployments()
        .get_deployment_hydration_status_with_lag(deploy_id, allowed_lag_secs)
        .await?;

    if initial_statuses.is_empty() {
        let json = serde_json::json!({"deploy_id": deploy_id, "clusters": [], "all_ready": true});
        log::output_json(&json);
        return Ok(());
    }

    let mut cluster_states: BTreeMap<String, ClusterStatusContext> = initial_statuses
        .into_iter()
        .map(|ctx| (ctx.cluster_name.clone(), ctx))
        .collect();

    let mut deployments = client.deployments_mut();
    let stream = deployments.subscribe_deployment_hydration(deploy_id, allowed_lag_secs);
    let mut stream = pin!(stream);

    while let Some(result) = stream.next().await {
        let update = result.map_err(CliError::Connection)?;
        let status = update_to_status(&update);

        cluster_states.insert(
            update.cluster_name.clone(),
            ClusterStatusContext {
                cluster_name: update.cluster_name,
                cluster_id: update.cluster_id,
                status,
                hydrated_count: update.hydrated_count,
                total_count: update.total_count,
                max_lag_secs: update.max_lag_secs,
                total_replicas: update.total_replicas,
                problematic_replicas: update.problematic_replicas,
            },
        );

        let all_ready = cluster_states
            .values()
            .all(|ctx| matches!(ctx.status, ClusterDeploymentStatus::Ready));

        // Emit one NDJSON line per update
        let statuses: Vec<_> = cluster_states.values().collect();
        let json = serde_json::json!({
            "deploy_id": deploy_id,
            "all_ready": all_ready,
            "clusters": statuses,
        });
        log::output_json(&json);

        if all_ready {
            return Ok(());
        }
    }

    Ok(())
}

/// Print status for a single cluster with visual formatting.
fn print_cluster_status(ctx: &ClusterStatusContext, allowed_lag_secs: i64) {
    // Cluster name header
    info!(
        "  {}",
        ctx.cluster_name
            .as_str()
            .if_supports_color(Stream::Stderr, |t| t.bold())
    );

    // Progress bar
    let bar = render_progress_bar(ctx.hydrated_count, ctx.total_count, 40);
    let progress_str = format!("{}/{} objects", ctx.hydrated_count, ctx.total_count);

    info_nonl!("  {} ", bar);
    match &ctx.status {
        ClusterDeploymentStatus::Ready => {
            info_nonl!(
                "{} {}",
                "✓".if_supports_color(Stream::Stderr, |t| t.green()),
                "ready".if_supports_color(Stream::Stderr, |t| t.green())
            )
        }
        ClusterDeploymentStatus::Hydrating { .. } => {
            info_nonl!(
                "{} {}",
                "◐".if_supports_color(Stream::Stderr, |t| t.yellow()),
                "hydrating".if_supports_color(Stream::Stderr, |t| t.yellow())
            )
        }
        ClusterDeploymentStatus::Lagging { .. } => {
            info_nonl!(
                "{} {}",
                "⚠".if_supports_color(Stream::Stderr, |t| t.yellow()),
                "lagging".if_supports_color(Stream::Stderr, |t| t.yellow())
            )
        }
        ClusterDeploymentStatus::Failing { .. } => {
            info_nonl!(
                "{} {}",
                "✗".if_supports_color(Stream::Stderr, |t| t.red()),
                "failing".if_supports_color(Stream::Stderr, |t| t.red())
            )
        }
    }
    info!(
        "  {}",
        progress_str.if_supports_color(Stream::Stderr, |t| t.dimmed())
    );

    // Additional context based on status
    match &ctx.status {
        ClusterDeploymentStatus::Ready => {
            info!(
                "  {} lag: {}s",
                "└".if_supports_color(Stream::Stderr, |t| t.dimmed()),
                ctx.max_lag_secs
                    .to_string()
                    .if_supports_color(Stream::Stderr, |t| t.green())
            );
        }
        ClusterDeploymentStatus::Hydrating { hydrated, total } => {
            #[allow(clippy::as_conversions)]
            let pct = if *total > 0 {
                (*hydrated as f64 / *total as f64 * 100.0) as u8
            } else {
                0
            };
            info!(
                "  {} {}% complete",
                "└".if_supports_color(Stream::Stderr, |t| t.dimmed()),
                pct.to_string()
                    .if_supports_color(Stream::Stderr, |t| t.yellow())
            );
        }
        ClusterDeploymentStatus::Lagging { max_lag_secs } => {
            let style = Style::new().yellow().bold();
            info!(
                "  {} lag: {}s (threshold: {}s)",
                "└".if_supports_color(Stream::Stderr, |t| t.dimmed()),
                max_lag_secs
                    .to_string()
                    .if_supports_color(Stream::Stderr, |t| style.style(t)),
                allowed_lag_secs
            );
        }
        ClusterDeploymentStatus::Failing { reason } => {
            info!(
                "  {} {}",
                "└".if_supports_color(Stream::Stderr, |t| t.dimmed()),
                reason
                    .to_string()
                    .if_supports_color(Stream::Stderr, |t| t.red())
            );
        }
    }
    info!();
}

/// Render a Unicode progress bar.
#[allow(clippy::as_conversions)]
fn render_progress_bar(current: i64, total: i64, width: usize) -> String {
    if total == 0 {
        return format!(
            "[{}]",
            "░"
                .repeat(width)
                .if_supports_color(Stream::Stderr, |t| t.dimmed())
        );
    }

    let filled = ((current as f64 / total as f64) * width as f64) as usize;
    let empty = width.saturating_sub(filled);

    format!(
        "[{}{}]",
        "█"
            .repeat(filled)
            .if_supports_color(Stream::Stderr, |t| t.cyan()),
        "░"
            .repeat(empty)
            .if_supports_color(Stream::Stderr, |t| t.dimmed())
    )
}

/// Print summary footer with counts.
fn print_summary(statuses: &[ClusterStatusContext]) {
    let mut ready = 0;
    let mut hydrating = 0;
    let mut lagging = 0;
    let mut failing = 0;

    for ctx in statuses {
        match ctx.status {
            ClusterDeploymentStatus::Ready => ready += 1,
            ClusterDeploymentStatus::Hydrating { .. } => hydrating += 1,
            ClusterDeploymentStatus::Lagging { .. } => lagging += 1,
            ClusterDeploymentStatus::Failing { .. } => failing += 1,
        }
    }

    info_nonl!("  ");
    let mut parts = Vec::new();
    if ready > 0 {
        parts.push(
            format!("{} ready", ready)
                .if_supports_color(Stream::Stderr, |t| t.green())
                .to_string(),
        );
    }
    if hydrating > 0 {
        parts.push(
            format!("{} hydrating", hydrating)
                .if_supports_color(Stream::Stderr, |t| t.yellow())
                .to_string(),
        );
    }
    if lagging > 0 {
        parts.push(
            format!("{} lagging", lagging)
                .if_supports_color(Stream::Stderr, |t| t.yellow())
                .to_string(),
        );
    }
    if failing > 0 {
        parts.push(
            format!("{} failing", failing)
                .if_supports_color(Stream::Stderr, |t| t.red())
                .to_string(),
        );
    }
    info!("{}", parts.join(" · "));
}

/// Run in continuous mode: subscribe to hydration updates and show live dashboard.
async fn run_continuous(
    deploy_id: &str,
    client: &mut Client,
    timeout: Option<u64>,
    allowed_lag_secs: i64,
) -> Result<(), CliError> {
    // Get initial hydration status
    let initial_statuses = client
        .deployments()
        .get_deployment_hydration_status_with_lag(deploy_id, allowed_lag_secs)
        .await?;

    if initial_statuses.is_empty() {
        info!("No clusters found in staging deployment '{}'", deploy_id);
        return Ok(());
    }

    let start_time = Instant::now();

    // Subscribe to hydration updates and monitor
    let monitor_future = monitor_hydration_live(
        deploy_id,
        client,
        initial_statuses,
        start_time,
        allowed_lag_secs,
    );

    if let Some(secs) = timeout {
        match tokio::time::timeout(Duration::from_secs(secs), monitor_future).await {
            Ok(result) => result,
            Err(_) => Err(CliError::ReadyTimeout {
                name: deploy_id.to_string(),
                seconds: secs,
            }),
        }
    } else {
        monitor_future.await
    }
}

/// Monitor hydration status via SUBSCRIBE and update live dashboard.
async fn monitor_hydration_live(
    deploy_id: &str,
    client: &mut Client,
    initial_statuses: Vec<ClusterStatusContext>,
    start_time: Instant,
    allowed_lag_secs: i64,
) -> Result<(), CliError> {
    let mut stdout = io::stdout();
    let num_clusters = initial_statuses.len();

    // Build initial state map
    let mut cluster_states: BTreeMap<String, ClusterStatusContext> = initial_statuses
        .into_iter()
        .map(|ctx| (ctx.cluster_name.clone(), ctx))
        .collect();

    // Calculate lines per render (header + per-cluster lines + summary)
    // Header: 3 lines, per cluster: 4 lines, summary: 2 lines
    let lines_per_render = 3 + (num_clusters * 4) + 2;

    // Initial render
    render_dashboard(
        &mut stdout,
        deploy_id,
        &cluster_states,
        start_time,
        allowed_lag_secs,
    )?;

    let mut deployments = client.deployments_mut();
    let stream = deployments.subscribe_deployment_hydration(deploy_id, allowed_lag_secs);
    let mut stream = pin!(stream);

    // Hide cursor during updates
    let _ = execute!(stdout, Hide);

    while let Some(result) = stream.next().await {
        let update = result.map_err(CliError::Connection)?;

        let status = update_to_status(&update);

        cluster_states.insert(
            update.cluster_name.clone(),
            ClusterStatusContext {
                cluster_name: update.cluster_name,
                cluster_id: update.cluster_id,
                status,
                hydrated_count: update.hydrated_count,
                total_count: update.total_count,
                max_lag_secs: update.max_lag_secs,
                total_replicas: update.total_replicas,
                problematic_replicas: update.problematic_replicas,
            },
        );

        // Move cursor up and clear, then re-render
        #[allow(clippy::as_conversions)]
        let lines = lines_per_render as u16;
        let _ = execute!(stdout, MoveUp(lines), MoveToColumn(0));
        for _ in 0..lines_per_render {
            let _ = execute!(stdout, Clear(ClearType::CurrentLine));
            info!();
        }
        let _ = execute!(stdout, MoveUp(lines), MoveToColumn(0));

        render_dashboard(
            &mut stdout,
            deploy_id,
            &cluster_states,
            start_time,
            allowed_lag_secs,
        )?;

        let all_ready = cluster_states
            .values()
            .all(|ctx| matches!(ctx.status, ClusterDeploymentStatus::Ready));

        if all_ready {
            let _ = execute!(stdout, Show);
            info!();
            let style = Style::new().green().bold();
            info!(
                "{}",
                "  All clusters are ready!".if_supports_color(Stream::Stderr, |t| style.style(t))
            );
            return Ok(());
        }
    }

    let _ = execute!(stdout, Show);
    Ok(())
}

/// Convert a HydrationStatusUpdate to a ClusterDeploymentStatus.
fn update_to_status(update: &HydrationStatusUpdate) -> ClusterDeploymentStatus {
    match update.status {
        ClusterDeploymentStatus::Ready => ClusterDeploymentStatus::Ready,
        ClusterDeploymentStatus::Hydrating { .. } => ClusterDeploymentStatus::Hydrating {
            hydrated: update.hydrated_count,
            total: update.total_count,
        },
        ClusterDeploymentStatus::Lagging { .. } => ClusterDeploymentStatus::Lagging {
            max_lag_secs: update.max_lag_secs,
        },
        ClusterDeploymentStatus::Failing { .. } => {
            let reason = match update.failure_reason {
                Some(FailureReason::NoReplicas) => FailureReason::NoReplicas,
                Some(FailureReason::AllReplicasProblematic { .. }) => {
                    FailureReason::AllReplicasProblematic {
                        problematic: update.problematic_replicas,
                        total: update.total_replicas,
                    }
                }
                None => FailureReason::NoReplicas,
            };
            ClusterDeploymentStatus::Failing { reason }
        }
    }
}

/// Render the live dashboard.
fn render_dashboard(
    stdout: &mut io::Stdout,
    deploy_id: &str,
    cluster_states: &BTreeMap<String, ClusterStatusContext>,
    start_time: Instant,
    allowed_lag_secs: i64,
) -> Result<(), CliError> {
    let elapsed = start_time.elapsed();
    let elapsed_str = format_duration(elapsed);

    // Header
    info!();
    let header_style = Style::new().cyan().bold();
    info!(
        "{}",
        format!("  mz-deploy wait · deployment: {}", deploy_id)
            .if_supports_color(Stream::Stderr, |t| header_style.style(t))
    );
    info!(
        "  {} {}",
        "elapsed:".if_supports_color(Stream::Stderr, |t| t.dimmed()),
        elapsed_str.if_supports_color(Stream::Stderr, |t| t.dimmed())
    );
    info!();

    // Sort clusters by name for consistent ordering
    let mut cluster_names: Vec<_> = cluster_states.keys().collect();
    cluster_names.sort();

    // Render each cluster
    for name in cluster_names {
        if let Some(ctx) = cluster_states.get(name) {
            print_cluster_status(ctx, allowed_lag_secs);
        }
    }

    // Summary
    let statuses: Vec<_> = cluster_states.values().cloned().collect();
    print_summary(&statuses);

    let _ = stdout.flush();
    Ok(())
}

/// Format a duration as human-readable string.
fn format_duration(duration: Duration) -> String {
    let secs = duration.as_secs();
    if secs < 60 {
        format!("{}s", secs)
    } else if secs < 3600 {
        let mins = secs / 60;
        let secs = secs % 60;
        format!("{}m {}s", mins, secs)
    } else {
        let hours = secs / 3600;
        let mins = (secs % 3600) / 60;
        format!("{}h {}m", hours, mins)
    }
}