smartdb/rust/crates/rustdb/src/management.rs

use anyhow::Result;
use serde::{Deserialize, Serialize};
use tokio::io::{AsyncBufReadExt, BufReader};
use tracing::{info, error};

use crate::RustDb;
use rustdb_config::RustDbOptions;

/// A management request from the TypeScript wrapper.
#[derive(Debug, Deserialize)]
pub struct ManagementRequest {
    pub id: String,
    pub method: String,
    #[serde(default)]
    pub params: serde_json::Value,
}

/// A management response back to the TypeScript wrapper.
#[derive(Debug, Serialize)]
pub struct ManagementResponse {
    pub id: String,
    pub success: bool,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub result: Option<serde_json::Value>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub error: Option<String>,
}

/// An unsolicited event from the server to the TypeScript wrapper.
#[derive(Debug, Serialize)]
pub struct ManagementEvent {
    pub event: String,
    pub data: serde_json::Value,
}

impl ManagementResponse {
    fn ok(id: String, result: serde_json::Value) -> Self {
        Self {
            id,
            success: true,
            result: Some(result),
            error: None,
        }
    }

    fn err(id: String, message: String) -> Self {
        Self {
            id,
            success: false,
            result: None,
            error: Some(message),
        }
    }
}

fn send_line(line: &str) {
    use std::io::Write;
    let stdout = std::io::stdout();
    let mut handle = stdout.lock();
    let _ = handle.write_all(line.as_bytes());
    let _ = handle.write_all(b"\n");
    let _ = handle.flush();
}

fn send_response(response: &ManagementResponse) {
    match serde_json::to_string(response) {
        Ok(json) => send_line(&json),
        Err(e) => error!("Failed to serialize management response: {}", e),
    }
}

fn send_event(event: &str, data: serde_json::Value) {
    let evt = ManagementEvent {
        event: event.to_string(),
        data,
    };
    match serde_json::to_string(&evt) {
        Ok(json) => send_line(&json),
        Err(e) => error!("Failed to serialize management event: {}", e),
    }
}

/// Run the management loop, reading JSON commands from stdin and writing responses to stdout.
pub async fn management_loop() -> Result<()> {
    let stdin = BufReader::new(tokio::io::stdin());
    let mut lines = stdin.lines();
    let mut db: Option<RustDb> = None;

    send_event("ready", serde_json::json!({}));

    loop {
        let line = match lines.next_line().await {
            Ok(Some(line)) => line,
            Ok(None) => {
                // stdin closed - parent process exited
                info!("Management stdin closed, shutting down");
                if let Some(ref mut d) = db {
                    let _ = d.stop().await;
                }
                break;
            }
            Err(e) => {
                error!("Error reading management stdin: {}", e);
                break;
            }
        };

        let line = line.trim().to_string();
        if line.is_empty() {
            continue;
        }

        let request: ManagementRequest = match serde_json::from_str(&line) {
            Ok(r) => r,
            Err(e) => {
                error!("Failed to parse management request: {}", e);
                send_response(&ManagementResponse::err(
                    "unknown".to_string(),
                    format!("Failed to parse request: {}", e),
                ));
                continue;
            }
        };

        let response = handle_request(&request, &mut db).await;
        send_response(&response);
    }

    Ok(())
}

async fn handle_request(
    request: &ManagementRequest,
    db: &mut Option<RustDb>,
) -> ManagementResponse {
    let id = request.id.clone();

    match request.method.as_str() {
        "start" => handle_start(&id, &request.params, db).await,
        "stop" => handle_stop(&id, db).await,
        "getStatus" => handle_get_status(&id, db),
        "getMetrics" => handle_get_metrics(&id, db),
        _ => ManagementResponse::err(id, format!("Unknown method: {}", request.method)),
    }
}

async fn handle_start(
    id: &str,
    params: &serde_json::Value,
    db: &mut Option<RustDb>,
) -> ManagementResponse {
    if db.is_some() {
        return ManagementResponse::err(id.to_string(), "Server is already running".to_string());
    }

    let config = match params.get("config") {
        Some(config) => config,
        None => return ManagementResponse::err(id.to_string(), "Missing 'config' parameter".to_string()),
    };

    let options: RustDbOptions = match serde_json::from_value(config.clone()) {
        Ok(o) => o,
        Err(e) => return ManagementResponse::err(id.to_string(), format!("Invalid config: {}", e)),
    };

    let connection_uri = options.connection_uri();

    match RustDb::new(options).await {
        Ok(mut d) => {
            match d.start().await {
                Ok(()) => {
                    send_event("started", serde_json::json!({}));
                    *db = Some(d);
                    ManagementResponse::ok(
                        id.to_string(),
                        serde_json::json!({ "connectionUri": connection_uri }),
                    )
                }
                Err(e) => {
                    send_event("error", serde_json::json!({"message": format!("{}", e)}));
                    ManagementResponse::err(id.to_string(), format!("Failed to start: {}", e))
                }
            }
        }
        Err(e) => ManagementResponse::err(id.to_string(), format!("Failed to create server: {}", e)),
    }
}

async fn handle_stop(
    id: &str,
    db: &mut Option<RustDb>,
) -> ManagementResponse {
    match db.as_mut() {
        Some(d) => {
            match d.stop().await {
                Ok(()) => {
                    *db = None;
                    send_event("stopped", serde_json::json!({}));
                    ManagementResponse::ok(id.to_string(), serde_json::json!({}))
                }
                Err(e) => ManagementResponse::err(id.to_string(), format!("Failed to stop: {}", e)),
            }
        }
        None => ManagementResponse::ok(id.to_string(), serde_json::json!({})),
    }
}

fn handle_get_status(
    id: &str,
    db: &Option<RustDb>,
) -> ManagementResponse {
    match db.as_ref() {
        Some(_d) => ManagementResponse::ok(
            id.to_string(),
            serde_json::json!({
                "running": true,
            }),
        ),
        None => ManagementResponse::ok(
            id.to_string(),
            serde_json::json!({ "running": false }),
        ),
    }
}

fn handle_get_metrics(
    id: &str,
    db: &Option<RustDb>,
) -> ManagementResponse {
    match db.as_ref() {
        Some(_d) => ManagementResponse::ok(
            id.to_string(),
            serde_json::json!({
                "connections": 0,
                "databases": 0,
            }),
        ),
        None => ManagementResponse::err(id.to_string(), "Server is not running".to_string()),
    }
}