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