tecton_compute/

execute.rs

//! One-shot dataset SQL execution with large-file optimizations.

use crate::error::ComputeError;
use crate::query::{batch_to_json_values, QueryResult};
use datafusion::dataframe::DataFrameWriteOptions;
use datafusion::prelude::*;
use futures::StreamExt;
use parquet::file::reader::{FileReader, SerializedFileReader};
use std::fs::File;
use std::path::{Path, PathBuf};
use std::time::Instant;
use tracing::{info, warn};

/// Default table name registered for the dataset (used in SQL as `FROM data`).
pub const DEFAULT_CSV_TABLE: &str = "data";

/// CSV files larger than this are auto-converted to Parquet when no cache exists.
pub const PARQUET_CONVERT_THRESHOLD_BYTES: u64 = 100 * 1024 * 1024;

/// Safety LIMIT appended to bare `SELECT *` queries.
pub const SELECT_STAR_DEFAULT_LIMIT: usize = 1000;

/// Hard cap when streaming result batches into memory for the UI/CLI.
const MAX_STREAM_ROWS: usize = 10_000;

/// Execute SQL against a CSV path (optionally via a Parquet cache) and return a rich result.
///
/// Optimizations:
/// - CSV > 100 MiB → convert/reuse sibling `.parquet` when present
/// - `SELECT *` without `LIMIT` → append `LIMIT 1000`
/// - Results are streamed via DataFusion `execute_stream` to limit RAM spikes
pub async fn execute_query(csv_path: &str, sql: &str) -> Result<QueryResult, ComputeError> {
    let started = Instant::now();
    let mut optimizations: Vec<String> = Vec::new();

    let (safe_sql, limit_note) = apply_select_star_safety(sql.trim())?;
    if let Some(note) = limit_note {
        info!(%note, "query safety applied");
        optimizations.push(note);
    }
    validate_sql(&safe_sql)?;

    let dataset_path = owned_absolute_path(csv_path)?;
    validate_dataset_path(Path::new(&dataset_path))?;

    let (exec_path, format_note) = ensure_query_source(Path::new(&dataset_path)).await?;
    if let Some(note) = format_note {
        info!(path = %exec_path.display(), %note, "format optimization applied");
        optimizations.push(note);
    }

    let table_name = DEFAULT_CSV_TABLE.to_owned();
    let ctx = SessionContext::new();
    register_table(&ctx, &table_name, &exec_path).await?;

    let df = ctx
        .sql(&safe_sql)
        .await
        .map_err(ComputeError::datafusion)?;

    let data = stream_json_rows(df, MAX_STREAM_ROWS).await?;

    let total_rows_affected = estimate_total_rows(
        &exec_path,
        &data,
        optimizations.iter().any(|o| o.contains("LIMIT")),
    )?;

    Ok(QueryResult {
        execution_time_ms: started.elapsed().as_millis() as u64,
        total_rows_affected,
        data,
        optimization_applied: merge_notes(optimizations),
    })
}

/// Convenience wrapper that returns pretty JSON (CLI / Tauri string bridge).
pub async fn execute_query_json(csv_path: &str, sql: &str) -> Result<String, ComputeError> {
    let result = execute_query(csv_path, sql).await?;
    Ok(result.to_pretty_json()?)
}

async fn ensure_query_source(path: &Path) -> Result<(PathBuf, Option<String>), ComputeError> {
    // Parquet file or directory of shards — use as-is.
    if path.is_dir() {
        return Ok((
            path.to_path_buf(),
            Some("Using Parquet dataset directory".to_string()),
        ));
    }

    let ext = path
        .extension()
        .and_then(|e| e.to_str())
        .unwrap_or("")
        .to_ascii_lowercase();

    if ext == "parquet" {
        return Ok((
            path.to_path_buf(),
            Some("Using Parquet dataset".to_string()),
        ));
    }

    // CSV path: prefer sibling .parquet cache, else convert when large.
    let parquet_path = path.with_extension("parquet");
    if parquet_path.exists() {
        return Ok((
            parquet_path,
            Some("Using existing Parquet cache".to_string()),
        ));
    }

    let meta = std::fs::metadata(path)?;
    if meta.len() <= PARQUET_CONVERT_THRESHOLD_BYTES {
        return Ok((path.to_path_buf(), None));
    }

    info!(
        csv = %path.display(),
        bytes = meta.len(),
        threshold = PARQUET_CONVERT_THRESHOLD_BYTES,
        "CSV exceeds threshold; converting to Parquet"
    );

    convert_csv_to_parquet(path, &parquet_path).await?;

    Ok((
        parquet_path,
        Some("Converted CSV to Parquet".to_string()),
    ))
}

