query_forge/

engine.rs

use std::{collections::HashSet, sync::OnceLock};

use anyhow::{Context, Result, anyhow, bail};
use regex::Regex;
use rusqlite::{
    Connection, params_from_iter,
    types::ValueRef,
};

use crate::input;
use crate::value::to_sql_value;
use crate::{
    ExtractionOptions, InputNormalizationOptions, QueryParam, QueryResult, QueryValue,
    TypeInferenceOptions, WorkbookInput,
};

const SQLITE_MAX_INSERT_VARIABLES: usize = 999;
const SQLITE_MAX_INSERT_ROWS: usize = 250;

pub fn run_query(
    workbook_path: &std::path::Path,
    sheet_name: Option<&str>,
    query: &str,
    has_headers: bool,
) -> Result<QueryResult> {
    run_query_with_params(workbook_path, sheet_name, query, &[], has_headers)
}

pub fn run_query_with_params(
    workbook_path: &std::path::Path,
    sheet_name: Option<&str>,
    query: &str,
    params: &[QueryParam],
    has_headers: bool,
) -> Result<QueryResult> {
    run_query_with_params_multi(&[workbook_path], sheet_name, query, params, has_headers)
}

pub fn run_query_with_params_multi(
    workbook_paths: &[&std::path::Path],
    sheet_name: Option<&str>,
    query: &str,
    params: &[QueryParam],
    has_headers: bool,
) -> Result<QueryResult> {
    let workbook_inputs = workbook_paths
        .iter()
        .map(|path| WorkbookInput {
            path,
            sheet_name,
            table_name: None,
            explicit_format: None,
        })
        .collect::<Vec<_>>();

    run_query_with_params_multi_inputs(&workbook_inputs, query, params, has_headers)
}

