csm/crates/server/src/db.rs

use sqlx::SqlitePool;
use anyhow::Result;

/// Database repository for device operations
pub struct DeviceRepo;

impl DeviceRepo {
    pub async fn upsert_status(pool: &SqlitePool, device_uid: &str, status: &csm_protocol::DeviceStatus) -> Result<()> {
        let top_procs_json = serde_json::to_string(&status.top_processes)?;

        // Update latest snapshot
        sqlx::query(
            "INSERT INTO device_status (device_uid, cpu_usage, memory_usage, memory_total_mb, disk_usage, disk_total_mb, network_rx_rate, network_tx_rate, running_procs, top_processes, reported_at)
             VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, datetime('now'))
             ON CONFLICT(device_uid) DO UPDATE SET
                cpu_usage = excluded.cpu_usage,
                memory_usage = excluded.memory_usage,
                memory_total_mb = excluded.memory_total_mb,
                disk_usage = excluded.disk_usage,
                disk_total_mb = excluded.disk_total_mb,
                network_rx_rate = excluded.network_rx_rate,
                network_tx_rate = excluded.network_tx_rate,
                running_procs = excluded.running_procs,
                top_processes = excluded.top_processes,
                reported_at = datetime('now'),
                updated_at = datetime('now')"
        )
        .bind(device_uid)
        .bind(status.cpu_usage)
        .bind(status.memory_usage)
        .bind(status.memory_total_mb as i64)
        .bind(status.disk_usage)
        .bind(status.disk_total_mb as i64)
        .bind(status.network_rx_rate as i64)
        .bind(status.network_tx_rate as i64)
        .bind(status.running_procs as i32)
        .bind(&top_procs_json)
        .execute(pool)
        .await?;

        // Insert into history
        sqlx::query(
            "INSERT INTO device_status_history (device_uid, cpu_usage, memory_usage, disk_usage, network_rx_rate, network_tx_rate, running_procs, reported_at)
             VALUES (?, ?, ?, ?, ?, ?, ?, datetime('now'))"
        )
        .bind(device_uid)
        .bind(status.cpu_usage)
        .bind(status.memory_usage)
        .bind(status.disk_usage)
        .bind(status.network_rx_rate as i64)
        .bind(status.network_tx_rate as i64)
        .bind(status.running_procs as i32)
        .execute(pool)
        .await?;

        // Update device heartbeat
        sqlx::query(
            "UPDATE devices SET status = 'online', last_heartbeat = datetime('now') WHERE device_uid = ?"
        )
        .bind(device_uid)
        .execute(pool)
        .await?;

        Ok(())
    }

    pub async fn insert_usb_event(pool: &SqlitePool, event: &csm_protocol::UsbEvent) -> Result<i64> {
        let result = sqlx::query(
            "INSERT INTO usb_events (device_uid, vendor_id, product_id, serial_number, device_name, event_type)
             VALUES (?, ?, ?, ?, ?, ?)"
        )
        .bind(&event.device_uid)
        .bind(&event.vendor_id)
        .bind(&event.product_id)
        .bind(&event.serial)
        .bind(&event.device_name)
        .bind(match event.event_type {
            csm_protocol::UsbEventType::Inserted => "inserted",
            csm_protocol::UsbEventType::Removed => "removed",
            csm_protocol::UsbEventType::Blocked => "blocked",
        })
        .execute(pool)
        .await?;

        Ok(result.last_insert_rowid())
    }

    pub async fn upsert_hardware(pool: &SqlitePool, asset: &csm_protocol::HardwareAsset) -> Result<()> {
        sqlx::query(
            "INSERT INTO hardware_assets (device_uid, cpu_model, cpu_cores, memory_total_mb, disk_model, disk_total_mb, gpu_model, motherboard, serial_number, reported_at)
             VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, datetime('now'))
             ON CONFLICT(device_uid) DO UPDATE SET
                cpu_model = excluded.cpu_model,
                cpu_cores = excluded.cpu_cores,
                memory_total_mb = excluded.memory_total_mb,
                disk_model = excluded.disk_model,
                disk_total_mb = excluded.disk_total_mb,
                gpu_model = excluded.gpu_model,
                motherboard = excluded.motherboard,
                serial_number = excluded.serial_number,
                reported_at = datetime('now'),
                updated_at = datetime('now')"
        )
        .bind(&asset.device_uid)
        .bind(&asset.cpu_model)
        .bind(asset.cpu_cores as i32)
        .bind(asset.memory_total_mb as i64)
        .bind(&asset.disk_model)
        .bind(asset.disk_total_mb as i64)
        .bind(&asset.gpu_model)
        .bind(&asset.motherboard)
        .bind(&asset.serial_number)
        .execute(pool)
        .await?;

        Ok(())
    }

    pub async fn upsert_software(pool: &SqlitePool, asset: &csm_protocol::SoftwareAsset) -> Result<()> {
        // Use INSERT OR REPLACE to handle the UNIQUE(device_uid, name, version) constraint
        // where version can be NULL (treated as distinct by SQLite)
        let version = asset.version.as_deref().unwrap_or("");
        sqlx::query(
            "INSERT OR REPLACE INTO software_assets (device_uid, name, version, publisher, install_date, install_path, updated_at)
             VALUES (?, ?, ?, ?, ?, ?, datetime('now'))"
        )
        .bind(&asset.device_uid)
        .bind(&asset.name)
        .bind(if version.is_empty() { None } else { Some(version) })
        .bind(&asset.publisher)
        .bind(&asset.install_date)
        .bind(&asset.install_path)
        .execute(pool)
        .await?;

        Ok(())
    }
}