async fn convert_csv_to_parquet(csv_path: &Path, parquet_path: &Path) -> Result<(), ComputeError> {
    let csv_str = path_to_utf8(csv_path)?;
    let tmp = parquet_path.with_extension("parquet.tmp");
    let tmp_str = path_to_utf8(&tmp)?;

    let ctx = SessionContext::new();
    let df = ctx
        .read_csv(&csv_str, CsvReadOptions::new())
        .await
        .map_err(ComputeError::datafusion)?;

    // Streamed write through DataFusion — avoids loading the full CSV as JSON.
    df.write_parquet(&tmp_str, DataFrameWriteOptions::new(), None)
        .await
        .map_err(ComputeError::datafusion)?;

    // DataFusion may write a directory for multi-file output; normalize to a single file path.
    finalize_parquet_output(&tmp, parquet_path)?;

    info!(
        from = %csv_path.display(),
        to = %parquet_path.display(),
        "Parquet conversion complete"
    );
    Ok(())
}

fn finalize_parquet_output(tmp: &Path, final_path: &Path) -> Result<(), ComputeError> {
    if tmp.is_file() {
        if final_path.exists() {
            std::fs::remove_file(final_path)?;
        }
        std::fs::rename(tmp, final_path)?;
        return Ok(());
    }

    if tmp.is_dir() {
        // Pick the first part file produced by the writer.
        let mut parts: Vec<PathBuf> = std::fs::read_dir(tmp)?
            .filter_map(|e| e.ok().map(|e| e.path()))
            .filter(|p| {
                p.extension()
                    .and_then(|e| e.to_str())
                    .is_some_and(|e| e.eq_ignore_ascii_case("parquet"))
            })
            .collect();
        parts.sort();
        let part = parts.into_iter().next().ok_or_else(|| {
            ComputeError::datafusion(format!(
                "parquet write produced no part files under {}",
                tmp.display()
            ))
        })?;
        if final_path.exists() {
            std::fs::remove_file(final_path)?;
        }
        std::fs::rename(&part, final_path)?;
        let _ = std::fs::remove_dir_all(tmp);
        return Ok(());
    }

    Err(ComputeError::datafusion(format!(
        "unexpected parquet write output: {}",
        tmp.display()
    )))
}

async fn register_table(
    ctx: &SessionContext,
    table_name: &str,
    path: &Path,
) -> Result<(), ComputeError> {
    let path_str = path_to_utf8(path)?;

    if path.is_dir() {
        // Multi-file Parquet dataset (e.g. bank_billion/part-*.parquet).
        ctx.register_parquet(table_name, &path_str, ParquetReadOptions::default())
            .await
            .map_err(ComputeError::datafusion)?;
        return Ok(());
    }

    let ext = path
        .extension()
        .and_then(|e| e.to_str())
        .unwrap_or("")
        .to_ascii_lowercase();

    match ext.as_str() {
        "parquet" => {
            ctx.register_parquet(table_name, &path_str, ParquetReadOptions::default())
                .await
                .map_err(ComputeError::datafusion)?;
        }
        "csv" => {
            let options = CsvReadOptions::new();
            ctx.register_csv(table_name, &path_str, options)
                .await
                .map_err(ComputeError::datafusion)?;
        }
        other => {
            return Err(ComputeError::invalid_input(format!(
                "unsupported dataset format '{other}'"
            )));
        }
    }
    Ok(())
}

async fn stream_json_rows(
    df: DataFrame,
    max_rows: usize,
) -> Result<Vec<serde_json::Value>, ComputeError> {
    let mut stream = df
        .execute_stream()
        .await
        .map_err(ComputeError::datafusion)?;

    let mut data = Vec::new();
    while let Some(batch) = stream.next().await {
        let batch = batch.map_err(ComputeError::datafusion)?;
        let mut rows = batch_to_json_values(&batch).map_err(|e| ComputeError::datafusion(e))?;
        for row in rows.drain(..) {
            if data.len() >= max_rows {
                warn!(max_rows, "result stream truncated at hard cap");
                return Ok(data);
            }
            data.push(row);
        }
    }
    Ok(data)
}

fn estimate_total_rows(
    exec_path: &Path,
    data: &[serde_json::Value],
    limit_applied: bool,
) -> Result<u64, ComputeError> {
    let returned = data.len() as u64;

    if !limit_applied {
        return Ok(returned);
    }

    if let Some(n) = parquet_dataset_num_rows(exec_path) {
        return Ok(n);
    }

    Ok(returned)
}

