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

ursa_sql/
lib.rs
