tecton_compute/
engine.rs

1//! DataFusion session management and dataset registration.
2
3use crate::query::{QueryRequest, QueryResponse, QueryResultSet};
4use datafusion::prelude::*;
5use std::path::{Path, PathBuf};
6use std::sync::Arc;
7use tecton_core::{ComputeConfig, Result, TectonError};
8use tracing::{debug, info, warn};
9
10/// SQL execution engine wrapping a DataFusion [`SessionContext`].
11pub struct ComputeEngine {
12    ctx: SessionContext,
13    data_dir: PathBuf,
14    max_rows: usize,
15}
16
17impl ComputeEngine {
18    /// Create a new engine and auto-register datasets found under `data_dir`.
19    pub async fn from_config(cfg: &ComputeConfig) -> Result<Self> {
20        std::fs::create_dir_all(&cfg.data_dir)?;
21
22        let ctx = SessionContext::new();
23        let mut engine = Self {
24            ctx,
25            data_dir: cfg.data_dir.clone(),
26            max_rows: cfg.max_rows,
27        };
28
29        engine.refresh_datasets().await?;
30        info!(dir = %engine.data_dir.display(), "compute engine ready");
31        Ok(engine)
32    }
33
34    pub fn data_dir(&self) -> &Path {
35        &self.data_dir
36    }
37
38    pub fn max_rows(&self) -> usize {
39        self.max_rows
40    }
41
42    /// Scan the data directory and register Parquet/CSV files as tables.
43    ///
44    /// Table names are derived from file stems (sanitized to SQL identifiers).
45    pub async fn refresh_datasets(&mut self) -> Result<Vec<String>> {
46        let mut registered = Vec::new();
47        let entries = std::fs::read_dir(&self.data_dir).map_err(|e| {
48            TectonError::compute(format!(
49                "cannot read data dir {}: {e}",
50                self.data_dir.display()
51            ))
52        })?;
53
54        for entry in entries {
55            let entry = entry.map_err(|e| TectonError::compute(e.to_string()))?;
56            let path = entry.path();
57            if !path.is_file() {
58                continue;
59            }
60            let ext = path
61                .extension()
62                .and_then(|e| e.to_str())
63                .unwrap_or("")
64                .to_ascii_lowercase();
65
66            let stem = path
67                .file_stem()
68                .and_then(|s| s.to_str())
69                .unwrap_or("table");
70            let table = sanitize_table_name(stem);
71            // Own the path before `.await` — `to_string_lossy().as_ref()` is a
72            // temporary `Cow` that must not be borrowed across suspension.
73            let path_owned = path.to_string_lossy().into_owned();
74
75            let result = match ext.as_str() {
76                "parquet" => {
77                    self.ctx
78                        .register_parquet(&table, &path_owned, ParquetReadOptions::default())
79                        .await
80                }
81                "csv" => {
82                    let csv_options = CsvReadOptions::new();
83                    self.ctx
84                        .register_csv(&table, &path_owned, csv_options)
85                        .await
86                }
87                _ => continue,
88            };
89
90            match result {
91                Ok(()) => {
92                    debug!(%table, path = %path.display(), "registered dataset");
93                    registered.push(table);
94                }
95                Err(e) => warn!(path = %path.display(), error = %e, "failed to register dataset"),
96            }
97        }
98
99        Ok(registered)
100    }
101
102    /// List currently registered table names.
103    pub async fn list_tables(&self) -> Result<Vec<String>> {
104        let names = self.ctx.catalog_names();
105        let mut tables = Vec::new();
106        for catalog_name in names {
107            let catalog = self
108                .ctx
109                .catalog(&catalog_name)
110                .ok_or_else(|| TectonError::compute("default catalog missing"))?;
111            for schema_name in catalog.schema_names() {
112                if let Some(schema) = catalog.schema(&schema_name) {
113                    tables.extend(schema.table_names());
114                }
115            }
116        }
117        tables.sort();
118        Ok(tables)
119    }
120
121    /// Execute a validated SQL query and return a JSON-friendly result set.
122    pub async fn sql(&self, request: QueryRequest) -> Result<QueryResponse> {
123        let sql = request.sql.trim();
124        validate_sql(sql)?;
125
126        let started = std::time::Instant::now();
127        let df = self
128            .ctx
129            .sql(sql)
130            .await
131            .map_err(|e| TectonError::compute(format!("SQL planning failed: {e}")))?;
132
133        let batches = df
134            .collect()
135            .await
136            .map_err(|e| TectonError::compute(format!("SQL execution failed: {e}")))?;
137
138        let result = QueryResultSet::from_batches(&batches, self.max_rows)?;
139        let elapsed_ms = started.elapsed().as_millis() as u64;
140
141        Ok(QueryResponse {
142            columns: result.columns,
143            rows: result.rows,
144            row_count: result.row_count,
145            truncated: result.truncated,
146            elapsed_ms,
147        })
148    }
149
150    /// Shared context for advanced callers.
151    pub fn context(&self) -> Arc<SessionContext> {
152        Arc::new(self.ctx.clone())
153    }
154}
155
156/// Allow only read-oriented statements for the MVP safety model.
157fn validate_sql(sql: &str) -> Result<()> {
158    if sql.is_empty() {
159        return Err(TectonError::invalid_input("SQL query must not be empty"));
160    }
161    if sql.len() > 32_768 {
162        return Err(TectonError::invalid_input("SQL query exceeds 32 KiB limit"));
163    }
164
165    let normalized = sql.trim_start().to_ascii_lowercase();
166    let forbidden = [
167        "insert ", "update ", "delete ", "drop ", "alter ", "create ", "truncate ",
168        "grant ", "revoke ", "copy ", "attach ", "detach ", "pragma ",
169    ];
170    for token in forbidden {
171        if normalized.starts_with(token.trim()) || normalized.contains(&format!(" {token}")) {
172            // Permit CREATE only? No — MVP is read-only.
173            return Err(TectonError::invalid_input(format!(
174                "mutating statement not allowed in MVP (found '{token}')"
175            )));
176        }
177    }
178
179    if !(normalized.starts_with("select")
180        || normalized.starts_with("with")
181        || normalized.starts_with("show")
182        || normalized.starts_with("describe")
183        || normalized.starts_with("explain"))
184    {
185        return Err(TectonError::invalid_input(
186            "only SELECT/WITH/SHOW/DESCRIBE/EXPLAIN statements are allowed",
187        ));
188    }
189
190    Ok(())
191}
192
193fn sanitize_table_name(raw: &str) -> String {
194    let mut name: String = raw
195        .chars()
196        .map(|c| if c.is_ascii_alphanumeric() { c } else { '_' })
197        .collect();
198    if name.is_empty() {
199        name = "table".into();
200    }
201    if name.chars().next().is_some_and(|c| c.is_ascii_digit()) {
202        name = format!("t_{name}");
203    }
204    name.to_ascii_lowercase()
205}
206
207#[cfg(test)]
208mod tests {
209    use super::*;
210
211    #[test]
212    fn sql_validation_blocks_mutations() {
213        assert!(validate_sql("SELECT 1").is_ok());
214        assert!(validate_sql("DROP TABLE x").is_err());
215        assert!(validate_sql("INSERT INTO t VALUES (1)").is_err());
216    }
217
218    #[test]
219    fn sanitize_table_names() {
220        assert_eq!(sanitize_table_name("My Data"), "my_data");
221        assert_eq!(sanitize_table_name("sales-2024"), "sales_2024");
222        assert_eq!(sanitize_table_name("9lives"), "t_9lives");
223    }
224}
tecton_compute/engine.rs

tecton_compute/
engine.rs