fn parquet_dataset_num_rows(path: &Path) -> Option<u64> {
    if path.is_file() {
        return parquet_file_num_rows(path);
    }
    if path.is_dir() {
        let mut total = 0u64;
        let mut any = false;
        for entry in std::fs::read_dir(path).ok()? {
            let entry = entry.ok()?;
            let p = entry.path();
            if p.extension()
                .and_then(|e| e.to_str())
                .is_some_and(|e| e.eq_ignore_ascii_case("parquet"))
            {
                total = total.saturating_add(parquet_file_num_rows(&p)?);
                any = true;
            }
        }
        return any.then_some(total);
    }
    None
}

fn parquet_file_num_rows(path: &Path) -> Option<u64> {
    let file = File::open(path).ok()?;
    let reader = SerializedFileReader::new(file).ok()?;
    let n = reader.metadata().file_metadata().num_rows();
    if n < 0 {
        None
    } else {
        Some(n as u64)
    }
}

/// If SQL is a `SELECT *` without LIMIT, append `LIMIT 1000`.
fn apply_select_star_safety(sql: &str) -> Result<(String, Option<String>), ComputeError> {
    if sql.is_empty() {
        return Err(ComputeError::invalid_input("SQL query must not be empty"));
    }

    let compact = collapse_ws_lower(sql);
    let is_select_star = compact.starts_with("select*")
        || compact.starts_with("select *")
        || compact.contains(" select *")
        || compact.contains("select*from")
        || compact.contains("select *from");

    if !is_select_star {
        return Ok((sql.to_owned(), None));
    }

    if has_limit_clause(&compact) {
        return Ok((sql.to_owned(), None));
    }

    let trimmed = sql.trim().trim_end_matches(';').trim_end();
    let rewritten = format!("{trimmed} LIMIT {SELECT_STAR_DEFAULT_LIMIT}");
    Ok((
        rewritten,
        Some(format!(
            "Appended LIMIT {SELECT_STAR_DEFAULT_LIMIT} to SELECT *"
        )),
    ))
}

fn has_limit_clause(sql_lower: &str) -> bool {
    // Token-based so identifiers like `limited_table` are not false positives.
    for token in sql_lower.split_whitespace().filter(|t| !t.is_empty()) {
        if token == "limit" {
            return true;
        }
        if let Some(rest) = token.strip_prefix("limit") {
            if !rest.is_empty() && rest.chars().all(|c| c.is_ascii_digit()) {
                return true;
            }
        }
    }
    false
}

fn collapse_ws_lower(sql: &str) -> String {
    let mut out = String::with_capacity(sql.len());
    let mut prev_space = false;
    for c in sql.chars() {
        if c.is_whitespace() {
            if !prev_space {
                out.push(' ');
                prev_space = true;
            }
        } else {
            out.push(c.to_ascii_lowercase());
            prev_space = false;
        }
    }
    out.trim().to_owned()
}

fn merge_notes(notes: Vec<String>) -> Option<String> {
    if notes.is_empty() {
        None
    } else {
        Some(notes.join("; "))
    }
}

fn owned_absolute_path(path: &str) -> Result<String, ComputeError> {
    if path.trim().is_empty() {
        return Err(ComputeError::invalid_input("csv path must not be empty"));
    }
    let path = Path::new(path);
    let absolute: PathBuf = if path.is_absolute() {
        path.to_path_buf()
    } else {
        std::env::current_dir()
            .map_err(ComputeError::Io)?
            .join(path)
    };
    Ok(absolute.to_string_lossy().into_owned())
}

fn path_to_utf8(path: &Path) -> Result<String, ComputeError> {
    path.to_str()
        .map(str::to_owned)
        .ok_or_else(|| ComputeError::invalid_input("path is not valid UTF-8"))
}

fn validate_dataset_path(path: &Path) -> Result<(), ComputeError> {
    if path.as_os_str().is_empty() {
        return Err(ComputeError::invalid_input("dataset path must not be empty"));
    }
    if !path.exists() {
        return Err(ComputeError::FileNotFound(path.display().to_string()));
    }

    if path.is_dir() {
        let has_parquet = std::fs::read_dir(path)
            .map_err(ComputeError::Io)?
            .filter_map(|e| e.ok())
            .any(|e| {
                e.path()
                    .extension()
                    .and_then(|ext| ext.to_str())
                    .is_some_and(|ext| ext.eq_ignore_ascii_case("parquet"))
            });
        if !has_parquet {
            return Err(ComputeError::invalid_input(format!(
                "directory has no .parquet shards: {}",
                path.display()
            )));
        }
        return Ok(());
    }

    if !path.is_file() {
        return Err(ComputeError::invalid_input(format!(
            "dataset path is not a file or directory: {}",
            path.display()
        )));
    }

    let ext = path
        .extension()
        .and_then(|e| e.to_str())
        .unwrap_or("")
        .to_ascii_lowercase();
    if ext != "csv" && ext != "parquet" {
        return Err(ComputeError::invalid_input(format!(
            "expected a .csv/.parquet file or parquet directory, got '{}'",
            path.display()
        )));
    }
    Ok(())
}

