From e94b8007c125a1022a01d46d85db7ada106be474 Mon Sep 17 00:00:00 2001
From: JT <547158+jntrnr@users.noreply.github.com>
Date: Tue, 28 Dec 2021 10:11:20 +1100
Subject: [PATCH] Allow update to also insert (#610)

---
 crates/nu-protocol/src/value/mod.rs | 43 +++++++++++++++++++++++++++--
 src/tests/test_table_operations.rs  |  5 ++++
 2 files changed, 46 insertions(+), 2 deletions(-)

diff --git a/crates/nu-protocol/src/value/mod.rs b/crates/nu-protocol/src/value/mod.rs
index 61f84127c5..5d75bea003 100644
--- a/crates/nu-protocol/src/value/mod.rs
+++ b/crates/nu-protocol/src/value/mod.rs
@@ -628,26 +628,65 @@ impl Value {
                         for val in vals.iter_mut() {
                             match val {
                                 Value::Record { cols, vals, .. } => {
-                                    for col in cols.iter().zip(vals) {
+                                    let mut found = false;
+                                    for col in cols.iter().zip(vals.iter_mut()) {
                                         if col.0 == col_name {
+                                            found = true;
                                             col.1.replace_data_at_cell_path(
                                                 &cell_path[1..],
                                                 new_val.clone(),
                                             )?
                                         }
                                     }
+                                    if !found {
+                                        if cell_path.len() == 1 {
+                                            cols.push(col_name.clone());
+                                            vals.push(new_val);
+                                            break;
+                                        } else {
+                                            let mut new_col = Value::Record {
+                                                cols: vec![],
+                                                vals: vec![],
+                                                span: new_val.span()?,
+                                            };
+                                            new_col.replace_data_at_cell_path(
+                                                &cell_path[1..],
+                                                new_val,
+                                            )?;
+                                            vals.push(new_col);
+                                            break;
+                                        }
+                                    }
                                 }
                                 v => return Err(ShellError::CantFindColumn(*span, v.span()?)),
                             }
                         }
                     }
                     Value::Record { cols, vals, .. } => {
-                        for col in cols.iter().zip(vals) {
+                        let mut found = false;
+
+                        for col in cols.iter().zip(vals.iter_mut()) {
                             if col.0 == col_name {
+                                found = true;
+
                                 col.1
                                     .replace_data_at_cell_path(&cell_path[1..], new_val.clone())?
                             }
                         }
+                        if !found {
+                            if cell_path.len() == 1 {
+                                cols.push(col_name.clone());
+                                vals.push(new_val);
+                            } else {
+                                let mut new_col = Value::Record {
+                                    cols: vec![],
+                                    vals: vec![],
+                                    span: new_val.span()?,
+                                };
+                                new_col.replace_data_at_cell_path(&cell_path[1..], new_val)?;
+                                vals.push(new_col);
+                            }
+                        }
                     }
                     v => return Err(ShellError::CantFindColumn(*span, v.span()?)),
                 },
diff --git a/src/tests/test_table_operations.rs b/src/tests/test_table_operations.rs
index bea9eb1ecb..8ce6f7355b 100644
--- a/src/tests/test_table_operations.rs
+++ b/src/tests/test_table_operations.rs
@@ -192,3 +192,8 @@ fn select() -> TestResult {
         "b",
     )
 }
+
+#[test]
+fn update_will_insert() -> TestResult {
+    run_test(r#"{} | update a b | get a"#, "b")
+}