krishiv_sql/
statement_completion.rs

1#![forbid(unsafe_code)]
2//! Spark-reference session/navigation statements that DataFusion's planner does
3//! not handle natively (Phase 60 statement completion).
4//!
5//! DataFusion 54 already plans `SET`/`RESET`, `SHOW TABLES|COLUMNS|FUNCTIONS`,
6//! `SHOW CREATE`, `TRUNCATE TABLE`, `DESCRIBE <table>`, and `EXPLAIN`. This
7//! module fills the two navigation gaps a SQL/BI client expects:
8//!
9//! - **`USE [CATALOG|SCHEMA|DATABASE|NAMESPACE] <name>`** — set the session's
10//!   current catalog/schema so subsequent unqualified table references resolve
11//!   there (Spark `USE`). `USE a.b` sets catalog `a`, schema `b`.
12//! - **`SHOW DATABASES` / `SHOW SCHEMAS`** — list schemas; rewritten to an
13//!   `information_schema.schemata` query returning a Spark-style `namespace`
14//!   column.
15//!
16//! `CACHE/UNCACHE TABLE` (session materialization), `SHOW PARTITIONS` (Iceberg
17//! metadata), and `DESCRIBE FUNCTION|DATABASE|QUERY` remain the itemized
18//! statement shortfall in the matrix.
19
20use datafusion::prelude::SessionContext;
21
22/// The parsed target of a `USE` statement.
23#[derive(Debug, Clone, PartialEq, Eq)]
24pub struct UseTarget {
25    /// New default catalog, if the statement set one.
26    pub catalog: Option<String>,
27    /// New default schema, if the statement set one.
28    pub schema: Option<String>,
29}
30
31/// Whether `ident` is wrapped in a matching pair of backticks or double quotes.
32fn is_quoted(ident: &str) -> bool {
33    let t = ident.trim();
34    ['`', '"']
35        .iter()
36        .any(|q| t.len() >= 2 && t.starts_with(*q) && t.ends_with(*q))
37}
38
39fn unquote(ident: &str) -> String {
40    let t = ident.trim();
41    for q in ['`', '"'] {
42        if let Some(inner) = t.strip_prefix(q).and_then(|s| s.strip_suffix(q)) {
43            return inner.to_string();
44        }
45    }
46    t.to_string()
47}
48
49/// Parse a `USE …` statement into its catalog/schema target, or `None` if the
50/// query is not a `USE` statement.
51///
52/// Forms (Spark/Databricks):
53/// - `USE db` / `USE SCHEMA db` / `USE DATABASE db` / `USE NAMESPACE ns` — schema
54/// - `USE CATALOG cat` — catalog
55/// - `USE cat.schema` — both
56pub fn parse_use(query: &str) -> Option<UseTarget> {
57    let q = query.trim().trim_end_matches(';').trim();
58    // First token must be USE (case-insensitive); the remainder is the target.
59    let (head, remainder) = q.split_once(char::is_whitespace)?;
60    if !head.eq_ignore_ascii_case("USE") {
61        return None;
62    }
63    let remainder = remainder.trim();
64    if remainder.is_empty() {
65        return None;
66    }
67    // Optional leading CATALOG / SCHEMA / DATABASE / NAMESPACE keyword; whatever
68    // follows it (which may be a back-quoted identifier containing spaces) is the
69    // whole name.
70    let (is_catalog, name_part) = match remainder.split_once(char::is_whitespace) {
71        Some((w, r)) if w.eq_ignore_ascii_case("CATALOG") => (true, r.trim()),
72        Some((w, r))
73            if w.eq_ignore_ascii_case("SCHEMA")
74                || w.eq_ignore_ascii_case("DATABASE")
75                || w.eq_ignore_ascii_case("NAMESPACE") =>
76        {
77            (false, r.trim())
78        }
79        _ => (false, remainder),
80    };
81    // A quoted identifier is one name even when it contains a dot: ``USE
82    // `my.schema` `` asks for the schema literally called `my.schema`, not for
83    // catalog `my` / schema `schema`. Only an *unquoted* name is qualified.
84    let was_quoted = is_quoted(name_part);
85    let name = unquote(name_part);
86    if name.is_empty() {
87        return None;
88    }
89    if is_catalog {
90        return Some(UseTarget {
91            catalog: Some(name),
92            schema: None,
93        });
94    }
95    // Schema form: allow a qualified `catalog.schema`.
96    if let Some((cat, sch)) = name.split_once('.').filter(|_| !was_quoted) {
97        return Some(UseTarget {
98            catalog: Some(unquote(cat)),
99            schema: Some(unquote(sch)),
100        });
101    }
102    Some(UseTarget {
103        catalog: None,
104        schema: Some(name),
105    })
106}
107
108/// Apply a `USE` target to the session by mutating the default catalog/schema.
109/// Returns `Some(Ok(()))` when the query was a `USE` statement (handled), or
110/// `None` when it was not.
111pub fn apply_use(ctx: &SessionContext, query: &str) -> Option<Result<(), String>> {
112    let target = parse_use(query)?;
113    let state_ref = ctx.state_ref();
114    let mut state = state_ref.write();
115    let opts = state.config_mut().options_mut();
116    if let Some(catalog) = target.catalog {
117        opts.catalog.default_catalog = catalog;
118    }
119    if let Some(schema) = target.schema {
120        opts.catalog.default_schema = schema;
121    }
122    Some(Ok(()))
123}
124
125/// If `query` is `SHOW DATABASES`/`SHOW SCHEMAS` (optionally `LIKE 'pat'`),
126/// return an equivalent `information_schema.schemata` SELECT; otherwise `None`.
127/// The result column is `namespace`, matching Spark's `SHOW DATABASES`.
128pub fn rewrite_show_databases(query: &str) -> Option<String> {
129    let q = query.trim().trim_end_matches(';').trim();
130    let upper = q.to_ascii_uppercase();
131    let is_show = upper.starts_with("SHOW DATABASES") || upper.starts_with("SHOW SCHEMAS");
132    if !is_show {
133        return None;
134    }
135
136    // Split the tail into the optional `{FROM|IN} <catalog>` qualifier and the
137    // optional `LIKE 'pattern'` filter.
138    //
139    // The catalog qualifier used to be ignored outright: `SHOW SCHEMAS IN prod`
140    // listed the schemas of *every* catalog, silently answering a broader
141    // question than the one asked. Spark's grammar is
142    // `SHOW SCHEMAS [ { FROM | IN } catalog ] [ LIKE pattern ]`.
143    let (before_like, like_pattern) = match upper.find(" LIKE ") {
144        Some(idx) => (
145            q.get(..idx).unwrap_or_default(),
146            q.get(idx + " LIKE ".len()..).map(str::trim),
147        ),
148        None => (q, None),
149    };
150
151    let mut predicates: Vec<String> = Vec::new();
152    if let Some(catalog) = catalog_qualifier(before_like) {
153        // Single-quote escaping so a catalog name containing `'` cannot end the
154        // literal early.
155        predicates.push(format!(
156            "catalog_name = '{}'",
157            catalog.replace('\'', "''")
158        ));
159    }
160    if let Some(pattern) = like_pattern.filter(|p| !p.is_empty()) {
161        predicates.push(format!("schema_name LIKE {pattern}"));
162    }
163
164    let where_clause = if predicates.is_empty() {
165        String::new()
166    } else {
167        format!(" WHERE {}", predicates.join(" AND "))
168    };
169    Some(format!(
170        "SELECT schema_name AS namespace FROM information_schema.schemata{where_clause} \
171         ORDER BY namespace"
172    ))
173}
174
175/// The catalog named by a trailing `{FROM|IN} <catalog>`, if present.
176///
177/// `head` is the statement up to any `LIKE`, e.g. `SHOW SCHEMAS IN prod`.
178fn catalog_qualifier(head: &str) -> Option<String> {
179    let mut tokens = head.split_whitespace().collect::<Vec<_>>();
180    let catalog = tokens.pop()?;
181    let keyword = tokens.pop()?;
182    if keyword.eq_ignore_ascii_case("FROM") || keyword.eq_ignore_ascii_case("IN") {
183        let name = unquote(catalog);
184        (!name.is_empty()).then_some(name)
185    } else {
186        None
187    }
188}
189
190#[cfg(test)]
191mod tests {
192    use super::*;
193
194    #[test]
195    fn parse_use_forms() {
196        assert_eq!(
197            parse_use("USE analytics"),
198            Some(UseTarget {
199                catalog: None,
200                schema: Some("analytics".into())
201            })
202        );
203        assert_eq!(
204            parse_use("USE SCHEMA sales"),
205            Some(UseTarget {
206                catalog: None,
207                schema: Some("sales".into())
208            })
209        );
210        assert_eq!(
211            parse_use("USE DATABASE sales;"),
212            Some(UseTarget {
213                catalog: None,
214                schema: Some("sales".into())
215            })
216        );
217        assert_eq!(
218            parse_use("USE CATALOG lakehouse"),
219            Some(UseTarget {
220                catalog: Some("lakehouse".into()),
221                schema: None
222            })
223        );
224        assert_eq!(
225            parse_use("USE lake.sales"),
226            Some(UseTarget {
227                catalog: Some("lake".into()),
228                schema: Some("sales".into())
229            })
230        );
231        assert_eq!(
232            parse_use("USE `my schema`").unwrap().schema.as_deref(),
233            Some("my schema")
234        );
235        assert_eq!(parse_use("SELECT 1"), None);
236    }
237
238    #[test]
239    fn show_databases_rewrite() {
240        assert!(
241            rewrite_show_databases("SHOW DATABASES")
242                .unwrap()
243                .contains("information_schema.schemata")
244        );
245        assert!(
246            rewrite_show_databases("SHOW SCHEMAS")
247                .unwrap()
248                .contains("AS namespace")
249        );
250        let with_like = rewrite_show_databases("SHOW DATABASES LIKE 'sal%'").unwrap();
251        assert!(with_like.contains("LIKE 'sal%'"));
252        assert_eq!(rewrite_show_databases("SHOW TABLES"), None);
253    }
254
255    /// `SHOW SCHEMAS { FROM | IN } <catalog>` must scope to that catalog.
256    ///
257    /// The qualifier used to be ignored outright, so the statement listed the
258    /// schemas of *every* catalog — a silently broader answer than the one
259    /// asked for.
260    #[test]
261    fn show_databases_honours_the_catalog_qualifier() {
262        for sql in ["SHOW SCHEMAS IN prod", "SHOW DATABASES FROM prod"] {
263            let rewritten = rewrite_show_databases(sql).expect("recognised");
264            assert!(
265                rewritten.contains("catalog_name = 'prod'"),
266                "{sql} must filter by catalog: {rewritten}"
267            );
268        }
269    }
270
271    /// Both filters compose rather than one displacing the other.
272    #[test]
273    fn catalog_qualifier_and_like_compose() {
274        let rewritten =
275            rewrite_show_databases("SHOW SCHEMAS IN prod LIKE 'sal%'").expect("recognised");
276        assert!(rewritten.contains("catalog_name = 'prod'"), "{rewritten}");
277        assert!(rewritten.contains("schema_name LIKE 'sal%'"), "{rewritten}");
278        assert!(rewritten.contains(" AND "), "{rewritten}");
279    }
280
281    /// A bare `SHOW DATABASES` still has no WHERE at all.
282    #[test]
283    fn bare_show_databases_is_unfiltered() {
284        let rewritten = rewrite_show_databases("SHOW DATABASES").expect("recognised");
285        assert!(!rewritten.contains("WHERE"), "{rewritten}");
286    }
287
288    /// A quoted identifier is one name even when it contains a dot.
289    #[test]
290    fn a_quoted_use_target_is_not_split_on_its_dot() {
291        assert_eq!(
292            parse_use("USE `my.schema`"),
293            Some(UseTarget {
294                catalog: None,
295                schema: Some("my.schema".into())
296            }),
297            "a back-quoted name is a single identifier, not catalog.schema"
298        );
299        // ...while an unquoted one still qualifies.
300        assert_eq!(
301            parse_use("USE lake.sales"),
302            Some(UseTarget {
303                catalog: Some("lake".into()),
304                schema: Some("sales".into())
305            })
306        );
307    }
308
309    #[tokio::test]
310    async fn use_changes_default_schema_end_to_end() {
311        let engine = crate::SqlEngine::new();
312        // USE mutates the session default schema, so a subsequent *unqualified*
313        // reference to an information_schema relation now resolves.
314        engine
315            .sql("USE information_schema")
316            .await
317            .expect("USE runs");
318        let batches = engine
319            .sql("SELECT count(*) AS c FROM tables")
320            .await
321            .expect("unqualified `tables` resolves via the new default schema")
322            .collect()
323            .await
324            .expect("collect");
325        let total: i64 = {
326            use arrow::array::Int64Array;
327            batches[0]
328                .column(0)
329                .as_any()
330                .downcast_ref::<Int64Array>()
331                .unwrap()
332                .value(0)
333        };
334        assert!(total > 0, "information_schema.tables should be non-empty");
335    }
336
337    #[tokio::test]
338    async fn show_databases_lists_schemas() {
339        let engine = crate::SqlEngine::new();
340        let batches = engine
341            .sql("SHOW DATABASES")
342            .await
343            .expect("SHOW DATABASES runs")
344            .collect()
345            .await
346            .expect("collect");
347        // At least the default schemas are present, under a `namespace` column.
348        let total: usize = batches.iter().map(|b| b.num_rows()).sum();
349        assert!(total >= 1);
350        assert_eq!(batches[0].schema().field(0).name(), "namespace");
351    }
352}
krishiv_sql/statement_completion.rs

krishiv_sql/
statement_completion.rs
