Skip to content

Commit decd291

Browse files
authored
feat(rules): read custom-rule SQL syntax trees from SQLQuack parsers (#1358)
SqlDocument is native and parses with the SQLQuack parser chosen per dialect: the native PostgreSQL parser, and the shared parser for DuckDB, MotherDuck, Snowflake and SQLite. ast() returns a SqlTreeNode tree and replaces polyglot_ast(). Other dialects, syntax errors, depth limits and statements kept as unparsed text raise RuleUsageError.
1 parent 3c5db67 commit decd291

42 files changed

Lines changed: 1044 additions & 180 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎crates/sqlbuild-analysis/src/assembly/analysis_session/main/prove_dynamic_contracts.rs‎

Lines changed: 4 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -5,15 +5,11 @@ use sqlbuild_core::constants::PANIC_MESSAGE;
55
use crate::assembly::analysis_session::_helpers::dynamic_pivot::pivot_outcomes;
66
use crate::assembly::analysis_session::constants::PIVOT_WORKER_STACK_BYTES;
77
use crate::assembly::analysis_session::models::{PivotBatchRequest, PivotOutcome};
8+
use crate::worker_stack::main::run_on_worker_stack::run_on_worker_stack;
89

910
/// Each model's proof in order; an internal failure of any proof fails them all.
1011
pub fn prove_dynamic_contracts(request: &PivotBatchRequest) -> Result<Vec<PivotOutcome>, String> {
11-
std::thread::scope(|scope| {
12-
std::thread::Builder::new()
13-
.stack_size(PIVOT_WORKER_STACK_BYTES)
14-
.spawn_scoped(scope, || pivot_outcomes(&request.models, &request.tables))
15-
.map_err(|error| error.to_string())?
16-
.join()
17-
.map_err(|_| PANIC_MESSAGE.to_owned())?
18-
})
12+
run_on_worker_stack(PIVOT_WORKER_STACK_BYTES, PANIC_MESSAGE, || {
13+
pivot_outcomes(&request.models, &request.tables)
14+
})?
1915
}

‎crates/sqlbuild-analysis/src/assembly/analysis_session/main/recover_cte_facts.rs‎

Lines changed: 6 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -5,17 +5,15 @@ use crate::assembly::analysis_session::_helpers::cte_facts::{
55
};
66
use crate::assembly::analysis_session::constants::CTE_FACT_WORKER_STACK_BYTES;
77
use crate::assembly::analysis_session::models::{CteFactRequest, CteFacts};
8+
use crate::worker_stack::main::run_on_worker_stack::run_on_worker_stack;
89

910
/// The facts enrichment recovers for the query, or why they are left to Python.
1011
pub fn recover_cte_facts(request: &CteFactRequest) -> Result<CteFacts, String> {
11-
std::thread::scope(|scope| {
12-
std::thread::Builder::new()
13-
.stack_size(CTE_FACT_WORKER_STACK_BYTES)
14-
.spawn_scoped(scope, || recovered_facts(request))
15-
.map_err(|error| error.to_string())?
16-
.join()
17-
.map_err(|_| "native CTE fact recovery panicked".to_owned())?
18-
})
12+
run_on_worker_stack(
13+
CTE_FACT_WORKER_STACK_BYTES,
14+
"native CTE fact recovery panicked",
15+
|| recovered_facts(request),
16+
)?
1917
}
2018

2119
fn recovered_facts(request: &CteFactRequest) -> Result<CteFacts, String> {

‎crates/sqlbuild-analysis/src/lib.rs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,5 +48,7 @@ pub mod snowflake_cost;
4848
pub mod sql_macros;
4949
pub mod sql_test_attachment;
5050
pub mod sql_tokens;
51+
pub mod sql_tree;
5152
pub mod target_paths;
5253
pub mod type_system;
54+
pub mod worker_stack;
Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,67 @@
1+
//! Join and CTE clauses, which SQLQuack keeps as statement fields rather than expressions.
2+
3+
use sqlquack_syntax::Expression;
4+
use sqlquack_syntax::expressions::queries::models::{Cte, Join, JoinKind, With};
5+
use sqlquack_syntax::tokens::models::Span;
6+
7+
use crate::sql_tree::constants::{CTE_KIND, JOIN_KIND};
8+
9+
/// One clause of a statement: its kind, its span, and the expressions it owns.
10+
pub(crate) struct Clause<'ast> {
11+
pub(crate) kind: &'static str,
12+
pub(crate) span: Option<Span>,
13+
pub(crate) children: Vec<&'ast Expression>,
14+
}
15+
16+
/// The CTEs and explicit joins `expression` holds directly, CTEs first.
17+
pub(crate) fn clauses(expression: &Expression) -> Vec<Clause<'_>> {
18+
let (with, joins): (Option<&With>, Vec<&[Join]>) = match expression {
19+
Expression::Select(select) => (select.with.as_ref(), vec![select.joins.as_slice()]),
20+
Expression::Union(union) => (union.with.as_ref(), Vec::new()),
21+
Expression::Intersect(intersect) => (intersect.with.as_ref(), Vec::new()),
22+
Expression::Except(except) => (except.with.as_ref(), Vec::new()),
23+
Expression::Insert(insert) => (insert.with.as_ref(), Vec::new()),
24+
Expression::Update(update) => (
25+
update.with.as_ref(),
26+
vec![update.table_joins.as_slice(), update.from_joins.as_slice()],
27+
),
28+
Expression::Delete(delete) => (delete.with.as_ref(), vec![delete.joins.as_slice()]),
29+
Expression::JoinedTable(joined) => (None, vec![joined.joins.as_slice()]),
30+
_ => (None, Vec::new()),
31+
};
32+
let ctes = with
33+
.into_iter()
34+
.flat_map(|with| with.ctes.iter())
35+
.map(cte_clause);
36+
let joins = joins
37+
.into_iter()
38+
.flatten()
39+
.filter(|join| !is_comma_item(join))
40+
.map(join_clause);
41+
ctes.chain(joins).collect()
42+
}
43+
44+
fn cte_clause(cte: &Cte) -> Clause<'_> {
45+
Clause {
46+
kind: CTE_KIND,
47+
span: cte.span.0,
48+
children: vec![&cte.this],
49+
}
50+
}
51+
52+
fn join_clause(join: &Join) -> Clause<'_> {
53+
let mut children: Vec<&Expression> = vec![&join.this];
54+
children.extend(join.on.as_ref());
55+
children.extend(join.match_condition.as_ref());
56+
children.extend(join.pivots.iter());
57+
Clause {
58+
kind: JOIN_KIND,
59+
span: join.span.0,
60+
children,
61+
}
62+
}
63+
64+
/// A comma-separated FROM item after an explicit join, which the shared parser keeps as a join.
65+
fn is_comma_item(join: &Join) -> bool {
66+
join.kind == JoinKind::Implicit
67+
}
Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
//! Parser selection, parsing, the tree walk, clause nodes and source positions.
2+
3+
pub(crate) mod clauses;
4+
pub(crate) mod parse;
5+
pub(crate) mod positions;
6+
pub(crate) mod rules_tree_parser;
7+
pub(crate) mod walk;
Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
//! Parse SQL into statements with node spans.
2+
3+
use sqlquack_syntax::{CallMode, Expression, ParseOptions};
4+
5+
use crate::sql_tree::constants::GUARD_CODE_PREFIX;
6+
use crate::sql_tree::errors::SqlTreeError;
7+
use crate::sql_tree::models::TreeParser;
8+
9+
/// The statements of `sql` under `parser`; a non-syntax PostgreSQL failure or a skipped statement is unanalysable.
10+
pub(crate) fn parse(sql: &str, parser: TreeParser) -> Result<Vec<Expression>, SqlTreeError> {
11+
match parser {
12+
TreeParser::Shared(dialect) => sqlquack_syntax::parse_statements_with(
13+
sql,
14+
dialect,
15+
&ParseOptions::default(),
16+
CallMode::FaithfulLocated,
17+
)
18+
.map_err(|error| shared_error(&error.to_string())),
19+
TreeParser::Postgres => {
20+
let parsed = sqlquack_syntax_postgres::parse_strict(sql).map_err(|failure| {
21+
match failure.parse_error() {
22+
Some(error) if error.syntax => SqlTreeError::Parse(error.message.clone()),
23+
Some(error) => SqlTreeError::Unanalysable(error.message.clone()),
24+
None => SqlTreeError::Parse("SQL has text outside any token".to_owned()),
25+
}
26+
})?;
27+
if !parsed.gaps.is_empty() {
28+
return Err(SqlTreeError::Unanalysable(
29+
"a statement uses syntax the PostgreSQL parser skips".to_owned(),
30+
));
31+
}
32+
Ok(parsed.statements)
33+
}
34+
}
35+
}
36+
37+
/// A shared-parser failure: a guard limit, such as nesting depth, is unanalysable, not a syntax error.
38+
fn shared_error(message: &str) -> SqlTreeError {
39+
if message.contains(GUARD_CODE_PREFIX) {
40+
SqlTreeError::Unanalysable(message.to_owned())
41+
} else {
42+
SqlTreeError::Parse(message.to_owned())
43+
}
44+
}
Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
//! Character offsets as 1-based lines and columns.
2+
3+
/// The character offset at which each line of `sql` starts.
4+
pub(crate) fn line_starts(sql: &str) -> Vec<usize> {
5+
let mut starts: Vec<usize> = vec![0];
6+
for (offset, character) in sql.chars().enumerate() {
7+
if character == '\n' {
8+
starts.push(offset + 1);
9+
}
10+
}
11+
starts
12+
}
13+
14+
/// The 1-based line and column of the character at `offset`.
15+
pub(crate) fn line_column(starts: &[usize], offset: usize) -> (usize, usize) {
16+
let line: usize = starts.partition_point(|start| *start <= offset);
17+
let start: usize = starts.get(line.saturating_sub(1)).copied().unwrap_or(0);
18+
(line.max(1), offset - start + 1)
19+
}
Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,23 @@
1+
//! `rules_tree_parser`: the parser that reads a dialect's SQL for rules.
2+
3+
use sqlquack_syntax::Dialect;
4+
5+
use crate::sql_tree::constants::{
6+
DUCKDB_DIALECTS, POSTGRES_DIALECTS, SNOWFLAKE_DIALECT, SQLITE_DIALECT,
7+
};
8+
use crate::sql_tree::models::TreeParser;
9+
10+
/// The parser for `dialect`, or `None`; add a dialect (one line) only once its parser passes the depth probe with an empty allowlist.
11+
pub(crate) fn rules_tree_parser(dialect: &str) -> Option<TreeParser> {
12+
if POSTGRES_DIALECTS.contains(&dialect) {
13+
return Some(TreeParser::Postgres);
14+
}
15+
if DUCKDB_DIALECTS.contains(&dialect) {
16+
return Some(TreeParser::Shared(Dialect::DuckDB));
17+
}
18+
match dialect {
19+
SNOWFLAKE_DIALECT => Some(TreeParser::Shared(Dialect::Snowflake)),
20+
SQLITE_DIALECT => Some(TreeParser::Shared(Dialect::SQLite)),
21+
_ => None,
22+
}
23+
}
Lines changed: 102 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,102 @@
1+
//! The iterative pre-order walk that turns parsed statements into rule tree nodes.
2+
3+
use std::collections::HashMap;
4+
5+
use sqlquack_syntax::tokens::models::Span;
6+
use sqlquack_syntax::{Expression, ExpressionWalk};
7+
8+
use crate::sql_tree::_helpers::clauses::clauses;
9+
use crate::sql_tree::_helpers::positions::line_column;
10+
use crate::sql_tree::constants::OPAQUE_KINDS;
11+
use crate::sql_tree::errors::SqlTreeError;
12+
use crate::sql_tree::models::SqlTreeNode;
13+
14+
/// Child expressions mapped to the index of the clause node that owns them.
15+
type Owners = HashMap<*const Expression, usize>;
16+
17+
/// The nodes of `statements` in pre-order, walked with an explicit stack so depth cannot overflow.
18+
pub(crate) fn walk(
19+
statements: &[Expression],
20+
starts: &[usize],
21+
) -> Result<Vec<SqlTreeNode>, SqlTreeError> {
22+
let mut nodes: Vec<SqlTreeNode> = Vec::new();
23+
let mut stack: Vec<(&Expression, Option<usize>)> = statements
24+
.iter()
25+
.rev()
26+
.map(|statement| (statement, None))
27+
.collect();
28+
while let Some((expression, parent)) = stack.pop() {
29+
let index: usize = nodes.len();
30+
let node: SqlTreeNode = expression_node(&nodes, expression, parent, starts)?;
31+
let (clause_nodes, owners) = clause_nodes(expression, &node, index, starts);
32+
nodes.push(node);
33+
nodes.extend(clause_nodes);
34+
for child in expression.children().into_iter().rev() {
35+
let owner: usize = owners
36+
.get(&std::ptr::from_ref(child))
37+
.copied()
38+
.unwrap_or(index);
39+
stack.push((child, Some(owner)));
40+
}
41+
}
42+
Ok(nodes)
43+
}
44+
45+
/// `expression`'s node, at its own start or else at its parent's.
46+
fn expression_node(
47+
nodes: &[SqlTreeNode],
48+
expression: &Expression,
49+
parent: Option<usize>,
50+
starts: &[usize],
51+
) -> Result<SqlTreeNode, SqlTreeError> {
52+
let kind: &'static str = expression.variant_name();
53+
if OPAQUE_KINDS.contains(&kind) {
54+
return Err(SqlTreeError::Unanalysable(
55+
"a statement uses syntax the parser keeps as unparsed text".to_owned(),
56+
));
57+
}
58+
let inherited = parent.map_or((1, 1), |parent| (nodes[parent].line, nodes[parent].column));
59+
let (line, column) =
60+
own_span(expression).map_or(inherited, |span| line_column(starts, span.start));
61+
Ok(SqlTreeNode {
62+
kind,
63+
line,
64+
column,
65+
parent,
66+
})
67+
}
68+
69+
/// The clause nodes `expression` holds, numbered after `index`, and the children each one owns.
70+
fn clause_nodes(
71+
expression: &Expression,
72+
node: &SqlTreeNode,
73+
index: usize,
74+
starts: &[usize],
75+
) -> (Vec<SqlTreeNode>, Owners) {
76+
let mut built: Vec<SqlTreeNode> = Vec::new();
77+
let mut owners: Owners = HashMap::new();
78+
for clause in clauses(expression) {
79+
let clause_index: usize = index + 1 + built.len();
80+
let (line, column) = clause.span.map_or((node.line, node.column), |span| {
81+
line_column(starts, span.start)
82+
});
83+
built.push(SqlTreeNode {
84+
kind: clause.kind,
85+
line,
86+
column,
87+
parent: Some(index),
88+
});
89+
for child in clause.children {
90+
let _ = owners.insert(std::ptr::from_ref(child), clause_index);
91+
}
92+
}
93+
(built, owners)
94+
}
95+
96+
/// The node's span; a qualified star without one starts at its qualifier.
97+
fn own_span(expression: &Expression) -> Option<Span> {
98+
expression.source_span().or_else(|| match expression {
99+
Expression::Star(star) => star.table.as_ref().and_then(|table| table.span),
100+
_ => None,
101+
})
102+
}
Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
//! Node kinds the rule tree adds to SQLQuack's expression variants, and kinds it refuses.
2+
3+
/// The kind of a join clause, which SQLQuack keeps as a field of its statement.
4+
pub(crate) const JOIN_KIND: &str = "join";
5+
/// The kind of a common table expression, which SQLQuack keeps as a field of its statement.
6+
pub(crate) const CTE_KIND: &str = "cte";
7+
/// Kinds that hold a statement the parser kept as unparsed text, so no rule can see inside it.
8+
pub(crate) const OPAQUE_KINDS: [&str; 2] = ["command", "raw"];
9+
/// Dialect names read with the shared DuckDB grammar; `generic` matches the fingerprint tokenizer.
10+
pub(crate) const DUCKDB_DIALECTS: [&str; 3] = ["duckdb", "motherduck", "generic"];
11+
/// Dialect names read with the native PostgreSQL parser.
12+
pub(crate) const POSTGRES_DIALECTS: [&str; 2] = ["postgres", "postgresql"];
13+
/// The Snowflake dialect name.
14+
pub(crate) const SNOWFLAKE_DIALECT: &str = "snowflake";
15+
/// The SQLite dialect name.
16+
pub(crate) const SQLITE_DIALECT: &str = "sqlite";
17+
/// Stack for the parse-and-walk worker, so deep SQL reaches the parser's depth limit, not a stack overflow.
18+
pub(crate) const TREE_WORKER_STACK_BYTES: usize = 16 * 1024 * 1024;
19+
/// The prefix of the shared parser's guard codes, such as `E_GUARD_NESTING_DEPTH_EXCEEDED`.
20+
pub(crate) const GUARD_CODE_PREFIX: &str = "E_GUARD_";

0 commit comments

Comments
 (0)