pub fn run_query_with_params_multi_inputs(
    workbook_inputs: &[WorkbookInput<'_>],
    query: &str,
    params: &[QueryParam],
    has_headers: bool,
) -> Result<QueryResult> {
    run_query_with_params_multi_inputs_and_options(
        workbook_inputs,
        query,
        params,
        &TypeInferenceOptions::default(),
        has_headers,
    )
}

pub fn run_query_with_params_multi_inputs_and_options(
    workbook_inputs: &[WorkbookInput<'_>],
    query: &str,
    params: &[QueryParam],
    inference_options: &TypeInferenceOptions,
    has_headers: bool,
) -> Result<QueryResult> {
    run_query_with_params_multi_inputs_and_options_and_normalization(
        workbook_inputs,
        query,
        params,
        inference_options,
        &InputNormalizationOptions::default(),
        &ExtractionOptions::default(),
        has_headers,
    )
}

pub fn run_query_with_params_multi_inputs_and_options_and_normalization(
    workbook_inputs: &[WorkbookInput<'_>],
    query: &str,
    params: &[QueryParam],
    inference_options: &TypeInferenceOptions,
    normalization_options: &InputNormalizationOptions,
    extraction_options: &ExtractionOptions,
    has_headers: bool,
) -> Result<QueryResult> {
    if workbook_inputs.is_empty() {
        bail!("at least one workbook input is required");
    }

    let connection =
        Connection::open_in_memory().context("failed to create in-memory SQLite database")?;
    let mut registered_sheet_views = HashSet::new();

    for (index, workbook_input) in workbook_inputs.iter().enumerate() {
        let mut sheet = input::load_input(
            workbook_input.path,
            workbook_input.sheet_name,
            inference_options,
            extraction_options,
            has_headers,
            workbook_input.explicit_format,
        )?;
        input::apply_input_normalization(&mut sheet, normalization_options);
        let table_name = workbook_input
            .table_name
            .map(str::to_owned)
            .unwrap_or_else(|| {
                if index == 0 {
                    "table".to_owned()
                } else {
                    format!("table{}", index + 1)
                }
            });

        register_sheet(
            &connection,
            &sheet,
            &table_name,
            &mut registered_sheet_views,
        )?;
    }

    execute_statements(&connection, query, params)
}

/// Splits a SQL string into individual statements delimited by `;`.
/// Handles single-quoted string literals (including `''` escapes) and
/// both `--` line comments and `/* */` block comments so that semicolons
/// inside them are not treated as statement separators.
fn split_sql_statements(sql: &str) -> Vec<String> {
    let mut statements = Vec::new();
    let mut current = String::new();
    let mut chars = sql.chars().peekable();

    while let Some(ch) = chars.next() {
        match ch {
            // Single-quoted string: consume until the closing quote,
            // treating '' as an escaped quote.
            '\'' => {
                current.push(ch);
                loop {
                    match chars.next() {
                        None => break,
                        Some('\'') => {
                            current.push('\'');
                            if chars.peek() == Some(&'\'') {
                                current.push(chars.next().unwrap());
                            } else {
                                break;
                            }
                        }
                        Some(c) => current.push(c),
                    }
                }
            }
            // -- line comment: consume until newline.
            '-' if chars.peek() == Some(&'-') => {
                current.push(ch);
                current.push(chars.next().unwrap()); // second '-'
                for c in chars.by_ref() {
                    current.push(c);
                    if c == '\n' {
                        break;
                    }
                }
            }
            // /* block comment: consume until */.
            '/' if chars.peek() == Some(&'*') => {
                current.push(ch);
                current.push(chars.next().unwrap()); // '*'
                let mut prev = '\0';
                for c in chars.by_ref() {
                    current.push(c);
                    if prev == '*' && c == '/' {
                        break;
                    }
                    prev = c;
                }
            }
            ';' => {
                let trimmed = current.trim().to_owned();
                if !trimmed.is_empty() {
                    statements.push(trimmed);
                }
                current.clear();
            }
            _ => current.push(ch),
        }
    }

    let trimmed = current.trim().to_owned();
    if !trimmed.is_empty() {
        statements.push(trimmed);
    }

    statements
}

/// Executes one or more `;`-separated SQL statements against `connection`.
/// All statements except the last are executed as DDL/DML via `execute_batch`.
/// The last statement is executed as a SELECT and its result is returned.
fn execute_statements(
    connection: &Connection,
    sql: &str,
    params: &[QueryParam],
) -> Result<QueryResult> {
    let statements = split_sql_statements(sql);

    match statements.as_slice() {
        [] => bail!("no SQL statements provided"),
        [single] => execute_query(connection, single, params),
        _ => {
            let last = statements.last().expect("non-empty slice always has a last element");
            let init = &statements[..statements.len() - 1];
            for stmt in init {
                connection
                    .execute_batch(stmt)
                    .with_context(|| format!("failed to execute statement: {stmt}"))?;
            }
            execute_query(connection, last, params)
        }
    }
}

fn register_sheet(
    connection: &Connection,
    sheet: &input::SheetData,
    table_name: &str,
    registered_views: &mut HashSet<String>,
) -> Result<()> {
    let columns = sheet
        .columns
        .iter()
        .map(|column| quote_identifier(column))
        .collect::<Vec<_>>()
        .join(", ");

    connection
        .execute(
            &format!("CREATE TABLE {} ({columns})", quote_identifier(table_name)),
            [],
        )
        .context("failed to create sheet table")?;

    let table1_key = normalize_view_key("table1");
    if table_name == "table" && !registered_views.contains(&table1_key) {
        connection
            .execute(
                &format!(
                    "CREATE VIEW {} AS SELECT * FROM {}",
                    quote_identifier("table1"),
                    quote_identifier("table")
                ),
                [],
            )
            .context("failed to register alias view table1 for first input")?;
        registered_views.insert(table1_key);
    }

    let sanitized_sheet_name = sanitize_table_name(&sheet.original_name);
    let sanitized_sheet_key = normalize_view_key(&sanitized_sheet_name);
    if sanitized_sheet_name != table_name && !registered_views.contains(&sanitized_sheet_key) {
        connection
            .execute(
                &format!(
                    "CREATE VIEW {} AS SELECT * FROM {}",
                    quote_identifier(&sanitized_sheet_name),
                    quote_identifier(table_name)
                ),
                [],
            )
            .with_context(|| {
                format!("failed to register view for sheet {}", sheet.original_name)
            })?;
        registered_views.insert(sanitized_sheet_key);
    }

    if sheet.rows.is_empty() {
        return Ok(());
    }

    insert_sheet_rows(connection, table_name, sheet)
}

fn insert_sheet_rows(connection: &Connection, table_name: &str, sheet: &input::SheetData) -> Result<()> {
    connection
        .execute_batch("BEGIN IMMEDIATE TRANSACTION")
        .context("failed to begin sheet insert transaction")?;

    let insert_result = (|| {
        let batch_size = calculate_insert_batch_size(sheet.columns.len());
        for rows_chunk in sheet.rows.chunks(batch_size) {
            let insert_sql =
                build_batched_insert_sql(table_name, sheet.columns.len(), rows_chunk.len());
            let mut values = Vec::with_capacity(rows_chunk.len() * sheet.columns.len());
            for row in rows_chunk {
                for index in 0..sheet.columns.len() {
                    let value = row.get(index).unwrap_or(&QueryValue::Null);
                    values.push(to_sql_value(value));
                }
            }

            connection
                .execute(&insert_sql, params_from_iter(values))
                .context("failed to insert sheet rows")?;
        }

        Ok(())
    })();

    match insert_result {
        Ok(()) => connection
            .execute_batch("COMMIT")
            .context("failed to commit sheet insert transaction"),
        Err(error) => {
            let _ = connection.execute_batch("ROLLBACK");
            Err(error)
        }
    }
}

fn calculate_insert_batch_size(column_count: usize) -> usize {
    let clamped_column_count = column_count.max(1);
    (SQLITE_MAX_INSERT_VARIABLES / clamped_column_count)
        .max(1)
        .min(SQLITE_MAX_INSERT_ROWS)
}

fn build_batched_insert_sql(table_name: &str, column_count: usize, row_count: usize) -> String {
    let row_placeholders = format!("({})", vec!["?"; column_count].join(", "));
    let values = vec![row_placeholders; row_count].join(", ");
    format!("INSERT INTO {} VALUES {values}", quote_identifier(table_name))
}

fn execute_query(
    connection: &Connection,
    query: &str,
    params: &[QueryParam],
) -> Result<QueryResult> {
    let normalized_query = normalize_query_for_reserved_identifiers(query);
    let mut statement = connection
        .prepare(&normalized_query)
        .map_err(|error| enrich_query_prepare_error(connection, query, error))?;

    if statement.column_count() == 0 {
        bail!("query must return rows");
    }

    let columns = statement
        .column_names()
        .into_iter()
        .map(str::to_owned)
        .collect::<Vec<_>>();

    bind_query_params(&mut statement, params)?;

    let column_count = statement.column_count();
    let mut rows = statement.raw_query();
    let mut result_rows = Vec::new();

    while let Some(row) = rows.next().context("failed to fetch query row")? {
        let mut values = Vec::with_capacity(column_count);

        for index in 0..column_count {
            values.push(match row.get_ref(index)? {
                ValueRef::Null => QueryValue::Null,
                ValueRef::Integer(value) => QueryValue::Integer(value),
                ValueRef::Real(value) => QueryValue::Real(value),
                ValueRef::Text(value) => {
                    QueryValue::Text(String::from_utf8_lossy(value).into_owned())
                }
                ValueRef::Blob(value) => QueryValue::Text(format_blob(value)),
            });
        }

        result_rows.push(values);
    }

    Ok(QueryResult {
        columns,
        rows: result_rows,
    })
}

fn enrich_query_prepare_error(
    connection: &Connection,
    query: &str,
    error: rusqlite::Error,
) -> anyhow::Error {
    let raw_error = error.to_string();

    if let Some(table) = raw_error.strip_prefix("no such table: ").map(str::trim) {
        let table = simplify_sqlite_missing_identifier(table);
        let available_tables = list_loaded_table_names(connection);
        let available_suffix = if available_tables.is_empty() {
            String::new()
        } else {
            format!(" Available tables/views: {}.", available_tables.join(", "))
        };

        return anyhow!(
            "query references unknown table '{table}'.{available_suffix} Check your table names or run `qf tables --input ...` to inspect available tables.\nQuery: {query}"
        );
    }

    if let Some(column) = raw_error.strip_prefix("no such column: ").map(str::trim) {
        let column = simplify_sqlite_missing_identifier(column);
        return anyhow!(
            "query references unknown column '{column}'. Check your column names or run `qf schema --input ...` to inspect available columns.\nQuery: {query}"
        );
    }

    anyhow!(error).context(format!("failed to prepare query: {query}"))
}

fn simplify_sqlite_missing_identifier(identifier: &str) -> &str {
    identifier
        .split_once(" in ")
        .map(|(name, _)| name)
        .unwrap_or(identifier)
        .trim()
}

fn list_loaded_table_names(connection: &Connection) -> Vec<String> {
    let mut statement = match connection.prepare(
        "SELECT name FROM sqlite_master WHERE type IN ('table', 'view') AND name NOT LIKE 'sqlite_%' ORDER BY name",
    ) {
        Ok(statement) => statement,
        Err(_) => return Vec::new(),
    };

    let names = match statement.query_map([], |row| row.get::<_, String>(0)) {
        Ok(names) => names,
        Err(_) => return Vec::new(),
    };

    names.filter_map(|name| name.ok()).collect()
}

fn normalize_query_for_reserved_identifiers(query: &str) -> String {
    static RESERVED_TABLE_RE: OnceLock<Regex> = OnceLock::new();

    let re = RESERVED_TABLE_RE.get_or_init(|| {
        Regex::new(r"(?i)\b(from|join|update|into)\s+table\b")
            .expect("reserved table regex should compile")
    });

    re.replace_all(query, |captures: &regex::Captures<'_>| {
        format!("{} \"table\"", &captures[1])
    })
    .into_owned()
}

fn bind_query_params(statement: &mut rusqlite::Statement<'_>, params: &[QueryParam]) -> Result<()> {
    for param in params {
        let mut bound = false;

        for prefix in [":", "@", "$"] {
            let parameter_name = format!("{prefix}{}", param.name);
            if let Some(index) = statement
                .parameter_index(&parameter_name)
                .with_context(|| format!("failed to inspect parameter {parameter_name}"))?
            {
                statement
                    .raw_bind_parameter(index, to_sql_value(&param.value))
                    .with_context(|| format!("failed to bind parameter {parameter_name}"))?;
                bound = true;
                break;
            }
        }

        if !bound {
            bail!("query does not contain parameter :{}", param.name);
        }
    }

    Ok(())
}

fn format_blob(value: &[u8]) -> String {
    use std::fmt::Write as _;

    let mut formatted = String::from("0x");
    for byte in value {
        let _ = write!(&mut formatted, "{byte:02x}");
    }

    formatted
}

fn quote_identifier(identifier: &str) -> String {
    format!("\"{}\"", identifier.replace('"', "\"\""))
}

fn sanitize_table_name(name: &str) -> String {
    let sanitized = name
        .chars()
        .map(|character| {
            if character.is_alphanumeric() || character == '_' {
                character
            } else {
                '_'
            }
        })
        .collect::<String>()
        .trim_matches('_')
        .to_string();

    if sanitized.is_empty() {
        "table_view".to_owned()
    } else {
        sanitized
    }
}

fn normalize_view_key(name: &str) -> String {
    name.to_ascii_lowercase()
}

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

    #[test]
    fn calculates_safe_insert_batch_size() {
        assert_eq!(calculate_insert_batch_size(1), SQLITE_MAX_INSERT_ROWS);
        assert_eq!(calculate_insert_batch_size(3), SQLITE_MAX_INSERT_ROWS);
        assert_eq!(calculate_insert_batch_size(400), 2);
        assert_eq!(calculate_insert_batch_size(2_000), 1);
    }

    #[test]
    fn builds_batched_insert_sql_for_multiple_rows() {
        assert_eq!(
            build_batched_insert_sql("table", 2, 3),
            "INSERT INTO \"table\" VALUES (?, ?), (?, ?), (?, ?)"
        );
    }
}