smartdb/rust/crates/rustdb-commands/src/handlers/insert_handler.rs

use std::collections::HashMap;

use bson::{doc, oid::ObjectId, Bson, Document};
use rustdb_index::IndexEngine;
use rustdb_storage::OpType;
use tracing::{debug, warn};

use crate::context::CommandContext;
use crate::error::{CommandError, CommandResult};

/// Handle the `insert` command.
pub async fn handle(
    cmd: &Document,
    db: &str,
    ctx: &CommandContext,
    document_sequences: Option<&HashMap<String, Vec<Document>>>,
) -> CommandResult<Document> {
    let coll = cmd
        .get_str("insert")
        .map_err(|_| CommandError::InvalidArgument("missing 'insert' field".into()))?;

    // Determine whether writes are ordered (default: true).
    let ordered = match cmd.get("ordered") {
        Some(Bson::Boolean(b)) => *b,
        _ => true,
    };

    // Collect documents from either the command body or OP_MSG document sequences.
    let docs: Vec<Document> = if let Some(seqs) = document_sequences {
        if let Some(seq_docs) = seqs.get("documents") {
            seq_docs.clone()
        } else {
            extract_docs_from_array(cmd)?
        }
    } else {
        extract_docs_from_array(cmd)?
    };

    if docs.is_empty() {
        return Err(CommandError::InvalidArgument(
            "no documents to insert".into(),
        ));
    }

    debug!(
        db = db,
        collection = coll,
        count = docs.len(),
        "insert command"
    );

    // Auto-create database and collection if they don't exist.
    ensure_collection_exists(db, coll, ctx).await?;

    let ns_key = format!("{}.{}", db, coll);
    let mut inserted_count: i32 = 0;
    let mut write_errors: Vec<Document> = Vec::new();

    for (idx, mut doc) in docs.into_iter().enumerate() {
        // Auto-generate _id if not present.
        if !doc.contains_key("_id") {
            doc.insert("_id", ObjectId::new());
        }

        // Attempt storage insert.
        match ctx.storage.insert_one(db, coll, doc.clone()).await {
            Ok(id_str) => {
                // Record in oplog.
                ctx.oplog.append(
                    OpType::Insert,
                    db,
                    coll,
                    &id_str,
                    Some(doc.clone()),
                    None,
                );

                // Update index engine.
                let mut engine = ctx
                    .indexes
                    .entry(ns_key.clone())
                    .or_insert_with(IndexEngine::new);
                if let Err(e) = engine.on_insert(&doc) {
                    warn!(
                        namespace = %ns_key,
                        error = %e,
                        "index update failed after successful insert"
                    );
                }
                inserted_count += 1;
            }
            Err(e) => {
                let err_msg = e.to_string();
                let (code, code_name) = if err_msg.contains("AlreadyExists")
                    || err_msg.contains("duplicate")
                {
                    (11000_i32, "DuplicateKey")
                } else {
                    (1_i32, "InternalError")
                };

                write_errors.push(doc! {
                    "index": idx as i32,
                    "code": code,
                    "codeName": code_name,
                    "errmsg": &err_msg,
                });

                if ordered {
                    // Stop on first error when ordered.
                    break;
                }
            }
        }
    }

    // Build response document.
    let mut response = doc! {
        "n": inserted_count,
        "ok": 1.0,
    };

    if !write_errors.is_empty() {
        response.insert(
            "writeErrors",
            write_errors
                .into_iter()
                .map(Bson::Document)
                .collect::<Vec<_>>(),
        );
    }

    Ok(response)
}

/// Extract documents from the `documents` array field in the command BSON.
fn extract_docs_from_array(cmd: &Document) -> CommandResult<Vec<Document>> {
    match cmd.get_array("documents") {
        Ok(arr) => {
            let mut docs = Vec::with_capacity(arr.len());
            for item in arr {
                match item {
                    Bson::Document(d) => docs.push(d.clone()),
                    _ => {
                        return Err(CommandError::InvalidArgument(
                            "documents array contains non-document element".into(),
                        ));
                    }
                }
            }
            Ok(docs)
        }
        Err(_) => Ok(Vec::new()),
    }
}

/// Ensure the target database and collection exist, creating them if needed.
async fn ensure_collection_exists(
    db: &str,
    coll: &str,
    ctx: &CommandContext,
) -> CommandResult<()> {
    // Create database (no-op if it already exists in most backends).
    if let Err(e) = ctx.storage.create_database(db).await {
        let msg = e.to_string();
        if !msg.contains("AlreadyExists") && !msg.contains("already exists") {
            return Err(CommandError::StorageError(msg));
        }
    }

    // Create collection if it doesn't exist.
    match ctx.storage.collection_exists(db, coll).await {
        Ok(true) => {}
        Ok(false) => {
            if let Err(e) = ctx.storage.create_collection(db, coll).await {
                let msg = e.to_string();
                if !msg.contains("AlreadyExists") && !msg.contains("already exists") {
                    return Err(CommandError::StorageError(msg));
                }
            }
        }
        Err(e) => {
            // Database might not exist yet; try creating collection anyway.
            if let Err(e2) = ctx.storage.create_collection(db, coll).await {
                let msg = e2.to_string();
                if !msg.contains("AlreadyExists") && !msg.contains("already exists") {
                    return Err(CommandError::StorageError(format!(
                        "collection_exists failed: {e}; create_collection failed: {msg}"
                    )));
                }
            }
        }
    }

    Ok(())
}