ursa_sql/
parser.rs

1//! Ursa's statement parser.
2//!
3//! It wraps DataFusion's [`DFParser`]: before each statement it peeks at the
4//! leading words, parses Ursa's own statements with sqlparser's public helpers,
5//! and hands everything else to `DFParser::parse_statement`. sqlparser knows no
6//! SQL/PGQ, so `PROPERTY`, `GRAPH` and `GRAPHS` are not keywords; they are
7//! matched as unquoted words, case-insensitively.
8
9use std::fmt;
10
11use datafusion::error::Result;
12use datafusion::sql::parser::{DFParser, DFParserBuilder, Statement as DFStatement};
13use datafusion::sql::sqlparser::ast::{Ident, ObjectName};
14use datafusion::sql::sqlparser::dialect::{Dialect, GenericDialect};
15use datafusion::sql::sqlparser::parser::ParserError;
16use datafusion::sql::sqlparser::tokenizer::{Location, Token};
17
18use crate::property_graph::{parse_tables, EdgeTableDef, VertexTableDef};
19
20/// Same as sqlparser's and DataFusion's default.
21pub const DEFAULT_RECURSION_LIMIT: usize = 50;
22
23/// One statement of Ursa SQL.
24#[derive(Debug, Clone, PartialEq)]
25pub enum UrsaStatement {
26    /// Anything DataFusion parses: queries, DDL, `COPY … TO`, `SET`, `EXPLAIN`.
27    DataFusion(Box<DFStatement>),
28    /// `CREATE [OR REPLACE] PROPERTY GRAPH [IF NOT EXISTS] name …` (SQL/PGQ).
29    CreatePropertyGraph(Box<CreatePropertyGraph>),
30    /// `DROP PROPERTY GRAPH [IF EXISTS] name`.
31    DropPropertyGraph(DropPropertyGraph),
32    /// `DESCRIBE GRAPH name` (or `DESC GRAPH name`).
33    DescribeGraph(DescribeGraph),
34    /// `SHOW GRAPHS`.
35    ShowGraphs,
36    /// `ATTACH [DATABASE] 'file.ursa' [AS name]`.
37    Attach(Attach),
38}
39
40/// `CREATE PROPERTY GRAPH`: Ursa's subset of SQL/PGQ (see
41/// [`crate::property_graph`]).
42#[derive(Debug, Clone, PartialEq, Eq)]
43pub struct CreatePropertyGraph {
44    pub name: ObjectName,
45    pub or_replace: bool,
46    pub if_not_exists: bool,
47    /// `VERTEX TABLES (…)`, which may be omitted.
48    pub vertices: Option<VertexTableDef>,
49    /// `EDGE TABLES (…)`.
50    pub edges: EdgeTableDef,
51}
52
53#[derive(Debug, Clone, PartialEq, Eq)]
54pub struct DropPropertyGraph {
55    pub name: ObjectName,
56    pub if_exists: bool,
57}
58
59#[derive(Debug, Clone, PartialEq, Eq)]
60pub struct DescribeGraph {
61    pub name: ObjectName,
62}
63
64#[derive(Debug, Clone, PartialEq, Eq)]
65pub struct Attach {
66    pub path: String,
67    pub alias: Option<Ident>,
68}
69
70impl UrsaStatement {
71    /// The statement's leading words, as an error message names it.
72    pub fn kind(&self) -> &'static str {
73        match self {
74            UrsaStatement::DataFusion(_) => "SQL statement",
75            UrsaStatement::CreatePropertyGraph(_) => "CREATE PROPERTY GRAPH",
76            UrsaStatement::DropPropertyGraph(_) => "DROP PROPERTY GRAPH",
77            UrsaStatement::DescribeGraph(_) => "DESCRIBE GRAPH",
78            UrsaStatement::ShowGraphs => "SHOW GRAPHS",
79            UrsaStatement::Attach(_) => "ATTACH",
80        }
81    }
82}
83
84impl fmt::Display for UrsaStatement {
85    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
86        match self {
87            UrsaStatement::DataFusion(s) => write!(f, "{s}"),
88            UrsaStatement::CreatePropertyGraph(s) => {
89                write!(f, "CREATE ")?;
90                if s.or_replace {
91                    write!(f, "OR REPLACE ")?;
92                }
93                write!(f, "PROPERTY GRAPH ")?;
94                if s.if_not_exists {
95                    write!(f, "IF NOT EXISTS ")?;
96                }
97                write!(f, "{}", s.name)?;
98                if let Some(vertices) = &s.vertices {
99                    write!(f, " {vertices}")?;
100                }
101                write!(f, " {}", s.edges)
102            }
103            UrsaStatement::DropPropertyGraph(s) => {
104                write!(f, "DROP PROPERTY GRAPH ")?;
105                if s.if_exists {
106                    write!(f, "IF EXISTS ")?;
107                }
108                write!(f, "{}", s.name)
109            }
110            UrsaStatement::DescribeGraph(s) => write!(f, "DESCRIBE GRAPH {}", s.name),
111            UrsaStatement::ShowGraphs => write!(f, "SHOW GRAPHS"),
112            UrsaStatement::Attach(s) => {
113                write!(f, "ATTACH '{}'", s.path.replace('\'', "''"))?;
114                if let Some(alias) = &s.alias {
115                    write!(f, " AS {alias}")?;
116                }
117                Ok(())
118            }
119        }
120    }
121}
122
123/// Parse `sql` (one or more statements separated by `;`) with the generic dialect.
124pub fn parse_sql(sql: &str) -> Result<Vec<UrsaStatement>> {
125    UrsaParser::new(sql, &GenericDialect {}, DEFAULT_RECURSION_LIMIT)?.parse_statements()
126}
127
128/// Ursa's parser: DataFusion's, with Ursa's statements recognised first.
129pub struct UrsaParser<'a> {
130    df: DFParser<'a>,
131    /// The text being parsed, for [`Self::parse_statements_with_text`].
132    sql: String,
133}
134
135impl<'a> UrsaParser<'a> {
136    pub fn new(sql: &str, dialect: &'a dyn Dialect, recursion_limit: usize) -> Result<Self> {
137        let df = DFParserBuilder::new(sql)
138            .with_dialect(dialect)
139            .with_recursion_limit(recursion_limit)
140            .build()?;
141        Ok(Self {
142            df,
143            sql: sql.to_string(),
144        })
145    }
146
147    /// Parse every statement, separated by `;`. Empty statements are skipped.
148    /// (The loop is `DFParser::parse_statements`, calling our `parse_statement`.)
149    pub fn parse_statements(&mut self) -> Result<Vec<UrsaStatement>> {
150        Ok(self
151            .parse_statements_with_text()?
152            .into_iter()
153            .map(|(statement, _)| statement)
154            .collect())
155    }
156
157    /// [`Self::parse_statements`], with each statement's text as written
158    /// (from its first token to its last, without the `;`). A database file
159    /// stores this text when a statement's `Display` does not round-trip.
160    pub fn parse_statements_with_text(&mut self) -> Result<Vec<(UrsaStatement, String)>> {
161        let mut statements = Vec::new();
162        let mut expecting_delimiter = false;
163        loop {
164            while self.df.parser.consume_token(&Token::SemiColon) {
165                expecting_delimiter = false;
166            }
167            let next = self.df.parser.peek_token();
168            if next.token == Token::EOF {
169                break;
170            }
171            if expecting_delimiter {
172                return Err(ParserError::ParserError(format!(
173                    "Expected: end of statement, found: {next}{}",
174                    next.span.start
175                ))
176                .into());
177            }
178            let start = next.span.start;
179            let statement = self.parse_statement()?;
180            let end = self.df.parser.get_current_token().span.end;
181            let text = source_text(&self.sql, start, end).unwrap_or_else(|| statement.to_string());
182            statements.push((statement, text));
183            expecting_delimiter = true;
184        }
185        Ok(statements)
186    }
187
188    /// Parse one statement.
189    pub fn parse_statement(&mut self) -> Result<UrsaStatement> {
190        if self.word(0, "CREATE") {
191            let at = if self.word(1, "OR") && self.word(2, "REPLACE") {
192                3
193            } else {
194                1
195            };
196            if self.word(at, "PROPERTY") && self.word(at + 1, "GRAPH") {
197                return self.parse_create_property_graph();
198            }
199        } else if self.word(0, "DROP") && self.word(1, "PROPERTY") && self.word(2, "GRAPH") {
200            return Ok(self.parse_drop_property_graph()?);
201        } else if (self.word(0, "DESCRIBE") || self.word(0, "DESC"))
202            && self.word(1, "GRAPH")
203            && matches!(self.df.parser.peek_nth_token_ref(2).token, Token::Word(_))
204        {
205            // `DESCRIBE graph` alone still describes a table named `graph`.
206            for _ in 0..2 {
207                self.df.parser.next_token();
208            }
209            let name = self.df.parser.parse_object_name(false)?;
210            return Ok(UrsaStatement::DescribeGraph(DescribeGraph { name }));
211        } else if self.word(0, "SHOW") && self.word(1, "GRAPHS") {
212            for _ in 0..2 {
213                self.df.parser.next_token();
214            }
215            return Ok(UrsaStatement::ShowGraphs);
216        } else if self.word(0, "ATTACH") {
217            return Ok(self.parse_attach()?);
218        }
219        Ok(UrsaStatement::DataFusion(Box::new(
220            self.df.parse_statement()?,
221        )))
222    }
223
224    /// Whether the `n`th token ahead is the unquoted word `word`, in any case.
225    fn word(&self, n: usize, word: &str) -> bool {
226        match &self.df.parser.peek_nth_token_ref(n).token {
227            Token::Word(w) => w.quote_style.is_none() && w.value.eq_ignore_ascii_case(word),
228            _ => false,
229        }
230    }
231
232    /// Consume the next token if it is the unquoted word `word`.
233    fn consume_word(&mut self, word: &str) -> bool {
234        let found = self.word(0, word);
235        if found {
236            self.df.parser.next_token();
237        }
238        found
239    }
240
241    fn consume_words(&mut self, words: &[&str]) -> bool {
242        let found = words.iter().enumerate().all(|(i, w)| self.word(i, w));
243        if found {
244            for _ in words {
245                self.df.parser.next_token();
246            }
247        }
248        found
249    }
250
251    fn parse_create_property_graph(&mut self) -> Result<UrsaStatement> {
252        self.df.parser.next_token(); // CREATE
253        let or_replace = self.consume_words(&["OR", "REPLACE"]);
254        self.consume_words(&["PROPERTY", "GRAPH"]);
255        let if_not_exists = self.consume_words(&["IF", "NOT", "EXISTS"]);
256        let name = self.df.parser.parse_object_name(false)?;
257        let (vertices, edges) = parse_tables(&mut self.df.parser)?;
258        Ok(UrsaStatement::CreatePropertyGraph(Box::new(
259            CreatePropertyGraph {
260                name,
261                or_replace,
262                if_not_exists,
263                vertices,
264                edges,
265            },
266        )))
267    }
268
269    fn parse_drop_property_graph(&mut self) -> Result<UrsaStatement, ParserError> {
270        self.consume_words(&["DROP", "PROPERTY", "GRAPH"]);
271        let if_exists = self.consume_words(&["IF", "EXISTS"]);
272        let name = self.df.parser.parse_object_name(false)?;
273        Ok(UrsaStatement::DropPropertyGraph(DropPropertyGraph {
274            name,
275            if_exists,
276        }))
277    }
278
279    fn parse_attach(&mut self) -> Result<UrsaStatement, ParserError> {
280        self.consume_word("ATTACH");
281        self.consume_word("DATABASE");
282        let path = self.df.parser.parse_literal_string()?;
283        let alias = if self.consume_word("AS") {
284            Some(self.df.parser.parse_identifier()?)
285        } else {
286            None
287        };
288        Ok(UrsaStatement::Attach(Attach { path, alias }))
289    }
290}
291
292/// The text of `sql` between two token locations, or `None` if a location is
293/// unknown (line 0) or outside the text.
294fn source_text(sql: &str, start: Location, end: Location) -> Option<String> {
295    let from = byte_offset(sql, start)?;
296    let to = byte_offset(sql, end)?;
297    (from < to).then(|| sql[from..to].trim().to_string())
298}
299
300/// The byte offset of a tokenizer location (1-based line and column, columns
301/// counted in characters).
302pub(crate) fn byte_offset(sql: &str, at: Location) -> Option<usize> {
303    if at.line == 0 || at.column == 0 {
304        return None;
305    }
306    let mut line_start = 0;
307    for _ in 1..at.line {
308        line_start += sql[line_start..].find('\n')? + 1;
309    }
310    let line = &sql[line_start..];
311    let column = usize::try_from(at.column - 1).ok()?;
312    match line.char_indices().nth(column) {
313        Some((offset, _)) => Some(line_start + offset),
314        None if line.chars().count() == column => Some(sql.len()),
315        None => None,
316    }
317}
318
319#[cfg(test)]
320mod tests {
321    use super::*;
322
323    fn one(sql: &str) -> UrsaStatement {
324        let mut statements = parse_sql(sql).unwrap();
325        assert_eq!(statements.len(), 1, "{sql}");
326        statements.pop().unwrap()
327    }
328
329    #[test]
330    fn create_property_graph() {
331        let sql = "CREATE PROPERTY GRAPH karate
332              VERTEX TABLES (members KEY (id))
333              EDGE TABLES (
334                ties SOURCE KEY (src) REFERENCES members (id)
335                     DESTINATION KEY (dst) REFERENCES members (id)
336              )";
337        let UrsaStatement::CreatePropertyGraph(g) = one(sql) else {
338            panic!("not a CREATE PROPERTY GRAPH")
339        };
340        assert_eq!(g.name.to_string(), "karate");
341        assert!(!g.or_replace && !g.if_not_exists);
342
343        let s = one("create or replace property graph if not exists db.social edge tables (e source key (a) destination key (b))");
344        let UrsaStatement::CreatePropertyGraph(g) = s else {
345            panic!("not a CREATE PROPERTY GRAPH")
346        };
347        assert_eq!(g.name.to_string(), "db.social");
348        assert!(g.or_replace && g.if_not_exists);
349    }
350
351    #[test]
352    fn drop_describe_show_attach() {
353        assert_eq!(
354            one("DROP PROPERTY GRAPH IF EXISTS karate"),
355            UrsaStatement::DropPropertyGraph(DropPropertyGraph {
356                name: ObjectName::from(vec![Ident::new("karate")]),
357                if_exists: true,
358            })
359        );
360        assert_eq!(
361            one("describe graph karate").to_string(),
362            "DESCRIBE GRAPH karate"
363        );
364        assert_eq!(one("DESC GRAPH karate").kind(), "DESCRIBE GRAPH");
365        assert_eq!(one("SHOW GRAPHS"), UrsaStatement::ShowGraphs);
366        assert_eq!(
367            one("ATTACH 'club.ursa' AS club"),
368            UrsaStatement::Attach(Attach {
369                path: "club.ursa".into(),
370                alias: Some(Ident::new("club")),
371            })
372        );
373        assert_eq!(
374            one("attach database 'club.ursa'").to_string(),
375            "ATTACH 'club.ursa'"
376        );
377    }
378
379    #[test]
380    fn everything_else_is_datafusion() {
381        for sql in [
382            "SELECT * FROM 'edges.csv' LIMIT 3",
383            "COPY (SELECT 1) TO 'out.parquet'",
384            "CREATE VIEW ties AS SELECT * FROM 'karate_edges.csv'",
385            "DESCRIBE graph",
386            "SHOW TABLES",
387            "CREATE TABLE property (graph INT)",
388        ] {
389            assert!(
390                matches!(one(sql), UrsaStatement::DataFusion(_)),
391                "{sql} should go to DataFusion"
392            );
393        }
394    }
395
396    #[test]
397    fn statements_split_on_semicolons() {
398        let statements = parse_sql(
399            "SELECT 1; ; CREATE PROPERTY GRAPH g EDGE TABLES (e SOURCE KEY (a) DESTINATION KEY (b));
400             SHOW GRAPHS; SELECT ';' AS semi;",
401        )
402        .unwrap();
403        let kinds: Vec<_> = statements.iter().map(UrsaStatement::kind).collect();
404        assert_eq!(
405            kinds,
406            [
407                "SQL statement",
408                "CREATE PROPERTY GRAPH",
409                "SHOW GRAPHS",
410                "SQL statement"
411            ]
412        );
413    }
414
415    #[test]
416    fn statements_keep_their_text() {
417        let sql =
418            "SELECT 1;\n  CREATE EXTERNAL TABLE t (a INT)\n  STORED AS CSV LOCATION 'é.csv'  ;\
419                   CREATE PROPERTY GRAPH g EDGE TABLES (e SOURCE KEY (a) DESTINATION KEY (b))";
420        let statements = UrsaParser::new(sql, &GenericDialect {}, DEFAULT_RECURSION_LIMIT)
421            .unwrap()
422            .parse_statements_with_text()
423            .unwrap();
424        let texts: Vec<_> = statements.iter().map(|(_, t)| t.as_str()).collect();
425        assert_eq!(
426            texts,
427            [
428                "SELECT 1",
429                "CREATE EXTERNAL TABLE t (a INT)\n  STORED AS CSV LOCATION 'é.csv'",
430                "CREATE PROPERTY GRAPH g EDGE TABLES (e SOURCE KEY (a) DESTINATION KEY (b))",
431            ]
432        );
433    }
434
435    #[test]
436    fn missing_delimiter_is_a_syntax_error() {
437        let err = parse_sql("SHOW GRAPHS SELECT 1").unwrap_err();
438        assert!(err.to_string().contains("end of statement"), "{err}");
439    }
440}
ursa_sql/parser.rs

ursa_sql/
parser.rs
