krishiv_sql/
introspection_sql.rs

1//! `DESCRIBE`, `SHOW COLUMNS`, and `EXPLAIN` SQL intercepts.
2
3use std::sync::Arc;
4
5use arrow::array::{BooleanArray, StringArray};
6use arrow::datatypes::{DataType, Field, Schema};
7use arrow::record_batch::RecordBatch;
8use datafusion::prelude::SessionContext;
9
10use crate::{SqlError, SqlResult};
11
12/// Parsed introspection statement.
13#[derive(Debug, Clone, PartialEq, Eq)]
14pub enum IntrospectionStatement {
15    Describe { table: String },
16    Explain { mode: ExplainSqlMode, query: String },
17}
18
19/// Detail level for `EXPLAIN` SQL.
20#[derive(Debug, Clone, Copy, PartialEq, Eq)]
21pub enum ExplainSqlMode {
22    Logical,
23    Physical,
24    Analyze,
25}
26
27/// Return `Some(stmt)` when `sql` is a Krishiv introspection statement.
28pub fn parse_introspection_statement(sql: &str) -> SqlResult<Option<IntrospectionStatement>> {
29    let trimmed = sql.trim();
30    if trimmed.is_empty() {
31        return Ok(None);
32    }
33    let upper = trimmed.to_ascii_uppercase();
34
35    if let Some(table) = parse_describe_target(trimmed, &upper) {
36        return Ok(Some(IntrospectionStatement::Describe { table }));
37    }
38
39    if let Some(table) = parse_show_columns_target(trimmed, &upper) {
40        return Ok(Some(IntrospectionStatement::Describe { table }));
41    }
42
43    if let Some((mode, query)) = parse_explain_target(trimmed, &upper) {
44        return Ok(Some(IntrospectionStatement::Explain { mode, query }));
45    }
46
47    Ok(None)
48}
49
50fn parse_describe_target(trimmed: &str, upper: &str) -> Option<String> {
51    const PREFIXES: &[&str] = &["DESCRIBE ", "DESC TABLE ", "DESC "];
52    for prefix in PREFIXES {
53        if upper.starts_with(prefix) {
54            let table = trimmed[prefix.len()..].trim().trim_end_matches(';').trim();
55            if !table.is_empty() {
56                return Some(table.to_string());
57            }
58        }
59    }
60    None
61}
62
63fn parse_show_columns_target(trimmed: &str, upper: &str) -> Option<String> {
64    if !upper.starts_with("SHOW COLUMNS ") {
65        return None;
66    }
67    let rest = trimmed["SHOW COLUMNS ".len()..].trim();
68    let upper_rest = rest.to_ascii_uppercase();
69    let table = if let Some(after) = upper_rest.strip_prefix("FROM ") {
70        rest[rest.len() - after.len()..].trim()
71    } else if let Some(after) = upper_rest.strip_prefix("IN ") {
72        rest[rest.len() - after.len()..].trim()
73    } else {
74        rest
75    };
76    let table = table.trim_end_matches(';').trim();
77    if table.is_empty() {
78        None
79    } else {
80        Some(table.to_string())
81    }
82}
83
84fn parse_explain_target(trimmed: &str, upper: &str) -> Option<(ExplainSqlMode, String)> {
85    if !upper.starts_with("EXPLAIN ") {
86        return None;
87    }
88    let rest = trimmed["EXPLAIN ".len()..].trim();
89    let upper_rest = rest.to_ascii_uppercase();
90    let (mode, query) = if let Some(query) = upper_rest.strip_prefix("LOGICAL ") {
91        (
92            ExplainSqlMode::Logical,
93            rest[rest.len() - query.len()..].trim(),
94        )
95    } else if let Some(query) = upper_rest.strip_prefix("PHYSICAL ") {
96        (
97            ExplainSqlMode::Physical,
98            rest[rest.len() - query.len()..].trim(),
99        )
100    } else if let Some(query) = upper_rest.strip_prefix("ANALYZE ") {
101        (
102            ExplainSqlMode::Analyze,
103            rest[rest.len() - query.len()..].trim(),
104        )
105    } else {
106        (ExplainSqlMode::Physical, rest)
107    };
108    let query = query.trim_end_matches(';').trim();
109    if query.is_empty() {
110        return None;
111    }
112    Some((mode, query.to_string()))
113}
114
115/// Build a `DESCRIBE` result batch for `table` using the DataFusion catalog.
116pub async fn describe_table(context: &SessionContext, table: &str) -> SqlResult<RecordBatch> {
117    let provider = context
118        .table_provider(table)
119        .await
120        .map_err(|error| SqlError::DataFusion {
121            message: format!("DESCRIBE: table '{table}' not found: {error}"),
122        })?;
123    let schema = provider.schema();
124    let col_name = Arc::new(StringArray::from(
125        schema
126            .fields()
127            .iter()
128            .map(|field| field.name().as_str())
129            .collect::<Vec<_>>(),
130    ));
131    let data_type = Arc::new(StringArray::from(
132        schema
133            .fields()
134            .iter()
135            .map(|field| field.data_type().to_string())
136            .collect::<Vec<_>>(),
137    ));
138    let nullable = Arc::new(BooleanArray::from(
139        schema
140            .fields()
141            .iter()
142            .map(|field| field.is_nullable())
143            .collect::<Vec<_>>(),
144    ));
145    let out_schema = Arc::new(Schema::new(vec![
146        Field::new("col_name", DataType::Utf8, false),
147        Field::new("data_type", DataType::Utf8, false),
148        Field::new("nullable", DataType::Boolean, false),
149    ]));
150    RecordBatch::try_new(out_schema, vec![col_name, data_type, nullable]).map_err(|error| {
151        SqlError::DataFusion {
152            message: format!("DESCRIBE: failed to build result batch: {error}"),
153        }
154    })
155}
156
157/// Render explain text for the embedded query.
158pub fn explain_query(query: &str, mode: ExplainSqlMode) -> SqlResult<String> {
159    match mode {
160        ExplainSqlMode::Logical => crate::explain_sql(query),
161        ExplainSqlMode::Physical => {
162            crate::explain_sql_optimized(query, &krishiv_plan::optimizer::Optimizer::default())
163        }
164        ExplainSqlMode::Analyze => {
165            let mut output = explain_query(query, ExplainSqlMode::Physical)?;
166            output.push_str(
167                "\n\nANALYZE: execute the query and call DataFrame::explain_with(Analyze) \
168                 for runtime statistics.",
169            );
170            Ok(output)
171        }
172    }
173}
174
175/// Build a single-row batch containing explain text.
176pub fn explain_result_batch(text: &str) -> SqlResult<RecordBatch> {
177    let schema = Arc::new(Schema::new(vec![Field::new("plan", DataType::Utf8, false)]));
178    let values = Arc::new(StringArray::from(vec![text]));
179    RecordBatch::try_new(schema, vec![values]).map_err(|error| SqlError::DataFusion {
180        message: format!("EXPLAIN: failed to build result batch: {error}"),
181    })
182}
183
184#[cfg(test)]
185mod tests {
186    use super::*;
187
188    #[test]
189    fn parse_describe_variants() {
190        let stmt = parse_introspection_statement("DESCRIBE orders")
191            .unwrap()
192            .unwrap();
193        assert_eq!(
194            stmt,
195            IntrospectionStatement::Describe {
196                table: "orders".into()
197            }
198        );
199        let stmt = parse_introspection_statement("DESC TABLE people")
200            .unwrap()
201            .unwrap();
202        assert!(matches!(stmt, IntrospectionStatement::Describe { .. }));
203    }
204
205    #[test]
206    fn parse_show_columns() {
207        let stmt = parse_introspection_statement("SHOW COLUMNS FROM events")
208            .unwrap()
209            .unwrap();
210        assert_eq!(
211            stmt,
212            IntrospectionStatement::Describe {
213                table: "events".into()
214            }
215        );
216    }
217
218    #[test]
219    fn parse_explain_modes() {
220        let stmt = parse_introspection_statement("EXPLAIN SELECT 1")
221            .unwrap()
222            .unwrap();
223        assert!(matches!(
224            stmt,
225            IntrospectionStatement::Explain {
226                mode: ExplainSqlMode::Physical,
227                ..
228            }
229        ));
230        let stmt = parse_introspection_statement("EXPLAIN LOGICAL SELECT 1")
231            .unwrap()
232            .unwrap();
233        assert!(matches!(
234            stmt,
235            IntrospectionStatement::Explain {
236                mode: ExplainSqlMode::Logical,
237                ..
238            }
239        ));
240    }
241
242    #[test]
243    fn explain_query_logical_renders_plan() {
244        let text = explain_query("SELECT 1 AS n", ExplainSqlMode::Logical).unwrap();
245        assert!(text.contains("SELECT") || text.contains("select") || !text.is_empty());
246    }
247}
krishiv_sql/introspection_sql.rs

krishiv_sql/
introspection_sql.rs
