query_forge/input/formats/

mod.rs

use std::path::Path;

use anyhow::Result;
use chrono::{NaiveDate, NaiveDateTime};

use crate::{
    ExtractionOptions, HeaderCase, InputNormalizationOptions, QueryValue, TypeInferenceOptions,
};

pub(crate) mod csv;
pub(crate) mod html;
pub(crate) mod json;
pub(crate) mod jsonl;
pub(crate) mod markdown;
pub(crate) mod parquet;
pub(crate) mod xlsx;
pub(crate) mod xml;

/// Raw tabular data loaded from a single input sheet or file.
#[derive(Debug)]
pub struct SheetData {
    /// Original name of the sheet, file, or tag as returned by the format loader.
    pub original_name: String,
    /// Column names in declaration order.
    pub columns: Vec<String>,
    /// Data rows; each row is a parallel `Vec` aligned with `columns`.
    pub rows: Vec<Vec<QueryValue>>,
}

pub fn load_input(
    input_path: &Path,
    requested_sheet: Option<&str>,
    inference_options: &TypeInferenceOptions,
    extraction_options: &ExtractionOptions,
    has_headers: bool,
) -> Result<SheetData> {
    if input_path
        .extension()
        .map(|extension| extension.to_string_lossy())
        .is_some_and(|extension| extension.eq_ignore_ascii_case("parquet"))
    {
        return parquet::load_parquet_sheet(input_path, requested_sheet, inference_options);
    }

    if input_path
        .extension()
        .map(|extension| extension.to_string_lossy())
        .is_some_and(|extension| extension.eq_ignore_ascii_case("xml"))
    {
        return xml::load_xml_sheet(
            input_path,
            requested_sheet,
            inference_options,
            extraction_options,
        );
    }

    if input_path
        .extension()
        .map(|extension| extension.to_string_lossy())
        .is_some_and(|extension| extension.eq_ignore_ascii_case("csv"))
    {
        return csv::load_csv_sheet(input_path, requested_sheet, inference_options, has_headers);
    }

    if input_path
        .extension()
        .map(|extension| extension.to_string_lossy())
        .is_some_and(|extension| extension.eq_ignore_ascii_case("jsonl"))
    {
        return jsonl::load_jsonl_sheet(input_path, requested_sheet, inference_options);
    }

    if input_path
        .extension()
        .map(|extension| extension.to_string_lossy())
        .is_some_and(|extension| extension.eq_ignore_ascii_case("json"))
    {
        return json::load_json_sheet(
            input_path,
            requested_sheet,
            inference_options,
            extraction_options,
        );
    }

    if input_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 markdown::load_markdown_sheet(
            input_path,
            requested_sheet,
            inference_options,
            has_headers,
        );
    }

    if input_path
        .extension()
        .map(|extension| extension.to_string_lossy())
        .is_some_and(|extension| {
            extension.eq_ignore_ascii_case("html") || extension.eq_ignore_ascii_case("htm")
        })
    {
        return html::load_html_sheet(
            input_path,
            requested_sheet,
            inference_options,
            has_headers,
        );
    }

    xlsx::load_xlsx_sheet(input_path, requested_sheet, inference_options, has_headers)
}

pub(crate) fn apply_input_normalization(sheet: &mut SheetData, options: &InputNormalizationOptions) {
    if options.trim {
        for column in &mut sheet.columns {
            *column = column.trim().to_owned();
        }

        for row in &mut sheet.rows {
            for value in row {
                if let QueryValue::Text(text) = value {
                    let trimmed = text.trim();
                    if trimmed.is_empty() {
                        *value = QueryValue::Null;
                    } else {
                        *text = trimmed.to_owned();
                    }
                }
            }
        }
    }

    if options.normalize_headers {
        sheet.columns = sheet
            .columns
            .iter()
            .enumerate()
            .map(|(index, header)| {
                normalize_header_value(header, index + 1, options.header_case.as_ref())
            })
            .collect();
    }

    if options.dedupe_headers {
        dedupe_headers(&mut sheet.columns);
    }

    if options.skip_empty_rows {
        sheet.rows.retain(|row| !is_row_empty(row));
    }
}

pub(crate) fn normalize_header_value(
    header: &str,
    index: usize,
    header_case: Option<&HeaderCase>,
) -> String {
    let normalized = header
        .trim()
        .chars()
        .map(|character| {
            if character.is_alphanumeric() || character == '_' {
                character
            } else {
                '_'
            }
        })
        .collect::<String>();
    let collapsed = collapse_underscores(&normalized)
        .trim_matches('_')
        .to_owned();

    if collapsed.is_empty() {
        return format!("column{index}");
    }

    let mut value = collapsed;

    if matches!(header_case, Some(HeaderCase::Snake)) {
        value = value.to_ascii_lowercase();
    } else if matches!(header_case, Some(HeaderCase::ScreamingSnake)) {
        value = value.to_ascii_uppercase();
    } else if matches!(header_case, Some(HeaderCase::Camel)) {
        value = to_camel_case(&value);
    } else if matches!(header_case, Some(HeaderCase::Pascal)) {
        value = to_pascal_case(&value);
    }

    value
}

