ursa_sql/
lib.rs

1//! Ursa SQL: DataFusion's dialect plus Ursa's own statements.
2//!
3//! [`parse_sql`] splits a string into [`UrsaStatement`]s. [`execute`] runs one
4//! on an engine session (see `ursa_engine::new_session`). Statements DataFusion
5//! knows are planned and run by DataFusion, except that one may not drop a
6//! table a graph uses, or replace it by one the graph no longer fits
7//! (`ursa_engine::graph::check_ddl`). `CREATE PROPERTY GRAPH` and `DROP
8//! PROPERTY GRAPH` define graphs (see [`property_graph`] for the subset of
9//! SQL/PGQ supported); `SHOW GRAPHS` and `DESCRIBE GRAPH` describe them.
10//! `SHOW ursa.<setting>` shows one of Ursa's settings, and `SHOW ALL` lists
11//! them with DataFusion's (see `settings`). `SHOW TABLES` lists the tables
12//! and views, without `information_schema`'s own views. `read_parquet(...)`,
13//! `read_csv(...)` and `read_json(...)` open their files before a statement
14//! is planned (see `ursa_engine::read_functions`).
15//!
16//! A session may have a `.ursa` database ([`database`]): [`open_session`]
17//! opens one, and `ATTACH 'club.ursa'` opens one in an in-memory session that
18//! has defined nothing yet. In such a session, statements that define or drop
19//! views, external tables and graphs are kept in the database's catalog file.
20//! Relative file paths resolve against the current directory when a statement
21//! runs, as in DuckDB; a stored definition's paths are rewritten relative to
22//! the directory holding the file (see [`database`]).
23
24use datafusion::dataframe::DataFrame;
25use datafusion::error::{DataFusionError, Result};
26use datafusion::prelude::SessionContext;
27use datafusion::sql::parser::Statement as DFStatement;
28use datafusion::sql::sqlparser::ast::Statement as SqlStatement;
29
30pub mod database;
31mod from_first;
32mod graph_ddl;
33pub mod graph_table;
34mod parser;
35mod paths;
36pub mod property_graph;
37mod settings;
38
39pub use database::open_session;
40pub use parser::{
41    parse_sql, Attach, CreatePropertyGraph, DescribeGraph, DropPropertyGraph, UrsaParser,
42    UrsaStatement, DEFAULT_RECURSION_LIMIT,
43};
44pub use ursa_engine::database::OpenOptions;
45
46/// Parse `sql` with the session's configured dialect and recursion limit.
47pub fn parse_sql_for(ctx: &SessionContext, sql: &str) -> Result<Vec<UrsaStatement>> {
48    let options = ctx.state().config().options().sql_parser.clone();
49    let dialect = datafusion::sql::sqlparser::dialect::dialect_from_str(options.dialect)
50        .ok_or_else(|| {
51            DataFusionError::Configuration(format!("unknown SQL dialect '{}'", options.dialect))
52        })?;
53    UrsaParser::new(sql, dialect.as_ref(), options.recursion_limit.get())?.parse_statements()
54}
55
56/// [`parse_sql_for`], with each statement's text as written (see
57/// [`execute_with_text`]).
58pub fn parse_sql_with_text_for(
59    ctx: &SessionContext,
60    sql: &str,
61) -> Result<Vec<(UrsaStatement, String)>> {
62    let options = ctx.state().config().options().sql_parser.clone();
63    let dialect = datafusion::sql::sqlparser::dialect::dialect_from_str(options.dialect)
64        .ok_or_else(|| {
65            DataFusionError::Configuration(format!("unknown SQL dialect '{}'", options.dialect))
66        })?;
67    UrsaParser::new(sql, dialect.as_ref(), options.recursion_limit.get())?
68        .parse_statements_with_text()
69}
70
71/// Run one statement on `ctx`. The returned frame is lazy for a query, and
72/// already applied for DDL (as with `SessionContext::sql`).
73///
74/// In a session with a database, a statement that defines an object is
75/// stored as its `Display` form; one whose `Display` does not parse back to
76/// the same statement (`CREATE EXTERNAL TABLE` with columns or options) needs
77/// [`execute_with_text`].
78pub async fn execute(ctx: &SessionContext, statement: UrsaStatement) -> Result<DataFrame> {
79    execute_inner(ctx, statement, None).await
80}
81
82/// [`execute`], with the statement's text as written (from
83/// [`parse_sql_with_text_for`]), which a database stores when the statement's
84/// `Display` does not round-trip.
85pub async fn execute_with_text(
86    ctx: &SessionContext,
87    statement: UrsaStatement,
88    text: &str,
89) -> Result<DataFrame> {
90    execute_inner(ctx, statement, Some(text)).await
91}
92
93async fn execute_inner(
94    ctx: &SessionContext,
95    statement: UrsaStatement,
96    text: Option<&str>,
97) -> Result<DataFrame> {
98    match ursa_engine::database::database(&ctx.state()) {
99        Some(db) => database::execute(ctx, &db, statement, text).await,
100        None => match statement {
101            UrsaStatement::Attach(attach) => database::attach(ctx, attach).await,
102            statement => run(ctx, statement).await,
103        },
104    }
105}
106
107/// Run one statement, leaving the database's catalog alone.
108pub(crate) async fn run(ctx: &SessionContext, statement: UrsaStatement) -> Result<DataFrame> {
109    match statement {
110        UrsaStatement::DataFusion(statement) => {
111            let state = ctx.state();
112            // A COPY target or external table location in object storage needs
113            // its store registered before DataFusion plans the write or read.
114            match statement.as_ref() {
115                DFStatement::CopyTo(copy) => {
116                    ursa_engine::ensure_object_store_for(state.runtime_env(), &copy.target).await?
117                }
118                DFStatement::CreateExternalTable(table) => {
119                    for location in &table.locations {
120                        ursa_engine::ensure_object_store_for(state.runtime_env(), location).await?
121                    }
122                }
123                _ => {}
124            }
125            // `SHOW ursa.<setting>` and `SHOW ALL` with Ursa's settings.
126            if let DFStatement::Statement(inner) = statement.as_ref() {
127                if let SqlStatement::ShowVariable { variable } = inner.as_ref() {
128                    if let Some(frame) = settings::show(ctx, variable).await {
129                        return frame;
130                    }
131                }
132                // `SHOW TABLES` without `information_schema`'s own views.
133                if is_plain_show_tables(inner) {
134                    return ctx
135                        .sql(
136                            "SELECT * FROM information_schema.tables \
137                             WHERE table_schema <> 'information_schema'",
138                        )
139                        .await;
140                }
141            }
142            // `read_csv(...)` and its siblings open their files first: the
143            // planner that plans them cannot wait on a listing.
144            let reads = ursa_engine::read_functions::open_reads(&state, &statement).await?;
145            let plan = state.statement_to_plan(*statement).await?;
146            drop(reads);
147            ursa_engine::graph::check_ddl(&state, &plan).await?;
148            ctx.execute_logical_plan(plan).await
149        }
150        UrsaStatement::CreatePropertyGraph(statement) => graph_ddl::create(ctx, *statement).await,
151        UrsaStatement::DropPropertyGraph(statement) => graph_ddl::drop(ctx, statement),
152        UrsaStatement::ShowGraphs => graph_ddl::show(ctx),
153        UrsaStatement::DescribeGraph(statement) => graph_ddl::describe(ctx, statement).await,
154        // `execute` handles ATTACH before it gets here.
155        UrsaStatement::Attach(_) => Err(DataFusionError::Plan(
156            "ATTACH cannot run while replaying a database's catalog".into(),
157        )),
158    }
159}
160
161/// Whether `statement` is `SHOW TABLES`, with no option DataFusion would
162/// refuse or treat differently.
163fn is_plain_show_tables(statement: &SqlStatement) -> bool {
164    let SqlStatement::ShowTables {
165        terse: false,
166        history: false,
167        extended: false,
168        full: false,
169        external: false,
170        show_options,
171    } = statement
172    else {
173        return false;
174    };
175    show_options.show_in.is_none()
176        && show_options.starts_with.is_none()
177        && show_options.limit.is_none()
178        && show_options.limit_from.is_none()
179        && show_options.filter_position.is_none()
180}
181
182#[cfg(test)]
183mod tests {
184    use super::*;
185
186    #[tokio::test]
187    async fn runs_datafusion_statements() {
188        let ctx = ursa_engine::new_session();
189        let mut statements = parse_sql_for(
190            &ctx,
191            "CREATE VIEW v AS SELECT 1 AS x UNION ALL SELECT 2; SELECT sum(x) AS s FROM v",
192        )
193        .unwrap();
194        let select = statements.pop().unwrap();
195        execute(&ctx, statements.pop().unwrap()).await.unwrap();
196        let batches = execute(&ctx, select)
197            .await
198            .unwrap()
199            .collect()
200            .await
201            .unwrap();
202        let formatted = datafusion::arrow::util::pretty::pretty_format_batches(&batches)
203            .unwrap()
204            .to_string();
205        assert!(formatted.contains("| 3 |"), "{formatted}");
206    }
207
208    #[tokio::test]
209    async fn show_tables_leaves_out_information_schema() {
210        let ctx = ursa_engine::new_session();
211        for sql in ["CREATE VIEW v AS SELECT 1 AS x", "SHOW TABLES"] {
212            let statement = parse_sql_for(&ctx, sql).unwrap().remove(0);
213            let batches = execute(&ctx, statement)
214                .await
215                .unwrap()
216                .collect()
217                .await
218                .unwrap();
219            if sql == "SHOW TABLES" {
220                let formatted = datafusion::arrow::util::pretty::pretty_format_batches(&batches)
221                    .unwrap()
222                    .to_string();
223                assert!(
224                    formatted.contains("| public       | v          | VIEW"),
225                    "{formatted}"
226                );
227                assert!(!formatted.contains("information_schema"), "{formatted}");
228            }
229        }
230    }
231}
ursa_sql/lib.rs

ursa_sql/
lib.rs
