Skip to main content

nu_command/stor/
update.rs

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                // FIXME Type::Any input added to disable pipeline input type checking, as run-time checks can raise undesirable type errors
21                // which aren't caught by the parser. see https://github.com/nushell/nushell/pull/14922 for more details
22                (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        // Check if the record is being passed as input or using the update record parameter
90        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            // Since input is being used, check if the data record parameter is used too
121            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    // Yup, this is a bit janky, but I'm not sure a better way to do this without having
182    // --and and --or flags as well as supporting ==, !=, <>, is null, is not null, etc.
183    // and other sql syntax. So, for now, just type a sql where clause as a string.
184    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    // dbg!(&update_stmt);
189
190    // Get the params from the passed values
191    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}