query-forge 0.9.0

Run SQL queries and dataset diffs on XLSX/XML/CSV/JSON/JSONL/Markdown/HTML/Feather/Parquet inputs and export results as text, CSV, JSONL, Markdown, XML, HTML, XLSX, Feather, or Parquet
Documentation
use std::time::Instant;

use anyhow::{Context, Result, anyhow, bail};
use clap::{CommandFactory, Parser};
use clap_complete::generate;
use clap_mangen::Man;
use serde_json::json;

mod cli;
mod clipboard;
mod conv;
mod diff;
mod input_spec;
mod inspect;
#[path = "cli_output/mod.rs"]
mod output;
mod pivot;
mod query;

use crate::cli::{Cli, Commands};

fn main() {
    match run() {
        Ok(code) => std::process::exit(code),
        Err(error) => {
            print_structured_error(&error);
            std::process::exit(2);
        }
    }
}

fn run() -> Result<i32> {
    let cli = Cli::try_parse().map_err(|error| anyhow!(error.to_string()))?;

    match cli.command {
        Commands::Query(cmd) => {
            let sql = query::load_query(&cmd)?;
            let sql = query::apply_with_ctes(&sql, &cmd.with_ctes);
            let params = query::parse_query_params(&cmd.params)?;
            let type_inference_options = query::resolve_type_inference_options(&cmd);
            let input_normalization_options = query::resolve_input_normalization_options(&cmd);
            let extraction_options = query::resolve_extraction_options(&cmd);
            let input_specs = input_spec::parse_input_specs(&cmd.input, None)?;
            let workbook_inputs = input_specs
                .iter()
                .map(|input| query_forge::WorkbookInput {
                    path: input.path.as_path(),
                    sheet_name: input.sheet_name.as_deref(),
                    table_name: input.table_name.as_deref(),
                    explicit_format: input.explicit_format.as_deref(),
                })
                .collect::<Vec<_>>();

            let start = Instant::now();
            let result =
                query_forge::run_query_with_params_multi_inputs_and_options_and_normalization(
                    &workbook_inputs,
                    &sql,
                    &params,
                    &type_inference_options,
                    &input_normalization_options,
                    &extraction_options,
                    !cmd.no_headers,
                )?;
            let elapsed = start.elapsed();

            let (output_path, output_format_hint) =
                output::parse_output_path_spec(cmd.output.as_deref());
            let input_tables = input_specs
                .iter()
                .enumerate()
                .map(|(i, spec)| output::InputTableMeta {
                    table_name: spec.table_name.clone().unwrap_or_else(|| {
                        if i == 0 {
                            "table".to_owned()
                        } else {
                            format!("table{}", i + 1)
                        }
                    }),
                    source_path: spec.path.display().to_string(),
                    source_selector: spec.sheet_name.clone(),
                })
                .collect::<Vec<_>>();

            let meta = output::QueryMetadata {
                row_count: result.rows.len(),
                columns: result.columns.clone(),
                input_tables,
                execution_time_ms: elapsed.as_millis(),
            };

            if cli.ai_mode {
                if output_path.is_some() {
                    output::write_query_result(
                        &result,
                        output_path.as_deref(),
                        cmd.format.clone(),
                        output_format_hint.as_deref(),
                    )?;
                }

                let data = output::query_result_to_json(&result);
                let metadata = 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,
                    "output_target": output_path.as_ref().map(|p| p.display().to_string()),
                    "output_format_hint": output_format_hint,
                });
                let envelope = output::ai_success_envelope("query", data, metadata, vec![]);
                println!("{}", serde_json::to_string_pretty(&envelope)?);
            } else {
                output::write_query_result(
                    &result,
                    output_path.as_deref(),
                    cmd.format,
                    output_format_hint.as_deref(),
                )?;

                if cmd.meta {
                    output::write_query_metadata(&meta, cmd.meta_format.as_ref());
                }
            }
        }
        Commands::Tables(cmd) => {
            if cli.ai_mode {
                let input_specs = input_spec::parse_input_specs(&cmd.input, cmd.sheet.as_deref())?;
                let workbook_inputs = input_specs
                    .iter()
                    .map(|s| query_forge::WorkbookInput {
                        path: s.path.as_path(),
                        sheet_name: s.sheet_name.as_deref(),
                        table_name: s.table_name.as_deref(),
                        explicit_format: s.explicit_format.as_deref(),
                    })
                    .collect::<Vec<_>>();
                let summaries = query_forge::load_table_summaries(&workbook_inputs, 0, false)?;
                let data = json!(summaries
                    .iter()
                    .map(|s| json!({
                        "table": s.table_name,
                        "source": s.source_path,
                        "selector": s.source_selector,
                        "rows": s.row_count,
                        "columns": s.columns.len(),
                    }))
                    .collect::<Vec<_>>());
                let warnings = summaries
                    .iter()
                    .flat_map(|s| s.warnings.clone())
                    .collect::<Vec<_>>();
                let metadata = json!({
                    "table_count": summaries.len(),
                    "format_requested": inspection_format_name(&cmd.format),
                });
                let envelope = output::ai_success_envelope("tables", data, metadata, warnings);
                println!("{}", serde_json::to_string_pretty(&envelope)?);
            } else {
                inspect::run_tables_command(cmd)?;
            }
        }
        Commands::Schema(cmd) => {
            if cli.ai_mode {
                let input_specs = input_spec::parse_input_specs(&cmd.input, cmd.sheet.as_deref())?;
                let workbook_inputs = input_specs
                    .iter()
                    .map(|s| query_forge::WorkbookInput {
                        path: s.path.as_path(),
                        sheet_name: s.sheet_name.as_deref(),
                        table_name: s.table_name.as_deref(),
                        explicit_format: s.explicit_format.as_deref(),
                    })
                    .collect::<Vec<_>>();
                let summaries = query_forge::load_table_summaries(&workbook_inputs, 0, false)?;
                let data = output::table_summaries_to_json(&summaries, false, false);
                let warnings = summaries
                    .iter()
                    .flat_map(|s| s.warnings.clone())
                    .collect::<Vec<_>>();
                let metadata = json!({
                    "table_count": summaries.len(),
                    "format_requested": inspection_format_name(&cmd.format),
                });
                let envelope = output::ai_success_envelope("schema", data, metadata, warnings);
                println!("{}", serde_json::to_string_pretty(&envelope)?);
            } else {
                inspect::run_schema_command(cmd)?;
            }
        }
        Commands::Inspect(cmd) => {
            if cli.ai_mode {
                let input_specs = input_spec::parse_input_specs(&cmd.input, cmd.sheet.as_deref())?;
                let workbook_inputs = input_specs
                    .iter()
                    .map(|s| query_forge::WorkbookInput {
                        path: s.path.as_path(),
                        sheet_name: s.sheet_name.as_deref(),
                        table_name: s.table_name.as_deref(),
                        explicit_format: s.explicit_format.as_deref(),
                    })
                    .collect::<Vec<_>>();
                let summaries = query_forge::load_table_summaries(&workbook_inputs, cmd.sample, cmd.stats)?;
                let data = output::table_summaries_to_json(&summaries, true, cmd.stats);
                let warnings = summaries
                    .iter()
                    .flat_map(|s| s.warnings.clone())
                    .collect::<Vec<_>>();
                let metadata = json!({
                    "table_count": summaries.len(),
                    "sample": cmd.sample,
                    "stats": cmd.stats,
                    "format_requested": inspection_format_name(&cmd.format),
                });
                let envelope = output::ai_success_envelope("inspect", data, metadata, warnings);
                println!("{}", serde_json::to_string_pretty(&envelope)?);
            } else {
                inspect::run_inspect_command(cmd)?;
            }
        }
        Commands::Diff(cmd) => {
            if cli.ai_mode {
                let execution = diff::execute_diff_command(&cmd)?;
                let has_differences = execution.has_differences;
                let data = output::query_result_to_json(&execution.result);
                let metadata = json!({
                    "has_differences": has_differences,
                    "exit_code": if has_differences { 1 } else { 0 },
                    "schema_only": cmd.schema_only,
                    "side_by_side": cmd.side_by_side,
                    "keys": cmd.key,
                    "ignored_columns": cmd.ignore_columns,
                });
                let envelope = output::ai_success_envelope("diff", data, metadata, vec![]);
                println!("{}", serde_json::to_string_pretty(&envelope)?);
                return Ok(if has_differences { 1 } else { 0 });
            }
            return diff::run_diff_command(cmd);
        }
        Commands::Pivot(cmd) => {
            if cli.ai_mode {
                let execution = pivot::execute_pivot_command(&cmd)?;
                let data = output::query_result_to_json(&execution.result);
                let metadata = json!({
                    "row_count": execution.result.rows.len(),
                    "column_count": execution.result.columns.len(),
                    "generated_sql": execution.pivot_sql,
                    "distinct_column_values": execution.distinct_column_values,
                    "show_sql_requested": cmd.show_sql,
                });
                let envelope = output::ai_success_envelope("pivot", data, metadata, vec![]);
                println!("{}", serde_json::to_string_pretty(&envelope)?);
                return Ok(0);
            }
            return pivot::run_pivot_command(cmd);
        }
        Commands::Conv(cmd) => {
            if cli.ai_mode {
                let result = conv::execute_conv_command(&cmd)?;
                let data = output::query_result_to_json(&result);
                let metadata = json!({
                    "row_count": result.rows.len(),
                    "column_count": result.columns.len(),
                    "input_count": cmd.input.len(),
                });
                let envelope = output::ai_success_envelope("conv", data, metadata, vec![]);
                println!("{}", serde_json::to_string_pretty(&envelope)?);
                return Ok(0);
            }
            conv::run_conv_command(cmd)?;
        }
        Commands::Completions(cmd) => {
            if cli.ai_mode {
                bail!("--ai-mode is currently supported for query, tables, schema, and inspect commands");
            }
            let mut app = Cli::command();
            let bin_name = app.get_name().to_owned();
            generate(cmd.shell, &mut app, bin_name, &mut std::io::stdout());
        }
        Commands::ManPage(cmd) => {
            if cli.ai_mode {
                bail!("--ai-mode is currently supported for query, tables, schema, and inspect commands");
            }
            let app = Cli::command();
            let man = Man::new(app);
            match cmd.output {
                Some(path) => {
                    let mut file = std::fs::File::create(&path).with_context(|| {
                        format!("Failed to create man page file: {}", path.display())
                    })?;
                    man.render(&mut file)?;
                }
                None => {
                    man.render(&mut std::io::stdout())?;
                }
            }
        }
        Commands::Capabilities(_) => {
            println!("{}", serde_json::to_string_pretty(&output::capabilities_json())?);
        }
    }

    Ok(0)
}

fn print_structured_error(error: &anyhow::Error) {
    let message = format!("{error:#}");
    let code = output::classify_error_code(&message);

    if std::env::args().any(|arg| arg == "--ai-mode") {
        let envelope = output::ai_error_envelope(code, &message);
        let rendered = serde_json::to_string_pretty(&envelope).unwrap_or_else(|_| {
            format!(
                "{{\"ok\":false,\"error\":{{\"code\":\"{}\",\"message\":\"{}\"}}}}",
                code,
                message.replace('"', "\\\"")
            )
        });
        eprintln!("{rendered}");
    } else {
        eprintln!("[{code}] {message}");
    }
}

fn inspection_format_name(format: &cli::InspectionFormat) -> &'static str {
    match format {
        cli::InspectionFormat::Text => "text",
        cli::InspectionFormat::Json => "json",
        cli::InspectionFormat::Markdown => "markdown",
    }
}