forward.rs

//! SQL forwarding to leader node in cluster mode.

use actix_web::{HttpRequest, HttpResponse};
use kalamdb_commons::models::{NamespaceId, NodeId};
use kalamdb_commons::Role;
use kalamdb_core::app_context::AppContext;
use kalamdb_core::error::KalamDbError;
use kalamdb_raft::{ClusterClient, ForwardSqlRequest, GroupId, RaftExecutor};
use std::sync::Arc;
use std::time::Instant;

use super::helpers::parse_forward_params;
use super::models::{ErrorCode, QueryRequest, SqlResponse};

fn header_to_string(req: &HttpRequest, name: &str) -> Option<String> {
    req.headers().get(name).and_then(|v| v.to_str().ok()).map(|v| v.to_string())
}

fn normalize_addr(addr: &str) -> String {
    addr.trim()
        .trim_start_matches("http://")
        .trim_start_matches("https://")
        .trim_end_matches('/')
        .to_string()
}

fn resolve_node_id_from_cluster_addr(app_context: &AppContext, addr: &str) -> Option<NodeId> {
    let target = normalize_addr(addr);
    let cluster_info = app_context.executor().get_cluster_info();
    cluster_info
        .nodes
        .iter()
        .find(|n| normalize_addr(&n.api_addr) == target || normalize_addr(&n.rpc_addr) == target)
        .map(|n| n.node_id)
}

fn cluster_client_for(app_context: &AppContext) -> Result<ClusterClient, HttpResponse> {
    let executor = app_context.executor();
    let raft_executor = executor.as_any().downcast_ref::<RaftExecutor>().ok_or_else(|| {
        HttpResponse::ServiceUnavailable().json(SqlResponse::error(
            ErrorCode::ClusterUnavailable,
            "Cluster forwarding requires Raft executor",
            0.0,
        ))
    })?;
    Ok(ClusterClient::new(Arc::clone(raft_executor.manager())))
}

/// Forward target: either the Meta leader or a specific node.
enum ForwardTarget {
    Leader,
    Node(NodeId),
}

async fn forward_sql_grpc(
    target: ForwardTarget,
    http_req: &HttpRequest,
    req: &QueryRequest,
    app_context: &AppContext,
    request_id: Option<&str>,
    start_time: Instant,
) -> Option<HttpResponse> {
    let client = match cluster_client_for(app_context) {
        Ok(c) => c,
        Err(resp) => return Some(resp),
    };

    let params = match parse_forward_params(&req.params) {
        Ok(v) => v,
        Err(e) => {
            return Some(HttpResponse::BadRequest().json(SqlResponse::error(
                ErrorCode::InvalidParameter,
                &e,
                start_time.elapsed().as_secs_f64() * 1000.0,
            )));
        },
    };

    let grpc_req = ForwardSqlRequest {
        sql: req.sql.clone(),
        namespace_id: req.namespace_id.as_ref().map(|ns| ns.to_string()),
        authorization_header: header_to_string(http_req, "Authorization"),
        request_id: header_to_string(http_req, "X-Request-ID")
            .or_else(|| request_id.map(ToOwned::to_owned)),
        params,
    };

    let response = match target {
        ForwardTarget::Leader => client.forward_sql_to_leader(grpc_req).await,
        ForwardTarget::Node(node_id) => client.forward_sql_to_node(node_id, grpc_req).await,
    };

    let response = match response {
        Ok(resp) => resp,
        Err(err) => {
            log::warn!("Failed to forward SQL over gRPC: {}", err);
            return Some(HttpResponse::ServiceUnavailable().json(SqlResponse::error(
                ErrorCode::ForwardFailed,
                "Failed to forward request to cluster leader",
                start_time.elapsed().as_secs_f64() * 1000.0,
            )));
        },
    };

    if !response.error.is_empty() && response.body.is_empty() {
        return Some(HttpResponse::BadGateway().json(SqlResponse::error(
            ErrorCode::ForwardFailed,
            &response.error,
            start_time.elapsed().as_secs_f64() * 1000.0,
        )));
    }

    let status = actix_web::http::StatusCode::from_u16(response.status_code as u16)
        .unwrap_or(actix_web::http::StatusCode::BAD_GATEWAY);
    Some(HttpResponse::build(status).content_type("application/json").body(response.body))
}

/// Forwards write operations to the leader node in cluster mode.
pub async fn forward_sql_if_follower(
    http_req: &HttpRequest,
    req: &QueryRequest,
    app_context: &Arc<AppContext>,
    default_namespace: &NamespaceId,
    request_id: Option<&str>,
) -> Option<HttpResponse> {
    let start_time = Instant::now();
    let executor = app_context.executor();

    if executor.is_leader(GroupId::Meta).await {
        return None;
    }

    let statements = match kalamdb_sql::split_statements(&req.sql) {
        Ok(stmts) => stmts,
        Err(_) => {
            return forward_sql_grpc(
                ForwardTarget::Leader,
                http_req,
                req,
                app_context.as_ref(),
                request_id,
                start_time,
            )
            .await
        },
    };

    let has_write = statements.iter().any(|sql| {
        let classify_sql =
            kalamdb_sql::execute_as::extract_inner_sql(sql).unwrap_or_else(|| sql.to_string());
        let stmt = kalamdb_sql::classifier::SqlStatement::classify_and_parse(
            &classify_sql,
            default_namespace,
            Role::System,
        )
        .unwrap_or_else(|_| {
            kalamdb_sql::classifier::SqlStatement::new(
                classify_sql,
                kalamdb_sql::classifier::SqlStatementKind::Unknown,
            )
        });
        stmt.is_write_operation()
    });

    if has_write {
        return forward_sql_grpc(
            ForwardTarget::Leader,
            http_req,
            req,
            app_context.as_ref(),
            request_id,
            start_time,
        )
        .await;
    }

    None
}

/// Handle typed NOT_LEADER errors by forwarding to the known shard leader over gRPC.
pub async fn handle_not_leader_error(
    err: &KalamDbError,
    http_req: &HttpRequest,
    req: &QueryRequest,
    app_context: &AppContext,
    request_id: Option<&str>,
    start_time: Instant,
) -> Option<HttpResponse> {
    if !app_context.is_cluster_mode() {
        return None;
    }

    let leader_addr = match err {
        KalamDbError::NotLeader { leader_addr } => leader_addr.as_ref(),
        _ => return None,
    };

    if let Some(addr) = leader_addr {
        if let Some(target_node_id) = resolve_node_id_from_cluster_addr(app_context, addr) {
            return forward_sql_grpc(
                ForwardTarget::Node(target_node_id),
                http_req,
                req,
                app_context,
                request_id,
                start_time,
            )
            .await;
        }
        log::debug!(
            target: "sql::forward",
            "NOT_LEADER: leader addr '{}' not found in cluster info, falling back to generic leader forward",
            addr
        );
    }

    forward_sql_grpc(ForwardTarget::Leader, http_req, req, app_context, request_id, start_time)
        .await
}