1use crate::database::{MEMORY_DB, SQLiteDatabase, get_shared_mem_conn, values_to_sql};
2use nu_engine::command_prelude::*;
3use nu_protocol::shell_error::generic::GenericError;
4use rusqlite::{Connection, params_from_iter};
5use std::fmt::Write;
6
7#[derive(Clone)]
8pub struct StorUpdate;
9
10impl Command for StorUpdate {
11 fn name(&self) -> &str {
12 "stor update"
13 }
14
15 fn signature(&self) -> Signature {
16 Signature::build("stor update")
17 .input_output_types(vec![
18 (Type::Nothing, Type::table()),
19 (Type::record(), Type::table()),
20 (Type::Any, Type::table()),
23 ])
24 .required_named(
25 "table-name",
26 SyntaxShape::String,
27 "Name of the table you want to insert into.",
28 Some('t'),
29 )
30 .named(
31 "update-record",
32 SyntaxShape::record(),
33 "A record of column names and column values to update in the specified table.",
34 Some('u'),
35 )
36 .named(
37 "where-clause",
38 SyntaxShape::String,
39 "A sql string to use as a where clause without the WHERE keyword.",
40 Some('w'),
41 )
42 .allow_variants_without_examples(true)
43 .category(Category::Database)
44 }
45
46 fn description(&self) -> &str {
47 "Update information in a specified table in the in-memory sqlite database."
48 }
49
50 fn search_terms(&self) -> Vec<&str> {
51 vec!["sqlite", "storing", "table", "saving", "changing"]
52 }
53
54 fn examples(&self) -> Vec<Example<'_>> {
55 vec![
56 Example {
57 description: "Update the in-memory sqlite database",
58 example: "stor update --table-name nudb --update-record {str1: nushell datetime1: 2020-04-17}",
59 result: None,
60 },
61 Example {
62 description: "Update the in-memory sqlite database with a where clause",
63 example: "stor update --table-name nudb --update-record {str1: nushell datetime1: 2020-04-17} --where-clause \"bool1 = 1\"",
64 result: None,
65 },
66 Example {
67 description: "Update the in-memory sqlite database through pipeline input",
68 example: "{str1: nushell datetime1: 2020-04-17} | stor update --table-name nudb",
69 result: None,
70 },
71 ]
72 }
73
74 fn run(
75 &self,
76 engine_state: &EngineState,
77 stack: &mut Stack,
78 call: &Call,
79 input: PipelineData,
80 ) -> Result<PipelineData, ShellError> {
81 let span = call.head;
82 let table_name: Option<String> = call.get_flag(engine_state, stack, "table-name")?;
83 let update_record: Option<Record> = call.get_flag(engine_state, stack, "update-record")?;
84 let where_clause_opt: Option<Spanned<String>> =
85 call.get_flag(engine_state, stack, "where-clause")?;
86
87 let conn = get_shared_mem_conn()?;
88
89 let columns = handle(span, update_record, input)?;
91
92 process(
93 engine_state,
94 table_name,
95 span,
96 &conn,
97 columns,
98 where_clause_opt,
99 )?;
100
101 let db = Box::new(SQLiteDatabase::new(
102 std::path::Path::new(MEMORY_DB),
103 engine_state.signals().clone(),
104 ));
105 Ok(Value::custom(db, span).into_pipeline_data())
106 }
107}
108
109fn handle(
110 span: Span,
111 update_record: Option<Record>,
112 input: PipelineData,
113) -> Result<Record, ShellError> {
114 match input {
115 PipelineData::Empty => update_record.ok_or_else(|| ShellError::MissingParameter {
116 param_name: "requires a record".into(),
117 span,
118 }),
119 PipelineData::Value(value, ..) => {
120 if update_record.is_some() {
122 return Err(ShellError::Generic(GenericError::new(
123 "Pipeline and Flag both being used",
124 "Use either pipeline input or '--update-record' parameter",
125 span,
126 )));
127 }
128 match value {
129 Value::Record { val, .. } => Ok(val.into_owned()),
130 val => Err(ShellError::OnlySupportsThisInputType {
131 exp_input_type: "record".into(),
132 wrong_type: val.get_type().to_string(),
133 dst_span: span,
134 src_span: val.span(),
135 }),
136 }
137 }
138 _ => {
139 if update_record.is_some() {
140 return Err(ShellError::Generic(GenericError::new(
141 "Pipeline and Flag both being used",
142 "Use either pipeline input or '--update-record' parameter",
143 span,
144 )));
145 }
146 Err(ShellError::OnlySupportsThisInputType {
147 exp_input_type: "record".into(),
148 wrong_type: "".into(),
149 dst_span: span,
150 src_span: span,
151 })
152 }
153 }
154}
155
156fn process(
157 engine_state: &EngineState,
158 table_name: Option<String>,
159 span: Span,
160 conn: &Connection,
161 record: Record,
162 where_clause_opt: Option<Spanned<String>>,
163) -> Result<(), ShellError> {
164 if table_name.is_none() {
165 return Err(ShellError::MissingParameter {
166 param_name: "requires at table name".into(),
167 span,
168 });
169 }
170 let new_table_name = table_name.unwrap_or("table".into());
171 let mut update_stmt = format!("UPDATE {new_table_name} ");
172
173 update_stmt.push_str("SET ");
174 let mut placeholders: Vec<String> = Vec::new();
175
176 for (index, (key, _)) in record.iter().enumerate() {
177 placeholders.push(format!("{} = ?{}", key, index + 1));
178 }
179 update_stmt.push_str(&placeholders.join(", "));
180
181 if let Some(where_clause) = where_clause_opt {
185 write!(update_stmt, " WHERE {}", where_clause.item)
186 .expect("writing to a String is infallible");
187 }
188 let params = values_to_sql(engine_state, record.values().cloned(), span)?;
192
193 conn.execute(&update_stmt, params_from_iter(params))
194 .map_err(|err| {
195 ShellError::Generic(GenericError::new_internal(
196 "Failed to open SQLite connection to the in-memory database from update",
197 err.to_string(),
198 ))
199 })?;
200 Ok(())
201}
202
203#[cfg(test)]
204mod test {
205 use super::*;
206
207 #[test]
208 fn test_examples() -> nu_test_support::Result {
209 nu_test_support::test().examples(StorUpdate)
210 }
211}