1use std::borrow::Cow;
2
3use nu_engine::{ClosureEval, ClosureEvalOnce, command_prelude::*};
4use nu_protocol::ast::PathMember;
5
6#[derive(Clone)]
7pub struct Upsert;
8
9impl Command for Upsert {
10 fn name(&self) -> &str {
11 "upsert"
12 }
13
14 fn signature(&self) -> Signature {
15 Signature::build("upsert")
16 .input_output_types(vec![
17 (Type::record(), Type::record()),
18 (Type::table(), Type::table()),
19 (
20 Type::List(Box::new(Type::Any)),
21 Type::List(Box::new(Type::Any)),
22 ),
23 ])
24 .required(
25 "field",
26 SyntaxShape::CellPath,
27 "The name of the column to update or insert.",
28 )
29 .required(
30 "replacement value",
31 SyntaxShape::Any,
32 "The new value to give the cell(s), or a closure to create the value.",
33 )
34 .allow_variants_without_examples(true)
35 .category(Category::Filters)
36 }
37
38 fn description(&self) -> &str {
39 "Update an existing column to have a new value, or insert a new column."
40 }
41
42 fn extra_description(&self) -> &str {
43 "When updating or inserting a column, the closure will be run for each row, and the current row will be passed as the first argument. \
44Referencing `$in` inside the closure will provide the value at the column for the current row or null if the column does not exist.
45
46When updating a specific index, the closure will instead be run once. The first argument to the closure and the `$in` value will both be the current value at the index. \
47If the command is inserting at the end of a list or table, then both of these values will be null."
48 }
49
50 fn search_terms(&self) -> Vec<&str> {
51 vec!["add"]
52 }
53
54 fn run(
55 &self,
56 engine_state: &EngineState,
57 stack: &mut Stack,
58 call: &Call,
59 input: PipelineData,
60 ) -> Result<PipelineData, ShellError> {
61 upsert(engine_state, stack, call, input)
62 }
63
64 fn examples(&self) -> Vec<Example<'_>> {
65 vec![
66 Example {
67 description: "Update a record's value.",
68 example: "{'name': 'nu', 'stars': 5} | upsert name 'Nushell'",
69 result: Some(Value::test_record(record! {
70 "name" => Value::test_string("Nushell"),
71 "stars" => Value::test_int(5),
72 })),
73 },
74 Example {
75 description: "Insert a new entry into a record.",
76 example: "{'name': 'nu', 'stars': 5} | upsert language 'Rust'",
77 result: Some(Value::test_record(record! {
78 "name" => Value::test_string("nu"),
79 "stars" => Value::test_int(5),
80 "language" => Value::test_string("Rust"),
81 })),
82 },
83 Example {
84 description: "Update each row of a table.",
85 example: "[[name lang]; [Nushell ''] [Reedline '']] | upsert lang 'Rust'",
86 result: Some(Value::test_list(vec![
87 Value::test_record(record! {
88 "name" => Value::test_string("Nushell"),
89 "lang" => Value::test_string("Rust"),
90 }),
91 Value::test_record(record! {
92 "name" => Value::test_string("Reedline"),
93 "lang" => Value::test_string("Rust"),
94 }),
95 ])),
96 },
97 Example {
98 description: "Insert a new column with values computed based off the other columns.",
99 example: "[[foo]; [7] [8] [9]] | upsert bar {|row| $row.foo * 2 }",
100 result: Some(Value::test_list(vec![
101 Value::test_record(record! {
102 "foo" => Value::test_int(7),
103 "bar" => Value::test_int(14),
104 }),
105 Value::test_record(record! {
106 "foo" => Value::test_int(8),
107 "bar" => Value::test_int(16),
108 }),
109 Value::test_record(record! {
110 "foo" => Value::test_int(9),
111 "bar" => Value::test_int(18),
112 }),
113 ])),
114 },
115 Example {
116 description: "Update null values in a column to a default value.",
117 example: "[[foo]; [2] [null] [4]] | upsert foo { default 0 }",
118 result: Some(Value::test_list(vec![
119 Value::test_record(record! {
120 "foo" => Value::test_int(2),
121 }),
122 Value::test_record(record! {
123 "foo" => Value::test_int(0),
124 }),
125 Value::test_record(record! {
126 "foo" => Value::test_int(4),
127 }),
128 ])),
129 },
130 Example {
131 description: "Upsert into a list, updating an existing value at an index.",
132 example: "[1 2 3] | upsert 0 2",
133 result: Some(Value::test_list(vec![
134 Value::test_int(2),
135 Value::test_int(2),
136 Value::test_int(3),
137 ])),
138 },
139 Example {
140 description: "Upsert into a list, inserting a new value at the end.",
141 example: "[1 2 3] | upsert 3 4",
142 result: Some(Value::test_list(vec![
143 Value::test_int(1),
144 Value::test_int(2),
145 Value::test_int(3),
146 Value::test_int(4),
147 ])),
148 },
149 Example {
150 description: "Upsert into a nested path, creating new values as needed.",
151 example: "[{} {a: [{}]}] | upsert a.0.b \"value\"",
152 result: Some(Value::test_list(vec![
153 Value::test_record(record!(
154 "a" => Value::test_record(record!(
155 "b" => Value::test_string("value"),
156 )),
157 )),
158 Value::test_record(record!(
159 "a" => Value::test_list(vec![Value::test_record(record!(
160 ))]),
161 )),
162 ])),
163 },
164 ]
165 }
166}
167
168fn upsert_recursive(
169 engine_state: &EngineState,
170 stack: &mut Stack,
171 head_span: Span,
172 replacement: Value,
173 input: PipelineData,
174 cell_paths: &[PathMember],
175) -> Result<PipelineData, ShellError> {
176 match input {
177 PipelineData::Value(mut value, metadata) => {
178 if let Value::Closure { val, .. } = replacement {
179 upsert_single_value_by_closure(
180 &mut value,
181 ClosureEvalOnce::new(engine_state, stack, *val),
182 head_span,
183 cell_paths,
184 false,
185 )?;
186 } else {
187 value.upsert_data_at_cell_path(cell_paths, replacement)?;
188 }
189 Ok(value.into_pipeline_data_with_metadata(metadata))
190 }
191 PipelineData::ListStream(stream, metadata) => {
192 if let Some((
193 &PathMember::Int {
194 val,
195 span: path_span,
196 ..
197 },
198 path,
199 )) = cell_paths.split_first()
200 {
201 let mut stream = stream.into_iter();
202 let mut pre_elems = vec![];
203
204 for idx in 0..val {
205 if let Some(v) = stream.next() {
206 pre_elems.push(v);
207 } else {
208 return Err(ShellError::InsertAfterNextFreeIndex {
209 available_idx: idx,
210 span: path_span,
211 });
212 }
213 }
214
215 let value = if path.is_empty() {
216 let value = stream.next().unwrap_or(Value::nothing(head_span));
217 if let Value::Closure { val, .. } = replacement {
218 ClosureEvalOnce::new(engine_state, stack, *val)
219 .run_with_value(value)?
220 .into_value(head_span)?
221 } else {
222 replacement
223 }
224 } else if let Some(mut value) = stream.next() {
225 if let Value::Closure { val, .. } = replacement {
226 upsert_single_value_by_closure(
227 &mut value,
228 ClosureEvalOnce::new(engine_state, stack, *val),
229 head_span,
230 path,
231 true,
232 )?;
233 } else {
234 value.upsert_data_at_cell_path(path, replacement)?;
235 }
236 value
237 } else {
238 return Err(ShellError::AccessBeyondEnd {
239 max_idx: pre_elems.len() - 1,
240 span: path_span,
241 });
242 };
243
244 pre_elems.push(value);
245
246 Ok(pre_elems
247 .into_iter()
248 .chain(stream)
249 .into_pipeline_data_with_metadata(
250 head_span,
251 engine_state.signals().clone(),
252 metadata,
253 ))
254 } else if let Some(new_cell_paths) = Value::try_put_int_path_member_on_top(cell_paths) {
255 upsert_recursive(
256 engine_state,
257 stack,
258 head_span,
259 replacement,
260 PipelineData::ListStream(stream, metadata),
261 &new_cell_paths,
262 )
263 } else if let Value::Closure { val, .. } = replacement {
264 let mut closure = ClosureEval::new(engine_state, stack, *val);
265 let cell_paths = cell_paths.to_vec();
266 let stream = stream.map(move |mut value| {
267 let err =
268 upsert_value_by_closure(&mut value, &mut closure, head_span, &cell_paths);
269
270 if let Err(e) = err {
271 Value::error(e, head_span)
272 } else {
273 value
274 }
275 });
276
277 Ok(PipelineData::list_stream(stream, metadata))
278 } else {
279 let cell_paths = cell_paths.to_vec();
280 let stream = stream.map(move |mut value| {
281 if let Err(e) = value.upsert_data_at_cell_path(&cell_paths, replacement.clone())
282 {
283 Value::error(e, head_span)
284 } else {
285 value
286 }
287 });
288
289 Ok(PipelineData::list_stream(stream, metadata))
290 }
291 }
292 PipelineData::Empty => Err(ShellError::IncompatiblePathAccess {
293 type_name: "empty pipeline".to_string(),
294 span: head_span,
295 }),
296 PipelineData::ByteStream(stream, ..) => Err(ShellError::IncompatiblePathAccess {
297 type_name: stream.type_().describe().into(),
298 span: head_span,
299 }),
300 }
301}
302
303fn upsert(
304 engine_state: &EngineState,
305 stack: &mut Stack,
306 call: &Call,
307 input: PipelineData,
308) -> Result<PipelineData, ShellError> {
309 let head = call.head;
310 let cell_path: CellPath = call.req(engine_state, stack, 0)?;
311 let replacement: Value = call.req(engine_state, stack, 1)?;
312 let input = input.into_stream_or_original(engine_state);
313
314 upsert_recursive(
315 engine_state,
316 stack,
317 head,
318 replacement,
319 input,
320 &cell_path.members,
321 )
322}
323
324fn upsert_value_by_closure(
325 value: &mut Value,
326 closure: &mut ClosureEval,
327 span: Span,
328 cell_path: &[PathMember],
329) -> Result<(), ShellError> {
330 let value_at_path = value.follow_cell_path(cell_path);
331
332 let input = value_at_path
333 .map(Cow::into_owned)
334 .map(IntoPipelineData::into_pipeline_data)
335 .unwrap_or(PipelineData::empty());
336
337 let new_value = closure
338 .add_arg(value.clone())
339 .run_with_input(input)?
340 .into_value(span)?;
341
342 value.upsert_data_at_cell_path(cell_path, new_value)
343}
344
345fn upsert_single_value_by_closure(
346 value: &mut Value,
347 closure: ClosureEvalOnce,
348 span: Span,
349 cell_path: &[PathMember],
350 cell_value_as_arg: bool,
351) -> Result<(), ShellError> {
352 let value_at_path = value.follow_cell_path(cell_path);
353
354 let arg = if cell_value_as_arg {
358 value_at_path
359 .as_deref()
360 .cloned()
361 .unwrap_or(Value::nothing(span))
362 } else {
363 value.clone()
364 };
365
366 let input = value_at_path
367 .map(Cow::into_owned)
368 .map(IntoPipelineData::into_pipeline_data)
369 .unwrap_or(PipelineData::empty());
370
371 let new_value = closure
372 .add_arg(arg)
373 .run_with_input(input)?
374 .into_value(span)?;
375
376 value.upsert_data_at_cell_path(cell_path, new_value)
377}
378
379#[cfg(test)]
380mod test {
381 use super::*;
382
383 #[test]
384 fn test_examples() -> nu_test_support::Result {
385 nu_test_support::test().examples(Upsert)
386 }
387}