claudetools/projects/msp-tools/guru-rmm/server/src/db/updates.rs

//! Database operations for agent updates

use anyhow::Result;
use sqlx::{PgPool, Row};
use uuid::Uuid;

/// Agent update record
#[derive(Debug, Clone, sqlx::FromRow)]
pub struct AgentUpdateRecord {
    pub id: Uuid,
    pub agent_id: Uuid,
    pub update_id: Uuid,
    pub old_version: String,
    pub target_version: String,
    pub download_url: Option<String>,
    pub checksum_sha256: Option<String>,
    pub status: Option<String>,
    pub started_at: Option<chrono::DateTime<chrono::Utc>>,
    pub completed_at: Option<chrono::DateTime<chrono::Utc>>,
    pub error_message: Option<String>,
}

/// Create a new agent update record
pub async fn create_agent_update(
    pool: &PgPool,
    agent_id: Uuid,
    update_id: Uuid,
    old_version: &str,
    target_version: &str,
    download_url: &str,
    checksum_sha256: &str,
) -> Result<()> {
    sqlx::query(
        r#"
        INSERT INTO agent_updates (agent_id, update_id, old_version, target_version, download_url, checksum_sha256, status)
        VALUES ($1, $2, $3, $4, $5, $6, 'pending')
        "#,
    )
    .bind(agent_id)
    .bind(update_id)
    .bind(old_version)
    .bind(target_version)
    .bind(download_url)
    .bind(checksum_sha256)
    .execute(pool)
    .await?;

    Ok(())
}

/// Mark an agent update as completed
pub async fn complete_agent_update(
    pool: &PgPool,
    update_id: Uuid,
    new_version: Option<&str>,
) -> Result<()> {
    sqlx::query(
        r#"
        UPDATE agent_updates
        SET status = 'completed', completed_at = NOW()
        WHERE update_id = $1
        "#,
    )
    .bind(update_id)
    .execute(pool)
    .await?;

    // If new_version provided, update the agent's version
    if let Some(version) = new_version {
        sqlx::query(
            r#"
            UPDATE agents
            SET agent_version = $2, updated_at = NOW()
            WHERE id = (SELECT agent_id FROM agent_updates WHERE update_id = $1)
            "#,
        )
        .bind(update_id)
        .bind(version)
        .execute(pool)
        .await?;
    }

    Ok(())
}

/// Mark an agent update as failed
pub async fn fail_agent_update(
    pool: &PgPool,
    update_id: Uuid,
    error_message: Option<&str>,
) -> Result<()> {
    sqlx::query(
        r#"
        UPDATE agent_updates
        SET status = 'failed', completed_at = NOW(), error_message = $2
        WHERE update_id = $1
        "#,
    )
    .bind(update_id)
    .bind(error_message)
    .execute(pool)
    .await?;

    Ok(())
}

/// Update the status of an agent update (for progress tracking)
pub async fn update_agent_update_status(
    pool: &PgPool,
    update_id: Uuid,
    status: &str,
) -> Result<()> {
    sqlx::query(
        r#"
        UPDATE agent_updates
        SET status = $2
        WHERE update_id = $1
        "#,
    )
    .bind(update_id)
    .bind(status)
    .execute(pool)
    .await?;

    Ok(())
}

/// Get pending update for an agent (if any)
pub async fn get_pending_update(
    pool: &PgPool,
    agent_id: Uuid,
) -> Result<Option<AgentUpdateRecord>> {
    let record = sqlx::query_as::<_, AgentUpdateRecord>(
        r#"
        SELECT id, agent_id, update_id, old_version, target_version,
               download_url, checksum_sha256, status, started_at, completed_at, error_message
        FROM agent_updates
        WHERE agent_id = $1 AND status IN ('pending', 'downloading', 'installing')
        ORDER BY started_at DESC
        LIMIT 1
        "#,
    )
    .bind(agent_id)
    .fetch_optional(pool)
    .await?;

    Ok(record)
}

/// Get stale updates (started but not completed within timeout)
pub async fn get_stale_updates(
    pool: &PgPool,
    timeout_secs: i64,
) -> Result<Vec<AgentUpdateRecord>> {
    let records = sqlx::query_as::<_, AgentUpdateRecord>(
        r#"
        SELECT id, agent_id, update_id, old_version, target_version,
               download_url, checksum_sha256, status, started_at, completed_at, error_message
        FROM agent_updates
        WHERE status IN ('pending', 'downloading', 'installing')
          AND started_at < NOW() - INTERVAL '1 second' * $1
        "#,
    )
    .bind(timeout_secs as f64)
    .fetch_all(pool)
    .await?;

    Ok(records)
}

/// Complete a pending update by matching agent reconnection
/// Called when an agent reconnects with a previous_version different from agent_version
pub async fn complete_update_by_agent(
    pool: &PgPool,
    agent_id: Uuid,
    pending_update_id: Option<Uuid>,
    old_version: &str,
    new_version: &str,
) -> Result<bool> {
    // First try by update_id if provided
    if let Some(update_id) = pending_update_id {
        let result = sqlx::query(
            r#"
            UPDATE agent_updates
            SET status = 'completed', completed_at = NOW()
            WHERE update_id = $1 AND agent_id = $2 AND status IN ('pending', 'downloading', 'installing')
            "#,
        )
        .bind(update_id)
        .bind(agent_id)
        .execute(pool)
        .await?;

        if result.rows_affected() > 0 {
            return Ok(true);
        }
    }

    // Fall back to finding by old_version match
    let result = sqlx::query(
        r#"
        UPDATE agent_updates
        SET status = 'completed', completed_at = NOW()
        WHERE agent_id = $1
          AND old_version = $2
          AND target_version = $3
          AND status IN ('pending', 'downloading', 'installing')
        "#,
    )
    .bind(agent_id)
    .bind(old_version)
    .bind(new_version)
    .execute(pool)
    .await?;

    Ok(result.rows_affected() > 0)
}