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