linkrs-server 0.1.0

//! Batch processing of HTTP handlers

use axum::{
    extract::{Extension, Json, Path, State},
    response::Json as JsonResponse,
};
use serde_json;

use crate::batch::{
    AddBatchItemsRequest, AddBatchItemsResponse, BatchStatusResponse, CreateBatchRequest,
    CreateBatchResponse, ExecuteBatchResponse,
};
use crate::http::{error::HttpError, state::AppState};
use crate::storage::{
    StorageClient, StorageOperationContextOps, StorageSchemaContextOps, StorageSyncContextOps,
};

#[utoipa::path(
    post,
    operation_id = "post_v1_batch",
    path = "/v1/batch",
    tag = "Batch",
    request_body = CreateBatchRequest,
    responses(
        (status = 200, body = CreateBatchResponse, description = "Batch task created"),
        (status = 500, description = "Internal error")
    )
)]
/// Create batch tasks
pub async fn create<
    S: StorageClient
        + StorageSchemaContextOps
        + StorageSyncContextOps
        + StorageOperationContextOps
        + Clone
        + Send
        + Sync
        + 'static,
>(
    State(state): State<AppState<S>>,
    Extension(session_id): Extension<i64>,
    Json(request): Json<CreateBatchRequest>,
) -> Result<JsonResponse<CreateBatchResponse>, HttpError> {
    super::authz::require_admin_session(&state, session_id)?;
    let batch_manager = state.server.get_batch_manager();

    match batch_manager.create_task(request.space_id, request.batch_type, request.batch_size) {
        Ok(task) => Ok(JsonResponse(CreateBatchResponse {
            batch_id: task.id,
            status: task.status,
            created_at: task.created_at.to_rfc3339(),
        })),
        Err(e) => Err(HttpError::InternalError(format!(
            "Failed to create batch task: {}",
            e
        ))),
    }
}

#[utoipa::path(
    get,
    operation_id = "get_v1_batch_id",
    path = "/v1/batch/{id}",
    tag = "Batch",
    params(("id" = String, Path, description = "Batch task id")),
    responses(
        (status = 200, body = BatchStatusResponse, description = "Batch task status"),
        (status = 404, description = "Not found"),
        (status = 500, description = "Internal error")
    )
)]
/// Obtaining the status of batch tasks
pub async fn status<
    S: StorageClient
        + StorageSchemaContextOps
        + StorageSyncContextOps
        + StorageOperationContextOps
        + Clone
        + Send
        + Sync
        + 'static,
>(
    State(state): State<AppState<S>>,
    Extension(session_id): Extension<i64>,
    Path(batch_id): Path<String>,
) -> Result<JsonResponse<BatchStatusResponse>, HttpError> {
    super::authz::require_admin_session(&state, session_id)?;
    let batch_manager = state.server.get_batch_manager();

    match batch_manager.get_task(&batch_id) {
        Some(task) => Ok(JsonResponse(BatchStatusResponse {
            batch_id: task.id,
            status: task.status,
            progress: task.progress,
            created_at: task.created_at.to_rfc3339(),
            updated_at: task.updated_at.to_rfc3339(),
        })),
        None => Err(HttpError::NotFound(format!(
            "Batch task does not exist: {}",
            batch_id
        ))),
    }
}

#[utoipa::path(
    post,
    path = "/v1/batch/{id}/items",
    tag = "Batch",
    params(("id" = String, Path, description = "Batch task id")),
    request_body = AddBatchItemsRequest,
    responses(
        (status = 200, body = AddBatchItemsResponse, description = "Items accepted"),
        (status = 404, description = "Not found"),
        (status = 500, description = "Internal error")
    )
)]
/// Add multiple items in batches
pub async fn add_items<
    S: StorageClient
        + StorageSchemaContextOps
        + StorageSyncContextOps
        + StorageOperationContextOps
        + Clone
        + Send
        + Sync
        + 'static,
>(
    State(state): State<AppState<S>>,
    Extension(session_id): Extension<i64>,
    Path(batch_id): Path<String>,
    Json(request): Json<AddBatchItemsRequest>,
) -> Result<JsonResponse<AddBatchItemsResponse>, HttpError> {
    super::authz::require_admin_session(&state, session_id)?;
    let batch_manager = state.server.get_batch_manager();

    match batch_manager.add_items(&batch_id, request.items) {
        Ok(accepted) => {
            let task = batch_manager.get_task(&batch_id).ok_or_else(|| {
                HttpError::NotFound(format!("Batch task does not exist: {}", batch_id))
            })?;

            Ok(JsonResponse(AddBatchItemsResponse {
                accepted,
                buffered: accepted,
                total_buffered: task.progress.buffered,
            }))
        }
        Err(e) => Err(HttpError::BadRequest(format!(
            "Adding batch items failed: {}",
            e
        ))),
    }
}

