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
31fn unquote(ident: &str) -> String {
32    let t = ident.trim();
33    for q in ['`', '"'] {
34        if let Some(inner) = t.strip_prefix(q).and_then(|s| s.strip_suffix(q)) {
35            return inner.to_string();
36        }
37    }
38    t.to_string()
39}
40
41/// Parse a `USE …` statement into its catalog/schema target, or `None` if the
42/// query is not a `USE` statement.
43///
44/// Forms (Spark/Databricks):
45/// - `USE db` / `USE SCHEMA db` / `USE DATABASE db` / `USE NAMESPACE ns` — schema
46/// - `USE CATALOG cat` — catalog
47/// - `USE cat.schema` — both
48pub fn parse_use(query: &str) -> Option<UseTarget> {
49    let q = query.trim().trim_end_matches(';').trim();
50    // First token must be USE (case-insensitive); the remainder is the target.
51    let (head, remainder) = q.split_once(char::is_whitespace)?;
52    if !head.eq_ignore_ascii_case("USE") {
53        return None;
54    }
55    let remainder = remainder.trim();
56    if remainder.is_empty() {
57        return None;
58    }
59    // Optional leading CATALOG / SCHEMA / DATABASE / NAMESPACE keyword; whatever
60    // follows it (which may be a back-quoted identifier containing spaces) is the
61    // whole name.
62    let (is_catalog, name_part) = match remainder.split_once(char::is_whitespace) {
63        Some((w, r)) if w.eq_ignore_ascii_case("CATALOG") => (true, r.trim()),
64        Some((w, r))
65            if w.eq_ignore_ascii_case("SCHEMA")
66                || w.eq_ignore_ascii_case("DATABASE")
67                || w.eq_ignore_ascii_case("NAMESPACE") =>
68        {
69            (false, r.trim())
70        }
71        _ => (false, remainder),
72    };
73    let name = unquote(name_part);
74    if name.is_empty() {
75        return None;
76    }
77    if is_catalog {
78        return Some(UseTarget {
79            catalog: Some(name),
80            schema: None,
81        });
82    }
83    // Schema form: allow a qualified `catalog.schema`.
84    if let Some((cat, sch)) = name.split_once('.') {
85        return Some(UseTarget {
86            catalog: Some(unquote(cat)),
87            schema: Some(unquote(sch)),
88        });
89    }
90    Some(UseTarget {
91        catalog: None,
92        schema: Some(name),
93    })
94}
95
96/// Apply a `USE` target to the session by mutating the default catalog/schema.
97/// Returns `Some(Ok(()))` when the query was a `USE` statement (handled), or
98/// `None` when it was not.
99pub fn apply_use(ctx: &SessionContext, query: &str) -> Option<Result<(), String>> {
100    let target = parse_use(query)?;
101    let state_ref = ctx.state_ref();
102    let mut state = state_ref.write();
103    let opts = state.config_mut().options_mut();
104    if let Some(catalog) = target.catalog {
105        opts.catalog.default_catalog = catalog;
106    }
107    if let Some(schema) = target.schema {
108        opts.catalog.default_schema = schema;
109    }
110    Some(Ok(()))
111}
112
113/// If `query` is `SHOW DATABASES`/`SHOW SCHEMAS` (optionally `LIKE 'pat'`),
114/// return an equivalent `information_schema.schemata` SELECT; otherwise `None`.
115/// The result column is `namespace`, matching Spark's `SHOW DATABASES`.
116pub fn rewrite_show_databases(query: &str) -> Option<String> {
117    let q = query.trim().trim_end_matches(';').trim();
118    let upper = q.to_ascii_uppercase();
119    let is_show = upper.starts_with("SHOW DATABASES") || upper.starts_with("SHOW SCHEMAS");
120    if !is_show {
121        return None;
122    }
123    // Optional `LIKE 'pattern'` filter.
124    let like_clause = if let Some(idx) = upper.find(" LIKE ") {
125        let pat = q[idx + 6..].trim().trim_end_matches(';').trim();
126        Some(format!(" WHERE schema_name LIKE {pat}"))
127    } else {
128        None
129    };
130    Some(format!(
131        "SELECT schema_name AS namespace FROM information_schema.schemata{} ORDER BY namespace",
132        like_clause.unwrap_or_default()
133    ))
134}
135
136#[cfg(test)]
137mod tests {
138    use super::*;
139
140    #[test]
141    fn parse_use_forms() {
142        assert_eq!(
143            parse_use("USE analytics"),
144            Some(UseTarget {
145                catalog: None,
146                schema: Some("analytics".into())
147            })
148        );
149        assert_eq!(
150            parse_use("USE SCHEMA sales"),
151            Some(UseTarget {
152                catalog: None,
153                schema: Some("sales".into())
154            })
155        );
156        assert_eq!(
157            parse_use("USE DATABASE sales;"),
158            Some(UseTarget {
159                catalog: None,
160                schema: Some("sales".into())
161            })
162        );
163        assert_eq!(
164            parse_use("USE CATALOG lakehouse"),
165            Some(UseTarget {
166                catalog: Some("lakehouse".into()),
167                schema: None
168            })
169        );
170        assert_eq!(
171            parse_use("USE lake.sales"),
172            Some(UseTarget {
173                catalog: Some("lake".into()),
174                schema: Some("sales".into())
175            })
176        );
177        assert_eq!(
178            parse_use("USE `my schema`").unwrap().schema.as_deref(),
179            Some("my schema")
180        );
181        assert_eq!(parse_use("SELECT 1"), None);
182    }
183
184    #[test]
185    fn show_databases_rewrite() {
186        assert!(
187            rewrite_show_databases("SHOW DATABASES")
188                .unwrap()
189                .contains("information_schema.schemata")
190        );
191        assert!(
192            rewrite_show_databases("SHOW SCHEMAS")
193                .unwrap()
194                .contains("AS namespace")
195        );
196        let with_like = rewrite_show_databases("SHOW DATABASES LIKE 'sal%'").unwrap();
197        assert!(with_like.contains("LIKE 'sal%'"));
198        assert_eq!(rewrite_show_databases("SHOW TABLES"), None);
199    }
200
201    #[tokio::test]
202    async fn use_changes_default_schema_end_to_end() {
203        let engine = crate::SqlEngine::new();
204        // USE mutates the session default schema, so a subsequent *unqualified*
205        // reference to an information_schema relation now resolves.
206        engine
207            .sql("USE information_schema")
208            .await
209            .expect("USE runs");
210        let batches = engine
211            .sql("SELECT count(*) AS c FROM tables")
212            .await
213            .expect("unqualified `tables` resolves via the new default schema")
214            .collect()
215            .await
216            .expect("collect");
217        let total: i64 = {
218            use arrow::array::Int64Array;
219            batches[0]
220                .column(0)
221                .as_any()
222                .downcast_ref::<Int64Array>()
223                .unwrap()
224                .value(0)
225        };
226        assert!(total > 0, "information_schema.tables should be non-empty");
227    }
228
229    #[tokio::test]
230    async fn show_databases_lists_schemas() {
231        let engine = crate::SqlEngine::new();
232        let batches = engine
233            .sql("SHOW DATABASES")
234            .await
235            .expect("SHOW DATABASES runs")
236            .collect()
237            .await
238            .expect("collect");
239        // At least the default schemas are present, under a `namespace` column.
240        let total: usize = batches.iter().map(|b| b.num_rows()).sum();
241        assert!(total >= 1);
242        assert_eq!(batches[0].schema().field(0).name(), "namespace");
243    }
244}
krishiv_sql/statement_completion.rs

krishiv_sql/
statement_completion.rs
