use crate::{helpers, parsing};
use bson;
use bson::Bson;
use std::fmt;
enum UpdateOperator {
Set(bson::Document),
Unset(Vec<String>),
}
impl UpdateOperator {
pub(crate) fn apply_inplace(&self, value: &mut bson::Document) {
match self {
UpdateOperator::Set(updates) => {
for (key, new_value) in updates.iter() {
value.insert(key, new_value.clone());
}
}
UpdateOperator::Unset(keys) => {
for key in keys {
value.remove(key);
}
}
}
}
fn to_bson(&self) -> bson::Document {
match self {
UpdateOperator::Set(updates) => bson::doc! {
"$set": updates.clone()
},
UpdateOperator::Unset(names_str) => {
let names_bson = Bson::Array(
names_str
.iter()
.map(|name| Bson::String(name.clone()))
.collect(),
);
bson::doc! {
"$unset": names_bson
}
}
}
}
}
fn parse_update_operator(name: &str, argument: &Bson) -> Result<UpdateOperator, String> {
match name {
"$set" => {
let updates = parsing::ensure_is_object(argument, "argument to $set")?;
Ok(UpdateOperator::Set(updates.clone()))
}
"$unset" => {
let names: Vec<String> = match argument {
Bson::Array(names_bson) => {
let mut names_str: Vec<String> = Vec::with_capacity(names_bson.len());
for name in names_bson {
let name = parsing::ensure_is_str(name, "argument to $unset")?.to_string();
names_str.push(name);
}
names_str
}
Bson::String(name) => vec![name.clone()],
_ => {
return Err(format!(
"argument to $unset must be a string or array of strings, not {}",
argument
));
}
};
Ok(UpdateOperator::Unset(names))
}
name => Err(format!("Unsupported update operator `{}`", name)),
}
}
pub struct UpdateQuery {
operators: Vec<UpdateOperator>,
}
impl UpdateQuery {
pub fn new_nop() -> UpdateQuery {
UpdateQuery {
operators: Vec::new(),
}
}
pub fn parse_bson<B: Into<Bson>>(query: B) -> Result<UpdateQuery, String> {
let query: Bson = query.into();
let query = parsing::ensure_is_object(&query, "update query")?;
let mut result = UpdateQuery::new_nop();
for (op_name, op_arg) in query.iter() {
result
.operators
.push(parse_update_operator(op_name, op_arg)?);
}
Ok(result)
}
pub fn to_bson(&self) -> bson::Document {
let mut result = bson::Document::new();
for operator in self.operators.iter() {
for (op_name, op_arg) in operator.to_bson().iter() {
result.insert(op_name, op_arg);
}
}
result
}
pub(crate) fn apply_inplace(&self, value: &mut bson::Document) {
for op in &self.operators {
op.apply_inplace(value);
}
}
pub(crate) fn apply_to_empty(&self) -> bson::Document {
let mut result = bson::Document::new();
self.apply_inplace(&mut result);
result
}
}
impl fmt::Display for UpdateQuery {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::result::Result<(), std::fmt::Error> {
let as_bson = self.to_bson();
write!(
f,
"{}",
helpers::dump_to_json_string(as_bson.into(), true, true).unwrap()
)
}
}