ursa_sql/
graph_table.rs

1//! SQL/PGQ's `GRAPH_TABLE ( … )`, rewritten before sqlparser sees it (#296).
2//!
3//! sqlparser has no SQL/PGQ grammar, and forking it is ruled out. So, as
4//! DuckPGQ puts a grammar extension in front of DuckDB's parser, Ursa rewrites
5//! each `GRAPH_TABLE ( body )` (or `GRAPH TABLE ( body )`) in a statement's
6//! text into an ordinary table function call, `graph_table('body')`, before
7//! the statement is parsed. The engine's relation planner parses the body (the
8//! graph's name, `MATCH`, the pattern, `WHERE` and `COLUMNS`) and plans it.
9//!
10//! The rewrite works on sqlparser's own tokens, so a `GRAPH_TABLE` inside a
11//! string or a comment is left alone, and the matching `)` is found by counting
12//! brackets over tokens, which strings and comments cannot unbalance. Text
13//! outside each `GRAPH_TABLE ( … )` is kept byte for byte.
14
15use datafusion::sql::sqlparser::dialect::Dialect;
16use datafusion::sql::sqlparser::tokenizer::{Token, TokenWithSpan, Tokenizer};
17
18use crate::parser::byte_offset;
19use crate::source_map::Rewritten;
20
21/// The table function a `GRAPH_TABLE ( … )` becomes.
22pub const GRAPH_TABLE_FUNCTION: &str = "graph_table";
23
24/// `sql` with every `GRAPH_TABLE ( body )` rewritten to `graph_table('body')`.
25/// Text that does not tokenize is returned as it is, for the parser to report.
26pub fn rewrite_graph_tables(sql: &str, dialect: &dyn Dialect) -> String {
27    rewrite_graph_tables_mapped(sql, dialect).text
28}
29
30/// [`rewrite_graph_tables`], with a record of what it replaced.
31pub(crate) fn rewrite_graph_tables_mapped(sql: &str, dialect: &dyn Dialect) -> Rewritten {
32    let mut out = Rewritten::new(sql.len());
33    let Ok(tokens) = Tokenizer::new(dialect, sql).tokenize_with_location() else {
34        return out.finish(sql);
35    };
36    // The tokens that are not whitespace or comments.
37    let tokens: Vec<&TokenWithSpan> = tokens
38        .iter()
39        .filter(|t| !matches!(t.token, Token::Whitespace(_)))
40        .collect();
41    let mut i = 0;
42    while i < tokens.len() {
43        let Some(open) = graph_table_at(&tokens, i) else {
44            i += 1;
45            continue;
46        };
47        let Some(close) = matching_paren(&tokens, open) else {
48            break; // unbalanced: let the parser report it
49        };
50        let (Some(start), Some(body_start), Some(body_end), Some(end)) = (
51            byte_offset(sql, tokens[i].span.start),
52            byte_offset(sql, tokens[open].span.end),
53            byte_offset(sql, tokens[close].span.start),
54            byte_offset(sql, tokens[close].span.end),
55        ) else {
56            break;
57        };
58        let body = sql[body_start..body_end].trim();
59        let call = format!("{GRAPH_TABLE_FUNCTION}('{}')", body.replace('\'', "''"));
60        out.replace(sql, start, end, &call);
61        i = close + 1;
62    }
63    out.finish(sql)
64}
65
66/// If a `GRAPH_TABLE (` or `GRAPH TABLE (` starts at token `i`, the index of
67/// its `(`.
68fn graph_table_at(tokens: &[&TokenWithSpan], i: usize) -> Option<usize> {
69    let word = |n: usize, w: &str| match tokens.get(n).map(|t| &t.token) {
70        Some(Token::Word(word)) => word.quote_style.is_none() && word.value.eq_ignore_ascii_case(w),
71        _ => false,
72    };
73    let paren = |n: usize| matches!(tokens.get(n).map(|t| &t.token), Some(Token::LParen));
74    if word(i, "GRAPH_TABLE") && paren(i + 1) {
75        Some(i + 1)
76    } else if word(i, "GRAPH") && word(i + 1, "TABLE") && paren(i + 2) {
77        Some(i + 2)
78    } else {
79        None
80    }
81}
82
83/// The index of the `)` that closes the `(` at `open`.
84fn matching_paren(tokens: &[&TokenWithSpan], open: usize) -> Option<usize> {
85    let mut depth = 0usize;
86    for (n, token) in tokens.iter().enumerate().skip(open) {
87        match token.token {
88            Token::LParen => depth += 1,
89            Token::RParen => {
90                depth -= 1;
91                if depth == 0 {
92                    return Some(n);
93                }
94            }
95            _ => {}
96        }
97    }
98    None
99}
100
101#[cfg(test)]
102mod tests {
103    use datafusion::sql::sqlparser::dialect::GenericDialect;
104
105    use super::*;
106
107    fn rewrite(sql: &str) -> String {
108        rewrite_graph_tables(sql, &GenericDialect {})
109    }
110
111    #[test]
112    fn a_graph_table_becomes_a_table_function_call() {
113        assert_eq!(
114            rewrite(
115                "SELECT * FROM GRAPH_TABLE (social MATCH (a:person)-[k:knows]->(b) COLUMNS (a.id, b.id)) t"
116            ),
117            "SELECT * FROM graph_table('social MATCH (a:person)-[k:knows]->(b) COLUMNS (a.id, b.id)') t"
118        );
119    }
120
121    #[test]
122    fn nested_brackets_quotes_and_two_word_spelling() {
123        assert_eq!(
124            rewrite(
125                "select x from graph table (g MATCH (a WHERE a.name = 'O''Brien (jr)')-[e]->(b) \
126                 WHERE (a.x + (b.y)) > 1 COLUMNS (concat(a.name, ')') AS n)) AS t WHERE x > 0"
127            ),
128            "select x from graph_table('g MATCH (a WHERE a.name = ''O''''Brien (jr)'')-[e]->(b) \
129             WHERE (a.x + (b.y)) > 1 COLUMNS (concat(a.name, '')'') AS n)') AS t WHERE x > 0"
130        );
131    }
132
133    #[test]
134    fn strings_comments_and_other_words_are_left_alone() {
135        for sql in [
136            "SELECT 'GRAPH_TABLE (g MATCH (a))' AS s",
137            "SELECT 1 -- GRAPH_TABLE (g MATCH (a))\n",
138            "SELECT 1 /* GRAPH_TABLE (g) */",
139            r#"SELECT * FROM "GRAPH_TABLE"(1)"#,
140            "SELECT graph_table FROM t",
141            "SELECT * FROM my_graph_table(1)",
142        ] {
143            assert_eq!(rewrite(sql), sql);
144        }
145    }
146
147    #[test]
148    fn every_graph_table_in_a_statement_and_in_ctes_subqueries_and_views() {
149        assert_eq!(
150            rewrite(
151                "WITH p AS (SELECT * FROM GRAPH_TABLE (g MATCH (a)-[e]->(b) COLUMNS (a.id AS a)))
152                 SELECT * FROM p JOIN (SELECT * FROM GRAPH_TABLE(g MATCH (c) COLUMNS (c.id AS a))) q USING (a)"
153            ),
154            "WITH p AS (SELECT * FROM graph_table('g MATCH (a)-[e]->(b) COLUMNS (a.id AS a)'))
155                 SELECT * FROM p JOIN (SELECT * FROM graph_table('g MATCH (c) COLUMNS (c.id AS a)')) q USING (a)"
156        );
157        assert_eq!(
158            rewrite("CREATE VIEW v AS SELECT * FROM GRAPH_TABLE (g MATCH (a) COLUMNS (a.id))"),
159            "CREATE VIEW v AS SELECT * FROM graph_table('g MATCH (a) COLUMNS (a.id)')"
160        );
161    }
162
163    #[test]
164    fn unbalanced_or_untokenizable_text_is_kept_for_the_parser() {
165        for sql in [
166            "SELECT * FROM GRAPH_TABLE (g MATCH (a)",
167            "SELECT 'unterminated FROM GRAPH_TABLE (g)",
168        ] {
169            assert_eq!(rewrite(sql), sql);
170        }
171    }
172
173    #[test]
174    fn multi_line_statements_keep_their_text() {
175        let sql = "SELECT a, b\nFROM GRAPH_TABLE (\n  g\n  MATCH (a)-[e]->(b)\n  COLUMNS (a.id AS a, b.id AS b)\n)\nORDER BY a;";
176        assert_eq!(
177            rewrite(sql),
178            "SELECT a, b\nFROM graph_table('g\n  MATCH (a)-[e]->(b)\n  COLUMNS (a.id AS a, b.id AS b)')\nORDER BY a;"
179        );
180    }
181}
ursa_sql/graph_table.rs

ursa_sql/
graph_table.rs