#[cfg(test)]
fn validate_csv_path(path: &Path) -> Result<(), ComputeError> {
    validate_dataset_path(path)
}

fn validate_sql(sql: &str) -> Result<(), ComputeError> {
    let trimmed = sql.trim();
    if trimmed.is_empty() {
        return Err(ComputeError::invalid_input("SQL query must not be empty"));
    }
    if trimmed.len() > 32_768 {
        return Err(ComputeError::invalid_input(
            "SQL query exceeds 32 KiB limit",
        ));
    }

    let normalized = trimmed.to_ascii_lowercase();
    let forbidden = [
        "insert ", "update ", "delete ", "drop ", "alter ", "create ", "truncate ", "grant ",
        "revoke ", "copy ", "attach ", "detach ", "pragma ",
    ];
    for token in forbidden {
        if normalized.starts_with(token.trim()) || normalized.contains(&format!(" {token}")) {
            return Err(ComputeError::invalid_input(format!(
                "mutating statement not allowed (found '{token}')"
            )));
        }
    }

    if !(normalized.starts_with("select")
        || normalized.starts_with("with")
        || normalized.starts_with("show")
        || normalized.starts_with("describe")
        || normalized.starts_with("explain"))
    {
        return Err(ComputeError::invalid_input(
            "only SELECT/WITH/SHOW/DESCRIBE/EXPLAIN statements are allowed",
        ));
    }

    Ok(())
}

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

    #[test]
    fn select_star_gets_limit() {
        let (sql, note) = apply_select_star_safety("SELECT * FROM data").unwrap();
        assert!(sql.to_ascii_uppercase().contains("LIMIT 1000"));
        assert!(note.unwrap().contains("LIMIT"));
    }

    #[test]
    fn select_star_with_limit_untouched() {
        let (sql, note) = apply_select_star_safety("SELECT * FROM data LIMIT 5").unwrap();
        assert_eq!(sql, "SELECT * FROM data LIMIT 5");
        assert!(note.is_none());
    }

    #[test]
    fn projected_select_untouched() {
        let (sql, note) = apply_select_star_safety("SELECT id FROM data").unwrap();
        assert_eq!(sql, "SELECT id FROM data");
        assert!(note.is_none());
    }

    #[tokio::test]
    async fn execute_query_returns_structured_result() {
        let dir = tempfile::tempdir().unwrap();
        let csv = dir.path().join("sample.csv");
        {
            let mut f = std::fs::File::create(&csv).unwrap();
            writeln!(f, "id,name").unwrap();
            writeln!(f, "1,alpha").unwrap();
            writeln!(f, "2,beta").unwrap();
        }

        let result = execute_query(
            csv.to_str().unwrap(),
            "SELECT id, name FROM data ORDER BY id LIMIT 1",
        )
        .await
        .unwrap();

        assert_eq!(result.data.len(), 1);
        assert!(result.execution_time_ms < 60_000);
        assert_eq!(result.total_rows_affected, 1);
        let row = result.data[0].as_object().unwrap();
        assert_eq!(row.get("name").and_then(|v| v.as_str()), Some("alpha"));
    }

    #[tokio::test]
    async fn select_star_safety_limits_rows() {
        let dir = tempfile::tempdir().unwrap();
        let csv = dir.path().join("many.csv");
        {
            let mut f = std::fs::File::create(&csv).unwrap();
            writeln!(f, "id").unwrap();
            for i in 0..1500 {
                writeln!(f, "{i}").unwrap();
            }
        }

        let result = execute_query(csv.to_str().unwrap(), "SELECT * FROM data")
            .await
            .unwrap();

        assert!(result.data.len() <= SELECT_STAR_DEFAULT_LIMIT);
        assert!(result
            .optimization_applied
            .as_deref()
            .unwrap_or("")
            .contains("LIMIT"));
    }

    #[test]
    fn rejects_non_csv() {
        assert!(validate_csv_path(Path::new("note.txt")).is_err());
    }
}