1use std::fmt;
10
11use datafusion::error::{DataFusionError, 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, Tokenizer};
17
18use crate::property_graph::{parse_tables, write_tables, EdgeTableDef, VertexTableDef};
19
20pub const DEFAULT_RECURSION_LIMIT: usize = 50;
22
23#[derive(Debug, Clone, PartialEq)]
25pub enum UrsaStatement {
26 DataFusion(Box<DFStatement>),
28 CreatePropertyGraph(Box<CreatePropertyGraph>),
30 DropPropertyGraph(DropPropertyGraph),
32 DescribeGraph(DescribeGraph),
34 ShowGraphs,
36 Attach(Attach),
38}
39
40#[derive(Debug, Clone, PartialEq, Eq)]
43pub struct CreatePropertyGraph {
44 pub name: ObjectName,
45 pub or_replace: bool,
46 pub if_not_exists: bool,
47 pub vertices: Vec<VertexTableDef>,
49 pub edges: Vec<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 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 write_tables(f, &s.vertices, &s.edges)
99 }
100 UrsaStatement::DropPropertyGraph(s) => {
101 write!(f, "DROP PROPERTY GRAPH ")?;
102 if s.if_exists {
103 write!(f, "IF EXISTS ")?;
104 }
105 write!(f, "{}", s.name)
106 }
107 UrsaStatement::DescribeGraph(s) => write!(f, "DESCRIBE GRAPH {}", s.name),
108 UrsaStatement::ShowGraphs => write!(f, "SHOW GRAPHS"),
109 UrsaStatement::Attach(s) => {
110 write!(f, "ATTACH '{}'", s.path.replace('\'', "''"))?;
111 if let Some(alias) = &s.alias {
112 write!(f, " AS {alias}")?;
113 }
114 Ok(())
115 }
116 }
117 }
118}
119
120pub fn parse_sql(sql: &str) -> Result<Vec<UrsaStatement>> {
122 UrsaParser::new(sql, &GenericDialect {}, DEFAULT_RECURSION_LIMIT)?.parse_statements()
123}
124
125pub struct UrsaParser<'a> {
127 df: DFParser<'a>,
128 dialect: &'a dyn Dialect,
129 sql: String,
131}
132
133impl<'a> UrsaParser<'a> {
134 pub fn new(sql: &str, dialect: &'a dyn Dialect, recursion_limit: usize) -> Result<Self> {
135 let sql = crate::graph_table::rewrite_graph_tables(sql, dialect);
138 let sql = crate::from_first::rewrite_from_first(&sql, dialect);
140 let df = DFParserBuilder::new(sql.as_str())
141 .with_dialect(dialect)
142 .with_recursion_limit(recursion_limit)
143 .build()?;
144 Ok(Self { df, dialect, sql })
145 }
146
147 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 pub fn parse_statements_with_text(&mut self) -> Result<Vec<(UrsaStatement, String)>> {
161 let statements = self.parse_all().map_err(|e| match e {
162 DataFusionError::SQL(..) if self.uses_sample_clause() => {
165 DataFusionError::NotImplemented(USING_SAMPLE.to_string())
166 }
167 e => e,
168 })?;
169 if self.uses_tablesample() {
172 return Err(DataFusionError::NotImplemented(TABLESAMPLE.to_string()));
173 }
174 Ok(statements)
175 }
176
177 fn uses_tablesample(&self) -> bool {
180 if !self.sql.to_ascii_uppercase().contains("TABLESAMPLE") {
181 return false;
182 }
183 let Ok(tokens) = Tokenizer::new(self.dialect, &self.sql).tokenize() else {
184 return false;
185 };
186 tokens.iter().any(|t| {
187 matches!(t, Token::Word(w) if w.quote_style.is_none()
188 && w.value.eq_ignore_ascii_case("TABLESAMPLE"))
189 })
190 }
191
192 fn parse_all(&mut self) -> Result<Vec<(UrsaStatement, String)>> {
193 let mut statements = Vec::new();
194 let mut expecting_delimiter = false;
195 loop {
196 while self.df.parser.consume_token(&Token::SemiColon) {
197 expecting_delimiter = false;
198 }
199 let next = self.df.parser.peek_token();
200 if next.token == Token::EOF {
201 break;
202 }
203 if expecting_delimiter {
204 return Err(ParserError::ParserError(format!(
205 "Expected: end of statement, found: {next}{}",
206 next.span.start
207 ))
208 .into());
209 }
210 let start = next.span.start;
211 let statement = self.parse_statement()?;
212 let end = self.df.parser.get_current_token().span.end;
213 let text = source_text(&self.sql, start, end).unwrap_or_else(|| statement.to_string());
214 statements.push((statement, text));
215 expecting_delimiter = true;
216 }
217 Ok(statements)
218 }
219
220 pub fn parse_statement(&mut self) -> Result<UrsaStatement> {
222 if let Some(message) = self.duckdb_statement() {
223 return Err(DataFusionError::NotImplemented(message.to_string()));
224 }
225 if self.word(0, "CREATE") {
226 let at = if self.word(1, "OR") && self.word(2, "REPLACE") {
227 3
228 } else {
229 1
230 };
231 if self.word(at, "PROPERTY") && self.word(at + 1, "GRAPH") {
232 return self.parse_create_property_graph();
233 }
234 } else if self.word(0, "DROP") && self.word(1, "PROPERTY") && self.word(2, "GRAPH") {
235 return Ok(self.parse_drop_property_graph()?);
236 } else if (self.word(0, "DESCRIBE") || self.word(0, "DESC"))
237 && self.word(1, "GRAPH")
238 && matches!(self.df.parser.peek_nth_token_ref(2).token, Token::Word(_))
239 {
240 for _ in 0..2 {
242 self.df.parser.next_token();
243 }
244 let name = self.df.parser.parse_object_name(false)?;
245 return Ok(UrsaStatement::DescribeGraph(DescribeGraph { name }));
246 } else if (self.word(0, "DESCRIBE") || self.word(0, "DESC"))
247 && self.word(1, "PROPERTY")
248 && self.word(2, "GRAPH")
249 {
250 for _ in 0..3 {
252 self.df.parser.next_token();
253 }
254 let name = self.df.parser.parse_object_name(false)?;
255 return Ok(UrsaStatement::DescribeGraph(DescribeGraph { name }));
256 } else if self.word(0, "SHOW") && self.word(1, "GRAPHS") {
257 for _ in 0..2 {
258 self.df.parser.next_token();
259 }
260 return Ok(UrsaStatement::ShowGraphs);
261 } else if self.word(0, "ATTACH") {
262 return Ok(self.parse_attach()?);
263 }
264 Ok(UrsaStatement::DataFusion(Box::new(
265 self.df.parse_statement()?,
266 )))
267 }
268
269 fn duckdb_statement(&self) -> Option<&'static str> {
272 DUCKDB_STATEMENTS
273 .iter()
274 .find(|(word, _)| self.word(0, word))
275 .filter(|(word, _)| !(*word == "LOAD" && self.word(1, "DATA")))
276 .map(|(_, message)| *message)
277 }
278
279 fn uses_sample_clause(&self) -> bool {
283 let Ok(tokens) = Tokenizer::new(self.dialect, &self.sql).tokenize() else {
284 return false;
285 };
286 let is = |token: &Token, word: &str| {
287 matches!(token, Token::Word(w) if w.quote_style.is_none()
288 && w.value.eq_ignore_ascii_case(word))
289 };
290 let mut words = tokens
291 .iter()
292 .filter(|t| !matches!(t, Token::Whitespace(_)))
293 .peekable();
294 while let Some(token) = words.next() {
295 if is(token, "USING") && words.peek().is_some_and(|next| is(next, "SAMPLE")) {
296 return true;
297 }
298 }
299 false
300 }
301
302 fn word(&self, n: usize, word: &str) -> bool {
304 match &self.df.parser.peek_nth_token_ref(n).token {
305 Token::Word(w) => w.quote_style.is_none() && w.value.eq_ignore_ascii_case(word),
306 _ => false,
307 }
308 }
309
310 fn consume_word(&mut self, word: &str) -> bool {
312 let found = self.word(0, word);
313 if found {
314 self.df.parser.next_token();
315 }
316 found
317 }
318
319 fn consume_words(&mut self, words: &[&str]) -> bool {
320 let found = words.iter().enumerate().all(|(i, w)| self.word(i, w));
321 if found {
322 for _ in words {
323 self.df.parser.next_token();
324 }
325 }
326 found
327 }
328
329 fn parse_create_property_graph(&mut self) -> Result<UrsaStatement> {
330 self.df.parser.next_token(); let or_replace = self.consume_words(&["OR", "REPLACE"]);
332 self.consume_words(&["PROPERTY", "GRAPH"]);
333 let if_not_exists = self.consume_words(&["IF", "NOT", "EXISTS"]);
334 let name = self.df.parser.parse_object_name(false)?;
335 let (vertices, edges) = parse_tables(&mut self.df.parser)?;
336 Ok(UrsaStatement::CreatePropertyGraph(Box::new(
337 CreatePropertyGraph {
338 name,
339 or_replace,
340 if_not_exists,
341 vertices,
342 edges,
343 },
344 )))
345 }
346
347 fn parse_drop_property_graph(&mut self) -> Result<UrsaStatement, ParserError> {
348 self.consume_words(&["DROP", "PROPERTY", "GRAPH"]);
349 let if_exists = self.consume_words(&["IF", "EXISTS"]);
350 let name = self.df.parser.parse_object_name(false)?;
351 Ok(UrsaStatement::DropPropertyGraph(DropPropertyGraph {
352 name,
353 if_exists,
354 }))
355 }
356
357 fn parse_attach(&mut self) -> Result<UrsaStatement, ParserError> {
358 self.consume_word("ATTACH");
359 self.consume_word("DATABASE");
360 let path = self.df.parser.parse_literal_string()?;
361 let alias = if self.consume_word("AS") {
362 Some(self.df.parser.parse_identifier()?)
363 } else {
364 None
365 };
366 Ok(UrsaStatement::Attach(Attach { path, alias }))
367 }
368}
369
370const DUCKDB_STATEMENTS: &[(&str, &str)] = &[
375 (
376 "SUMMARIZE",
377 "SUMMARIZE is not supported; use DESCRIBE t for the columns and their types, and \
378 SELECT count(*), min(c), max(c), avg(c) FROM t for a column's statistics",
379 ),
380 (
381 "PIVOT",
382 "PIVOT is not supported; use GROUP BY with one aggregate per column, as \
383 sum(x) FILTER (WHERE k = 'a') AS a",
384 ),
385 (
386 "UNPIVOT",
387 "UNPIVOT is not supported; use UNION ALL of one SELECT per column",
388 ),
389 (
390 "INSTALL",
391 "INSTALL is not supported: Ursa has no extensions; Parquet, CSV, JSON and object \
392 storage are built in",
393 ),
394 (
395 "LOAD",
396 "LOAD is not supported: Ursa has no extensions; Parquet, CSV, JSON and object \
397 storage are built in",
398 ),
399 (
400 "PRAGMA",
401 "PRAGMA is not supported; use SHOW ALL, SHOW name or SET name = value for settings, \
402 and SELECT * FROM ursa_functions() for the graph functions",
403 ),
404];
405
406const TABLESAMPLE: &str = "TABLESAMPLE is not supported (it would be ignored and every row read); \
408 use ORDER BY random() LIMIT n for a sample of n rows";
409
410const USING_SAMPLE: &str =
412 "USING SAMPLE is not supported; use ORDER BY random() LIMIT n for a sample of n rows";
413
414fn source_text(sql: &str, start: Location, end: Location) -> Option<String> {
417 let from = byte_offset(sql, start)?;
418 let to = byte_offset(sql, end)?;
419 (from < to).then(|| sql[from..to].trim().to_string())
420}
421
422pub(crate) fn byte_offset(sql: &str, at: Location) -> Option<usize> {
425 if at.line == 0 || at.column == 0 {
426 return None;
427 }
428 let mut line_start = 0;
429 for _ in 1..at.line {
430 line_start += sql[line_start..].find('\n')? + 1;
431 }
432 let line = &sql[line_start..];
433 let column = usize::try_from(at.column - 1).ok()?;
434 match line.char_indices().nth(column) {
435 Some((offset, _)) => Some(line_start + offset),
436 None if line.chars().count() == column => Some(sql.len()),
437 None => None,
438 }
439}
440
441#[cfg(test)]
442mod tests {
443 use super::*;
444
445 fn one(sql: &str) -> UrsaStatement {
446 let mut statements = parse_sql(sql).unwrap();
447 assert_eq!(statements.len(), 1, "{sql}");
448 statements.pop().unwrap()
449 }
450
451 #[test]
452 fn create_property_graph() {
453 let sql = "CREATE PROPERTY GRAPH karate
454 VERTEX TABLES (members KEY (id))
455 EDGE TABLES (
456 ties SOURCE KEY (src) REFERENCES members (id)
457 DESTINATION KEY (dst) REFERENCES members (id)
458 )";
459 let UrsaStatement::CreatePropertyGraph(g) = one(sql) else {
460 panic!("not a CREATE PROPERTY GRAPH")
461 };
462 assert_eq!(g.name.to_string(), "karate");
463 assert!(!g.or_replace && !g.if_not_exists);
464
465 let s = one("create or replace property graph if not exists db.social edge tables (e source key (a) destination key (b))");
466 let UrsaStatement::CreatePropertyGraph(g) = s else {
467 panic!("not a CREATE PROPERTY GRAPH")
468 };
469 assert_eq!(g.name.to_string(), "db.social");
470 assert!(g.or_replace && g.if_not_exists);
471 }
472
473 #[test]
474 fn drop_describe_show_attach() {
475 assert_eq!(
476 one("DROP PROPERTY GRAPH IF EXISTS karate"),
477 UrsaStatement::DropPropertyGraph(DropPropertyGraph {
478 name: ObjectName::from(vec![Ident::new("karate")]),
479 if_exists: true,
480 })
481 );
482 assert_eq!(
483 one("describe graph karate").to_string(),
484 "DESCRIBE GRAPH karate"
485 );
486 assert_eq!(one("DESC GRAPH karate").kind(), "DESCRIBE GRAPH");
487 assert_eq!(one("SHOW GRAPHS"), UrsaStatement::ShowGraphs);
488 assert_eq!(
489 one("ATTACH 'club.ursa' AS club"),
490 UrsaStatement::Attach(Attach {
491 path: "club.ursa".into(),
492 alias: Some(Ident::new("club")),
493 })
494 );
495 assert_eq!(
496 one("attach database 'club.ursa'").to_string(),
497 "ATTACH 'club.ursa'"
498 );
499 }
500
501 #[test]
502 fn everything_else_is_datafusion() {
503 for sql in [
504 "SELECT * FROM 'edges.csv' LIMIT 3",
505 "COPY (SELECT 1) TO 'out.parquet'",
506 "CREATE VIEW ties AS SELECT * FROM 'karate_edges.csv'",
507 "DESCRIBE graph",
508 "SHOW TABLES",
509 "CREATE TABLE property (graph INT)",
510 ] {
511 assert!(
512 matches!(one(sql), UrsaStatement::DataFusion(_)),
513 "{sql} should go to DataFusion"
514 );
515 }
516 }
517
518 #[test]
519 fn statements_split_on_semicolons() {
520 let statements = parse_sql(
521 "SELECT 1; ; CREATE PROPERTY GRAPH g EDGE TABLES (e SOURCE KEY (a) DESTINATION KEY (b));
522 SHOW GRAPHS; SELECT ';' AS semi;",
523 )
524 .unwrap();
525 let kinds: Vec<_> = statements.iter().map(UrsaStatement::kind).collect();
526 assert_eq!(
527 kinds,
528 [
529 "SQL statement",
530 "CREATE PROPERTY GRAPH",
531 "SHOW GRAPHS",
532 "SQL statement"
533 ]
534 );
535 }
536
537 #[test]
538 fn statements_keep_their_text() {
539 let sql =
540 "SELECT 1;\n CREATE EXTERNAL TABLE t (a INT)\n STORED AS CSV LOCATION 'é.csv' ;\
541 CREATE PROPERTY GRAPH g EDGE TABLES (e SOURCE KEY (a) DESTINATION KEY (b))";
542 let statements = UrsaParser::new(sql, &GenericDialect {}, DEFAULT_RECURSION_LIMIT)
543 .unwrap()
544 .parse_statements_with_text()
545 .unwrap();
546 let texts: Vec<_> = statements.iter().map(|(_, t)| t.as_str()).collect();
547 assert_eq!(
548 texts,
549 [
550 "SELECT 1",
551 "CREATE EXTERNAL TABLE t (a INT)\n STORED AS CSV LOCATION 'é.csv'",
552 "CREATE PROPERTY GRAPH g EDGE TABLES (e SOURCE KEY (a) DESTINATION KEY (b))",
553 ]
554 );
555 }
556
557 #[test]
558 fn missing_delimiter_is_a_syntax_error() {
559 let err = parse_sql("SHOW GRAPHS SELECT 1").unwrap_err();
560 assert!(err.to_string().contains("end of statement"), "{err}");
561 }
562
563 #[test]
564 fn tablesample_is_refused_not_ignored() {
565 for sql in [
566 "SELECT * FROM t TABLESAMPLE SYSTEM (10)",
567 "select * from t tablesample bernoulli (5 percent)",
568 ] {
569 let err = parse_sql(sql).unwrap_err().to_string();
570 assert!(err.contains("TABLESAMPLE is not supported"), "{sql}: {err}");
571 }
572 parse_sql("SELECT 'tablesample' AS x").unwrap();
574 parse_sql("SELECT \"TABLESAMPLE\" FROM t").unwrap();
575 }
576
577 #[test]
578 fn duckdb_statements_say_what_to_write_instead() {
579 let message = |sql: &str| parse_sql(sql).unwrap_err().to_string();
580 for (sql, start) in [
581 ("SUMMARIZE t", "SUMMARIZE is not supported; use DESCRIBE t"),
582 ("summarize SELECT 1", "SUMMARIZE is not supported;"),
583 (
584 "PIVOT t ON k USING sum(x)",
585 "PIVOT is not supported; use GROUP BY",
586 ),
587 ("UNPIVOT t ON a, b", "UNPIVOT is not supported;"),
588 (
589 "INSTALL httpfs",
590 "INSTALL is not supported: Ursa has no extensions",
591 ),
592 (
593 "LOAD httpfs",
594 "LOAD is not supported: Ursa has no extensions",
595 ),
596 ("PRAGMA version", "PRAGMA is not supported; use SHOW ALL"),
597 (
598 "SELECT * FROM t USING SAMPLE 10",
599 "USING SAMPLE is not supported;",
600 ),
601 (
602 "SELECT * FROM t USING SAMPLE 10%",
603 "USING SAMPLE is not supported;",
604 ),
605 (
606 "SELECT 1; SELECT * FROM (SELECT * FROM t USING SAMPLE 5 ROWS)",
607 "USING SAMPLE is not supported; use ORDER BY random() LIMIT n",
608 ),
609 ] {
610 let message = message(sql);
611 assert!(message.contains(start), "{sql}\n{message}");
612 }
613 assert!(parse_sql("SELECT 'USING SAMPLE' AS summarize").is_ok());
615 let message = message("SELEC 'USING SAMPLE'");
616 assert!(message.contains("Expected: an SQL statement"), "{message}");
617 }
618}