#[utoipa::path(
    post,
    operation_id = "post_v1_batch_id_execute",
    path = "/v1/batch/{id}/execute",
    tag = "Batch",
    params(("id" = String, Path, description = "Batch task id")),
    responses(
        (status = 200, body = ExecuteBatchResponse, description = "Batch task executed"),
        (status = 404, description = "Not found"),
        (status = 500, description = "Internal error")
    )
)]
/// Perform batch tasks
pub async fn execute<
    S: StorageClient
        + StorageSchemaContextOps
        + StorageSyncContextOps
        + StorageOperationContextOps
        + Clone
        + Send
        + Sync
        + 'static,
>(
    State(state): State<AppState<S>>,
    Extension(session_id): Extension<i64>,
    Path(batch_id): Path<String>,
) -> Result<JsonResponse<ExecuteBatchResponse>, HttpError> {
    super::authz::require_admin_session(&state, session_id)?;
    let batch_manager = state.server.get_batch_manager();

    // Retrieve task information in order to obtain the space_id.
    let task = batch_manager.get_task(&batch_id).ok_or_else(|| {
        HttpError::NotFound(format!("The batch task does not exist: {}", batch_id))
    })?;

    // Query the `space_name` using the `space_id`.
    let space_name = {
        let storage = state.server.get_storage();
        let storage = storage.write();
        match storage.get_space_by_id(task.space_id) {
            Ok(Some(space_info)) => space_info.space_name,
            Ok(None) => {
                return Err(HttpError::NotFound(format!(
                    "The graph space does not exist: {}",
                    task.space_id
                )))
            }
            Err(e) => {
                return Err(HttpError::InternalError(format!(
                    "Querying the graph space failed: {}",
                    e
                )))
            }
        }
    };

    match batch_manager.execute_task(&batch_id, &space_name).await {
        Ok(result) => {
            let task = batch_manager.get_task(&batch_id).ok_or_else(|| {
                HttpError::NotFound(format!("Batch task does not exist: {}", batch_id))
            })?;

            Ok(JsonResponse(ExecuteBatchResponse {
                batch_id: batch_id.clone(),
                status: task.status,
                result,
                completed_at: Some(task.updated_at.to_rfc3339()),
            }))
        }
        Err(e) => Err(HttpError::InternalError(format!(
            "Failed to execute batch tasks: {}",
            e
        ))),
    }
}

#[utoipa::path(
    post,
    path = "/v1/batch/{id}/cancel",
    tag = "Batch",
    params(("id" = String, Path, description = "Batch task id")),
    responses(
        (status = 200, body = serde_json::Value, description = "Batch task cancelled"),
        (status = 500, description = "Internal error")
    )
)]
/// Cancel the batch task.
pub async fn cancel<
    S: StorageClient
        + StorageSchemaContextOps
        + StorageSyncContextOps
        + StorageOperationContextOps
        + Clone
        + Send
        + Sync
        + 'static,
>(
    State(state): State<AppState<S>>,
    Extension(session_id): Extension<i64>,
    Path(batch_id): Path<String>,
) -> Result<JsonResponse<serde_json::Value>, HttpError> {
    super::authz::require_admin_session(&state, session_id)?;
    let batch_manager = state.server.get_batch_manager();

    match batch_manager.cancel_task(&batch_id) {
        Ok(()) => Ok(JsonResponse(serde_json::json!({
            "message": "Batch task cancelled",
            "batch_id": batch_id,
        }))),
        Err(e) => Err(HttpError::BadRequest(format!(
            "Failed to cancel the batch task: {}",
            e
        ))),
    }
}

#[utoipa::path(
    delete,
    path = "/v1/batch/{id}",
    tag = "Batch",
    params(("id" = String, Path, description = "Batch task id")),
    responses(
        (status = 200, body = serde_json::Value, description = "Batch task deleted"),
        (status = 404, description = "Not found"),
        (status = 500, description = "Internal error")
    )
)]
/// Delete batch tasks.
pub async fn delete<
    S: StorageClient
        + StorageSchemaContextOps
        + StorageSyncContextOps
        + StorageOperationContextOps
        + Clone
        + Send
        + Sync
        + 'static,
>(
    State(state): State<AppState<S>>,
    Extension(session_id): Extension<i64>,
    Path(batch_id): Path<String>,
) -> Result<JsonResponse<serde_json::Value>, HttpError> {
    super::authz::require_admin_session(&state, session_id)?;
    let batch_manager = state.server.get_batch_manager();

    match batch_manager.remove_task(&batch_id) {
        Ok(()) => Ok(JsonResponse(serde_json::json!({
            "message": "Batch task deleted",
            "batch_id": batch_id,
        }))),
        Err(e) => Err(HttpError::NotFound(format!(
            "Failed to delete batch tasks: {}",
            e
        ))),
    }
}