use std::{fs, io::Write as _, path::Path};
use anyhow::{Result, anyhow, bail};
use serde_json::json;
use crate::cli::{MetadataFormat, OutputFormat};
pub(crate) fn write_query_result(
result: &query_forge::QueryResult,
output_path: Option<&Path>,
explicit_format: Option<OutputFormat>,
) -> Result<()> {
match resolve_output_format(output_path, explicit_format)? {
OutputFormat::Text => {
let rendered = query_forge::render_text(result);
write_or_print(&rendered, output_path)?;
}
OutputFormat::Csv => {
let rendered = query_forge::render_csv(result);
write_or_print(&rendered, output_path)?;
}
OutputFormat::Json => {
let rendered = query_forge::render_json(result);
write_or_print(&rendered, output_path)?;
}
OutputFormat::Jsonl => {
let rendered = query_forge::render_jsonl(result);
write_or_print(&rendered, output_path)?;
}
OutputFormat::Markdown => {
let rendered = query_forge::render_markdown(result);
write_or_print(&rendered, output_path)?;
}
OutputFormat::Xml => {
let rendered = query_forge::render_xml(result);
write_or_print(&rendered, output_path)?;
}
OutputFormat::Xlsx => {
let output_path = output_path
.ok_or_else(|| anyhow!("--output is required when --format xlsx is selected"))?;
query_forge::write_xlsx(result, output_path)?;
}
OutputFormat::Parquet => {
let output_path = output_path
.ok_or_else(|| anyhow!("--output is required when --format parquet is selected"))?;
query_forge::write_parquet(result, output_path)?;
}
}
Ok(())
}
fn write_or_print(rendered: &str, output_path: Option<&Path>) -> Result<()> {
if let Some(output_path) = output_path {
fs::write(output_path, rendered)?;
} else {
println!("{rendered}");
}
Ok(())
}
pub(crate) fn resolve_output_format(
output: Option<&Path>,
format: Option<OutputFormat>,
) -> Result<OutputFormat> {
if let Some(format) = format {
if matches!(format, OutputFormat::Xlsx | OutputFormat::Parquet) && output.is_none() {
let name = if matches!(format, OutputFormat::Xlsx) {
"xlsx"
} else {
"parquet"
};
bail!("--output is required when --format {name} is selected");
}
return Ok(format);
}
if output
.and_then(|path| path.extension())
.map(|extension| extension.to_string_lossy())
.is_some_and(|extension| extension.eq_ignore_ascii_case("xlsx"))
{
return Ok(OutputFormat::Xlsx);
}
if output
.and_then(|path| path.extension())
.map(|extension| extension.to_string_lossy())
.is_some_and(|extension| extension.eq_ignore_ascii_case("parquet"))
{
return Ok(OutputFormat::Parquet);
}
if output
.and_then(|path| path.extension())
.map(|extension| extension.to_string_lossy())
.is_some_and(|extension| extension.eq_ignore_ascii_case("csv"))
{
return Ok(OutputFormat::Csv);
}
if output
.and_then(|path| path.extension())
.map(|extension| extension.to_string_lossy())
.is_some_and(|extension| extension.eq_ignore_ascii_case("json"))
{
return Ok(OutputFormat::Json);
}
if output
.and_then(|path| path.extension())
.map(|extension| extension.to_string_lossy())
.is_some_and(|extension| extension.eq_ignore_ascii_case("jsonl"))
{
return Ok(OutputFormat::Jsonl);
}
if output
.and_then(|path| path.extension())
.map(|extension| extension.to_string_lossy())
.is_some_and(|extension| {
extension.eq_ignore_ascii_case("md") || extension.eq_ignore_ascii_case("markdown")
})
{
return Ok(OutputFormat::Markdown);
}
if output
.and_then(|path| path.extension())
.map(|extension| extension.to_string_lossy())
.is_some_and(|extension| extension.eq_ignore_ascii_case("xml"))
{
return Ok(OutputFormat::Xml);
}
Ok(OutputFormat::Text)
}
pub(crate) struct InputTableMeta {
pub(crate) table_name: String,
pub(crate) source_path: String,
pub(crate) source_selector: Option<String>,
}
pub(crate) struct QueryMetadata {
pub(crate) row_count: usize,
pub(crate) columns: Vec<String>,
pub(crate) input_tables: Vec<InputTableMeta>,
pub(crate) execution_time_ms: u128,
}
pub(crate) fn write_query_metadata(meta: &QueryMetadata, format: Option<&MetadataFormat>) {
let rendered = match format {
Some(MetadataFormat::Json) => render_meta_json(meta),
_ => render_meta_text(meta),
};
let stderr = std::io::stderr();
let mut handle = stderr.lock();
let _ = writeln!(handle, "{rendered}");
}
fn render_meta_text(meta: &QueryMetadata) -> String {
let mut out = String::new();
out.push_str(&format!("rows: {}\n", meta.row_count));
out.push_str(&format!(
"columns: {}\n",
if meta.columns.is_empty() {
"-".to_owned()
} else {
meta.columns.join(", ")
}
));
let w_table = meta
.input_tables
.iter()
.map(|t| t.table_name.len())
.max()
.unwrap_or(0)
.max("table".len());
let w_source = meta
.input_tables
.iter()
.map(|t| t.source_path.len())
.max()
.unwrap_or(0)
.max("source".len());
for (i, t) in meta.input_tables.iter().enumerate() {
let label = if i == 0 { "tables: " } else { " " };
out.push_str(&format!(
"{}{:<w_table$} {:<w_source$} {}\n",
label,
t.table_name,
t.source_path,
t.source_selector.as_deref().unwrap_or("-"),
));
}
if meta.input_tables.is_empty() {
out.push_str("tables: -\n");
}
out.push_str(&format!("time: {}ms", meta.execution_time_ms));
out
}
fn render_meta_json(meta: &QueryMetadata) -> String {
let value = json!({
"row_count": meta.row_count,
"columns": meta.columns,
"input_tables": meta.input_tables.iter().map(|t| json!({
"table_name": t.table_name,
"source_path": t.source_path,
"source_selector": t.source_selector,
})).collect::<Vec<_>>(),
"execution_time_ms": meta.execution_time_ms,
});
serde_json::to_string_pretty(&value).unwrap_or_else(|_| "{}".to_owned())
}
#[cfg(test)]
mod tests {
use std::path::Path;
use crate::cli::OutputFormat;
use super::resolve_output_format;
#[test]
fn infers_output_format_from_extension() {
assert!(matches!(
resolve_output_format(Some(Path::new("result.xlsx")), None).expect("xlsx format"),
OutputFormat::Xlsx
));
assert!(matches!(
resolve_output_format(Some(Path::new("result.parquet")), None).expect("parquet format"),
OutputFormat::Parquet
));
assert!(matches!(
resolve_output_format(Some(Path::new("result.csv")), None).expect("csv format"),
OutputFormat::Csv
));
assert!(matches!(
resolve_output_format(Some(Path::new("result.jsonl")), None).expect("jsonl format"),
OutputFormat::Jsonl
));
assert!(matches!(
resolve_output_format(Some(Path::new("result.json")), None).expect("json format"),
OutputFormat::Json
));
assert!(matches!(
resolve_output_format(Some(Path::new("result.md")), None).expect("md format"),
OutputFormat::Markdown
));
assert!(matches!(
resolve_output_format(Some(Path::new("result.txt")), None).expect("text format"),
OutputFormat::Text
));
}
fn make_meta() -> super::QueryMetadata {
super::QueryMetadata {
row_count: 3,
columns: vec!["id".to_owned(), "name".to_owned()],
input_tables: vec![
super::InputTableMeta {
table_name: "table".to_owned(),
source_path: "./data.csv".to_owned(),
source_selector: None,
},
super::InputTableMeta {
table_name: "sales".to_owned(),
source_path: "./sales.xlsx".to_owned(),
source_selector: Some("Q1".to_owned()),
},
],
execution_time_ms: 42,
}
}
#[test]
fn meta_text_contains_row_count_and_columns() {
let meta = make_meta();
let rendered = super::render_meta_text(&meta);
assert!(rendered.contains("rows: 3"), "missing row count: {rendered}");
assert!(rendered.contains("id, name"), "missing columns: {rendered}");
assert!(rendered.contains("42ms"), "missing timing: {rendered}");
}
#[test]
fn meta_text_contains_input_tables() {
let meta = make_meta();
let rendered = super::render_meta_text(&meta);
assert!(rendered.contains("table"), "missing table name: {rendered}");
assert!(rendered.contains("./data.csv"), "missing source path: {rendered}");
assert!(rendered.contains("sales"), "missing table name: {rendered}");
assert!(rendered.contains("./sales.xlsx"), "missing source path: {rendered}");
assert!(rendered.contains("Q1"), "missing selector: {rendered}");
}
#[test]
fn meta_text_shows_dash_for_missing_selector() {
let meta = make_meta();
let rendered = super::render_meta_text(&meta);
assert!(rendered.contains(" -\n") || rendered.contains(" -"), "missing dash for no selector: {rendered}");
}
#[test]
fn meta_json_is_valid_structure() {
let meta = make_meta();
let rendered = super::render_meta_json(&meta);
let parsed: serde_json::Value =
serde_json::from_str(&rendered).expect("metadata JSON should be valid");
assert_eq!(parsed["row_count"], 3);
assert_eq!(parsed["execution_time_ms"], 42);
assert_eq!(parsed["columns"][0], "id");
assert_eq!(parsed["columns"][1], "name");
assert_eq!(parsed["input_tables"][0]["table_name"], "table");
assert_eq!(parsed["input_tables"][0]["source_selector"], serde_json::Value::Null);
assert_eq!(parsed["input_tables"][1]["table_name"], "sales");
assert_eq!(parsed["input_tables"][1]["source_selector"], "Q1");
}
}