/// Split `value` on `_` boundaries, lowercase every word, then join with the
/// first word unchanged and all subsequent words title-cased.  Produces camelCase.
fn to_camel_case(value: &str) -> String {
    let mut result = String::with_capacity(value.len());
    for (i, word) in value.split('_').enumerate() {
        if word.is_empty() {
            continue;
        }
        if i == 0 {
            result.push_str(&word.to_ascii_lowercase());
        } else {
            let mut chars = word.chars();
            if let Some(first) = chars.next() {
                result.push(first.to_ascii_uppercase());
                result.push_str(&chars.as_str().to_ascii_lowercase());
            }
        }
    }
    result
}

/// Split `value` on `_` boundaries, lowercase every word, then join with all
/// words title-cased.  Produces PascalCase.
fn to_pascal_case(value: &str) -> String {
    let mut result = String::with_capacity(value.len());
    for word in value.split('_') {
        if word.is_empty() {
            continue;
        }
        let mut chars = word.chars();
        if let Some(first) = chars.next() {
            result.push(first.to_ascii_uppercase());
            result.push_str(&chars.as_str().to_ascii_lowercase());
        }
    }
    result
}

fn collapse_underscores(value: &str) -> String {
    let mut output = String::with_capacity(value.len());
    let mut previous_was_underscore = false;

    for character in value.chars() {
        if character == '_' {
            if !previous_was_underscore {
                output.push(character);
            }
            previous_was_underscore = true;
        } else {
            output.push(character);
            previous_was_underscore = false;
        }
    }

    output
}

pub(crate) fn dedupe_headers(headers: &mut [String]) {
    let mut seen = std::collections::HashMap::new();
    let mut assigned = std::collections::HashSet::new();

    for header in headers {
        let base = header.clone();
        let count = seen
            .entry(base.clone())
            .and_modify(|count| *count += 1)
            .or_insert(1usize);

        let mut candidate = if *count == 1 {
            base.clone()
        } else {
            format!("{base}_{count}")
        };

        while assigned.contains(&candidate) {
            *count += 1;
            candidate = format!("{base}_{count}");
        }

        assigned.insert(candidate.clone());
        *header = candidate;
    }
}

fn is_row_empty(row: &[QueryValue]) -> bool {
    row.iter().all(|value| match value {
        QueryValue::Null => true,
        QueryValue::Text(text) => text.trim().is_empty(),
        QueryValue::Integer(_) | QueryValue::Real(_) => false,
    })
}

pub(super) fn normalize_text_headers(headers: &[String]) -> Vec<String> {
    let mut seen = std::collections::HashMap::new();

    headers
        .iter()
        .enumerate()
        .map(|(index, header)| {
            let base = if header.trim().is_empty() {
                format!("column{}", index + 1)
            } else {
                header.clone()
            };

            let count = seen
                .entry(base.clone())
                .and_modify(|count| *count += 1)
                .or_insert(1usize);

            if *count == 1 {
                base
            } else {
                format!("{base}_{count}")
            }
        })
        .collect()
}

pub(super) fn rows_maps_to_sheet_data(
    rows_maps: Vec<Vec<(String, QueryValue)>>,
    original_name: String,
) -> SheetData {
    let mut columns = Vec::<String>::new();
    for row in &rows_maps {
        for (column, _) in row {
            if !columns.iter().any(|existing| existing == column) {
                columns.push(column.clone());
            }
        }
    }

    let rows = rows_maps
        .into_iter()
        .map(|row| {
            columns
                .iter()
                .map(|column| {
                    row.iter()
                        .find(|(key, _)| key == column)
                        .map(|(_, value)| value.clone())
                        .unwrap_or(QueryValue::Null)
                })
                .collect::<Vec<_>>()
        })
        .collect::<Vec<_>>();

    SheetData {
        original_name,
        columns,
        rows,
    }
}

pub(super) fn parse_scalar_value(raw: &str, inference_options: &TypeInferenceOptions) -> QueryValue {
    let trimmed = raw.trim();
    if matches_token(trimmed, &inference_options.null_values) {
        return QueryValue::Null;
    }

    if !inference_options.infer_types {
        return QueryValue::Text(raw.to_owned());
    }

    if matches_token(trimmed, &inference_options.true_values) {
        return QueryValue::Integer(1);
    }

    if matches_token(trimmed, &inference_options.false_values) {
        return QueryValue::Integer(0);
    }

    if let Ok(value) = trimmed.parse::<i64>() {
        return QueryValue::Integer(value);
    }

    if let Some(value) = parse_decimal_value(trimmed, inference_options.decimal_comma) {
        return QueryValue::Real(value);
    }

    if let Some(value) = parse_date_value(trimmed, inference_options.date_format.as_deref()) {
        return QueryValue::Text(value);
    }

    QueryValue::Text(raw.to_owned())
}

