From 4b33e9e848744b545f17894401fbef86adceb1c3 Mon Sep 17 00:00:00 2001 From: Ziyang Hu Date: Wed, 4 Jan 2023 23:44:47 +0800 Subject: [PATCH] import data from file or URL --- Cargo.lock | 1 + cozo-core/src/lib.rs | 8 +- cozoserver/Cargo.toml | 1 + cozoserver/README-zh.md | 1 + cozoserver/README.md | 1 + cozoserver/src/main.rs | 2 +- cozoserver/src/repl.rs | 292 +++++++++++++++++++--------------------- 7 files changed, 149 insertions(+), 157 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 266a8ec8..77e00260 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -670,6 +670,7 @@ dependencies = [ "env_logger", "log", "miette", + "minreq", "prettytable", "rand 0.8.5", "rouille", diff --git a/cozo-core/src/lib.rs b/cozo-core/src/lib.rs index 4fcff6b8..b30a1927 100644 --- a/cozo-core/src/lib.rs +++ b/cozo-core/src/lib.rs @@ -285,7 +285,7 @@ impl DbInstance { /// Note that triggers are _not_ run for the relations, if any exists. /// If you need to activate triggers, use queries with parameters. pub fn import_relations_str(&self, data: &str) -> String { - match self.import_relations_str_inner(data) { + match self.import_relations_str_with_err(data) { Ok(()) => { format!("{}", json!({"ok": true})) } @@ -294,7 +294,11 @@ impl DbInstance { } } } - fn import_relations_str_inner(&self, data: &str) -> Result<()> { + /// Import a relation, the data is given as a JSON string + /// + /// Note that triggers are _not_ run for the relations, if any exists. + /// If you need to activate triggers, use queries with parameters. + pub fn import_relations_str_with_err(&self, data: &str) -> Result<()> { let j_obj: BTreeMap = serde_json::from_str(data).into_diagnostic()?; self.import_relations(j_obj) } diff --git a/cozoserver/Cargo.toml b/cozoserver/Cargo.toml index 47e31974..1adeb5f9 100644 --- a/cozoserver/Cargo.toml +++ b/cozoserver/Cargo.toml @@ -54,4 +54,5 @@ chrono = "0.4.19" serde_json = "1.0.81" prettytable = "0.10.0" rustyline = "10.0.0" +minreq = { version = "2.6.0", features = ["https-rustls"] } miette = { version = "5.5.0", features = ["fancy"] } diff --git a/cozoserver/README-zh.md b/cozoserver/README-zh.md index 2e2d2c6a..47ed3cef 100644 --- a/cozoserver/README-zh.md +++ b/cozoserver/README-zh.md @@ -30,6 +30,7 @@ * `%unset <键>`：删除已设置的参数值。 * `%clear`：清空所有已设置的参数。 * `%params`：显示当前所有参数。 +* `%import <文件或 URL>`：将文件或 URL 里的 JSON 数据导入至数据库。 * `%save <文件>`：下一个成功查询的结果将会以 JSON 格式存储在指定的文件中。如果文件参数未给出，则清除上次的文件设置。 * `%backup <文件>`：备份全部数据至指定的文件。 * `%restore <文件>`：将指定的备份文件中的数据加载到当前数据库中。当前数据库必须为空。 diff --git a/cozoserver/README.md b/cozoserver/README.md index cae78ada..37f6e432 100644 --- a/cozoserver/README.md +++ b/cozoserver/README.md @@ -38,6 +38,7 @@ You can use the following meta ops in the REPL: * `%unset `: unset a parameter. * `%clear`: unset all parameters. * `%params`: print all set parameters. +* `%import `: import data in JSON format from the file or URL. * `%save `: the result of the next successful query will be saved in JSON format in a file instead of printed on screen. If `` is omitted, then the effect of any previous `%save` command is nullified. * `%backup `: the current database will be backed up into the file. * `%restore `: restore the data in the backup to the current database. The current database must be empty. diff --git a/cozoserver/src/main.rs b/cozoserver/src/main.rs index bda0ce31..e1eb3e34 100644 --- a/cozoserver/src/main.rs +++ b/cozoserver/src/main.rs @@ -72,7 +72,6 @@ macro_rules! check_auth { } fn main() { - env_logger::Builder::from_env(Env::default().default_filter_or("info")).init(); let args = Args::parse(); if args.bind != "127.0.0.1" { eprintln!("{}", SECURITY_WARNING); @@ -95,6 +94,7 @@ fn main() { exit(-1); } } else { + env_logger::Builder::from_env(Env::default().default_filter_or("info")).init(); server_main(args, db) } } diff --git a/cozoserver/src/repl.rs b/cozoserver/src/repl.rs index 242af1d4..06b16c79 100644 --- a/cozoserver/src/repl.rs +++ b/cozoserver/src/repl.rs @@ -10,8 +10,10 @@ use std::collections::BTreeMap; use std::error::Error; -use std::io::Write; +use std::fs::File; +use std::io::{Read, Write}; +use miette::{bail, miette, IntoDiagnostic}; use prettytable; use rustyline; use serde_json::{json, Value}; @@ -72,159 +74,8 @@ pub(crate) fn repl_main(db: DbInstance) -> Result<(), Box> { let readline = rl.readline("=> "); match readline { Ok(line) => { - let line = line.trim(); - if line.is_empty() { - continue; - } - if let Some(remaining) = line.strip_prefix("%") { - let remaining = remaining.trim(); - let (op, payload) = remaining - .split_once(|c: char| c.is_whitespace()) - .unwrap_or((remaining, "")); - match op { - "set" => { - if let Some((key, v_str)) = - payload.trim().split_once(|c: char| c.is_whitespace()) - { - match serde_json::from_str(v_str) { - Ok(val) => { - params.insert(key.to_string(), val); - } - Err(e) => { - eprintln!("{:?}", e) - } - } - } else { - eprintln!("Bad set syntax. Should be '%set '.") - } - } - "unset" => { - let key = remaining.trim(); - if params.remove(key).is_none() { - eprintln!("Key not found: '{}'", key) - } - } - "clear" => { - params.clear(); - } - "params" => match serde_json::to_string_pretty(&json!(¶ms)) { - Ok(display) => { - println!("{}", display) - } - Err(err) => { - eprintln!("{:?}", err) - } - }, - "backup" => { - let path = remaining.trim(); - if path.is_empty() { - eprintln!("Backup requires a path"); - } else { - match db.backup_db(path.to_string()) { - Ok(_) => { - println!("Backup written successfully to {}", path) - } - Err(err) => { - eprintln!("{:?}", err) - } - } - } - } - "restore" => { - let path = remaining.trim(); - if path.is_empty() { - eprintln!("Restore requires a path"); - } else { - match db.restore_backup(path) { - Ok(_) => { - println!("Backup successfully loaded from {}", path) - } - Err(err) => { - eprintln!("{:?}", err) - } - } - } - } - "save" => { - let next_path = remaining.trim(); - if next_path.is_empty() { - eprintln!("Next result will NOT be saved to file"); - } else { - eprintln!("Next result will be saved to file: {}", next_path); - save_next = Some(next_path.to_string()) - } - } - op => eprintln!("Unknown op: {}", op), - } - } else { - match db.run_script(&line, params.clone()) { - Ok(out) => { - if let Some(path) = save_next.as_ref() { - println!( - "Query has returned {} rows, saving to file {}", - out.rows.len(), - path - ); - - let to_save = out - .rows - .iter() - .map(|row| -> Value { - row.iter() - .zip(out.headers.iter()) - .map(|(v, k)| (k.to_string(), v.clone())) - .collect() - }) - .collect(); - - let j_payload = Value::Array(to_save); - - match std::fs::File::create(path) { - Ok(mut file) => { - match file.write_all(j_payload.to_string().as_bytes()) { - Ok(_) => { - save_next = None; - } - Err(err) => { - eprintln!("{:?}", err); - } - } - } - Err(err) => { - eprintln!("{:?}", err); - } - } - } else { - use prettytable::format; - let mut table = prettytable::Table::new(); - let headers = out - .headers - .iter() - .map(prettytable::Cell::from) - .collect::>(); - table.set_titles(prettytable::Row::new(headers)); - let rows = out - .rows - .iter() - .map(|r| r.iter().map(|c| format!("{}", c)).collect::>()) - .collect::>(); - let rows = rows.iter().map(|r| { - r.iter().map(prettytable::Cell::from).collect::>() - }); - for row in rows { - table.add_row(prettytable::Row::new(row)); - } - table.set_format(*format::consts::FORMAT_NO_BORDER_LINE_SEPARATOR); - table.printstd(); - } - } - Err(mut err) => { - if err.source_code().is_none() { - err = err.with_source_code(line.to_string()); - } - eprintln!("{:?}", err); - } - }; + if let Err(err) = process_line(&line, &db, &mut params, &mut save_next) { + eprintln!("{:?}", err); } rl.add_history_entry(line); } @@ -242,3 +93,136 @@ pub(crate) fn repl_main(db: DbInstance) -> Result<(), Box> { } Ok(()) } + +fn process_line( + line: &str, + db: &DbInstance, + params: &mut BTreeMap, + save_next: &mut Option, +) -> miette::Result<()> { + let line = line.trim(); + if line.is_empty() { + return Ok(()); + } + if let Some(remaining) = line.strip_prefix("%") { + let remaining = remaining.trim(); + let (op, payload) = remaining + .split_once(|c: char| c.is_whitespace()) + .unwrap_or((remaining, "")); + match op { + "set" => { + let (key, v_str) = payload + .trim() + .split_once(|c: char| c.is_whitespace()) + .ok_or_else(|| miette!("Bad set syntax. Should be '%set '."))?; + let val = serde_json::from_str(v_str).into_diagnostic()?; + params.insert(key.to_string(), val); + } + "unset" => { + let key = payload.trim(); + if params.remove(key).is_none() { + bail!("Key not found: '{}'", key) + } + } + "clear" => { + params.clear(); + } + "params" => { + let display = serde_json::to_string_pretty(&json!(¶ms)).into_diagnostic()?; + println!("{}", display); + } + "backup" => { + let path = payload.trim(); + if path.is_empty() { + bail!("Backup requires a path"); + }; + db.backup_db(path.to_string())?; + println!("Backup written successfully to {}", path) + } + "restore" => { + let path = payload.trim(); + if path.is_empty() { + bail!("Restore requires a path"); + }; + db.restore_backup(path)?; + println!("Backup successfully loaded from {}", path) + } + "save" => { + let next_path = payload.trim(); + if next_path.is_empty() { + println!("Next result will NOT be saved to file"); + } else { + println!("Next result will be saved to file: {}", next_path); + *save_next = Some(next_path.to_string()) + } + } + "import" => { + let url = payload.trim(); + if url.starts_with("http://") || url.starts_with("https://") { + let data = minreq::get(url).send().into_diagnostic()?; + let data = data.as_str().into_diagnostic()?; + db.import_relations_str_with_err(data)?; + println!("Imported data from {}", url) + } else { + let file_path = url.strip_prefix("file://").unwrap_or(url); + let mut file = File::open(file_path).into_diagnostic()?; + let mut content = String::new(); + file.read_to_string(&mut content).into_diagnostic()?; + db.import_relations_str_with_err(&content)?; + println!("Imported data from {}", url); + } + } + op => bail!("Unknown op: {}", op), + } + } else { + let out = db.run_script(&line, params.clone())?; + if let Some(path) = save_next.as_ref() { + println!( + "Query has returned {} rows, saving to file {}", + out.rows.len(), + path + ); + + let to_save = out + .rows + .iter() + .map(|row| -> Value { + row.iter() + .zip(out.headers.iter()) + .map(|(v, k)| (k.to_string(), v.clone())) + .collect() + }) + .collect(); + + let j_payload = Value::Array(to_save); + + let mut file = File::create(path).into_diagnostic()?; + file.write_all(j_payload.to_string().as_bytes()) + .into_diagnostic()?; + *save_next = None; + } else { + use prettytable::format; + let mut table = prettytable::Table::new(); + let headers = out + .headers + .iter() + .map(prettytable::Cell::from) + .collect::>(); + table.set_titles(prettytable::Row::new(headers)); + let rows = out + .rows + .iter() + .map(|r| r.iter().map(|c| format!("{}", c)).collect::>()) + .collect::>(); + let rows = rows + .iter() + .map(|r| r.iter().map(prettytable::Cell::from).collect::>()); + for row in rows { + table.add_row(prettytable::Row::new(row)); + } + table.set_format(*format::consts::FORMAT_NO_BORDER_LINE_SEPARATOR); + table.printstd(); + } + } + Ok(()) +}