tecton-compute 0.1.1

//! Query request/response types and Arrow → JSON conversion.

use datafusion::arrow::array::Array;
use datafusion::arrow::record_batch::RecordBatch;
use datafusion::arrow::util::display::array_value_to_string;
use serde::{Deserialize, Serialize};
use serde_json::{json, Map, Value};
use tecton_core::{Result, TectonError};

/// Incoming SQL query from CLI or UI.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct QueryRequest {
    pub sql: String,
}

impl QueryRequest {
    pub fn new(sql: impl Into<String>) -> Self {
        Self { sql: sql.into() }
    }
}

/// Rich result returned by [`crate::execute_query`] for large-dataset workloads.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct QueryResult {
    /// Wall-clock execution time in milliseconds.
    pub execution_time_ms: u64,
    /// Rows matched / estimated (may exceed `data.len()` when a safety LIMIT applied).
    pub total_rows_affected: u64,
    /// Limited result set as JSON objects (one per row).
    pub data: Vec<Value>,
    /// Human-readable note about automatic optimizations (Parquet conversion, LIMIT, …).
    pub optimization_applied: Option<String>,
}

impl QueryResult {
    /// Serialize as pretty JSON for CLI / Tauri string bridges.
    pub fn to_pretty_json(&self) -> std::result::Result<String, serde_json::Error> {
        serde_json::to_string_pretty(self)
    }
}

/// Tabular result returned to callers (legacy engine path).
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct QueryResponse {
    pub columns: Vec<String>,
    pub rows: Vec<Vec<Option<String>>>,
    pub row_count: usize,
    pub truncated: bool,
    pub elapsed_ms: u64,
}

/// Intermediate conversion helper.
pub struct QueryResultSet {
    pub columns: Vec<String>,
    pub rows: Vec<Vec<Option<String>>>,
    pub row_count: usize,
    pub truncated: bool,
}

impl QueryResultSet {
    pub fn from_batches(batches: &[RecordBatch], max_rows: usize) -> Result<Self> {
        if batches.is_empty() {
            return Ok(Self {
                columns: vec![],
                rows: vec![],
                row_count: 0,
                truncated: false,
            });
        }

        let columns: Vec<String> = batches[0]
            .schema()
            .fields()
            .iter()
            .map(|f| f.name().clone())
            .collect();

        let mut rows = Vec::new();
        let mut truncated = false;

        'outer: for batch in batches {
            let n = batch.num_rows();
            for row_idx in 0..n {
                if rows.len() >= max_rows {
                    truncated = true;
                    break 'outer;
                }
                let mut row = Vec::with_capacity(batch.num_columns());
                for col_idx in 0..batch.num_columns() {
                    let array = batch.column(col_idx);
                    if array.is_null(row_idx) {
                        row.push(None);
                    } else {
                        let value = array_value_to_string(array.as_ref(), row_idx)
                            .map_err(|e| TectonError::compute(e.to_string()))?;
                        row.push(Some(value));
                    }
                }
                rows.push(row);
            }
        }

        let row_count = rows.len();
        Ok(Self {
            columns,
            rows,
            row_count,
            truncated,
        })
    }
}

/// Convert a record batch into JSON objects (column name → value).
pub fn batch_to_json_values(batch: &RecordBatch) -> std::result::Result<Vec<Value>, TectonError> {
    let columns: Vec<String> = batch
        .schema()
        .fields()
        .iter()
        .map(|f| f.name().clone())
        .collect();

    let mut rows = Vec::with_capacity(batch.num_rows());
    for row_idx in 0..batch.num_rows() {
        let mut map = Map::with_capacity(batch.num_columns());
        for (col_idx, name) in columns.iter().enumerate() {
            let array = batch.column(col_idx);
            if array.is_null(row_idx) {
                map.insert(name.clone(), Value::Null);
                continue;
            }
            let text = array_value_to_string(array.as_ref(), row_idx)
                .map_err(|e| TectonError::compute(e.to_string()))?;
            map.insert(name.clone(), json_primitive_from_text(&text));
        }
        rows.push(Value::Object(map));
    }
    Ok(rows)
}

fn json_primitive_from_text(text: &str) -> Value {
    if let Ok(n) = text.parse::<i64>() {
        return json!(n);
    }
    if let Ok(n) = text.parse::<f64>() {
        if n.is_finite() {
            return json!(n);
        }
    }
    if text.eq_ignore_ascii_case("true") {
        return Value::Bool(true);
    }
    if text.eq_ignore_ascii_case("false") {
        return Value::Bool(false);
    }
    Value::String(text.to_owned())
}