From 171d4a5f796dad6f02bcf33a673c0e740659573d Mon Sep 17 00:00:00 2001 From: Ziyang Hu Date: Wed, 25 Jan 2023 17:33:19 +0800 Subject: [PATCH] REPL can run files stored on the FS --- cozo-bin/README-zh.md | 1 + cozo-bin/README.md | 1 + cozo-bin/src/repl.rs | 112 ++++++++++++++++++++++++------------------ 3 files changed, 66 insertions(+), 48 deletions(-) diff --git a/cozo-bin/README-zh.md b/cozo-bin/README-zh.md index b451335a..66d51580 100644 --- a/cozo-bin/README-zh.md +++ b/cozo-bin/README-zh.md @@ -30,6 +30,7 @@ * `%unset <键>`：删除已设置的参数值。 * `%clear`：清空所有已设置的参数。 * `%params`：显示当前所有参数。 +* `%run <文件>`: 运行 `<文件>` 中包含的查询。 * `%import <文件或 URL>`：将文件或 URL 里的 JSON 数据导入至数据库。 * `%save <文件>`：下一个成功查询的结果将会以 JSON 格式存储在指定的文件中。如果文件参数未给出，则清除上次的文件设置。 * `%backup <文件>`：备份全部数据至指定的文件。 diff --git a/cozo-bin/README.md b/cozo-bin/README.md index 6df2d145..d85aec7f 100644 --- a/cozo-bin/README.md +++ b/cozo-bin/README.md @@ -37,6 +37,7 @@ You can use the following meta ops in the REPL: * `%unset `: unset a parameter. * `%clear`: unset all parameters. * `%params`: print all set parameters. +* `%run `: run the script contained in ``. * `%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. diff --git a/cozo-bin/src/repl.rs b/cozo-bin/src/repl.rs index 88203e04..d83fde14 100644 --- a/cozo-bin/src/repl.rs +++ b/cozo-bin/src/repl.rs @@ -10,6 +10,7 @@ use std::collections::BTreeMap; use std::error::Error; +use std::fs; use std::fs::File; use std::io::{Read, Write}; @@ -17,7 +18,7 @@ use clap::Args; use miette::{bail, miette, IntoDiagnostic}; use serde_json::{json, Value}; -use cozo::{DataValue, DbInstance}; +use cozo::{DataValue, DbInstance, NamedRows}; struct Indented; @@ -144,6 +145,58 @@ fn process_line( if line.is_empty() { return Ok(()); } + + let mut process_out = |out: NamedRows| -> miette::Result<()> { + 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(()) + }; + if let Some(remaining) = line.strip_prefix('%') { let remaining = remaining.trim(); let (op, payload) = remaining @@ -179,6 +232,15 @@ fn process_line( db.backup_db(path)?; println!("Backup written successfully to {path}") } + "run" => { + let path = payload.trim(); + if path.is_empty() { + bail!("Run requires path to a script"); + } + let content = fs::read_to_string(path).into_diagnostic()?; + let out = db.run_script(&content, params.clone())?; + process_out(out)?; + } "restore" => { let path = payload.trim(); if path.is_empty() { @@ -216,53 +278,7 @@ fn process_line( } } 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(); - } + process_out(out)?; } Ok(()) }