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_list(vec![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 "b" => Value::test_string("value"),
161 ))]),
162 )),
163 ])),
164 },
165 ]
166 }
167}
168
169fn upsert(
170 engine_state: &EngineState,
171 stack: &mut Stack,
172 call: &Call,
173 input: PipelineData,
174) -> Result<PipelineData, ShellError> {
175 let head = call.head;
176 let cell_path: CellPath = call.req(engine_state, stack, 0)?;
177 let replacement: Value = call.req(engine_state, stack, 1)?;
178
179 match input {
180 PipelineData::Value(mut value, metadata) => {
181 if let Value::Closure { val, .. } = replacement {
182 match (cell_path.members.first(), &mut value) {
183 (Some(PathMember::String { .. }), Value::List { vals, .. }) => {
184 let mut closure = ClosureEval::new(engine_state, stack, *val);
185 for val in vals {
186 upsert_value_by_closure(
187 val,
188 &mut closure,
189 head,
190 &cell_path.members,
191 false,
192 )?;
193 }
194 }
195 (first, _) => {
196 upsert_single_value_by_closure(
197 &mut value,
198 ClosureEvalOnce::new(engine_state, stack, *val),
199 head,
200 &cell_path.members,
201 matches!(first, Some(PathMember::Int { .. })),
202 )?;
203 }
204 }
205 } else {
206 value.upsert_data_at_cell_path(&cell_path.members, replacement)?;
207 }
208 Ok(value.into_pipeline_data_with_metadata(metadata))
209 }
210 PipelineData::ListStream(stream, metadata) => {
211 if let Some((
212 &PathMember::Int {
213 val,
214 span: path_span,
215 ..
216 },
217 path,
218 )) = cell_path.members.split_first()
219 {
220 let mut stream = stream.into_iter();
221 let mut pre_elems = vec![];
222
223 for idx in 0..val {
224 if let Some(v) = stream.next() {
225 pre_elems.push(v);
226 } else {
227 return Err(ShellError::InsertAfterNextFreeIndex {
228 available_idx: idx,
229 span: path_span,
230 });
231 }
232 }
233
234 let value = if path.is_empty() {
235 let value = stream.next().unwrap_or(Value::nothing(head));
236 if let Value::Closure { val, .. } = replacement {
237 ClosureEvalOnce::new(engine_state, stack, *val)
238 .run_with_value(value)?
239 .into_value(head)?
240 } else {
241 replacement
242 }
243 } else if let Some(mut value) = stream.next() {
244 if let Value::Closure { val, .. } = replacement {
245 upsert_single_value_by_closure(
246 &mut value,
247 ClosureEvalOnce::new(engine_state, stack, *val),
248 head,
249 path,
250 true,
251 )?;
252 } else {
253 value.upsert_data_at_cell_path(path, replacement)?;
254 }
255 value
256 } else {
257 return Err(ShellError::AccessBeyondEnd {
258 max_idx: pre_elems.len() - 1,
259 span: path_span,
260 });
261 };
262
263 pre_elems.push(value);
264
265 Ok(pre_elems
266 .into_iter()
267 .chain(stream)
268 .into_pipeline_data_with_metadata(
269 head,
270 engine_state.signals().clone(),
271 metadata,
272 ))
273 } else if let Value::Closure { val, .. } = replacement {
274 let mut closure = ClosureEval::new(engine_state, stack, *val);
275 let stream = stream.map(move |mut value| {
276 let err = upsert_value_by_closure(
277 &mut value,
278 &mut closure,
279 head,
280 &cell_path.members,
281 false,
282 );
283
284 if let Err(e) = err {
285 Value::error(e, head)
286 } else {
287 value
288 }
289 });
290
291 Ok(PipelineData::list_stream(stream, metadata))
292 } else {
293 let stream = stream.map(move |mut value| {
294 if let Err(e) =
295 value.upsert_data_at_cell_path(&cell_path.members, replacement.clone())
296 {
297 Value::error(e, head)
298 } else {
299 value
300 }
301 });
302
303 Ok(PipelineData::list_stream(stream, metadata))
304 }
305 }
306 PipelineData::Empty => Err(ShellError::IncompatiblePathAccess {
307 type_name: "empty pipeline".to_string(),
308 span: head,
309 }),
310 PipelineData::ByteStream(stream, ..) => Err(ShellError::IncompatiblePathAccess {
311 type_name: stream.type_().describe().into(),
312 span: head,
313 }),
314 }
315}
316
317fn upsert_value_by_closure(
318 value: &mut Value,
319 closure: &mut ClosureEval,
320 span: Span,
321 cell_path: &[PathMember],
322 first_path_member_int: bool,
323) -> Result<(), ShellError> {
324 let value_at_path = value.follow_cell_path(cell_path);
325
326 let arg = if first_path_member_int {
327 value_at_path
328 .as_deref()
329 .cloned()
330 .unwrap_or(Value::nothing(span))
331 } else {
332 value.clone()
333 };
334
335 let input = value_at_path
336 .map(Cow::into_owned)
337 .map(IntoPipelineData::into_pipeline_data)
338 .unwrap_or(PipelineData::empty());
339
340 let new_value = closure
341 .add_arg(arg)
342 .run_with_input(input)?
343 .into_value(span)?;
344
345 value.upsert_data_at_cell_path(cell_path, new_value)
346}
347
348fn upsert_single_value_by_closure(
349 value: &mut Value,
350 closure: ClosureEvalOnce,
351 span: Span,
352 cell_path: &[PathMember],
353 first_path_member_int: bool,
354) -> Result<(), ShellError> {
355 let value_at_path = value.follow_cell_path(cell_path);
356
357 let arg = if first_path_member_int {
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() {
385 use crate::test_examples;
386
387 test_examples(Upsert {})
388 }
389}