tecton-compute 0.1.0

//! DataFusion session management and dataset registration.

use crate::query::{QueryRequest, QueryResponse, QueryResultSet};
use datafusion::prelude::*;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use tecton_core::{ComputeConfig, Result, TectonError};
use tracing::{debug, info, warn};

/// SQL execution engine wrapping a DataFusion [`SessionContext`].
pub struct ComputeEngine {
    ctx: SessionContext,
    data_dir: PathBuf,
    max_rows: usize,
}

impl ComputeEngine {
    /// Create a new engine and auto-register datasets found under `data_dir`.
    pub async fn from_config(cfg: &ComputeConfig) -> Result<Self> {
        std::fs::create_dir_all(&cfg.data_dir)?;

        let ctx = SessionContext::new();
        let mut engine = Self {
            ctx,
            data_dir: cfg.data_dir.clone(),
            max_rows: cfg.max_rows,
        };

        engine.refresh_datasets().await?;
        info!(dir = %engine.data_dir.display(), "compute engine ready");
        Ok(engine)
    }

    pub fn data_dir(&self) -> &Path {
        &self.data_dir
    }

    pub fn max_rows(&self) -> usize {
        self.max_rows
    }

    /// Scan the data directory and register Parquet/CSV files as tables.
    ///
    /// Table names are derived from file stems (sanitized to SQL identifiers).
    pub async fn refresh_datasets(&mut self) -> Result<Vec<String>> {
        let mut registered = Vec::new();
        let entries = std::fs::read_dir(&self.data_dir).map_err(|e| {
            TectonError::compute(format!(
                "cannot read data dir {}: {e}",
                self.data_dir.display()
            ))
        })?;

        for entry in entries {
            let entry = entry.map_err(|e| TectonError::compute(e.to_string()))?;
            let path = entry.path();
            if !path.is_file() {
                continue;
            }
            let ext = path
                .extension()
                .and_then(|e| e.to_str())
                .unwrap_or("")
                .to_ascii_lowercase();

            let stem = path
                .file_stem()
                .and_then(|s| s.to_str())
                .unwrap_or("table");
            let table = sanitize_table_name(stem);
            // Own the path before `.await` — `to_string_lossy().as_ref()` is a
            // temporary `Cow` that must not be borrowed across suspension.
            let path_owned = path.to_string_lossy().into_owned();

            let result = match ext.as_str() {
                "parquet" => {
                    self.ctx
                        .register_parquet(&table, &path_owned, ParquetReadOptions::default())
                        .await
                }
                "csv" => {
                    let csv_options = CsvReadOptions::new();
                    self.ctx
                        .register_csv(&table, &path_owned, csv_options)
                        .await
                }
                _ => continue,
            };

            match result {
                Ok(()) => {
                    debug!(%table, path = %path.display(), "registered dataset");
                    registered.push(table);
                }
                Err(e) => warn!(path = %path.display(), error = %e, "failed to register dataset"),
            }
        }

        Ok(registered)
    }

    /// List currently registered table names.
    pub async fn list_tables(&self) -> Result<Vec<String>> {
        let names = self.ctx.catalog_names();
        let mut tables = Vec::new();
        for catalog_name in names {
            let catalog = self
                .ctx
                .catalog(&catalog_name)
                .ok_or_else(|| TectonError::compute("default catalog missing"))?;
            for schema_name in catalog.schema_names() {
                if let Some(schema) = catalog.schema(&schema_name) {
                    tables.extend(schema.table_names());
                }
            }
        }
        tables.sort();
        Ok(tables)
    }

    /// Execute a validated SQL query and return a JSON-friendly result set.
    pub async fn sql(&self, request: QueryRequest) -> Result<QueryResponse> {
        let sql = request.sql.trim();
        validate_sql(sql)?;

        let started = std::time::Instant::now();
        let df = self
            .ctx
            .sql(sql)
            .await
            .map_err(|e| TectonError::compute(format!("SQL planning failed: {e}")))?;

        let batches = df
            .collect()
            .await
            .map_err(|e| TectonError::compute(format!("SQL execution failed: {e}")))?;

        let result = QueryResultSet::from_batches(&batches, self.max_rows)?;
        let elapsed_ms = started.elapsed().as_millis() as u64;

        Ok(QueryResponse {
            columns: result.columns,
            rows: result.rows,
            row_count: result.row_count,
            truncated: result.truncated,
            elapsed_ms,
        })
    }

    /// Shared context for advanced callers.
    pub fn context(&self) -> Arc<SessionContext> {
        Arc::new(self.ctx.clone())
    }
}

/// Allow only read-oriented statements for the MVP safety model.
fn validate_sql(sql: &str) -> Result<()> {
    if sql.is_empty() {
        return Err(TectonError::invalid_input("SQL query must not be empty"));
    }
    if sql.len() > 32_768 {
        return Err(TectonError::invalid_input("SQL query exceeds 32 KiB limit"));
    }

    let normalized = sql.trim_start().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}")) {
            // Permit CREATE only? No — MVP is read-only.
            return Err(TectonError::invalid_input(format!(
                "mutating statement not allowed in MVP (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(TectonError::invalid_input(
            "only SELECT/WITH/SHOW/DESCRIBE/EXPLAIN statements are allowed",
        ));
    }

    Ok(())
}

fn sanitize_table_name(raw: &str) -> String {
    let mut name: String = raw
        .chars()
        .map(|c| if c.is_ascii_alphanumeric() { c } else { '_' })
        .collect();
    if name.is_empty() {
        name = "table".into();
    }
    if name.chars().next().is_some_and(|c| c.is_ascii_digit()) {
        name = format!("t_{name}");
    }
    name.to_ascii_lowercase()
}

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

    #[test]
    fn sql_validation_blocks_mutations() {
        assert!(validate_sql("SELECT 1").is_ok());
        assert!(validate_sql("DROP TABLE x").is_err());
        assert!(validate_sql("INSERT INTO t VALUES (1)").is_err());
    }

    #[test]
    fn sanitize_table_names() {
        assert_eq!(sanitize_table_name("My Data"), "my_data");
        assert_eq!(sanitize_table_name("sales-2024"), "sales_2024");
        assert_eq!(sanitize_table_name("9lives"), "t_9lives");
    }
}