use bson;
use num_traits::Num;
use std::collections::HashMap;
use crate::expr::Expr;
use crate::field_witnesses::{FieldName, HasField};
use crate::path::Path;
pub struct UpdateBuilder<T> {
pub prefix: Vec<String>,
clauses: HashMap<UpdateOperation, Vec<(String, bson::Bson)>>,
_marker: std::marker::PhantomData<T>,
}
impl<T> Default for UpdateBuilder<T> {
fn default() -> Self {
Self::new()
}
}
impl<T> UpdateBuilder<T> {
pub fn new() -> Self {
UpdateBuilder {
prefix: Vec::new(),
clauses: HashMap::new(),
_marker: std::marker::PhantomData,
}
}
fn field_path<F: FieldName>(&self) -> String {
if self.prefix.is_empty() {
F::field_name().to_string()
} else {
format!("{}.{}", self.prefix.join("."), F::field_name())
}
}
fn push_clause(&mut self, op: UpdateOperation, path: String, clause: bson::Bson) {
self.clauses.entry(op).or_default().push((path, clause));
}
pub fn set<F: FieldName, V: Into<bson::Bson>>(&mut self, value: V) -> &mut Self
where
T: HasField<F>,
{
let path = self.field_path::<F>();
self.push_clause(UpdateOperation::Set, path, value.into());
self
}
pub fn set_expr<F: FieldName, V>(&mut self, expr: Expr<T, V>) -> &mut Self
where
T: HasField<F>,
{
self.set::<F, _>(expr)
}
pub fn unset<F: FieldName>(&mut self) -> &mut Self
where
T: HasField<F>,
{
let path = self.field_path::<F>();
self.push_clause(UpdateOperation::Unset, path, bson::Bson::Null);
self
}
pub fn inc<F: FieldName, N: Num + Into<bson::Bson>>(&mut self, value: N) -> &mut Self
where
T: HasField<F>,
{
let path = self.field_path::<F>();
self.push_clause(UpdateOperation::Inc, path, value.into());
self
}
pub fn max<F: FieldName, N: Num + Into<bson::Bson>>(&mut self, value: N) -> &mut Self
where
T: HasField<F>,
{
let path = self.field_path::<F>();
self.push_clause(UpdateOperation::Max, path, value.into());
self
}
pub fn min<F: FieldName, N: Num + Into<bson::Bson>>(&mut self, value: N) -> &mut Self
where
T: HasField<F>,
{
let path = self.field_path::<F>();
self.push_clause(UpdateOperation::Min, path, value.into());
self
}
pub fn mul<F: FieldName, N: Num + Into<bson::Bson>>(&mut self, value: N) -> &mut Self
where
T: HasField<F>,
{
let path = self.field_path::<F>();
self.push_clause(UpdateOperation::Mul, path, value.into());
self
}
pub fn rename<F: FieldName>(&mut self, new_name: &str) -> &mut Self
where
T: HasField<F>,
{
let path = self.field_path::<F>();
self.push_clause(
UpdateOperation::Rename,
path,
bson::Bson::String(new_name.to_string()),
);
self
}
pub fn current_date<F: FieldName>(&mut self, date_type: CurrentDateType) -> &mut Self
where
T: HasField<F>,
{
let path = self.field_path::<F>();
self.push_clause(
UpdateOperation::CurrentDate,
path,
bson::Bson::String(date_type.to_string()),
);
self
}
pub fn add_to_set<F: FieldName, V: Into<bson::Bson>>(&mut self, value: V) -> &mut Self
where
T: HasField<F>,
T::Value: IntoIterator<Item = V>,
{
let path = self.field_path::<F>();
self.push_clause(UpdateOperation::AddToSet, path, value.into());
self
}
pub fn add_to_set_each<F: FieldName, I: IntoIterator<Item = V>, V: Into<bson::Bson>>(
&mut self,
values: I,
) -> &mut Self
where
T: HasField<F>,
T::Value: IntoIterator<Item = V>,
{
let path = self.field_path::<F>();
let values_vec: Vec<bson::Bson> = values.into_iter().map(|v| v.into()).collect();
self.push_clause(
UpdateOperation::AddToSet,
path,
bson::doc! { "$each": values_vec }.into(),
);
self
}
pub fn pop<F: FieldName>(&mut self, strategy: PopStrategy) -> &mut Self
where
T: HasField<F>,
T::Value: IntoIterator,
{
let path = self.field_path::<F>();
self.push_clause(UpdateOperation::Pop, path, strategy.into());
self
}
pub fn pull_expr<F: FieldName>(&mut self, cond: bson::Bson) -> &mut Self
where
T: HasField<F>,
T::Value: IntoIterator,
{
let path = self.field_path::<F>();
self.push_clause(UpdateOperation::Pull, path, cond);
self
}
pub fn pull<F: FieldName, V: Into<bson::Bson>>(&mut self, value: V) -> &mut Self
where
T: HasField<F>,
T::Value: IntoIterator<Item = V>,
{
let path = self.field_path::<F>();
self.push_clause(UpdateOperation::Pull, path, value.into());
self
}
pub fn pull_all<F: FieldName, I>(&mut self, values: I) -> &mut Self
where
T: HasField<F>,
I: Into<bson::Bson> + IntoIterator,
T::Value: IntoIterator<Item = I::Item>,
{
let path = self.field_path::<F>();
self.push_clause(UpdateOperation::PullAll, path, values.into());
self
}
pub fn push<F: FieldName, V: Into<bson::Bson>>(&mut self, value: V) -> &mut Self
where
T: HasField<F>,
T::Value: IntoIterator<Item = V>,
{
let path = self.field_path::<F>();
self.push_clause(UpdateOperation::Push, path, value.into());
self
}
pub fn push_each<
F: FieldName,
I: IntoIterator<Item = V>,
V: Into<bson::Bson>,
Clause: Into<PushEach<I, V>>,
>(
&mut self,
clause: Clause,
) -> &mut Self
where
T: HasField<F>,
T::Value: IntoIterator<Item = V>,
{
let path = self.field_path::<F>();
let c: PushEach<I, V> = clause.into();
self.push_clause(UpdateOperation::Push, path, c.into());
self
}
pub fn if_some<A, F>(&mut self, value: Option<A>, f: F) -> &mut Self
where
F: FnOnce(&mut Self, A) -> &mut Self,
{
if let Some(v) = value {
f(self, v);
}
self
}
pub fn untyped<F: FieldName>(&mut self, op: UpdateOperation, expr: bson::Bson) -> &mut Self
where
T: HasField<F>,
{
let path = self.field_path::<F>();
self.push_clause(op, path, expr);
self
}
pub fn with_lookup<F: FieldName, L, G: FieldName, U: HasField<G>, N>(
&mut self,
lookup: L,
f: N,
) -> &mut Self
where
T: HasField<F>,
L: FnOnce(&Path<F, T, T>) -> Path<G, U, T>,
N: FnOnce(&mut UpdateBuilder<U>),
{
let base_field: Path<F, T, T> = Path {
prefix: self.prefix.clone(),
_marker: std::marker::PhantomData,
};
let resolved_field = lookup(&base_field);
let mut nested_builder = UpdateBuilder::<U> {
prefix: resolved_field.prefix.clone(),
clauses: HashMap::new(),
_marker: std::marker::PhantomData,
};
f(&mut nested_builder);
for (operation, clauses_vec) in nested_builder.clauses {
self.clauses
.entry(operation)
.or_default()
.extend(clauses_vec);
}
self
}
pub fn with_field<F: FieldName, N>(&mut self, f: N) -> &mut Self
where
T: HasField<F>,
N: FnOnce(&mut UpdateBuilder<T>),
{
self.with_lookup::<F, _, F, T, _>(|path| path.clone(), f)
}
pub fn build(&mut self) -> bson::Document {
let mut doc = bson::Document::new();
for (op, op_clauses) in &self.clauses {
let operation = op.as_str();
let mut operation_doc = bson::Document::new();
for (field, clause) in op_clauses {
operation_doc.insert(field.clone(), clause.clone());
}
doc.insert(operation, operation_doc);
}
doc
}
}
#[derive(Eq, Hash, PartialEq)]
pub enum UpdateOperation {
Set,
Unset,
Inc,
Max,
Min,
Mul,
Rename,
CurrentDate,
AddToSet,
Pop,
Pull,
PullAll,
Push,
}
impl UpdateOperation {
pub const fn as_str(&self) -> &'static str {
match self {
UpdateOperation::Set => "$set",
UpdateOperation::Unset => "$unset",
UpdateOperation::Inc => "$inc",
UpdateOperation::Max => "$max",
UpdateOperation::Min => "$min",
UpdateOperation::Mul => "$mul",
UpdateOperation::Rename => "$rename",
UpdateOperation::CurrentDate => "$currentDate",
UpdateOperation::AddToSet => "$addToSet",
UpdateOperation::Pop => "$pop",
UpdateOperation::Pull => "$pull",
UpdateOperation::PullAll => "$pullAll",
UpdateOperation::Push => "$push",
}
}
}
pub enum PushEachSlice {
PushEmptySlice,
PushLastSlice(usize),
PushFirstSlice(usize),
}
impl From<PushEachSlice> for bson::Bson {
fn from(slice: PushEachSlice) -> Self {
match slice {
PushEachSlice::PushEmptySlice => bson::Bson::Int32(0),
PushEachSlice::PushLastSlice(n) => bson::Bson::Int32(-(n as i32)),
PushEachSlice::PushFirstSlice(n) => bson::Bson::Int32(n as i32),
}
}
}
pub enum PushEachSort {
PushSortAscending,
PushSortDescending,
PushSortExpression(bson::Document),
}
impl From<PushEachSort> for bson::Bson {
fn from(sort: PushEachSort) -> Self {
match sort {
PushEachSort::PushSortAscending => bson::Bson::Int32(1),
PushEachSort::PushSortDescending => bson::Bson::Int32(-1),
PushEachSort::PushSortExpression(expr) => expr.into(),
}
}
}
pub enum PushEachPosition {
PushTakeFirst(usize),
PushTakeLast(usize),
}
impl From<PushEachPosition> for bson::Bson {
fn from(position: PushEachPosition) -> Self {
match position {
PushEachPosition::PushTakeFirst(n) => bson::Bson::Int32(n as i32),
PushEachPosition::PushTakeLast(n) => bson::Bson::Int32(-(n as i32)),
}
}
}
pub struct PushEach<Values: IntoIterator<Item = V>, V: Into<bson::Bson>> {
pub values: Values,
pub slice: std::option::Option<PushEachSlice>,
pub sort: std::option::Option<PushEachSort>,
pub position: std::option::Option<PushEachPosition>,
}
impl<Values: IntoIterator> From<Values> for PushEach<Values, <Values as IntoIterator>::Item>
where
<Values as IntoIterator>::Item: Into<bson::Bson>,
{
fn from(values: Values) -> Self {
PushEach {
values,
slice: None,
sort: None,
position: None,
}
}
}
impl<Values: IntoIterator<Item = V>, V: Into<bson::Bson>> PushEach<Values, V> {
pub fn new(values: Values) -> Self {
Self {
values,
slice: None,
sort: None,
position: None,
}
}
pub fn with_slice(mut self, slice: PushEachSlice) -> Self {
self.slice = Some(slice);
self
}
pub fn with_sort(mut self, sort: PushEachSort) -> Self {
self.sort = Some(sort);
self
}
pub fn with_position(mut self, position: PushEachPosition) -> Self {
self.position = Some(position);
self
}
}
impl<Values: IntoIterator<Item = V>, V: Into<bson::Bson>> From<PushEach<Values, V>> for bson::Bson {
fn from(push_each: PushEach<Values, V>) -> Self {
let mut expr = bson::Document::new();
let values: Vec<bson::Bson> = push_each.values.into_iter().map(|v| v.into()).collect();
expr.insert("$each", bson::Bson::Array(values));
if let Some(slice) = push_each.slice {
expr.insert("$slice", bson::Bson::from(slice));
}
if let Some(sort) = push_each.sort {
expr.insert("$sort", bson::Bson::from(sort));
}
if let Some(position) = push_each.position {
expr.insert("$position", bson::Bson::from(position));
}
expr.into()
}
}
impl std::fmt::Display for UpdateOperation {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}", self.as_str())
}
}
pub enum CurrentDateType {
Date,
Timestamp,
}
impl std::fmt::Display for CurrentDateType {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let s = match self {
CurrentDateType::Date => "date",
CurrentDateType::Timestamp => "timestamp",
};
write!(f, "{s}")
}
}
pub enum PopStrategy {
First,
Last,
}
impl From<PopStrategy> for bson::Bson {
fn from(strategy: PopStrategy) -> Self {
match strategy {
PopStrategy::First => bson::Bson::Int32(-1),
PopStrategy::Last => bson::Bson::Int32(1),
}
}
}
pub fn empty<T>() -> UpdateBuilder<T> {
UpdateBuilder::new()
}
#[cfg(test)]
mod tests {
use super::*;
use crate::field_witnesses::FieldName;
struct TestFieldName;
impl FieldName for TestFieldName {
fn field_name() -> &'static str {
"test_field"
}
}
struct AnotherFieldName;
impl FieldName for AnotherFieldName {
fn field_name() -> &'static str {
"another_field"
}
}
struct NestedFieldName;
impl FieldName for NestedFieldName {
fn field_name() -> &'static str {
"nested.field"
}
}
#[test]
fn test_field_path_empty_prefix() {
let builder = UpdateBuilder::<()>::new();
let path = builder.field_path::<TestFieldName>();
assert_eq!(path, "test_field");
}
#[test]
fn test_field_path_single_prefix() {
let mut builder = UpdateBuilder::<()>::new();
builder.prefix.push("parent".to_string());
let path = builder.field_path::<TestFieldName>();
assert_eq!(path, "parent.test_field");
}
#[test]
fn test_field_path_multiple_prefix() {
let mut builder = UpdateBuilder::<()>::new();
builder.prefix.push("root".to_string());
builder.prefix.push("parent".to_string());
builder.prefix.push("child".to_string());
let path = builder.field_path::<TestFieldName>();
assert_eq!(path, "root.parent.child.test_field");
}
#[test]
fn test_field_path_different_field_types() {
let mut builder = UpdateBuilder::<()>::new();
builder.prefix.push("prefix".to_string());
let path1 = builder.field_path::<TestFieldName>();
assert_eq!(path1, "prefix.test_field");
let path2 = builder.field_path::<AnotherFieldName>();
assert_eq!(path2, "prefix.another_field");
}
#[test]
fn test_field_path_nested_field_name() {
let mut builder = UpdateBuilder::<()>::new();
builder.prefix.push("outer".to_string());
let path = builder.field_path::<NestedFieldName>();
assert_eq!(path, "outer.nested.field");
}
#[test]
fn test_field_path_empty_string_prefix() {
let mut builder = UpdateBuilder::<()>::new();
builder.prefix.push("".to_string());
let path = builder.field_path::<TestFieldName>();
assert_eq!(path, ".test_field");
}
#[test]
fn test_field_path_consistency_across_multiple_calls() {
let mut builder = UpdateBuilder::<()>::new();
builder.prefix.push("consistent".to_string());
let path1 = builder.field_path::<TestFieldName>();
let path2 = builder.field_path::<TestFieldName>();
assert_eq!(path1, path2);
assert_eq!(path1, "consistent.test_field");
}
#[test]
fn test_field_path_deeply_nested_prefix() {
let mut builder = UpdateBuilder::<()>::new();
for i in 0..10 {
builder.prefix.push(format!("level{i}"));
}
let path = builder.field_path::<TestFieldName>();
assert_eq!(
path,
"level0.level1.level2.level3.level4.level5.level6.level7.level8.level9.test_field"
);
}
#[test]
fn test_field_path_special_characters_in_prefix() {
let mut builder = UpdateBuilder::<()>::new();
builder.prefix.push("with-dash".to_string());
builder.prefix.push("with_underscore".to_string());
builder.prefix.push("with123numbers".to_string());
let path = builder.field_path::<TestFieldName>();
assert_eq!(path, "with-dash.with_underscore.with123numbers.test_field");
}
#[test]
fn test_field_path_with_numeric_string_prefix() {
let mut builder = UpdateBuilder::<()>::new();
builder.prefix.push("0".to_string());
builder.prefix.push("123".to_string());
let path = builder.field_path::<TestFieldName>();
assert_eq!(path, "0.123.test_field");
}
}