From 6e08f2e8f524693e474d3bbf3447cf0b2dc446d7 Mon Sep 17 00:00:00 2001
From: Jada <jada.lilleboe@mongodb.com>
Date: Tue, 19 Nov 2024 13:05:12 -0500
Subject: [PATCH] RUST-622 Add document square bracket indexing

---
 src/bson.rs                   | 15 ++++++++++++
 src/document.rs               | 12 ++++++++++
 src/tests/modules/document.rs | 45 +++++++++++++++++++++++++++++++++++
 3 files changed, 72 insertions(+)

diff --git a/src/bson.rs b/src/bson.rs
index 9c732c50..806fee0b 100644
--- a/src/bson.rs
+++ b/src/bson.rs
@@ -25,6 +25,7 @@ use std::{
     convert::{TryFrom, TryInto},
     fmt::{self, Debug, Display, Formatter},
     hash::Hash,
+    ops::Index,
 };
 
 use serde_json::{json, Value};
@@ -236,6 +237,20 @@ impl Debug for Bson {
     }
 }
 
+impl Index<&str> for Bson {
+    type Output = Bson;
+
+    fn index(&self, index: &str) -> &Self::Output {
+        match *self {
+            Bson::Document(ref doc) => match doc.get(index) {
+                Some(v) => v,
+                None => &Bson::Null,
+            },
+            _ => &Bson::Null,
+        }
+    }
+}
+
 impl From<f32> for Bson {
     fn from(a: f32) -> Bson {
         Bson::Double(a.into())
diff --git a/src/document.rs b/src/document.rs
index 18b57851..b3e24067 100644
--- a/src/document.rs
+++ b/src/document.rs
@@ -7,6 +7,7 @@ use std::{
     fmt::{self, Debug, Display, Formatter},
     io::{Read, Write},
     iter::{Extend, FromIterator, IntoIterator},
+    ops::Index,
 };
 
 use ahash::RandomState;
@@ -139,6 +140,17 @@ impl Debug for Document {
     }
 }
 
+impl Index<&str> for Document {
+    type Output = Bson;
+
+    fn index(&self, index: &str) -> &Self::Output {
+        match self.get(index) {
+            Some(v) => v,
+            None => &Bson::Null,
+        }
+    }
+}
+
 /// An iterator over Document entries.
 pub struct IntoIter {
     inner: indexmap::map::IntoIter<String, Bson>,
diff --git a/src/tests/modules/document.rs b/src/tests/modules/document.rs
index e8be3199..c6f5390e 100644
--- a/src/tests/modules/document.rs
+++ b/src/tests/modules/document.rs
@@ -370,3 +370,48 @@ fn test_pretty_printing() {
 }"#;
     assert_eq!(expected, format!("{d:#}"));
 }
+
+#[test]
+fn test_indexing() {
+    let d = doc! {"x": 1};
+    let val = d["x"].as_i32().unwrap();
+    assert_eq!(val, 1);
+
+    let d = doc! {"x": {"y": 100}};
+    let val = d["x"]["y"].as_i32().unwrap();
+    assert_eq!(val, 100);
+
+    let d = doc! {"x" : true};
+    let val = d["x"].as_bool().unwrap();
+    assert!(val);
+
+    let d = doc! {"x": "y"};
+    let val = d["x"].as_str().unwrap();
+    assert_eq!(val, "y");
+
+    let d = doc! {"x": 1.9};
+    let val = d["x"].as_f64().unwrap();
+    assert_eq!(val, 1.9);
+
+    let d = doc! {"x": [1, 2, 3]};
+    let val = d["x"].as_array().unwrap();
+    assert_eq!(val.len(), 3);
+}
+
+#[test]
+fn test_indexing_key_not_found() {
+    let d = doc! {"x": 1};
+    let val = &d["y"];
+    assert!(val.as_null().is_some());
+
+    let d = doc! {"x": {"y": 1}};
+    let val = &d["x"]["z"];
+    assert!(val.as_null().is_some());
+}
+
+#[test]
+fn test_indexing_on_wrong_bson_type() {
+    let d = doc! {"x": {"y": 1}};
+    let val = &d["x"]["y"]["z"];
+    assert!(val.as_null().is_some());
+}