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`).
12//!
13//! A session may have a `.ursa` database ([`database`]): [`open_session`]
14//! opens one, and `ATTACH 'club.ursa'` opens one in an in-memory session that
15//! has defined nothing yet. In such a session, statements that define or drop
16//! views, external tables and graphs are kept in the database's catalog file.
17//! Relative file paths resolve against the current directory when a statement
18//! runs, as in DuckDB; a stored definition's paths are rewritten relative to
19//! the directory holding the file (see [`database`]).
20
21use datafusion::dataframe::DataFrame;
22use datafusion::error::{DataFusionError, Result};
23use datafusion::prelude::SessionContext;
24use datafusion::sql::parser::Statement as DFStatement;
25use datafusion::sql::sqlparser::ast::Statement as SqlStatement;
26
27pub mod database;
28mod graph_ddl;
29mod parser;
30mod paths;
31pub mod property_graph;
32mod settings;
33
34pub use database::open_session;
35pub use parser::{
36    parse_sql, Attach, CreatePropertyGraph, DescribeGraph, DropPropertyGraph, UrsaParser,
37    UrsaStatement, DEFAULT_RECURSION_LIMIT,
38};
39pub use ursa_engine::database::OpenOptions;
40
41/// Parse `sql` with the session's configured dialect and recursion limit.
42pub fn parse_sql_for(ctx: &SessionContext, sql: &str) -> Result<Vec<UrsaStatement>> {
43    let options = ctx.state().config().options().sql_parser.clone();
44    let dialect = datafusion::sql::sqlparser::dialect::dialect_from_str(options.dialect)
45        .ok_or_else(|| {
46            DataFusionError::Configuration(format!("unknown SQL dialect '{}'", options.dialect))
47        })?;
48    UrsaParser::new(sql, dialect.as_ref(), options.recursion_limit.get())?.parse_statements()
49}
50
51/// [`parse_sql_for`], with each statement's text as written (see
52/// [`execute_with_text`]).
53pub fn parse_sql_with_text_for(
54    ctx: &SessionContext,
55    sql: &str,
56) -> Result<Vec<(UrsaStatement, String)>> {
57    let options = ctx.state().config().options().sql_parser.clone();
58    let dialect = datafusion::sql::sqlparser::dialect::dialect_from_str(options.dialect)
59        .ok_or_else(|| {
60            DataFusionError::Configuration(format!("unknown SQL dialect '{}'", options.dialect))
61        })?;
62    UrsaParser::new(sql, dialect.as_ref(), options.recursion_limit.get())?
63        .parse_statements_with_text()
64}
65
66/// Run one statement on `ctx`. The returned frame is lazy for a query, and
67/// already applied for DDL (as with `SessionContext::sql`).
68///
69/// In a session with a database, a statement that defines an object is
70/// stored as its `Display` form; one whose `Display` does not parse back to
71/// the same statement (`CREATE EXTERNAL TABLE` with columns or options) needs
72/// [`execute_with_text`].
73pub async fn execute(ctx: &SessionContext, statement: UrsaStatement) -> Result<DataFrame> {
74    execute_inner(ctx, statement, None).await
75}
76
77/// [`execute`], with the statement's text as written (from
78/// [`parse_sql_with_text_for`]), which a database stores when the statement's
79/// `Display` does not round-trip.
80pub async fn execute_with_text(
81    ctx: &SessionContext,
82    statement: UrsaStatement,
83    text: &str,
84) -> Result<DataFrame> {
85    execute_inner(ctx, statement, Some(text)).await
86}
87
88async fn execute_inner(
89    ctx: &SessionContext,
90    statement: UrsaStatement,
91    text: Option<&str>,
92) -> Result<DataFrame> {
93    match ursa_engine::database::database(&ctx.state()) {
94        Some(db) => database::execute(ctx, &db, statement, text).await,
95        None => match statement {
96            UrsaStatement::Attach(attach) => database::attach(ctx, attach).await,
97            statement => run(ctx, statement).await,
98        },
99    }
100}
101
102/// Run one statement, leaving the database's catalog alone.
103pub(crate) async fn run(ctx: &SessionContext, statement: UrsaStatement) -> Result<DataFrame> {
104    match statement {
105        UrsaStatement::DataFusion(statement) => {
106            let state = ctx.state();
107            // A COPY target or external table location in object storage needs
108            // its store registered before DataFusion plans the write or read.
109            match statement.as_ref() {
110                DFStatement::CopyTo(copy) => {
111                    ursa_engine::ensure_object_store_for(state.runtime_env(), &copy.target).await?
112                }
113                DFStatement::CreateExternalTable(table) => {
114                    for location in &table.locations {
115                        ursa_engine::ensure_object_store_for(state.runtime_env(), location).await?
116                    }
117                }
118                _ => {}
119            }
120            // `SHOW ursa.<setting>` and `SHOW ALL` with Ursa's settings.
121            if let DFStatement::Statement(inner) = statement.as_ref() {
122                if let SqlStatement::ShowVariable { variable } = inner.as_ref() {
123                    if let Some(frame) = settings::show(ctx, variable).await {
124                        return frame;
125                    }
126                }
127            }
128            let plan = state.statement_to_plan(*statement).await?;
129            ursa_engine::graph::check_ddl(&state, &plan).await?;
130            ctx.execute_logical_plan(plan).await
131        }
132        UrsaStatement::CreatePropertyGraph(statement) => graph_ddl::create(ctx, *statement).await,
133        UrsaStatement::DropPropertyGraph(statement) => graph_ddl::drop(ctx, statement),
134        UrsaStatement::ShowGraphs => graph_ddl::show(ctx),
135        UrsaStatement::DescribeGraph(statement) => graph_ddl::describe(ctx, statement).await,
136        // `execute` handles ATTACH before it gets here.
137        UrsaStatement::Attach(_) => Err(DataFusionError::Plan(
138            "ATTACH cannot run while replaying a database's catalog".into(),
139        )),
140    }
141}
142
143#[cfg(test)]
144mod tests {
145    use super::*;
146
147    #[tokio::test]
148    async fn runs_datafusion_statements() {
149        let ctx = ursa_engine::new_session();
150        let mut statements = parse_sql_for(
151            &ctx,
152            "CREATE VIEW v AS SELECT 1 AS x UNION ALL SELECT 2; SELECT sum(x) AS s FROM v",
153        )
154        .unwrap();
155        let select = statements.pop().unwrap();
156        execute(&ctx, statements.pop().unwrap()).await.unwrap();
157        let batches = execute(&ctx, select)
158            .await
159            .unwrap()
160            .collect()
161            .await
162            .unwrap();
163        let formatted = datafusion::arrow::util::pretty::pretty_format_batches(&batches)
164            .unwrap()
165            .to_string();
166        assert!(formatted.contains("| 3 |"), "{formatted}");
167    }
168}
ursa_sql/lib.rs

ursa_sql/
lib.rs
