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>>, ) -> CommandResult { 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 = 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 = 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::>(), ); } Ok(response) } /// Extract documents from the `documents` array field in the command BSON. fn extract_docs_from_array(cmd: &Document) -> CommandResult> { 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(()) }