use clap::{Arg, ArgMatches, Command};
use crate::domain::model::query::{
QueryIdentifier, QueryResult, QuerySchema, QueryScript, QueryValue,
};
use crate::domain::usecases::query::{execute_query, list_queries, query_schema, QueryRunner};
use crate::infra::driving::cli::output::{render_structured, OutputFormat};
use crate::infra::driving::cli::Context;
pub(in super::super) fn subcommand() -> Command {
Command::new("query")
.about("Run a CozoScript query over the workspace")
.subcommand_required(true)
.arg_required_else_help(true)
.subcommand(
Command::new("exec")
.about("Run an inline CozoScript expression")
.arg(Arg::new("script").required(true).value_name("COZOSCRIPT")),
)
.subcommand(
Command::new("run")
.about("Run a named CozoScript file from the queries directory")
.arg(Arg::new("name").required(true).value_name("NAME")),
)
.subcommand(
Command::new("list").about("List named queries available in the queries directory"),
)
.subcommand(
Command::new("schema").about("Print the relations exposed to CozoScript queries"),
)
}
pub(in super::super) fn execute(matches: &ArgMatches, ctx: &Context<'_>) {
match matches.subcommand() {
Some(("exec", sub)) => execute_exec(sub, ctx),
Some(("run", sub)) => execute_run(sub, ctx),
Some(("list", _)) => execute_list(ctx),
Some(("schema", _)) => execute_schema(ctx),
_ => unreachable!(),
}
}
fn execute_exec(matches: &ArgMatches, ctx: &Context<'_>) {
let script_text = crate::infra::driving::cli::helpers::required_str(matches, "script");
let runner = build_runner(ctx);
let script = QueryScript::new(script_text.to_string());
let result = match QueryRunner::run(&runner, &script) {
Ok(r) => r,
Err(e) => {
eprintln!(
"error: {}",
enrich_with_line_col(&format!("{e}"), script_text)
);
std::process::exit(1);
}
};
render(ctx, &result);
}
fn execute_run(matches: &ArgMatches, ctx: &Context<'_>) {
let raw = crate::infra::driving::cli::helpers::required_str(matches, "name");
let id = match QueryIdentifier::new(raw) {
Ok(i) => i,
Err(e) => {
eprintln!("error: {e}");
std::process::exit(1);
}
};
let store = ctx.query_store();
let runner = build_runner(ctx);
let result = match execute_query(&store, &runner, &id) {
Ok(r) => r,
Err(e) => {
eprintln!("error: {e}");
std::process::exit(1);
}
};
render(ctx, &result);
}
fn execute_list(ctx: &Context<'_>) {
let store = ctx.query_store();
let queries = match list_queries(&store) {
Ok(q) => q,
Err(e) => {
eprintln!("error: {e}");
std::process::exit(1);
}
};
for query in &queries {
match query.description.as_ref() {
Some(d) => println!("{}\t{}", query.name, d),
None => println!("{}", query.name),
}
}
}
fn execute_schema(ctx: &Context<'_>) {
let runner = build_runner(ctx);
let schema = match query_schema(&runner) {
Ok(s) => s,
Err(e) => {
eprintln!("error: {e}");
std::process::exit(1);
}
};
render_schema(&schema);
}
fn render_schema(schema: &QuerySchema) {
let mut first = true;
for relation in &schema.relations {
if !first {
println!();
}
first = false;
println!("{}", relation.name);
let name_width = relation
.columns
.iter()
.map(|c| c.name.len())
.max()
.unwrap_or(0);
let type_width = relation
.columns
.iter()
.map(|c| c.type_.len())
.max()
.unwrap_or(0);
for column in &relation.columns {
let key_marker = if column.is_key { " (key)" } else { "" };
println!(
" {n:<name_width$} {t:<type_width$}{key_marker}",
n = column.name,
t = column.type_,
);
}
}
}
fn build_runner(ctx: &Context<'_>) -> crate::infra::driven::cozo::adapter::CozoAdapter {
match ctx.query_runner() {
Ok(r) => r,
Err(e) => {
eprintln!("error: {e}");
std::process::exit(1);
}
}
}
fn render(ctx: &Context<'_>, result: &QueryResult) {
match ctx.output_fmt {
OutputFormat::Human => render_tsv(result),
OutputFormat::Json | OutputFormat::Yaml => {
let payload = rows_to_json(result);
render_structured(&payload, ctx.output_fmt);
}
}
}
fn enrich_with_line_col(message: &str, script: &str) -> String {
let Some(offset) = parse_at_offset(message) else {
return message.to_owned();
};
let (line, col) = byte_offset_to_line_col(script, offset);
format!("{message} (line {line}, col {col})")
}
fn parse_at_offset(message: &str) -> Option<usize> {
let needle = " at ";
let pos = message.rfind(needle)?;
let tail = &message[pos + needle.len()..];
let digits: String = tail.chars().take_while(|c| c.is_ascii_digit()).collect();
if digits.is_empty() {
return None;
}
digits.parse().ok()
}
fn byte_offset_to_line_col(source: &str, offset: usize) -> (usize, usize) {
let mut line = 1usize;
let mut col = 1usize;
for (i, ch) in source.char_indices() {
if i >= offset {
return (line, col);
}
if ch == '\n' {
line += 1;
col = 1;
} else {
col += 1;
}
}
(line, col)
}
fn render_tsv(r: &QueryResult) {
println!("{}", r.headers.join("\t"));
for row in &r.rows {
let cells: Vec<String> = row.iter().map(value_str).collect();
println!("{}", cells.join("\t"));
}
}
fn rows_to_json(r: &QueryResult) -> Vec<serde_json::Map<String, serde_json::Value>> {
r.rows
.iter()
.map(|row| {
let mut obj = serde_json::Map::new();
for (header, value) in r.headers.iter().zip(row.iter()) {
obj.insert(header.clone(), value_to_json(value));
}
obj
})
.collect()
}
fn value_to_json(v: &QueryValue) -> serde_json::Value {
use serde_json::Value;
match v {
QueryValue::Null => Value::Null,
QueryValue::Bool(b) => Value::Bool(*b),
QueryValue::Int(i) => Value::Number((*i).into()),
QueryValue::Float(f) => serde_json::Number::from_f64(*f)
.map(Value::Number)
.unwrap_or(Value::Null),
QueryValue::Str(s) => Value::String(s.clone()),
QueryValue::Other(s) => Value::String(s.clone()),
}
}
fn value_str(v: &QueryValue) -> String {
match v {
QueryValue::Null => "NULL".to_owned(),
QueryValue::Bool(b) => b.to_string(),
QueryValue::Int(i) => i.to_string(),
QueryValue::Float(f) => f.to_string(),
QueryValue::Str(s) => s.clone(),
QueryValue::Other(s) => s.clone(),
}
}