Загрузил(а) файлы в 'cozo-bin/src'
parent
528f4c8c22
commit
247a78c6e8
@ -0,0 +1,301 @@
|
||||
/*
|
||||
* Copyright 2022, The Cozo Project Authors.
|
||||
*
|
||||
* This Source Code Form is subject to the terms of the Mozilla Public License, v. 2.0.
|
||||
* If a copy of the MPL was not distributed with this file,
|
||||
* You can obtain one at https://mozilla.org/MPL/2.0/.
|
||||
*/
|
||||
|
||||
// This file is based on code contributed by https://github.com/rhn
|
||||
|
||||
use std::collections::BTreeMap;
|
||||
use std::error::Error;
|
||||
use std::fs;
|
||||
use std::fs::File;
|
||||
use std::io::{Read, Write};
|
||||
|
||||
use clap::Args;
|
||||
use miette::{bail, miette, IntoDiagnostic};
|
||||
use rustyline::history::DefaultHistory;
|
||||
use rustyline::Changeset;
|
||||
use serde_json::{json, Value};
|
||||
|
||||
use cozo::{evaluate_expressions, DataValue, DbInstance, NamedRows, ScriptMutability};
|
||||
|
||||
struct Indented;
|
||||
|
||||
impl rustyline::hint::Hinter for Indented {
|
||||
type Hint = String;
|
||||
}
|
||||
|
||||
impl rustyline::highlight::Highlighter for Indented {}
|
||||
impl rustyline::completion::Completer for Indented {
|
||||
type Candidate = String;
|
||||
|
||||
fn update(
|
||||
&self,
|
||||
_line: &mut rustyline::line_buffer::LineBuffer,
|
||||
_start: usize,
|
||||
_elected: &str,
|
||||
_cl: &mut Changeset,
|
||||
) {
|
||||
unreachable!();
|
||||
}
|
||||
}
|
||||
|
||||
impl rustyline::Helper for Indented {}
|
||||
|
||||
impl rustyline::validate::Validator for Indented {
|
||||
fn validate(
|
||||
&self,
|
||||
ctx: &mut rustyline::validate::ValidationContext<'_>,
|
||||
) -> rustyline::Result<rustyline::validate::ValidationResult> {
|
||||
Ok(if ctx.input().starts_with(' ') {
|
||||
if ctx.input().ends_with('\n') {
|
||||
rustyline::validate::ValidationResult::Valid(None)
|
||||
} else {
|
||||
rustyline::validate::ValidationResult::Incomplete
|
||||
}
|
||||
} else {
|
||||
rustyline::validate::ValidationResult::Valid(None)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Args, Debug)]
|
||||
pub(crate) struct ReplArgs {
|
||||
/// Database engine, can be `mem`, `sqlite`, `rocksdb` and others.
|
||||
#[clap(short, long, default_value_t = String::from("mem"))]
|
||||
engine: String,
|
||||
|
||||
/// Path to the directory to store the database
|
||||
#[clap(short, long, default_value_t = String::from("cozo.db"))]
|
||||
path: String,
|
||||
|
||||
/// Extra config in JSON format
|
||||
#[clap(short, long, default_value_t = String::from("{}"))]
|
||||
config: String,
|
||||
}
|
||||
|
||||
pub(crate) fn repl_main(args: ReplArgs) -> Result<(), Box<dyn Error>> {
|
||||
let db = DbInstance::new(&args.engine, args.path, &args.config).unwrap();
|
||||
|
||||
let db_copy = db.clone();
|
||||
ctrlc::set_handler(move || {
|
||||
let running = db_copy
|
||||
.run_default("::running")
|
||||
.expect("Cannot determine running queries");
|
||||
for row in running.rows {
|
||||
let id = row.into_iter().next().unwrap();
|
||||
eprintln!("Killing running query {id}");
|
||||
db_copy
|
||||
.run_script(
|
||||
"::kill $id",
|
||||
BTreeMap::from([("id".to_string(), id)]),
|
||||
ScriptMutability::Mutable,
|
||||
)
|
||||
.expect("Cannot kill process");
|
||||
}
|
||||
})
|
||||
.expect("Error setting Ctrl-C handler");
|
||||
|
||||
println!("Welcome to the fluidB REPL.");
|
||||
println!("Type a space followed by newline to enter multiline mode.");
|
||||
|
||||
let mut exit = false;
|
||||
let mut rl = rustyline::Editor::<Indented, DefaultHistory>::new()?;
|
||||
let mut params = BTreeMap::new();
|
||||
let mut save_next: Option<String> = None;
|
||||
rl.set_helper(Some(Indented));
|
||||
|
||||
let history_file = ".fluidB_repl_history";
|
||||
if rl.load_history(history_file).is_ok() {
|
||||
println!("Loaded history from {history_file}");
|
||||
}
|
||||
|
||||
loop {
|
||||
let readline = rl.readline("fluidB:~>");
|
||||
match readline {
|
||||
Ok(line) => {
|
||||
if let Err(err) = process_line(&line, &db, &mut params, &mut save_next) {
|
||||
eprintln!("{err:?}");
|
||||
}
|
||||
if let Err(err) = rl.add_history_entry(line) {
|
||||
eprintln!("{err:?}");
|
||||
}
|
||||
exit = false;
|
||||
}
|
||||
Err(rustyline::error::ReadlineError::Interrupted) => {
|
||||
if exit {
|
||||
break;
|
||||
} else {
|
||||
println!("Again to exit");
|
||||
exit = true;
|
||||
}
|
||||
}
|
||||
Err(rustyline::error::ReadlineError::Eof) => break,
|
||||
Err(e) => eprintln!("{e:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
if rl.save_history(history_file).is_ok() {
|
||||
eprintln!("Query history saved in {history_file}");
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn process_line(
|
||||
line: &str,
|
||||
db: &DbInstance,
|
||||
params: &mut BTreeMap<String, DataValue>,
|
||||
save_next: &mut Option<String>,
|
||||
) -> miette::Result<()> {
|
||||
let line = line.trim();
|
||||
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::<Vec<_>>();
|
||||
table.set_titles(prettytable::Row::new(headers));
|
||||
let rows = out
|
||||
.rows
|
||||
.iter()
|
||||
.map(|r| r.iter().map(|c| format!("{c}")).collect::<Vec<_>>())
|
||||
.collect::<Vec<_>>();
|
||||
let rows = rows
|
||||
.iter()
|
||||
.map(|r| r.iter().map(prettytable::Cell::from).collect::<Vec<_>>());
|
||||
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
|
||||
.split_once(|c: char| c.is_whitespace())
|
||||
.unwrap_or((remaining, ""));
|
||||
match op {
|
||||
"eval" => {
|
||||
let out = evaluate_expressions(payload, params, params)?;
|
||||
println!("{out}");
|
||||
}
|
||||
"set" => {
|
||||
let (key, v_str) = payload
|
||||
.trim()
|
||||
.split_once(|c: char| c.is_whitespace())
|
||||
.ok_or_else(|| miette!("Bad set syntax. Should be '%set <KEY> <VALUE>'."))?;
|
||||
let val: Value = serde_json::from_str(v_str).into_diagnostic()?;
|
||||
let val = DataValue::from(val);
|
||||
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)?;
|
||||
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(), ScriptMutability::Mutable)?;
|
||||
process_out(out)?;
|
||||
}
|
||||
"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}");
|
||||
}
|
||||
}
|
||||
_ => {
|
||||
let out = db.run_script(line, params.clone(), ScriptMutability::Mutable)?;
|
||||
process_out(out)?;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
let out = db.run_script(line, params.clone(), ScriptMutability::Mutable)?;
|
||||
process_out(out)?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
Loading…
Reference in New Issue