BREAKING CHANGE(core): replace the TypeScript database engine with a Rust-backed embedded server and bridge
This commit is contained in:
617
rust/crates/rustdb-commands/src/handlers/update_handler.rs
Normal file
617
rust/crates/rustdb-commands/src/handlers/update_handler.rs
Normal file
@@ -0,0 +1,617 @@
|
||||
use std::collections::HashSet;
|
||||
|
||||
use bson::{doc, oid::ObjectId, Bson, Document};
|
||||
use rustdb_index::IndexEngine;
|
||||
use rustdb_query::{QueryMatcher, UpdateEngine, sort_documents, apply_projection};
|
||||
use tracing::debug;
|
||||
|
||||
use crate::context::CommandContext;
|
||||
use crate::error::{CommandError, CommandResult};
|
||||
|
||||
/// Handle `update` and `findAndModify` commands.
|
||||
pub async fn handle(
|
||||
cmd: &Document,
|
||||
db: &str,
|
||||
ctx: &CommandContext,
|
||||
command_name: &str,
|
||||
) -> CommandResult<Document> {
|
||||
match command_name {
|
||||
"findAndModify" | "findandmodify" => handle_find_and_modify(cmd, db, ctx).await,
|
||||
_ => handle_update(cmd, db, ctx).await,
|
||||
}
|
||||
}
|
||||
|
||||
/// Handle the `update` command.
|
||||
async fn handle_update(
|
||||
cmd: &Document,
|
||||
db: &str,
|
||||
ctx: &CommandContext,
|
||||
) -> CommandResult<Document> {
|
||||
let coll = cmd
|
||||
.get_str("update")
|
||||
.map_err(|_| CommandError::InvalidArgument("missing 'update' field".into()))?;
|
||||
|
||||
let updates = cmd
|
||||
.get_array("updates")
|
||||
.map_err(|_| CommandError::InvalidArgument("missing 'updates' array".into()))?;
|
||||
|
||||
let ordered = match cmd.get("ordered") {
|
||||
Some(Bson::Boolean(b)) => *b,
|
||||
_ => true,
|
||||
};
|
||||
|
||||
debug!(db = db, collection = coll, count = updates.len(), "update command");
|
||||
|
||||
// Auto-create database and collection if needed.
|
||||
ensure_collection_exists(db, coll, ctx).await?;
|
||||
|
||||
let ns_key = format!("{}.{}", db, coll);
|
||||
let mut total_n: i32 = 0;
|
||||
let mut total_n_modified: i32 = 0;
|
||||
let mut upserted_list: Vec<Document> = Vec::new();
|
||||
let mut write_errors: Vec<Document> = Vec::new();
|
||||
|
||||
for (idx, update_bson) in updates.iter().enumerate() {
|
||||
let update_spec = match update_bson {
|
||||
Bson::Document(d) => d,
|
||||
_ => {
|
||||
write_errors.push(doc! {
|
||||
"index": idx as i32,
|
||||
"code": 14_i32,
|
||||
"codeName": "TypeMismatch",
|
||||
"errmsg": "update spec must be a document",
|
||||
});
|
||||
if ordered {
|
||||
break;
|
||||
}
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
let filter = match update_spec.get("q") {
|
||||
Some(Bson::Document(d)) => d.clone(),
|
||||
_ => Document::new(),
|
||||
};
|
||||
|
||||
let update = match update_spec.get("u") {
|
||||
Some(Bson::Document(d)) => d.clone(),
|
||||
Some(Bson::Array(_pipeline)) => {
|
||||
// Aggregation pipeline updates are not yet supported; treat as error.
|
||||
write_errors.push(doc! {
|
||||
"index": idx as i32,
|
||||
"code": 14_i32,
|
||||
"codeName": "TypeMismatch",
|
||||
"errmsg": "aggregation pipeline updates not yet supported",
|
||||
});
|
||||
if ordered {
|
||||
break;
|
||||
}
|
||||
continue;
|
||||
}
|
||||
_ => {
|
||||
write_errors.push(doc! {
|
||||
"index": idx as i32,
|
||||
"code": 14_i32,
|
||||
"codeName": "TypeMismatch",
|
||||
"errmsg": "missing or invalid 'u' field in update spec",
|
||||
});
|
||||
if ordered {
|
||||
break;
|
||||
}
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
let multi = match update_spec.get("multi") {
|
||||
Some(Bson::Boolean(b)) => *b,
|
||||
_ => false,
|
||||
};
|
||||
|
||||
let upsert = match update_spec.get("upsert") {
|
||||
Some(Bson::Boolean(b)) => *b,
|
||||
_ => false,
|
||||
};
|
||||
|
||||
let array_filters: Option<Vec<Document>> =
|
||||
update_spec.get_array("arrayFilters").ok().map(|arr| {
|
||||
arr.iter()
|
||||
.filter_map(|v| {
|
||||
if let Bson::Document(d) = v {
|
||||
Some(d.clone())
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
.collect()
|
||||
});
|
||||
|
||||
// Load all documents and filter.
|
||||
let all_docs = load_filtered_docs(db, coll, &filter, &ns_key, ctx).await?;
|
||||
|
||||
if all_docs.is_empty() && upsert {
|
||||
// Upsert: create a new document.
|
||||
let new_doc = build_upsert_doc(&filter);
|
||||
|
||||
// Apply update operators or replacement.
|
||||
match UpdateEngine::apply_update(&new_doc, &update, array_filters.as_deref()) {
|
||||
Ok(mut updated) => {
|
||||
// Apply $setOnInsert if present.
|
||||
if let Some(Bson::Document(soi)) = update.get("$setOnInsert") {
|
||||
UpdateEngine::apply_set_on_insert(&mut updated, soi);
|
||||
}
|
||||
|
||||
// Ensure _id exists.
|
||||
let new_id = if !updated.contains_key("_id") {
|
||||
let oid = ObjectId::new();
|
||||
updated.insert("_id", oid);
|
||||
Bson::ObjectId(oid)
|
||||
} else {
|
||||
updated.get("_id").unwrap().clone()
|
||||
};
|
||||
|
||||
// Insert the new document.
|
||||
match ctx.storage.insert_one(db, coll, updated.clone()).await {
|
||||
Ok(_) => {
|
||||
// Update index.
|
||||
let mut engine = ctx
|
||||
.indexes
|
||||
.entry(ns_key.clone())
|
||||
.or_insert_with(IndexEngine::new);
|
||||
let _ = engine.on_insert(&updated);
|
||||
|
||||
total_n += 1;
|
||||
upserted_list.push(doc! {
|
||||
"index": idx as i32,
|
||||
"_id": new_id,
|
||||
});
|
||||
}
|
||||
Err(e) => {
|
||||
write_errors.push(doc! {
|
||||
"index": idx as i32,
|
||||
"code": 1_i32,
|
||||
"codeName": "InternalError",
|
||||
"errmsg": e.to_string(),
|
||||
});
|
||||
if ordered {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
write_errors.push(doc! {
|
||||
"index": idx as i32,
|
||||
"code": 14_i32,
|
||||
"codeName": "TypeMismatch",
|
||||
"errmsg": e.to_string(),
|
||||
});
|
||||
if ordered {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// Update matched documents.
|
||||
let docs_to_update = if multi {
|
||||
all_docs
|
||||
} else {
|
||||
all_docs.into_iter().take(1).collect()
|
||||
};
|
||||
|
||||
for matched_doc in &docs_to_update {
|
||||
match UpdateEngine::apply_update(
|
||||
matched_doc,
|
||||
&update,
|
||||
array_filters.as_deref(),
|
||||
) {
|
||||
Ok(updated_doc) => {
|
||||
let id_str = extract_id_string(matched_doc);
|
||||
match ctx
|
||||
.storage
|
||||
.update_by_id(db, coll, &id_str, updated_doc.clone())
|
||||
.await
|
||||
{
|
||||
Ok(()) => {
|
||||
// Update index.
|
||||
if let Some(mut engine) = ctx.indexes.get_mut(&ns_key) {
|
||||
let _ = engine.on_update(matched_doc, &updated_doc);
|
||||
}
|
||||
total_n += 1;
|
||||
// Check if the document actually changed.
|
||||
if matched_doc != &updated_doc {
|
||||
total_n_modified += 1;
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
write_errors.push(doc! {
|
||||
"index": idx as i32,
|
||||
"code": 1_i32,
|
||||
"codeName": "InternalError",
|
||||
"errmsg": e.to_string(),
|
||||
});
|
||||
if ordered {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
write_errors.push(doc! {
|
||||
"index": idx as i32,
|
||||
"code": 14_i32,
|
||||
"codeName": "TypeMismatch",
|
||||
"errmsg": e.to_string(),
|
||||
});
|
||||
if ordered {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Build response.
|
||||
let mut response = doc! {
|
||||
"n": total_n,
|
||||
"nModified": total_n_modified,
|
||||
"ok": 1.0,
|
||||
};
|
||||
|
||||
if !upserted_list.is_empty() {
|
||||
response.insert(
|
||||
"upserted",
|
||||
upserted_list
|
||||
.into_iter()
|
||||
.map(Bson::Document)
|
||||
.collect::<Vec<_>>(),
|
||||
);
|
||||
}
|
||||
|
||||
if !write_errors.is_empty() {
|
||||
response.insert(
|
||||
"writeErrors",
|
||||
write_errors
|
||||
.into_iter()
|
||||
.map(Bson::Document)
|
||||
.collect::<Vec<_>>(),
|
||||
);
|
||||
}
|
||||
|
||||
Ok(response)
|
||||
}
|
||||
|
||||
/// Handle the `findAndModify` command.
|
||||
async fn handle_find_and_modify(
|
||||
cmd: &Document,
|
||||
db: &str,
|
||||
ctx: &CommandContext,
|
||||
) -> CommandResult<Document> {
|
||||
let coll = cmd
|
||||
.get_str("findAndModify")
|
||||
.or_else(|_| cmd.get_str("findandmodify"))
|
||||
.map_err(|_| CommandError::InvalidArgument("missing 'findAndModify' field".into()))?;
|
||||
|
||||
let query = match cmd.get("query") {
|
||||
Some(Bson::Document(d)) => d.clone(),
|
||||
_ => Document::new(),
|
||||
};
|
||||
|
||||
let sort = match cmd.get("sort") {
|
||||
Some(Bson::Document(d)) => Some(d.clone()),
|
||||
_ => None,
|
||||
};
|
||||
|
||||
let update_doc = match cmd.get("update") {
|
||||
Some(Bson::Document(d)) => Some(d.clone()),
|
||||
_ => None,
|
||||
};
|
||||
|
||||
let remove = match cmd.get("remove") {
|
||||
Some(Bson::Boolean(b)) => *b,
|
||||
_ => false,
|
||||
};
|
||||
|
||||
let return_new = match cmd.get("new") {
|
||||
Some(Bson::Boolean(b)) => *b,
|
||||
_ => false,
|
||||
};
|
||||
|
||||
let upsert = match cmd.get("upsert") {
|
||||
Some(Bson::Boolean(b)) => *b,
|
||||
_ => false,
|
||||
};
|
||||
|
||||
let fields = match cmd.get("fields") {
|
||||
Some(Bson::Document(d)) => Some(d.clone()),
|
||||
_ => None,
|
||||
};
|
||||
|
||||
let array_filters: Option<Vec<Document>> =
|
||||
cmd.get_array("arrayFilters").ok().map(|arr| {
|
||||
arr.iter()
|
||||
.filter_map(|v| {
|
||||
if let Bson::Document(d) = v {
|
||||
Some(d.clone())
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
.collect()
|
||||
});
|
||||
|
||||
// Auto-create database and collection.
|
||||
ensure_collection_exists(db, coll, ctx).await?;
|
||||
|
||||
let ns_key = format!("{}.{}", db, coll);
|
||||
|
||||
// Load and filter documents.
|
||||
let mut matched = load_filtered_docs(db, coll, &query, &ns_key, ctx).await?;
|
||||
|
||||
// Sort if specified.
|
||||
if let Some(ref sort_spec) = sort {
|
||||
sort_documents(&mut matched, sort_spec);
|
||||
}
|
||||
|
||||
// Take the first match.
|
||||
let target = matched.into_iter().next();
|
||||
|
||||
if remove {
|
||||
// Remove operation.
|
||||
if let Some(ref doc) = target {
|
||||
let id_str = extract_id_string(doc);
|
||||
ctx.storage.delete_by_id(db, coll, &id_str).await?;
|
||||
|
||||
// Update index.
|
||||
if let Some(mut engine) = ctx.indexes.get_mut(&ns_key) {
|
||||
engine.on_delete(doc);
|
||||
}
|
||||
|
||||
let value = apply_fields_projection(doc, &fields);
|
||||
|
||||
return Ok(doc! {
|
||||
"value": value,
|
||||
"lastErrorObject": {
|
||||
"n": 1_i32,
|
||||
"updatedExisting": false,
|
||||
},
|
||||
"ok": 1.0,
|
||||
});
|
||||
} else {
|
||||
return Ok(doc! {
|
||||
"value": Bson::Null,
|
||||
"lastErrorObject": {
|
||||
"n": 0_i32,
|
||||
"updatedExisting": false,
|
||||
},
|
||||
"ok": 1.0,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
// Update operation.
|
||||
let update = match update_doc {
|
||||
Some(u) => u,
|
||||
None => {
|
||||
return Ok(doc! {
|
||||
"value": Bson::Null,
|
||||
"lastErrorObject": {
|
||||
"n": 0_i32,
|
||||
"updatedExisting": false,
|
||||
},
|
||||
"ok": 1.0,
|
||||
});
|
||||
}
|
||||
};
|
||||
|
||||
if let Some(original_doc) = target {
|
||||
// Update the matched document.
|
||||
let updated_doc = UpdateEngine::apply_update(
|
||||
&original_doc,
|
||||
&update,
|
||||
array_filters.as_deref(),
|
||||
)
|
||||
.map_err(|e| CommandError::InternalError(e.to_string()))?;
|
||||
|
||||
let id_str = extract_id_string(&original_doc);
|
||||
ctx.storage
|
||||
.update_by_id(db, coll, &id_str, updated_doc.clone())
|
||||
.await?;
|
||||
|
||||
// Update index.
|
||||
if let Some(mut engine) = ctx.indexes.get_mut(&ns_key) {
|
||||
let _ = engine.on_update(&original_doc, &updated_doc);
|
||||
}
|
||||
|
||||
let return_doc = if return_new {
|
||||
&updated_doc
|
||||
} else {
|
||||
&original_doc
|
||||
};
|
||||
|
||||
let value = apply_fields_projection(return_doc, &fields);
|
||||
|
||||
Ok(doc! {
|
||||
"value": value,
|
||||
"lastErrorObject": {
|
||||
"n": 1_i32,
|
||||
"updatedExisting": true,
|
||||
},
|
||||
"ok": 1.0,
|
||||
})
|
||||
} else if upsert {
|
||||
// Upsert: create a new document.
|
||||
let new_doc = build_upsert_doc(&query);
|
||||
|
||||
let mut updated_doc = UpdateEngine::apply_update(
|
||||
&new_doc,
|
||||
&update,
|
||||
array_filters.as_deref(),
|
||||
)
|
||||
.map_err(|e| CommandError::InternalError(e.to_string()))?;
|
||||
|
||||
// Apply $setOnInsert if present.
|
||||
if let Some(Bson::Document(soi)) = update.get("$setOnInsert") {
|
||||
UpdateEngine::apply_set_on_insert(&mut updated_doc, soi);
|
||||
}
|
||||
|
||||
// Ensure _id.
|
||||
let upserted_id = if !updated_doc.contains_key("_id") {
|
||||
let oid = ObjectId::new();
|
||||
updated_doc.insert("_id", oid);
|
||||
Bson::ObjectId(oid)
|
||||
} else {
|
||||
updated_doc.get("_id").unwrap().clone()
|
||||
};
|
||||
|
||||
ctx.storage
|
||||
.insert_one(db, coll, updated_doc.clone())
|
||||
.await?;
|
||||
|
||||
// Update index.
|
||||
{
|
||||
let mut engine = ctx
|
||||
.indexes
|
||||
.entry(ns_key.clone())
|
||||
.or_insert_with(IndexEngine::new);
|
||||
let _ = engine.on_insert(&updated_doc);
|
||||
}
|
||||
|
||||
let value = if return_new {
|
||||
apply_fields_projection(&updated_doc, &fields)
|
||||
} else {
|
||||
Bson::Null
|
||||
};
|
||||
|
||||
Ok(doc! {
|
||||
"value": value,
|
||||
"lastErrorObject": {
|
||||
"n": 1_i32,
|
||||
"updatedExisting": false,
|
||||
"upserted": upserted_id,
|
||||
},
|
||||
"ok": 1.0,
|
||||
})
|
||||
} else {
|
||||
Ok(doc! {
|
||||
"value": Bson::Null,
|
||||
"lastErrorObject": {
|
||||
"n": 0_i32,
|
||||
"updatedExisting": false,
|
||||
},
|
||||
"ok": 1.0,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// ---- Helpers ----
|
||||
|
||||
/// Load documents from storage, optionally using index for candidate narrowing, then filter.
|
||||
async fn load_filtered_docs(
|
||||
db: &str,
|
||||
coll: &str,
|
||||
filter: &Document,
|
||||
ns_key: &str,
|
||||
ctx: &CommandContext,
|
||||
) -> CommandResult<Vec<Document>> {
|
||||
// Try to use index to narrow candidates.
|
||||
let candidate_ids: Option<HashSet<String>> = ctx
|
||||
.indexes
|
||||
.get(ns_key)
|
||||
.and_then(|engine| engine.find_candidate_ids(filter));
|
||||
|
||||
let docs = if let Some(ids) = candidate_ids {
|
||||
if ids.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
ctx.storage.find_by_ids(db, coll, ids).await?
|
||||
} else {
|
||||
ctx.storage.find_all(db, coll).await?
|
||||
};
|
||||
|
||||
// Apply filter.
|
||||
if filter.is_empty() {
|
||||
Ok(docs)
|
||||
} else {
|
||||
Ok(QueryMatcher::filter(&docs, filter))
|
||||
}
|
||||
}
|
||||
|
||||
/// Build a base document for an upsert from the filter's equality conditions.
|
||||
fn build_upsert_doc(filter: &Document) -> Document {
|
||||
let mut doc = Document::new();
|
||||
for (key, value) in filter {
|
||||
if key.starts_with('$') {
|
||||
// Skip top-level operators like $and, $or.
|
||||
continue;
|
||||
}
|
||||
match value {
|
||||
Bson::Document(d) if d.keys().any(|k| k.starts_with('$')) => {
|
||||
// If the value has operators (e.g., $gt), extract $eq if present.
|
||||
if let Some(eq_val) = d.get("$eq") {
|
||||
doc.insert(key.clone(), eq_val.clone());
|
||||
}
|
||||
}
|
||||
_ => {
|
||||
doc.insert(key.clone(), value.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
doc
|
||||
}
|
||||
|
||||
/// Extract _id as a string for storage operations.
|
||||
fn extract_id_string(doc: &Document) -> String {
|
||||
match doc.get("_id") {
|
||||
Some(Bson::ObjectId(oid)) => oid.to_hex(),
|
||||
Some(Bson::String(s)) => s.clone(),
|
||||
Some(other) => format!("{}", other),
|
||||
None => String::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Apply fields projection if specified, returning Bson.
|
||||
fn apply_fields_projection(doc: &Document, fields: &Option<Document>) -> Bson {
|
||||
match fields {
|
||||
Some(proj) if !proj.is_empty() => Bson::Document(apply_projection(doc, proj)),
|
||||
_ => Bson::Document(doc.clone()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Ensure the target database and collection exist, creating them if needed.
|
||||
async fn ensure_collection_exists(
|
||||
db: &str,
|
||||
coll: &str,
|
||||
ctx: &CommandContext,
|
||||
) -> CommandResult<()> {
|
||||
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));
|
||||
}
|
||||
}
|
||||
|
||||
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) => {
|
||||
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(())
|
||||
}
|
||||
Reference in New Issue
Block a user