fn matches_token(raw: &str, candidates: &[String]) -> bool {
    candidates
        .iter()
        .any(|candidate| raw.eq_ignore_ascii_case(candidate.trim()))
}

fn parse_decimal_value(raw: &str, decimal_comma: bool) -> Option<f64> {
    if let Ok(value) = raw.parse::<f64>() {
        return Some(value);
    }

    if !decimal_comma {
        return None;
    }

    let compact = raw.replace(' ', "");
    if !compact.contains(',') {
        return None;
    }

    let normalized = if compact.contains('.') {
        compact.replace('.', "").replace(',', ".")
    } else {
        compact.replace(',', ".")
    };

    normalized.parse::<f64>().ok()
}

fn parse_date_value(raw: &str, date_format: Option<&str>) -> Option<String> {
    let format = date_format?;

    if let Ok(value) = NaiveDate::parse_from_str(raw, format) {
        return Some(value.format("%Y-%m-%d").to_string());
    }

    if let Ok(value) = NaiveDateTime::parse_from_str(raw, format) {
        return Some(value.format("%Y-%m-%d %H:%M:%S").to_string());
    }

    None
}

pub(super) fn apply_inference_overrides(
    value: QueryValue,
    inference_options: &TypeInferenceOptions,
) -> QueryValue {
    if inference_options.infer_types {
        return value;
    }

    match value {
        QueryValue::Null => QueryValue::Null,
        QueryValue::Integer(number) => QueryValue::Text(number.to_string()),
        QueryValue::Real(number) => QueryValue::Text(number.to_string()),
        QueryValue::Text(text) => QueryValue::Text(text),
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn normalizes_special_character_headers_to_column_name() {
        assert_eq!(
            normalize_header_value(" !!! ", 3, Some(&HeaderCase::Snake)),
            "column3"
        );
    }

    #[test]
    fn header_case_snake_lowercases_all() {
        assert_eq!(
            normalize_header_value("First_Name", 1, Some(&HeaderCase::Snake)),
            "first_name"
        );
        assert_eq!(
            normalize_header_value("PRODUCT ID", 1, Some(&HeaderCase::Snake)),
            "product_id"
        );
    }

    #[test]
    fn header_case_camel_produces_camel_case() {
        assert_eq!(
            normalize_header_value("first_name", 1, Some(&HeaderCase::Camel)),
            "firstName"
        );
        assert_eq!(
            normalize_header_value("PRODUCT ID", 1, Some(&HeaderCase::Camel)),
            "productId"
        );
        assert_eq!(
            normalize_header_value("product", 1, Some(&HeaderCase::Camel)),
            "product"
        );
    }

    #[test]
    fn header_case_pascal_produces_pascal_case() {
        assert_eq!(
            normalize_header_value("first_name", 1, Some(&HeaderCase::Pascal)),
            "FirstName"
        );
        assert_eq!(
            normalize_header_value("PRODUCT ID", 1, Some(&HeaderCase::Pascal)),
            "ProductId"
        );
        assert_eq!(
            normalize_header_value("product", 1, Some(&HeaderCase::Pascal)),
            "Product"
        );
    }

    #[test]
    fn header_case_screaming_snake_uppercases_all() {
        assert_eq!(
            normalize_header_value("first_name", 1, Some(&HeaderCase::ScreamingSnake)),
            "FIRST_NAME"
        );
        assert_eq!(
            normalize_header_value("PRODUCT ID", 1, Some(&HeaderCase::ScreamingSnake)),
            "PRODUCT_ID"
        );
    }

    #[test]
    fn header_case_fallback_to_column_index_for_empty() {
        assert_eq!(
            normalize_header_value(" !!! ", 3, Some(&HeaderCase::Camel)),
            "column3"
        );
        assert_eq!(
            normalize_header_value(" !!! ", 5, Some(&HeaderCase::Pascal)),
            "column5"
        );
        assert_eq!(
            normalize_header_value(" !!! ", 2, Some(&HeaderCase::ScreamingSnake)),
            "column2"
        );
    }

    #[test]
    fn dedupe_headers_handles_generated_name_collisions() {
        let mut headers = vec![
            "name".to_owned(),
            "name".to_owned(),
            "name_2".to_owned(),
            "name".to_owned(),
        ];

        dedupe_headers(&mut headers);

        assert_eq!(headers, vec!["name", "name_2", "name_2_2", "name_3"]);
    }
}