yaak/crates-cli/yaak-cli/src/commands/send.rs

use crate::cli::SendArgs;
use crate::commands::request;
use crate::context::CliContext;
use futures::future::join_all;
use yaak_models::queries::any_request::AnyRequest;

enum ExecutionMode {
    Sequential,
    Parallel,
}

pub async fn run(
    ctx: &CliContext,
    args: SendArgs,
    environment: Option<&str>,
    cookie_jar_id: Option<&str>,
    verbose: bool,
) -> i32 {
    match send_target(ctx, args, environment, cookie_jar_id, verbose).await {
        Ok(()) => 0,
        Err(error) => {
            eprintln!("Error: {error}");
            1
        }
    }
}

async fn send_target(
    ctx: &CliContext,
    args: SendArgs,
    environment: Option<&str>,
    cookie_jar_id: Option<&str>,
    verbose: bool,
) -> Result<(), String> {
    let mode = if args.parallel { ExecutionMode::Parallel } else { ExecutionMode::Sequential };

    if let Ok(request) = ctx.db().get_any_request(&args.id) {
        let workspace_id = match &request {
            AnyRequest::HttpRequest(r) => r.workspace_id.clone(),
            AnyRequest::GrpcRequest(r) => r.workspace_id.clone(),
            AnyRequest::WebsocketRequest(r) => r.workspace_id.clone(),
        };
        let resolved_cookie_jar_id =
            request::resolve_cookie_jar_id(ctx, &workspace_id, cookie_jar_id)?;

        return request::send_request_by_id(
            ctx,
            &args.id,
            environment,
            resolved_cookie_jar_id.as_deref(),
            verbose,
        )
        .await;
    }

    if let Ok(folder) = ctx.db().get_folder(&args.id) {
        let resolved_cookie_jar_id =
            request::resolve_cookie_jar_id(ctx, &folder.workspace_id, cookie_jar_id)?;

        let request_ids = collect_folder_request_ids(ctx, &args.id)?;
        if request_ids.is_empty() {
            println!("No requests found in folder {}", args.id);
            return Ok(());
        }
        return send_many(
            ctx,
            request_ids,
            mode,
            args.fail_fast,
            environment,
            resolved_cookie_jar_id.as_deref(),
            verbose,
        )
        .await;
    }

    if let Ok(workspace) = ctx.db().get_workspace(&args.id) {
        let resolved_cookie_jar_id =
            request::resolve_cookie_jar_id(ctx, &workspace.id, cookie_jar_id)?;

        let request_ids = collect_workspace_request_ids(ctx, &args.id)?;
        if request_ids.is_empty() {
            println!("No requests found in workspace {}", args.id);
            return Ok(());
        }
        return send_many(
            ctx,
            request_ids,
            mode,
            args.fail_fast,
            environment,
            resolved_cookie_jar_id.as_deref(),
            verbose,
        )
        .await;
    }

    Err(format!("Could not resolve ID '{}' as request, folder, or workspace", args.id))
}

fn collect_folder_request_ids(ctx: &CliContext, folder_id: &str) -> Result<Vec<String>, String> {
    let mut ids = Vec::new();

    let mut http_ids = ctx
        .db()
        .list_http_requests_for_folder_recursive(folder_id)
        .map_err(|e| format!("Failed to list HTTP requests in folder: {e}"))?
        .into_iter()
        .map(|r| r.id)
        .collect::<Vec<_>>();
    ids.append(&mut http_ids);

    let mut grpc_ids = ctx
        .db()
        .list_grpc_requests_for_folder_recursive(folder_id)
        .map_err(|e| format!("Failed to list gRPC requests in folder: {e}"))?
        .into_iter()
        .map(|r| r.id)
        .collect::<Vec<_>>();
    ids.append(&mut grpc_ids);

    let mut websocket_ids = ctx
        .db()
        .list_websocket_requests_for_folder_recursive(folder_id)
        .map_err(|e| format!("Failed to list WebSocket requests in folder: {e}"))?
        .into_iter()
        .map(|r| r.id)
        .collect::<Vec<_>>();
    ids.append(&mut websocket_ids);

    Ok(ids)
}

fn collect_workspace_request_ids(
    ctx: &CliContext,
    workspace_id: &str,
) -> Result<Vec<String>, String> {
    let mut ids = Vec::new();

    let mut http_ids = ctx
        .db()
        .list_http_requests(workspace_id)
        .map_err(|e| format!("Failed to list HTTP requests in workspace: {e}"))?
        .into_iter()
        .map(|r| r.id)
        .collect::<Vec<_>>();
    ids.append(&mut http_ids);

    let mut grpc_ids = ctx
        .db()
        .list_grpc_requests(workspace_id)
        .map_err(|e| format!("Failed to list gRPC requests in workspace: {e}"))?
        .into_iter()
        .map(|r| r.id)
        .collect::<Vec<_>>();
    ids.append(&mut grpc_ids);

    let mut websocket_ids = ctx
        .db()
        .list_websocket_requests(workspace_id)
        .map_err(|e| format!("Failed to list WebSocket requests in workspace: {e}"))?
        .into_iter()
        .map(|r| r.id)
        .collect::<Vec<_>>();
    ids.append(&mut websocket_ids);

    Ok(ids)
}

async fn send_many(
    ctx: &CliContext,
    request_ids: Vec<String>,
    mode: ExecutionMode,
    fail_fast: bool,
    environment: Option<&str>,
    cookie_jar_id: Option<&str>,
    verbose: bool,
) -> Result<(), String> {
    let mut success_count = 0usize;
    let mut failures: Vec<(String, String)> = Vec::new();

    match mode {
        ExecutionMode::Sequential => {
            for request_id in request_ids {
                match request::send_request_by_id(
                    ctx,
                    &request_id,
                    environment,
                    cookie_jar_id,
                    verbose,
                )
                .await
                {
                    Ok(()) => success_count += 1,
                    Err(error) => {
                        failures.push((request_id, error));
                        if fail_fast {
                            break;
                        }
                    }
                }
            }
        }
        ExecutionMode::Parallel => {
            let tasks = request_ids
                .iter()
                .map(|request_id| async move {
                    (
                        request_id.clone(),
                        request::send_request_by_id(
                            ctx,
                            request_id,
                            environment,
                            cookie_jar_id,
                            verbose,
                        )
                        .await,
                    )
                })
                .collect::<Vec<_>>();

            for (request_id, result) in join_all(tasks).await {
                match result {
                    Ok(()) => success_count += 1,
                    Err(error) => failures.push((request_id, error)),
                }
            }
        }
    }

    let failure_count = failures.len();
    println!("Send summary: {success_count} succeeded, {failure_count} failed");

    if failure_count == 0 {
        return Ok(());
    }

    for (request_id, error) in failures {
        eprintln!("  {}: {}", request_id, error);
    }
    Err("One or more requests failed".to_string())
}