1use crate::database::{MEMORY_DB, SQLiteDatabase, get_shared_mem_conn};
2use nu_engine::command_prelude::*;
3use nu_protocol::shell_error::generic::GenericError;
4use rusqlite::Connection;
5use std::fmt::Write;
6
7#[derive(Clone)]
8pub struct StorCreate;
9
10impl Command for StorCreate {
11 fn name(&self) -> &str {
12 "stor create"
13 }
14
15 fn signature(&self) -> Signature {
16 Signature::build("stor create")
17 .input_output_types(vec![(Type::Nothing, Type::table())])
18 .required_named(
19 "table-name",
20 SyntaxShape::String,
21 "Name of the table you want to create.",
22 Some('t'),
23 )
24 .required_named(
25 "columns",
26 SyntaxShape::record(),
27 "A record of column names and datatypes.",
28 Some('c'),
29 )
30 .allow_variants_without_examples(true)
31 .category(Category::Database)
32 }
33
34 fn description(&self) -> &str {
35 "Create a table in the in-memory sqlite database."
36 }
37
38 fn search_terms(&self) -> Vec<&str> {
39 vec!["sqlite", "storing", "table"]
40 }
41
42 fn examples(&self) -> Vec<Example<'_>> {
43 vec![
44 Example {
45 description: "Create an in-memory sqlite database with specified table name, column names, and column data types",
46 example: "stor create --table-name nudb --columns {bool1: bool, int1: int, float1: float, str1: str, datetime1: datetime}",
47 result: None,
48 },
49 Example {
50 description: "Create an in-memory sqlite database with a json column",
51 example: "stor create --table-name files_with_md --columns {file: str, metadata: jsonb}",
52 result: None,
53 },
54 ]
55 }
56
57 fn run(
58 &self,
59 engine_state: &EngineState,
60 stack: &mut Stack,
61 call: &Call,
62 _input: PipelineData,
63 ) -> Result<PipelineData, ShellError> {
64 let span = call.head;
65 let table_name: Option<String> = call.get_flag(engine_state, stack, "table-name")?;
66 let columns: Option<Record> = call.get_flag(engine_state, stack, "columns")?;
67
68 let conn = get_shared_mem_conn()?;
69 process(table_name, span, &conn, columns)?;
70
71 let db = Box::new(SQLiteDatabase::new(
72 std::path::Path::new(MEMORY_DB),
73 engine_state.signals().clone(),
74 ));
75 Ok(Value::custom(db, span).into_pipeline_data())
76 }
77}
78
79fn process(
80 table_name: Option<String>,
81 span: Span,
82 conn: &Connection,
83 columns: Option<Record>,
84) -> Result<(), ShellError> {
85 if table_name.is_none() {
86 return Err(ShellError::MissingParameter {
87 param_name: "requires at table name".into(),
88 span,
89 });
90 }
91 let new_table_name = table_name.unwrap_or("table".into());
92 match columns {
93 Some(record) => {
94 let mut create_stmt = format!("CREATE TABLE {new_table_name} ( ");
95 for (column_name, column_datatype) in record {
96 match column_datatype.coerce_str()?.to_lowercase().as_ref() {
97 "int" => {
98 write!(create_stmt, "{column_name} INTEGER, ")
99 .expect("writing to a String is infallible");
100 }
101 "float" => {
102 write!(create_stmt, "{column_name} REAL, ")
103 .expect("writing to a String is infallible");
104 }
105 "str" => {
106 write!(create_stmt, "{column_name} VARCHAR(255), ")
107 .expect("writing to a String is infallible");
108 }
109
110 "bool" => {
111 write!(create_stmt, "{column_name} BOOLEAN, ")
112 .expect("writing to a String is infallible");
113 }
114 "datetime" => {
115 write!(
116 create_stmt,
117 "{column_name} DATETIME DEFAULT(STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')), "
118 )
119 .expect("writing to a String is infallible");
120 }
121 "json" => {
122 write!(create_stmt, "{column_name} JSON, ")
123 .expect("writing to a String is infallible");
124 }
125 "jsonb" => {
126 write!(create_stmt, "{column_name} JSONB, ")
127 .expect("writing to a String is infallible");
128 }
129
130 _ => {
131 return Err(ShellError::UnsupportedInput {
132 msg: "Unsupported column data type. Please use: int, float, str, bool, datetime, json, jsonb".into(),
133 input: format!("{column_datatype:?}"),
134 msg_span: column_datatype.span(),
135 input_span: column_datatype.span(),
136 });
137 }
138 }
139 }
140 if create_stmt.ends_with(", ") {
141 create_stmt.pop();
142 create_stmt.pop();
143 }
144 create_stmt.push_str(" )");
145
146 conn.execute(&create_stmt, []).map_err(|err| {
149 ShellError::Generic(GenericError::new_internal(
150 "Failed to open SQLite connection to the in-memory SQLite databasefrom create.",
151 err.to_string(),
152 ))
153 })?;
154 }
155 None => {
156 return Err(ShellError::MissingParameter {
157 param_name: "requires at least one column".into(),
158 span,
159 });
160 }
161 };
162 Ok(())
163}
164
165#[cfg(test)]
166mod test {
167 use super::*;
168
169 #[test]
170 fn test_examples() -> nu_test_support::Result {
171 nu_test_support::test().examples(StorCreate)
172 }
173
174 #[test]
175 fn test_process_with_valid_parameters() {
176 let table_name = Some("test_table".to_string());
177 let span = Span::test_data();
178 let conn = get_shared_mem_conn().expect("Test was unable to get shared connection.");
179 let mut columns = Record::new();
180 columns.insert(
181 "int_column".to_string(),
182 Value::test_string("int".to_string()),
183 );
184
185 let result = process(table_name, span, &conn, Some(columns));
186
187 assert!(result.is_ok());
188 }
189
190 #[test]
191 fn test_process_with_missing_table_name() {
192 let table_name = None;
193 let span = Span::test_data();
194 let conn = get_shared_mem_conn().expect("Test was unable to get shared connection.");
195 let mut columns = Record::new();
196 columns.insert(
197 "int_column".to_string(),
198 Value::test_string("int".to_string()),
199 );
200
201 let result = process(table_name, span, &conn, Some(columns));
202
203 assert!(result.is_err());
204 assert!(
205 result
206 .unwrap_err()
207 .to_string()
208 .contains("requires at table name")
209 );
210 }
211
212 #[test]
213 fn test_process_with_missing_columns() {
214 let table_name = Some("test_table".to_string());
215 let span = Span::test_data();
216 let conn = get_shared_mem_conn().expect("Test was unable to get shared connection.");
217
218 let result = process(table_name, span, &conn, None);
219
220 assert!(result.is_err());
221 assert!(
222 result
223 .unwrap_err()
224 .to_string()
225 .contains("requires at least one column")
226 );
227 }
228
229 #[test]
230 fn test_process_with_unsupported_column_data_type() {
231 let table_name = Some("test_table".to_string());
232 let span = Span::test_data();
233 let conn = get_shared_mem_conn().expect("Test was unable to get shared connection.");
234 let mut columns = Record::new();
235 let column_datatype = "bogus_data_type".to_string();
236 columns.insert(
237 "column0".to_string(),
238 Value::test_string(column_datatype.clone()),
239 );
240
241 let result = process(table_name, span, &conn, Some(columns));
242
243 assert!(result.is_err());
244
245 let expected_err = ShellError::UnsupportedInput {
246 msg: "Unsupported column data type. Please use: int, float, str, bool, datetime, json, jsonb".into(),
247 input: format!("{:?}", column_datatype.clone()),
248 msg_span: Span::test_data(),
249 input_span: Span::test_data(),
250 };
251 assert_eq!(result.unwrap_err().to_string(), expected_err.to_string());
252 }
253}