datafusion-query-builder 0.1.2__tar.gz → 0.3.1__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {datafusion_query_builder-0.1.2 → datafusion_query_builder-0.3.1}/Cargo.lock +1 -1
- {datafusion_query_builder-0.1.2 → datafusion_query_builder-0.3.1}/Cargo.toml +1 -1
- {datafusion_query_builder-0.1.2 → datafusion_query_builder-0.3.1}/PKG-INFO +12 -1
- {datafusion_query_builder-0.1.2 → datafusion_query_builder-0.3.1}/README.md +11 -0
- {datafusion_query_builder-0.1.2 → datafusion_query_builder-0.3.1}/pyproject.toml +1 -1
- {datafusion_query_builder-0.1.2 → datafusion_query_builder-0.3.1}/python/datafusion_query_builder/__init__.pyi +6 -0
- datafusion_query_builder-0.3.1/src/dialect.rs +54 -0
- {datafusion_query_builder-0.1.2 → datafusion_query_builder-0.3.1}/src/expr.rs +7 -0
- {datafusion_query_builder-0.1.2 → datafusion_query_builder-0.3.1}/src/lib.rs +1 -0
- {datafusion_query_builder-0.1.2 → datafusion_query_builder-0.3.1}/src/lower.rs +168 -56
- {datafusion_query_builder-0.1.2 → datafusion_query_builder-0.3.1}/src/python.rs +34 -4
- {datafusion_query_builder-0.1.2 → datafusion_query_builder-0.3.1}/src/query.rs +14 -0
- {datafusion_query_builder-0.1.2 → datafusion_query_builder-0.3.1}/src/render.rs +4 -32
- {datafusion_query_builder-0.1.2 → datafusion_query_builder-0.3.1}/tests/core.rs +242 -1
- {datafusion_query_builder-0.1.2 → datafusion_query_builder-0.3.1}/tests/properties.proptest-regressions +1 -0
- {datafusion_query_builder-0.1.2 → datafusion_query_builder-0.3.1}/tests/properties.rs +167 -0
- {datafusion_query_builder-0.1.2 → datafusion_query_builder-0.3.1}/tests/test_python.py +67 -0
- {datafusion_query_builder-0.1.2 → datafusion_query_builder-0.3.1}/.github/workflows/ci.yml +0 -0
- {datafusion_query_builder-0.1.2 → datafusion_query_builder-0.3.1}/.gitignore +0 -0
- {datafusion_query_builder-0.1.2 → datafusion_query_builder-0.3.1}/LICENSE +0 -0
- {datafusion_query_builder-0.1.2 → datafusion_query_builder-0.3.1}/python/datafusion_query_builder/__init__.py +0 -0
- {datafusion_query_builder-0.1.2 → datafusion_query_builder-0.3.1}/python/datafusion_query_builder/py.typed +0 -0
- {datafusion_query_builder-0.1.2 → datafusion_query_builder-0.3.1}/src/functions.rs +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: datafusion-query-builder
|
|
3
|
-
Version: 0.1
|
|
3
|
+
Version: 0.3.1
|
|
4
4
|
Classifier: Development Status :: 3 - Alpha
|
|
5
5
|
Classifier: Intended Audience :: Developers
|
|
6
6
|
Classifier: License :: OSI Approved :: MIT License
|
|
@@ -64,6 +64,17 @@ literals are escaped — values are injection-safe by construction. Reach for `r
|
|
|
64
64
|
fragment), `f.call("name", ...)` (any function), or `param("name")` (a `${name}` placeholder) when
|
|
65
65
|
you step outside the v1 grammar.
|
|
66
66
|
|
|
67
|
+
A `raw(...)` fragment must be one complete SQL expression. The builder parses it on its own, drops
|
|
68
|
+
any comments, and adds brackets where its top-level operator binds looser than the surrounding
|
|
69
|
+
one, as it does for native expressions. Pass a user-written predicate as `raw(predicate)`, not
|
|
70
|
+
`raw(f"({predicate})")`. Text wrapping breaks on a trailing `-- comment`, and a predicate such as
|
|
71
|
+
`1=1) OR (1=1` can close the bracket.
|
|
72
|
+
|
|
73
|
+
The JSONB key-exists operators are first-class: `col("attributes").has_key("gen_ai.input.messages")`
|
|
74
|
+
renders `attributes ? 'gen_ai.input.messages'` — a cheap presence check that never extracts the
|
|
75
|
+
value (contrast `raw("attributes ->> 'k'").is_not_null()`, which reads the whole value out just to
|
|
76
|
+
test it). `has_any_key` and `has_all_keys` render `?|` / `?&`.
|
|
77
|
+
|
|
67
78
|
## Architecture
|
|
68
79
|
|
|
69
80
|
```
|
|
@@ -44,6 +44,17 @@ literals are escaped — values are injection-safe by construction. Reach for `r
|
|
|
44
44
|
fragment), `f.call("name", ...)` (any function), or `param("name")` (a `${name}` placeholder) when
|
|
45
45
|
you step outside the v1 grammar.
|
|
46
46
|
|
|
47
|
+
A `raw(...)` fragment must be one complete SQL expression. The builder parses it on its own, drops
|
|
48
|
+
any comments, and adds brackets where its top-level operator binds looser than the surrounding
|
|
49
|
+
one, as it does for native expressions. Pass a user-written predicate as `raw(predicate)`, not
|
|
50
|
+
`raw(f"({predicate})")`. Text wrapping breaks on a trailing `-- comment`, and a predicate such as
|
|
51
|
+
`1=1) OR (1=1` can close the bracket.
|
|
52
|
+
|
|
53
|
+
The JSONB key-exists operators are first-class: `col("attributes").has_key("gen_ai.input.messages")`
|
|
54
|
+
renders `attributes ? 'gen_ai.input.messages'` — a cheap presence check that never extracts the
|
|
55
|
+
value (contrast `raw("attributes ->> 'k'").is_not_null()`, which reads the whole value out just to
|
|
56
|
+
test it). `has_any_key` and `has_all_keys` render `?|` / `?&`.
|
|
57
|
+
|
|
47
58
|
## Architecture
|
|
48
59
|
|
|
49
60
|
```
|
|
@@ -18,6 +18,11 @@ class Expr:
|
|
|
18
18
|
def is_null(self) -> Expr: ...
|
|
19
19
|
def is_not_null(self) -> Expr: ...
|
|
20
20
|
def between(self, low: IntoExpr, high: IntoExpr) -> Expr: ...
|
|
21
|
+
# JSONB key-exists operators (`?` / `?|` / `?&`). `has_key` is a cheap presence check that,
|
|
22
|
+
# unlike `col ->> 'k' is not None`, never extracts the value.
|
|
23
|
+
def has_key(self, key: IntoExpr) -> Expr: ...
|
|
24
|
+
def has_any_key(self, keys: Sequence[IntoExpr]) -> Expr: ...
|
|
25
|
+
def has_all_keys(self, keys: Sequence[IntoExpr]) -> Expr: ...
|
|
21
26
|
def asc(self, nulls_first: bool | None = ...) -> SortExpr: ...
|
|
22
27
|
def desc(self, nulls_first: bool | None = ...) -> SortExpr: ...
|
|
23
28
|
# Valid only on a function-call expression:
|
|
@@ -79,6 +84,7 @@ class Query:
|
|
|
79
84
|
) -> Query: ...
|
|
80
85
|
def cross_join(self, other: str | Query, alias: str | None = ...) -> Query: ...
|
|
81
86
|
def with_cte(self, name: str, query: Query) -> Query: ...
|
|
87
|
+
def with_recursive_cte(self, name: str, query: Query) -> Query: ...
|
|
82
88
|
def order_by(self, *sorts: IntoSort) -> Query: ...
|
|
83
89
|
def limit(self, n: int) -> Query: ...
|
|
84
90
|
def offset(self, n: int) -> Query: ...
|
|
@@ -0,0 +1,54 @@
|
|
|
1
|
+
//! The single SQL dialect the builder parses against.
|
|
2
|
+
//!
|
|
3
|
+
//! Two call sites re-parse SQL: [`crate::lower`] parses `raw(...)` fragments and cast-type strings
|
|
4
|
+
//! into AST nodes, and [`crate::render`] re-parses the fully rendered query as a self-test. Both use
|
|
5
|
+
//! this one dialect so a `raw` fragment round-trips through the *same* grammar the final query is
|
|
6
|
+
//! validated against — otherwise a fragment could parse one way going in and fail (or mean something
|
|
7
|
+
//! else) coming out.
|
|
8
|
+
//!
|
|
9
|
+
//! It is `GenericDialect` plus the handful of capabilities DataFusion's own parser accepts that the
|
|
10
|
+
//! generic dialect gates off by default.
|
|
11
|
+
|
|
12
|
+
use sqlparser::dialect::{Dialect, GenericDialect};
|
|
13
|
+
|
|
14
|
+
/// Generic SQL plus the DataFusion-flavoured extensions the builder relies on.
|
|
15
|
+
#[derive(Debug, Default)]
|
|
16
|
+
pub(crate) struct BuilderDialect;
|
|
17
|
+
|
|
18
|
+
impl Dialect for BuilderDialect {
|
|
19
|
+
fn is_identifier_start(&self, ch: char) -> bool {
|
|
20
|
+
GenericDialect {}.is_identifier_start(ch)
|
|
21
|
+
}
|
|
22
|
+
|
|
23
|
+
fn is_identifier_part(&self, ch: char) -> bool {
|
|
24
|
+
GenericDialect {}.is_identifier_part(ch)
|
|
25
|
+
}
|
|
26
|
+
|
|
27
|
+
/// `${var}` dollar-brace placeholders (the pydantic `sqlparser` fork extension the DataFusion
|
|
28
|
+
/// planner relies on).
|
|
29
|
+
fn supports_dollar_placeholder(&self) -> bool {
|
|
30
|
+
true
|
|
31
|
+
}
|
|
32
|
+
|
|
33
|
+
// Capabilities DataFusion's parser accepts that the generic dialect gates off by default.
|
|
34
|
+
fn supports_filter_during_aggregation(&self) -> bool {
|
|
35
|
+
true
|
|
36
|
+
}
|
|
37
|
+
|
|
38
|
+
fn supports_group_by_expr(&self) -> bool {
|
|
39
|
+
true
|
|
40
|
+
}
|
|
41
|
+
|
|
42
|
+
/// Parse `?` / `?|` / `?&` as the Postgres-style JSON key-exists operators (the same ones our
|
|
43
|
+
/// engine uses), rather than as prepared-statement `?` placeholders.
|
|
44
|
+
///
|
|
45
|
+
/// The method name is an `sqlparser` quirk, not our intent: the fork happens to gate `?`-family
|
|
46
|
+
/// tokenization behind `supports_geometric_types` — it's the single lever that flips `?` from a
|
|
47
|
+
/// placeholder to `Token::Question`, and the real Postgres dialect enables the operators through
|
|
48
|
+
/// the same switch. This has nothing to do with geometry; we're here only for the JSON operators.
|
|
49
|
+
/// With it off, `attributes ? 'key'` tokenizes `?` as a placeholder, `parse_expr` stops at
|
|
50
|
+
/// `attributes`, and the operator is silently dropped.
|
|
51
|
+
fn supports_geometric_types(&self) -> bool {
|
|
52
|
+
true
|
|
53
|
+
}
|
|
54
|
+
}
|
|
@@ -36,6 +36,13 @@ pub enum BinaryOp {
|
|
|
36
36
|
And,
|
|
37
37
|
Or,
|
|
38
38
|
StringConcat,
|
|
39
|
+
/// JSONB key-exists `?`: does the string on the right exist as a top-level key/element of the
|
|
40
|
+
/// JSON on the left. DataFusion's JSON extension maps it to `json_contains`.
|
|
41
|
+
JsonExists,
|
|
42
|
+
/// JSONB `?|`: does *any* string in the right-hand array exist as a top-level key/element.
|
|
43
|
+
JsonExistsAny,
|
|
44
|
+
/// JSONB `?&`: do *all* strings in the right-hand array exist as top-level keys/elements.
|
|
45
|
+
JsonExistsAll,
|
|
39
46
|
}
|
|
40
47
|
|
|
41
48
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
|
@@ -6,10 +6,10 @@
|
|
|
6
6
|
//! high-churn nodes are isolated into the small `lower_*` helpers below for exactly that reason.
|
|
7
7
|
|
|
8
8
|
use sqlparser::ast;
|
|
9
|
-
use sqlparser::dialect::GenericDialect;
|
|
10
9
|
use sqlparser::parser::Parser;
|
|
11
10
|
use sqlparser::tokenizer::Span;
|
|
12
11
|
|
|
12
|
+
use crate::dialect::BuilderDialect;
|
|
13
13
|
use crate::expr::{BinaryOp, Call, Expr, Scalar, SortExpr, UnaryOp, Window};
|
|
14
14
|
use crate::query::{Body, BuildError, Cte, Join, JoinKind, Query, Result, Select, SetOp, TableRef};
|
|
15
15
|
|
|
@@ -78,7 +78,7 @@ fn string_literal(s: &str) -> ast::Expr {
|
|
|
78
78
|
}
|
|
79
79
|
|
|
80
80
|
fn parse_data_type(text: &str) -> Result<ast::DataType> {
|
|
81
|
-
let dialect =
|
|
81
|
+
let dialect = BuilderDialect;
|
|
82
82
|
let mut parser = Parser::new(&dialect)
|
|
83
83
|
.try_with_sql(text)
|
|
84
84
|
.map_err(|e| BuildError::UnparsableSql(format!("invalid cast type {text:?}: {e}")))?;
|
|
@@ -95,37 +95,131 @@ fn parse_data_type(text: &str) -> Result<ast::DataType> {
|
|
|
95
95
|
Ok(data_type)
|
|
96
96
|
}
|
|
97
97
|
|
|
98
|
-
|
|
99
|
-
|
|
100
|
-
|
|
101
|
-
|
|
98
|
+
// Binding tightness tiers (higher binds tighter), used to decide where to insert parentheses.
|
|
99
|
+
// `sqlparser`'s `Display` is purely structural — it never adds precedence parens — so the builder
|
|
100
|
+
// must wrap operands itself or a tree like `a / (b - c)` would render as the semantically different
|
|
101
|
+
// `a / b - c`.
|
|
102
|
+
const OR_PREC: u8 = 10;
|
|
103
|
+
const AND_PREC: u8 = 20;
|
|
104
|
+
const NOT_PREC: u8 = 30;
|
|
105
|
+
/// `IS [NOT] NULL`, `IN`, `BETWEEN`, `LIKE` and friends. Dialects disagree on where these sit
|
|
106
|
+
/// relative to comparisons (Postgres binds `LIKE` tighter than `=`, `GenericDialect` looser), so
|
|
107
|
+
/// they rank just below comparisons: always wrapped inside a comparison or arithmetic operand,
|
|
108
|
+
/// never wrapped under `AND` / `OR` / `NOT`.
|
|
109
|
+
const PREDICATE_PREC: u8 = 35;
|
|
110
|
+
const COMPARISON_PREC: u8 = 40;
|
|
111
|
+
/// Operators outside the builder's own set (`->>`, `@>`, bitwise, custom). Postgres binds these
|
|
112
|
+
/// tighter than comparisons and looser than arithmetic.
|
|
113
|
+
const OTHER_OP_PREC: u8 = 45;
|
|
114
|
+
const ADDITIVE_PREC: u8 = 50;
|
|
115
|
+
const MULTIPLICATIVE_PREC: u8 = 60;
|
|
116
|
+
const AT_TIME_ZONE_PREC: u8 = 65;
|
|
117
|
+
const PREFIX_PREC: u8 = 70;
|
|
118
|
+
const ATOM_PREC: u8 = u8::MAX;
|
|
119
|
+
|
|
102
120
|
fn binary_op_prec(op: BinaryOp) -> u8 {
|
|
103
121
|
match op {
|
|
104
|
-
BinaryOp::Or =>
|
|
105
|
-
BinaryOp::And =>
|
|
122
|
+
BinaryOp::Or => OR_PREC,
|
|
123
|
+
BinaryOp::And => AND_PREC,
|
|
106
124
|
BinaryOp::Eq
|
|
107
125
|
| BinaryOp::NotEq
|
|
108
126
|
| BinaryOp::Lt
|
|
109
127
|
| BinaryOp::LtEq
|
|
110
128
|
| BinaryOp::Gt
|
|
111
|
-
| BinaryOp::GtEq =>
|
|
112
|
-
|
|
113
|
-
|
|
129
|
+
| BinaryOp::GtEq => COMPARISON_PREC,
|
|
130
|
+
// JSONB key-exists operators return a boolean and are typically combined with AND/OR/NOT;
|
|
131
|
+
// one comparison-level tier keeps `(a ? 'x') AND (b ? 'y')` and `NOT (a ? 'x')` grouped
|
|
132
|
+
// correctly, which is all realistic usage needs.
|
|
133
|
+
BinaryOp::JsonExists | BinaryOp::JsonExistsAny | BinaryOp::JsonExistsAll => COMPARISON_PREC,
|
|
134
|
+
BinaryOp::Plus | BinaryOp::Minus | BinaryOp::StringConcat => ADDITIVE_PREC,
|
|
135
|
+
BinaryOp::Multiply | BinaryOp::Divide | BinaryOp::Modulo => MULTIPLICATIVE_PREC,
|
|
114
136
|
}
|
|
115
137
|
}
|
|
116
138
|
|
|
117
|
-
/// Binding tightness of an expression
|
|
118
|
-
///
|
|
119
|
-
|
|
139
|
+
/// Binding tightness of an already-lowered expression.
|
|
140
|
+
///
|
|
141
|
+
/// This reads the `sqlparser` AST rather than the façade, because a `raw(...)` fragment is opaque
|
|
142
|
+
/// until it is parsed: `raw("a = 1 OR b = 2")` is an `OR` and must be wrapped like one, or
|
|
143
|
+
/// `raw("a = 1 OR b = 2") & c` renders as `a = 1 OR b = 2 AND c`. Anything not recognized as an
|
|
144
|
+
/// atom ranks lowest, so an unfamiliar node gets redundant parentheses rather than a wrong grouping.
|
|
145
|
+
fn ast_prec(expr: &ast::Expr) -> u8 {
|
|
146
|
+
use ast::BinaryOperator as B;
|
|
147
|
+
use ast::Expr as E;
|
|
120
148
|
match expr {
|
|
121
|
-
|
|
122
|
-
|
|
123
|
-
|
|
124
|
-
|
|
125
|
-
|
|
126
|
-
|
|
127
|
-
|
|
128
|
-
|
|
149
|
+
E::BinaryOp { op, .. } => match op {
|
|
150
|
+
B::Or => OR_PREC,
|
|
151
|
+
B::And => AND_PREC,
|
|
152
|
+
B::Eq
|
|
153
|
+
| B::NotEq
|
|
154
|
+
| B::Lt
|
|
155
|
+
| B::LtEq
|
|
156
|
+
| B::Gt
|
|
157
|
+
| B::GtEq
|
|
158
|
+
| B::Spaceship
|
|
159
|
+
| B::Question
|
|
160
|
+
| B::QuestionAnd
|
|
161
|
+
| B::QuestionPipe => COMPARISON_PREC,
|
|
162
|
+
B::Plus | B::Minus | B::StringConcat => ADDITIVE_PREC,
|
|
163
|
+
B::Multiply | B::Divide | B::Modulo => MULTIPLICATIVE_PREC,
|
|
164
|
+
_ => OTHER_OP_PREC,
|
|
165
|
+
},
|
|
166
|
+
E::UnaryOp {
|
|
167
|
+
op: ast::UnaryOperator::Not,
|
|
168
|
+
..
|
|
169
|
+
}
|
|
170
|
+
| E::Exists { negated: true, .. } => NOT_PREC,
|
|
171
|
+
E::UnaryOp { .. } => PREFIX_PREC,
|
|
172
|
+
E::IsFalse(_)
|
|
173
|
+
| E::IsNotFalse(_)
|
|
174
|
+
| E::IsTrue(_)
|
|
175
|
+
| E::IsNotTrue(_)
|
|
176
|
+
| E::IsNull(_)
|
|
177
|
+
| E::IsNotNull(_)
|
|
178
|
+
| E::IsUnknown(_)
|
|
179
|
+
| E::IsNotUnknown(_)
|
|
180
|
+
| E::IsDistinctFrom(..)
|
|
181
|
+
| E::IsNotDistinctFrom(..)
|
|
182
|
+
| E::IsNormalized { .. }
|
|
183
|
+
| E::InList { .. }
|
|
184
|
+
| E::InSubquery { .. }
|
|
185
|
+
| E::InUnnest { .. }
|
|
186
|
+
| E::Between { .. }
|
|
187
|
+
| E::Like { .. }
|
|
188
|
+
| E::ILike { .. }
|
|
189
|
+
| E::SimilarTo { .. }
|
|
190
|
+
| E::RLike { .. }
|
|
191
|
+
| E::AnyOp { .. }
|
|
192
|
+
| E::AllOp { .. }
|
|
193
|
+
| E::MemberOf(_) => PREDICATE_PREC,
|
|
194
|
+
E::AtTimeZone { .. } => AT_TIME_ZONE_PREC,
|
|
195
|
+
E::Identifier(_)
|
|
196
|
+
| E::CompoundIdentifier(_)
|
|
197
|
+
| E::CompoundFieldAccess { .. }
|
|
198
|
+
| E::JsonAccess { .. }
|
|
199
|
+
| E::Value(_)
|
|
200
|
+
| E::TypedString { .. }
|
|
201
|
+
| E::Function(_)
|
|
202
|
+
| E::Cast { .. }
|
|
203
|
+
| E::Convert { .. }
|
|
204
|
+
| E::Case { .. }
|
|
205
|
+
| E::Nested(_)
|
|
206
|
+
| E::Tuple(_)
|
|
207
|
+
| E::Array(_)
|
|
208
|
+
| E::Subquery(_)
|
|
209
|
+
| E::Exists { negated: false, .. }
|
|
210
|
+
| E::Interval(_)
|
|
211
|
+
| E::Extract { .. }
|
|
212
|
+
| E::Ceil { .. }
|
|
213
|
+
| E::Floor { .. }
|
|
214
|
+
| E::Position { .. }
|
|
215
|
+
| E::Substring { .. }
|
|
216
|
+
| E::Trim { .. }
|
|
217
|
+
| E::Overlay { .. }
|
|
218
|
+
| E::Collate { .. }
|
|
219
|
+
| E::Struct { .. }
|
|
220
|
+
| E::Map(_)
|
|
221
|
+
| E::Dictionary(_) => ATOM_PREC,
|
|
222
|
+
_ => 0,
|
|
129
223
|
}
|
|
130
224
|
}
|
|
131
225
|
|
|
@@ -154,6 +248,9 @@ fn lower_binary_op(op: BinaryOp) -> ast::BinaryOperator {
|
|
|
154
248
|
BinaryOp::And => B::And,
|
|
155
249
|
BinaryOp::Or => B::Or,
|
|
156
250
|
BinaryOp::StringConcat => B::StringConcat,
|
|
251
|
+
BinaryOp::JsonExists => B::Question,
|
|
252
|
+
BinaryOp::JsonExistsAny => B::QuestionPipe,
|
|
253
|
+
BinaryOp::JsonExistsAll => B::QuestionAnd,
|
|
157
254
|
}
|
|
158
255
|
}
|
|
159
256
|
|
|
@@ -229,13 +326,20 @@ fn lower_call(call: &Call) -> Result<ast::Expr> {
|
|
|
229
326
|
}))
|
|
230
327
|
}
|
|
231
328
|
|
|
329
|
+
/// Lower an operand, and parenthesize it when its precedence is `<= wrap_at_or_below`.
|
|
330
|
+
fn lower_operand(expr: &Expr, wrap_at_or_below: u8) -> Result<ast::Expr> {
|
|
331
|
+
let lowered = lower_expr(expr)?;
|
|
332
|
+
let wrap = ast_prec(&lowered) <= wrap_at_or_below;
|
|
333
|
+
Ok(maybe_paren(lowered, wrap))
|
|
334
|
+
}
|
|
335
|
+
|
|
232
336
|
fn lower_binary(left: &Expr, op: BinaryOp, right: &Expr) -> Result<ast::Expr> {
|
|
233
337
|
let parent = binary_op_prec(op);
|
|
234
338
|
// Left keeps its operator unwrapped while it binds at least as tight (left-associative); the
|
|
235
339
|
// right operand is wrapped even at equal precedence, so `a - (b - c)` and `a / (b - c)` keep
|
|
236
340
|
// their grouping.
|
|
237
|
-
let left_ast =
|
|
238
|
-
let right_ast =
|
|
341
|
+
let left_ast = lower_operand(left, parent - 1)?;
|
|
342
|
+
let right_ast = lower_operand(right, parent)?;
|
|
239
343
|
Ok(ast::Expr::BinaryOp {
|
|
240
344
|
left: Box::new(left_ast),
|
|
241
345
|
op: lower_binary_op(op),
|
|
@@ -243,31 +347,20 @@ fn lower_binary(left: &Expr, op: BinaryOp, right: &Expr) -> Result<ast::Expr> {
|
|
|
243
347
|
})
|
|
244
348
|
}
|
|
245
349
|
|
|
246
|
-
/// Whether an expression's `Display` begins with a `-` token (a negative numeric literal or a
|
|
247
|
-
/// nested unary negation). A unary minus placed directly in front of one emits `--`, which
|
|
248
|
-
/// DataFusion lexes as a line comment — silently swallowing the rest of the query — so such an
|
|
249
|
-
/// operand must be parenthesized.
|
|
250
|
-
fn renders_with_leading_minus(expr: &Expr) -> bool {
|
|
251
|
-
match expr {
|
|
252
|
-
Expr::Literal(Scalar::Int(i)) => *i < 0,
|
|
253
|
-
Expr::Literal(Scalar::Float(f)) => f.is_sign_negative(),
|
|
254
|
-
Expr::Unary {
|
|
255
|
-
op: UnaryOp::Neg, ..
|
|
256
|
-
} => true,
|
|
257
|
-
_ => false,
|
|
258
|
-
}
|
|
259
|
-
}
|
|
260
|
-
|
|
261
350
|
fn lower_unary(op: UnaryOp, operand: &Expr) -> Result<ast::Expr> {
|
|
262
351
|
let (parent, ast_op) = match op {
|
|
263
|
-
UnaryOp::Not => (
|
|
264
|
-
UnaryOp::Neg => (
|
|
352
|
+
UnaryOp::Not => (NOT_PREC, ast::UnaryOperator::Not),
|
|
353
|
+
UnaryOp::Neg => (PREFIX_PREC, ast::UnaryOperator::Minus),
|
|
265
354
|
};
|
|
266
|
-
let
|
|
267
|
-
|
|
355
|
+
let lowered = lower_expr(operand)?;
|
|
356
|
+
// A unary minus placed directly in front of an operand that renders with a leading `-` (a
|
|
357
|
+
// negative literal, a nested negation, a `raw("-x")`) emits `--`, which DataFusion lexes as a
|
|
358
|
+
// line comment that swallows the rest of the query.
|
|
359
|
+
let wrap = ast_prec(&lowered) < parent
|
|
360
|
+
|| (matches!(op, UnaryOp::Neg) && lowered.to_string().starts_with('-'));
|
|
268
361
|
Ok(ast::Expr::UnaryOp {
|
|
269
362
|
op: ast_op,
|
|
270
|
-
expr: Box::new(maybe_paren(
|
|
363
|
+
expr: Box::new(maybe_paren(lowered, wrap)),
|
|
271
364
|
})
|
|
272
365
|
}
|
|
273
366
|
|
|
@@ -288,8 +381,10 @@ pub fn lower_expr(expr: &Expr) -> Result<ast::Expr> {
|
|
|
288
381
|
Expr::Raw(sql) => parse_raw_expr(sql)?,
|
|
289
382
|
Expr::Binary { left, op, right } => lower_binary(left, *op, right)?,
|
|
290
383
|
Expr::Unary { op, expr } => lower_unary(*op, expr)?,
|
|
384
|
+
// The operand of a predicate form is wrapped when it is itself a comparison or looser, so
|
|
385
|
+
// `(a = b) IS NULL` and `(a OR b) IN (…)` keep their grouping in every dialect.
|
|
291
386
|
Expr::IsNull { expr, negated } => {
|
|
292
|
-
let inner = Box::new(
|
|
387
|
+
let inner = Box::new(lower_operand(expr, COMPARISON_PREC)?);
|
|
293
388
|
if *negated {
|
|
294
389
|
ast::Expr::IsNotNull(inner)
|
|
295
390
|
} else {
|
|
@@ -301,7 +396,7 @@ pub fn lower_expr(expr: &Expr) -> Result<ast::Expr> {
|
|
|
301
396
|
list,
|
|
302
397
|
negated,
|
|
303
398
|
} => ast::Expr::InList {
|
|
304
|
-
expr: Box::new(
|
|
399
|
+
expr: Box::new(lower_operand(expr, COMPARISON_PREC)?),
|
|
305
400
|
list: list.iter().map(lower_expr).collect::<Result<Vec<_>>>()?,
|
|
306
401
|
negated: *negated,
|
|
307
402
|
},
|
|
@@ -310,7 +405,7 @@ pub fn lower_expr(expr: &Expr) -> Result<ast::Expr> {
|
|
|
310
405
|
subquery,
|
|
311
406
|
negated,
|
|
312
407
|
} => ast::Expr::InSubquery {
|
|
313
|
-
expr: Box::new(
|
|
408
|
+
expr: Box::new(lower_operand(expr, COMPARISON_PREC)?),
|
|
314
409
|
subquery: Box::new(lower_query(subquery)?),
|
|
315
410
|
negated: *negated,
|
|
316
411
|
},
|
|
@@ -320,10 +415,10 @@ pub fn lower_expr(expr: &Expr) -> Result<ast::Expr> {
|
|
|
320
415
|
high,
|
|
321
416
|
negated,
|
|
322
417
|
} => ast::Expr::Between {
|
|
323
|
-
expr: Box::new(
|
|
418
|
+
expr: Box::new(lower_operand(expr, COMPARISON_PREC)?),
|
|
324
419
|
negated: *negated,
|
|
325
|
-
low: Box::new(
|
|
326
|
-
high: Box::new(
|
|
420
|
+
low: Box::new(lower_operand(low, COMPARISON_PREC)?),
|
|
421
|
+
high: Box::new(lower_operand(high, COMPARISON_PREC)?),
|
|
327
422
|
},
|
|
328
423
|
Expr::Case {
|
|
329
424
|
when_then,
|
|
@@ -361,14 +456,26 @@ pub fn lower_expr(expr: &Expr) -> Result<ast::Expr> {
|
|
|
361
456
|
}
|
|
362
457
|
|
|
363
458
|
/// Parse a raw fragment into an expression so it composes with the rest of the AST.
|
|
459
|
+
///
|
|
460
|
+
/// `parse_expr` stops at the first token it can't continue on and returns the partial expression it
|
|
461
|
+
/// has so far — so a fragment like `attributes ? 'k'` parsed under a dialect that doesn't know `?`
|
|
462
|
+
/// would yield just `attributes`, silently dropping the rest. We parse with [`BuilderDialect`] (which
|
|
463
|
+
/// *does* understand the JSONB operators) and then assert the parser reached EOF, turning any
|
|
464
|
+
/// leftover tokens into a loud error instead of a truncated query. Mirrors [`parse_data_type`].
|
|
364
465
|
fn parse_raw_expr(sql: &str) -> Result<ast::Expr> {
|
|
365
|
-
let dialect =
|
|
466
|
+
let dialect = BuilderDialect;
|
|
366
467
|
let mut parser = Parser::new(&dialect)
|
|
367
468
|
.try_with_sql(sql)
|
|
368
469
|
.map_err(|e| BuildError::UnparsableSql(format!("invalid raw expression {sql:?}: {e}")))?;
|
|
369
|
-
parser
|
|
470
|
+
let expr = parser
|
|
370
471
|
.parse_expr()
|
|
371
|
-
.map_err(|e| BuildError::UnparsableSql(format!("invalid raw expression {sql:?}: {e}")))
|
|
472
|
+
.map_err(|e| BuildError::UnparsableSql(format!("invalid raw expression {sql:?}: {e}")))?;
|
|
473
|
+
if parser.peek_token().token != sqlparser::tokenizer::Token::EOF {
|
|
474
|
+
return Err(BuildError::UnparsableSql(format!(
|
|
475
|
+
"invalid raw expression {sql:?}: unexpected trailing tokens"
|
|
476
|
+
)));
|
|
477
|
+
}
|
|
478
|
+
Ok(expr)
|
|
372
479
|
}
|
|
373
480
|
|
|
374
481
|
fn lower_select_item(expr: &Expr) -> Result<ast::SelectItem> {
|
|
@@ -428,12 +535,17 @@ fn and_combine(filters: &[Expr]) -> Result<Option<ast::Expr>> {
|
|
|
428
535
|
let Some(first) = iter.next() else {
|
|
429
536
|
return Ok(None);
|
|
430
537
|
};
|
|
431
|
-
|
|
538
|
+
// A lone filter needs no parentheses. Once filters are AND-combined, each one that binds looser
|
|
539
|
+
// than `AND` is wrapped, or `.filter(a | b).filter(c)` renders as `a OR b AND c`.
|
|
540
|
+
if filters.len() == 1 {
|
|
541
|
+
return Ok(Some(lower_expr(first)?));
|
|
542
|
+
}
|
|
543
|
+
let mut acc = lower_operand(first, AND_PREC - 1)?;
|
|
432
544
|
for f in iter {
|
|
433
545
|
acc = ast::Expr::BinaryOp {
|
|
434
546
|
left: Box::new(acc),
|
|
435
547
|
op: ast::BinaryOperator::And,
|
|
436
|
-
right: Box::new(
|
|
548
|
+
right: Box::new(lower_operand(f, AND_PREC - 1)?),
|
|
437
549
|
};
|
|
438
550
|
}
|
|
439
551
|
Ok(Some(acc))
|
|
@@ -621,7 +733,7 @@ pub fn lower_query(query: &Query) -> Result<ast::Query> {
|
|
|
621
733
|
} else {
|
|
622
734
|
Some(ast::With {
|
|
623
735
|
with_token: ast::helpers::attached_token::AttachedToken::empty(),
|
|
624
|
-
recursive:
|
|
736
|
+
recursive: query.recursive,
|
|
625
737
|
cte_tables: query
|
|
626
738
|
.ctes
|
|
627
739
|
.iter()
|
|
@@ -172,7 +172,7 @@ impl PyExpr {
|
|
|
172
172
|
self.inner.clone().in_subquery(q.inner.clone(), false),
|
|
173
173
|
));
|
|
174
174
|
}
|
|
175
|
-
let list = coerce_list(values)?;
|
|
175
|
+
let list = coerce_list(values, "expected a list/tuple of values or a subquery")?;
|
|
176
176
|
Ok(PyExpr::new(self.inner.clone().in_list(list, false)))
|
|
177
177
|
}
|
|
178
178
|
|
|
@@ -183,7 +183,7 @@ impl PyExpr {
|
|
|
183
183
|
self.inner.clone().in_subquery(q.inner.clone(), true),
|
|
184
184
|
));
|
|
185
185
|
}
|
|
186
|
-
let list = coerce_list(values)?;
|
|
186
|
+
let list = coerce_list(values, "expected a list/tuple of values or a subquery")?;
|
|
187
187
|
Ok(PyExpr::new(self.inner.clone().in_list(list, true)))
|
|
188
188
|
}
|
|
189
189
|
|
|
@@ -199,6 +199,28 @@ impl PyExpr {
|
|
|
199
199
|
PyExpr::new(self.inner.clone().between(low.0, high.0, false))
|
|
200
200
|
}
|
|
201
201
|
|
|
202
|
+
/// JSONB key-exists `?`: does `key` exist as a top-level key/element of this JSON value.
|
|
203
|
+
/// A cheap presence check that (unlike `->> key IS NOT NULL`) never extracts the value.
|
|
204
|
+
fn has_key(&self, key: ExprArg) -> PyExpr {
|
|
205
|
+
self.bin(BinaryOp::JsonExists, key)
|
|
206
|
+
}
|
|
207
|
+
|
|
208
|
+
/// JSONB `?|`: does *any* of `keys` exist as a top-level key/element.
|
|
209
|
+
fn has_any_key(&self, keys: &Bound<'_, PyAny>) -> PyResult<PyExpr> {
|
|
210
|
+
let array = Expr::Array(coerce_list(keys, "expected a list/tuple of keys")?);
|
|
211
|
+
Ok(PyExpr::new(
|
|
212
|
+
self.inner.clone().binary(BinaryOp::JsonExistsAny, array),
|
|
213
|
+
))
|
|
214
|
+
}
|
|
215
|
+
|
|
216
|
+
/// JSONB `?&`: do *all* of `keys` exist as top-level keys/elements.
|
|
217
|
+
fn has_all_keys(&self, keys: &Bound<'_, PyAny>) -> PyResult<PyExpr> {
|
|
218
|
+
let array = Expr::Array(coerce_list(keys, "expected a list/tuple of keys")?);
|
|
219
|
+
Ok(PyExpr::new(
|
|
220
|
+
self.inner.clone().binary(BinaryOp::JsonExistsAll, array),
|
|
221
|
+
))
|
|
222
|
+
}
|
|
223
|
+
|
|
202
224
|
#[pyo3(signature = (nulls_first=None))]
|
|
203
225
|
fn asc(&self, nulls_first: Option<bool>) -> PySortExpr {
|
|
204
226
|
PySortExpr {
|
|
@@ -334,10 +356,10 @@ impl PyExpr {
|
|
|
334
356
|
}
|
|
335
357
|
}
|
|
336
358
|
|
|
337
|
-
fn coerce_list(obj: &Bound<'_, PyAny
|
|
359
|
+
fn coerce_list(obj: &Bound<'_, PyAny>, expected: &'static str) -> PyResult<Vec<Expr>> {
|
|
338
360
|
let items: Vec<ExprArg> = obj
|
|
339
361
|
.extract()
|
|
340
|
-
.map_err(|_| QueryBuilderError::new_err(
|
|
362
|
+
.map_err(|_| QueryBuilderError::new_err(expected))?;
|
|
341
363
|
Ok(items.into_iter().map(|a| a.0).collect())
|
|
342
364
|
}
|
|
343
365
|
|
|
@@ -519,6 +541,14 @@ impl PyQuery {
|
|
|
519
541
|
PyQuery::new(self.inner.clone().with_cte(name, query.inner.clone()))
|
|
520
542
|
}
|
|
521
543
|
|
|
544
|
+
fn with_recursive_cte(&self, name: &str, query: PyRef<'_, PyQuery>) -> PyQuery {
|
|
545
|
+
PyQuery::new(
|
|
546
|
+
self.inner
|
|
547
|
+
.clone()
|
|
548
|
+
.with_recursive_cte(name, query.inner.clone()),
|
|
549
|
+
)
|
|
550
|
+
}
|
|
551
|
+
|
|
522
552
|
#[pyo3(signature = (*sorts))]
|
|
523
553
|
fn order_by(&self, sorts: Vec<SortArg>) -> PyQuery {
|
|
524
554
|
let sorts = sorts.into_iter().map(|s| s.0).collect();
|
|
@@ -110,6 +110,7 @@ pub struct Cte {
|
|
|
110
110
|
#[derive(Debug, Clone, PartialEq)]
|
|
111
111
|
pub struct Query {
|
|
112
112
|
pub ctes: Vec<Cte>,
|
|
113
|
+
pub recursive: bool,
|
|
113
114
|
pub body: Body,
|
|
114
115
|
pub order_by: Vec<SortExpr>,
|
|
115
116
|
pub limit: Option<i64>,
|
|
@@ -138,6 +139,7 @@ impl Query {
|
|
|
138
139
|
fn from_select(select: Select) -> Self {
|
|
139
140
|
Query {
|
|
140
141
|
ctes: vec![],
|
|
142
|
+
recursive: false,
|
|
141
143
|
body: Body::Select(Box::new(select)),
|
|
142
144
|
order_by: vec![],
|
|
143
145
|
limit: None,
|
|
@@ -232,6 +234,17 @@ impl Query {
|
|
|
232
234
|
self
|
|
233
235
|
}
|
|
234
236
|
|
|
237
|
+
/// Add a recursive CTE; reference it later by name in a `FROM` slot.
|
|
238
|
+
#[must_use]
|
|
239
|
+
pub fn with_recursive_cte(mut self, name: impl Into<String>, query: Query) -> Self {
|
|
240
|
+
self.recursive = true;
|
|
241
|
+
self.ctes.push(Cte {
|
|
242
|
+
name: name.into(),
|
|
243
|
+
query,
|
|
244
|
+
});
|
|
245
|
+
self
|
|
246
|
+
}
|
|
247
|
+
|
|
235
248
|
/// Combine with another query via a set operation. The new outer query starts with no
|
|
236
249
|
/// CTEs/ORDER BY/LIMIT (add them afterward to apply to the whole result); each operand keeps
|
|
237
250
|
/// its own, and any operand carrying them is wrapped in a derived subquery during lowering.
|
|
@@ -239,6 +252,7 @@ impl Query {
|
|
|
239
252
|
pub fn set_op(self, op: SetOp, all: bool, other: Query) -> Self {
|
|
240
253
|
Query {
|
|
241
254
|
ctes: vec![],
|
|
255
|
+
recursive: false,
|
|
242
256
|
body: Body::SetOp {
|
|
243
257
|
op,
|
|
244
258
|
all,
|
|
@@ -1,8 +1,8 @@
|
|
|
1
1
|
//! Rendering façade queries to SQL text, plus a cheap round-trip parse self-test.
|
|
2
2
|
|
|
3
|
-
use sqlparser::dialect::{Dialect, GenericDialect};
|
|
4
3
|
use sqlparser::parser::Parser;
|
|
5
4
|
|
|
5
|
+
use crate::dialect::BuilderDialect;
|
|
6
6
|
use crate::lower::lower_query;
|
|
7
7
|
use crate::query::{BuildError, Query, Result};
|
|
8
8
|
|
|
@@ -11,42 +11,14 @@ pub fn to_sql(query: &Query) -> Result<String> {
|
|
|
11
11
|
Ok(lower_query(query)?.to_string())
|
|
12
12
|
}
|
|
13
13
|
|
|
14
|
-
/// A parsing dialect that mirrors what DataFusion accepts closely enough to prove a generated
|
|
15
|
-
/// query is well-formed: generic SQL plus `${var}` dollar-brace placeholders (the pydantic
|
|
16
|
-
/// `sqlparser` fork extension).
|
|
17
|
-
#[derive(Debug, Default)]
|
|
18
|
-
struct ValidationDialect;
|
|
19
|
-
|
|
20
|
-
impl Dialect for ValidationDialect {
|
|
21
|
-
fn is_identifier_start(&self, ch: char) -> bool {
|
|
22
|
-
GenericDialect {}.is_identifier_start(ch)
|
|
23
|
-
}
|
|
24
|
-
|
|
25
|
-
fn is_identifier_part(&self, ch: char) -> bool {
|
|
26
|
-
GenericDialect {}.is_identifier_part(ch)
|
|
27
|
-
}
|
|
28
|
-
|
|
29
|
-
fn supports_dollar_placeholder(&self) -> bool {
|
|
30
|
-
true
|
|
31
|
-
}
|
|
32
|
-
|
|
33
|
-
// Capabilities DataFusion's parser accepts that the generic dialect gates off by default.
|
|
34
|
-
fn supports_filter_during_aggregation(&self) -> bool {
|
|
35
|
-
true
|
|
36
|
-
}
|
|
37
|
-
|
|
38
|
-
fn supports_group_by_expr(&self) -> bool {
|
|
39
|
-
true
|
|
40
|
-
}
|
|
41
|
-
}
|
|
42
|
-
|
|
43
14
|
/// Render and re-parse the query, returning the SQL on success. Proves *parseability* (catches a
|
|
44
15
|
/// malformed `raw(...)` fragment or a structural bug); it does not prove the query will *plan*
|
|
45
16
|
/// (unknown columns/functions/types are caught by DataFusion at plan time). Returns the SQL so callers can
|
|
46
|
-
/// validate-and-use in one step.
|
|
17
|
+
/// validate-and-use in one step. Re-parses with [`BuilderDialect`] — the same grammar `raw(...)`
|
|
18
|
+
/// fragments are parsed against — so a fragment that lowered cleanly always re-parses.
|
|
47
19
|
pub fn validate(query: &Query) -> Result<String> {
|
|
48
20
|
let sql = to_sql(query)?;
|
|
49
|
-
let dialect =
|
|
21
|
+
let dialect = BuilderDialect;
|
|
50
22
|
Parser::parse_sql(&dialect, &sql).map_err(|e| {
|
|
51
23
|
// The builder produced SQL that doesn't parse — a bug in the builder, not caller input.
|
|
52
24
|
BuildError::Misuse(format!(
|
|
@@ -6,7 +6,7 @@ use insta::assert_snapshot;
|
|
|
6
6
|
|
|
7
7
|
use datafusion_query_builder::expr::{BinaryOp, Expr, Scalar, UnaryOp};
|
|
8
8
|
use datafusion_query_builder::functions::{call, count_star};
|
|
9
|
-
use datafusion_query_builder::query::{JoinKind, Query, SetOp, TableRef};
|
|
9
|
+
use datafusion_query_builder::query::{BuildError, JoinKind, Query, SetOp, TableRef};
|
|
10
10
|
use datafusion_query_builder::{to_sql, validate};
|
|
11
11
|
|
|
12
12
|
fn lit_str(s: &str) -> Expr {
|
|
@@ -142,6 +142,26 @@ fn slo_burn_rate_ctes_cross_join() {
|
|
|
142
142
|
assert_snapshot!(validate(&q).unwrap(), @"WITH long_window AS (SELECT round(avg(bad), 2) AS burn_rate FROM records WHERE service_name = 'my-service' AND start_timestamp > now() - INTERVAL '1 hour'), short_window AS (SELECT round(avg(bad), 2) AS burn_rate FROM records WHERE service_name = 'my-service' AND start_timestamp > now() - INTERVAL '5 minutes') SELECT long_window.burn_rate AS long_burn, short_window.burn_rate AS short_burn FROM long_window CROSS JOIN short_window");
|
|
143
143
|
}
|
|
144
144
|
|
|
145
|
+
#[test]
|
|
146
|
+
fn recursive_cte() {
|
|
147
|
+
let seed = Query::empty()
|
|
148
|
+
.select(vec![Expr::lit(Scalar::Int(1)).alias("n")])
|
|
149
|
+
.unwrap();
|
|
150
|
+
let step = Query::table("numbers")
|
|
151
|
+
.select(vec![
|
|
152
|
+
Expr::column("n")
|
|
153
|
+
.binary(BinaryOp::Plus, Expr::lit(Scalar::Int(1)))
|
|
154
|
+
.alias("n"),
|
|
155
|
+
])
|
|
156
|
+
.unwrap()
|
|
157
|
+
.filter(Expr::column("n").binary(BinaryOp::Lt, Expr::lit(Scalar::Int(3))))
|
|
158
|
+
.unwrap();
|
|
159
|
+
let recursive = seed.set_op(SetOp::Union, true, step);
|
|
160
|
+
let q = Query::table("numbers").with_recursive_cte("numbers", recursive);
|
|
161
|
+
|
|
162
|
+
assert_snapshot!(validate(&q).unwrap(), @"WITH RECURSIVE numbers AS (SELECT 1 AS n UNION ALL SELECT n + 1 AS n FROM numbers WHERE n < 3) SELECT * FROM numbers");
|
|
163
|
+
}
|
|
164
|
+
|
|
145
165
|
#[test]
|
|
146
166
|
fn union_all_distinct_keys() {
|
|
147
167
|
let attr = Query::table("metrics")
|
|
@@ -224,6 +244,84 @@ fn invalid_raw_fragment_is_rejected() {
|
|
|
224
244
|
assert!(to_sql(&q).is_err());
|
|
225
245
|
}
|
|
226
246
|
|
|
247
|
+
#[test]
|
|
248
|
+
fn raw_json_key_exists_no_longer_truncates() {
|
|
249
|
+
// Regression: `raw("attributes ? '...'")` used to render as just `attributes` — the JSONB `?`
|
|
250
|
+
// key-exists operator was silently dropped because the raw fragment was parsed with a dialect
|
|
251
|
+
// that tokenized `?` as a prepared-statement placeholder and `parse_expr` stopped early. It now
|
|
252
|
+
// parses (and re-parses) faithfully.
|
|
253
|
+
let q = Query::table("records")
|
|
254
|
+
.select(vec![
|
|
255
|
+
Expr::raw("attributes ? 'gen_ai.input.messages'").alias("x"),
|
|
256
|
+
])
|
|
257
|
+
.unwrap();
|
|
258
|
+
assert_snapshot!(
|
|
259
|
+
validate(&q).unwrap(),
|
|
260
|
+
@"SELECT attributes ? 'gen_ai.input.messages' AS x FROM records"
|
|
261
|
+
);
|
|
262
|
+
}
|
|
263
|
+
|
|
264
|
+
#[test]
|
|
265
|
+
fn raw_fragment_with_trailing_tokens_is_rejected() {
|
|
266
|
+
// The other half of the truncation bug: a fragment that parses a valid leading expression but
|
|
267
|
+
// leaves tokens behind must error, not silently keep only the prefix. Pre-fix `1 + 1 oops`
|
|
268
|
+
// rendered as `1 + 1`.
|
|
269
|
+
let q = Query::table("t")
|
|
270
|
+
.select(vec![Expr::raw("1 + 1 oops").alias("x")])
|
|
271
|
+
.unwrap();
|
|
272
|
+
let err = to_sql(&q).unwrap_err();
|
|
273
|
+
assert!(
|
|
274
|
+
matches!(&err, BuildError::UnparsableSql(m) if m.contains("trailing")),
|
|
275
|
+
"unexpected error: {err:?}"
|
|
276
|
+
);
|
|
277
|
+
}
|
|
278
|
+
|
|
279
|
+
#[test]
|
|
280
|
+
fn json_key_exists_operators_render_natively() {
|
|
281
|
+
// The three JSONB key-exists operators as first-class `BinaryOp`s, so callers don't need `raw()`.
|
|
282
|
+
let exists = Expr::column("attributes")
|
|
283
|
+
.binary(BinaryOp::JsonExists, lit_str("gen_ai.input.messages"))
|
|
284
|
+
.alias("has_msgs");
|
|
285
|
+
let any = Expr::column("attributes")
|
|
286
|
+
.binary(
|
|
287
|
+
BinaryOp::JsonExistsAny,
|
|
288
|
+
Expr::Array(vec![lit_str("a"), lit_str("b")]),
|
|
289
|
+
)
|
|
290
|
+
.alias("has_any");
|
|
291
|
+
let all = Expr::column("attributes")
|
|
292
|
+
.binary(
|
|
293
|
+
BinaryOp::JsonExistsAll,
|
|
294
|
+
Expr::Array(vec![lit_str("a"), lit_str("b")]),
|
|
295
|
+
)
|
|
296
|
+
.alias("has_all");
|
|
297
|
+
let q = Query::table("records")
|
|
298
|
+
.select(vec![exists, any, all])
|
|
299
|
+
.unwrap();
|
|
300
|
+
assert_snapshot!(
|
|
301
|
+
validate(&q).unwrap(),
|
|
302
|
+
@"SELECT attributes ? 'gen_ai.input.messages' AS has_msgs, attributes ?| ARRAY['a', 'b'] AS has_any, attributes ?& ARRAY['a', 'b'] AS has_all FROM records"
|
|
303
|
+
);
|
|
304
|
+
}
|
|
305
|
+
|
|
306
|
+
#[test]
|
|
307
|
+
fn json_key_exists_combines_with_boolean_ops_without_losing_grouping() {
|
|
308
|
+
// Combined with AND / NOT: the key-exists operands must stay grouped so the meaning survives a
|
|
309
|
+
// round-trip through the parser.
|
|
310
|
+
let pred = Expr::unary(
|
|
311
|
+
UnaryOp::Not,
|
|
312
|
+
Expr::column("attributes").binary(BinaryOp::JsonExists, lit_str("a")),
|
|
313
|
+
)
|
|
314
|
+
.binary(
|
|
315
|
+
BinaryOp::And,
|
|
316
|
+
Expr::column("resource").binary(BinaryOp::JsonExists, lit_str("b")),
|
|
317
|
+
);
|
|
318
|
+
let q = Query::table("records").filter(pred).unwrap();
|
|
319
|
+
assert_snapshot!(
|
|
320
|
+
validate(&q).unwrap(),
|
|
321
|
+
@"SELECT * FROM records WHERE NOT attributes ? 'a' AND resource ? 'b'"
|
|
322
|
+
);
|
|
323
|
+
}
|
|
324
|
+
|
|
227
325
|
#[test]
|
|
228
326
|
fn string_literal_with_backslash_quote_round_trips() {
|
|
229
327
|
// Regression: a value of `\'` (backslash + single quote) must render as a properly escaped
|
|
@@ -251,3 +349,146 @@ fn string_literals_escape_quotes_and_keep_backslashes() {
|
|
|
251
349
|
@"SELECT * FROM records WHERE service_name = 'O''Brien' AND path = 'C:\\logs'"
|
|
252
350
|
);
|
|
253
351
|
}
|
|
352
|
+
|
|
353
|
+
#[test]
|
|
354
|
+
fn raw_fragments_keep_their_grouping_when_combined() {
|
|
355
|
+
// A `raw(...)` fragment is opaque until parsed, so it used to be treated as an atom and never
|
|
356
|
+
// wrapped: `raw("a = 1 OR b = 2") & c` rendered `a = 1 OR b = 2 AND c`, which means
|
|
357
|
+
// `a = 1 OR (b = 2 AND c)`. Callers had to wrap the text themselves as `({fragment})`, which
|
|
358
|
+
// breaks on a trailing `-- comment` and lets a fragment like `1=1) OR (1=1` close the bracket.
|
|
359
|
+
let or_fragment = || Expr::raw("a = 1 OR b = 2");
|
|
360
|
+
let c = || Expr::column("c");
|
|
361
|
+
let q = Query::table("t")
|
|
362
|
+
.select(vec![
|
|
363
|
+
or_fragment().binary(BinaryOp::And, c()).alias("and_right"),
|
|
364
|
+
c().binary(BinaryOp::And, or_fragment()).alias("and_left"),
|
|
365
|
+
Expr::unary(UnaryOp::Not, or_fragment()).alias("not_or"),
|
|
366
|
+
Expr::lit(Scalar::Bool(true))
|
|
367
|
+
.binary(BinaryOp::Eq, Expr::raw("x = 1"))
|
|
368
|
+
.alias("cmp_in_cmp"),
|
|
369
|
+
Expr::unary(UnaryOp::Neg, Expr::raw("x + 1")).alias("neg_sum"),
|
|
370
|
+
Expr::unary(UnaryOp::Neg, Expr::raw("-x")).alias("neg_neg"),
|
|
371
|
+
Expr::raw("x IS NULL")
|
|
372
|
+
.binary(BinaryOp::Eq, Expr::lit(Scalar::Bool(false)))
|
|
373
|
+
.alias("is_null_in_cmp"),
|
|
374
|
+
])
|
|
375
|
+
.unwrap();
|
|
376
|
+
assert_snapshot!(
|
|
377
|
+
validate(&q).unwrap(),
|
|
378
|
+
@"SELECT (a = 1 OR b = 2) AND c AS and_right, c AND (a = 1 OR b = 2) AS and_left, NOT (a = 1 OR b = 2) AS not_or, true = (x = 1) AS cmp_in_cmp, -(x + 1) AS neg_sum, -(-x) AS neg_neg, (x IS NULL) = false AS is_null_in_cmp FROM t"
|
|
379
|
+
);
|
|
380
|
+
}
|
|
381
|
+
|
|
382
|
+
#[test]
|
|
383
|
+
fn raw_fragments_that_bind_tightly_stay_unwrapped() {
|
|
384
|
+
// Output compatibility: a fragment that already binds tighter than its parent renders exactly
|
|
385
|
+
// as before, so existing callers see no change in the common shapes.
|
|
386
|
+
let q = Query::table("t")
|
|
387
|
+
.filter(Expr::raw("attributes ->> 'k'").binary(BinaryOp::Eq, lit_str("v")))
|
|
388
|
+
.unwrap()
|
|
389
|
+
.filter(Expr::raw("x IS NULL"))
|
|
390
|
+
.unwrap()
|
|
391
|
+
.filter(Expr::raw("(a OR b)"))
|
|
392
|
+
.unwrap()
|
|
393
|
+
.filter(Expr::column("ts").binary(BinaryOp::Gt, Expr::raw("now() - INTERVAL '1 hour'")))
|
|
394
|
+
.unwrap()
|
|
395
|
+
.select(vec![
|
|
396
|
+
Expr::raw("attributes ->> 'n'")
|
|
397
|
+
.binary(BinaryOp::Plus, Expr::lit(Scalar::Int(1)))
|
|
398
|
+
.alias("n"),
|
|
399
|
+
])
|
|
400
|
+
.unwrap();
|
|
401
|
+
assert_snapshot!(
|
|
402
|
+
validate(&q).unwrap(),
|
|
403
|
+
@"SELECT (attributes ->> 'n') + 1 AS n FROM t WHERE attributes ->> 'k' = 'v' AND x IS NULL AND (a OR b) AND ts > now() - INTERVAL '1 hour'"
|
|
404
|
+
);
|
|
405
|
+
}
|
|
406
|
+
|
|
407
|
+
#[test]
|
|
408
|
+
fn raw_fragment_with_trailing_line_comment_combines_safely() {
|
|
409
|
+
// The fragment is parsed on its own, so a trailing `-- comment` ends at the end of the
|
|
410
|
+
// fragment and is dropped. It can no longer swallow a bracket or operator that follows it.
|
|
411
|
+
let q = Query::table("t")
|
|
412
|
+
.filter(Expr::column("service_name").binary(BinaryOp::Eq, lit_str("svc")))
|
|
413
|
+
.unwrap()
|
|
414
|
+
.filter(Expr::raw("a = 1 OR b = 2 -- either marker"))
|
|
415
|
+
.unwrap()
|
|
416
|
+
.select(vec![count_star().alias("n")])
|
|
417
|
+
.unwrap();
|
|
418
|
+
assert_snapshot!(
|
|
419
|
+
validate(&q).unwrap(),
|
|
420
|
+
@"SELECT count(*) AS n FROM t WHERE service_name = 'svc' AND (a = 1 OR b = 2)"
|
|
421
|
+
);
|
|
422
|
+
}
|
|
423
|
+
|
|
424
|
+
#[test]
|
|
425
|
+
fn chained_filters_keep_disjunctions_grouped() {
|
|
426
|
+
// `filter()` calls are AND-combined. A disjunction passed to one of them used to render bare:
|
|
427
|
+
// `.filter(a | b).filter(c)` became `a OR b AND c`.
|
|
428
|
+
let a_or_b = || eq("a", "1").binary(BinaryOp::Or, eq("b", "1"));
|
|
429
|
+
let q = Query::table("t")
|
|
430
|
+
.filter(a_or_b())
|
|
431
|
+
.unwrap()
|
|
432
|
+
.filter(eq("c", "1"))
|
|
433
|
+
.unwrap()
|
|
434
|
+
.filter(a_or_b())
|
|
435
|
+
.unwrap()
|
|
436
|
+
.select(vec![Expr::column("x")])
|
|
437
|
+
.unwrap();
|
|
438
|
+
assert_snapshot!(
|
|
439
|
+
validate(&q).unwrap(),
|
|
440
|
+
@"SELECT x FROM t WHERE (a = '1' OR b = '1') AND c = '1' AND (a = '1' OR b = '1')"
|
|
441
|
+
);
|
|
442
|
+
|
|
443
|
+
// A single filter is the whole `WHERE` clause and needs no brackets.
|
|
444
|
+
let single = Query::table("t")
|
|
445
|
+
.filter(a_or_b())
|
|
446
|
+
.unwrap()
|
|
447
|
+
.select(vec![Expr::column("x")])
|
|
448
|
+
.unwrap();
|
|
449
|
+
assert_snapshot!(
|
|
450
|
+
validate(&single).unwrap(),
|
|
451
|
+
@"SELECT x FROM t WHERE a = '1' OR b = '1'"
|
|
452
|
+
);
|
|
453
|
+
}
|
|
454
|
+
|
|
455
|
+
#[test]
|
|
456
|
+
fn predicate_operands_keep_their_grouping() {
|
|
457
|
+
// `IS NULL`, `IN` and `BETWEEN` bind tighter than `AND` / `OR` and, in some dialects, looser
|
|
458
|
+
// than comparisons, so a boolean operand must be wrapped.
|
|
459
|
+
let a_or_b = || Expr::column("a").binary(BinaryOp::Or, Expr::column("b"));
|
|
460
|
+
let q = Query::table("t")
|
|
461
|
+
.select(vec![
|
|
462
|
+
a_or_b().is_null(false).alias("n"),
|
|
463
|
+
Expr::raw("x = 1")
|
|
464
|
+
.in_list(vec![Expr::lit(Scalar::Bool(true))], false)
|
|
465
|
+
.alias("i"),
|
|
466
|
+
Expr::column("x")
|
|
467
|
+
.between(
|
|
468
|
+
Expr::raw("lo OR hi"),
|
|
469
|
+
Expr::column("y").binary(BinaryOp::Plus, Expr::lit(Scalar::Int(1))),
|
|
470
|
+
false,
|
|
471
|
+
)
|
|
472
|
+
.alias("b"),
|
|
473
|
+
Expr::raw("attributes ->> 'k'")
|
|
474
|
+
.is_null(true)
|
|
475
|
+
.alias("unchanged"),
|
|
476
|
+
])
|
|
477
|
+
.unwrap();
|
|
478
|
+
assert_snapshot!(
|
|
479
|
+
validate(&q).unwrap(),
|
|
480
|
+
@"SELECT (a OR b) IS NULL AS n, (x = 1) IN (true) AS i, x BETWEEN (lo OR hi) AND y + 1 AS b, attributes ->> 'k' IS NOT NULL AS unchanged FROM t"
|
|
481
|
+
);
|
|
482
|
+
}
|
|
483
|
+
|
|
484
|
+
#[test]
|
|
485
|
+
fn raw_fragment_cannot_close_a_bracket_it_did_not_open() {
|
|
486
|
+
// Without text wrapping there is no caller bracket to close, so an unbalanced fragment is a
|
|
487
|
+
// parse error instead of an expression that escapes its scope.
|
|
488
|
+
let q = Query::table("t")
|
|
489
|
+
.filter(eq("service_name", "svc"))
|
|
490
|
+
.unwrap()
|
|
491
|
+
.filter(Expr::raw("1=1) OR (1=1"))
|
|
492
|
+
.unwrap();
|
|
493
|
+
assert!(matches!(to_sql(&q), Err(BuildError::UnparsableSql(_))));
|
|
494
|
+
}
|
|
@@ -5,3 +5,4 @@
|
|
|
5
5
|
# It is recommended to check this file in to source control so that
|
|
6
6
|
# everyone who runs the test benefits from these saved cases.
|
|
7
7
|
cc 7917b1826878898c3c20c818b8243678c7bedb47511b9b8d9f829083e4248df2 # shrinks to expr = Binary { left: Literal(Int(0)), op: Plus, right: Binary { left: Literal(Int(0)), op: Plus, right: Unary { op: Neg, expr: Literal(Int(-1)) } } }
|
|
8
|
+
cc 537d1664bbc206dafb87318db35f12f0d466a9c8ffdc6bcb99cdcdf082a8a489 # shrinks to (native, mixed) = (Binary { left: Binary { left: Binary { left: Literal(Bool(false)), op: Eq, right: Literal(Bool(false)) }, op: Or, right: Literal(Bool(false)) }, op: And, right: Binary { left: Literal(Bool(false)), op: And, right: Literal(Bool(false)) } }, Binary { left: Raw("false = false OR false"), op: And, right: Binary { left: Literal(Bool(false)), op: And, right: Literal(Bool(false)) } })
|
|
@@ -176,6 +176,56 @@ fn known_nasty_scalars_round_trip() {
|
|
|
176
176
|
}
|
|
177
177
|
}
|
|
178
178
|
|
|
179
|
+
#[test]
|
|
180
|
+
fn recursive_cte_executes_in_datafusion() {
|
|
181
|
+
let seed = Query::empty()
|
|
182
|
+
.select(vec![Expr::lit(Scalar::Int(1)).alias("n")])
|
|
183
|
+
.expect("select seed");
|
|
184
|
+
let step = Query::table("numbers")
|
|
185
|
+
.select(vec![
|
|
186
|
+
Expr::column("n")
|
|
187
|
+
.binary(BinaryOp::Plus, Expr::lit(Scalar::Int(1)))
|
|
188
|
+
.alias("n"),
|
|
189
|
+
])
|
|
190
|
+
.expect("select recursive step")
|
|
191
|
+
.filter(Expr::column("n").binary(BinaryOp::Lt, Expr::lit(Scalar::Int(3))))
|
|
192
|
+
.expect("filter recursive step");
|
|
193
|
+
let query = Query::table("numbers").with_recursive_cte(
|
|
194
|
+
"numbers",
|
|
195
|
+
seed.set_op(datafusion_query_builder::query::SetOp::Union, true, step),
|
|
196
|
+
);
|
|
197
|
+
let sql = to_sql(&query).expect("render recursive query");
|
|
198
|
+
|
|
199
|
+
let values = runtime().block_on(async move {
|
|
200
|
+
let ctx = SessionContext::new();
|
|
201
|
+
let batches = ctx
|
|
202
|
+
.sql(&sql)
|
|
203
|
+
.await
|
|
204
|
+
.expect("plan recursive query")
|
|
205
|
+
.collect()
|
|
206
|
+
.await
|
|
207
|
+
.expect("execute recursive query");
|
|
208
|
+
batches
|
|
209
|
+
.iter()
|
|
210
|
+
.flat_map(|batch| {
|
|
211
|
+
(0..batch.num_rows()).map(|row| {
|
|
212
|
+
ScalarValue::try_from_array(batch.column(0), row)
|
|
213
|
+
.expect("extract recursive value")
|
|
214
|
+
})
|
|
215
|
+
})
|
|
216
|
+
.collect::<Vec<_>>()
|
|
217
|
+
});
|
|
218
|
+
|
|
219
|
+
assert_eq!(
|
|
220
|
+
values,
|
|
221
|
+
vec![
|
|
222
|
+
ScalarValue::Int64(Some(1)),
|
|
223
|
+
ScalarValue::Int64(Some(2)),
|
|
224
|
+
ScalarValue::Int64(Some(3)),
|
|
225
|
+
]
|
|
226
|
+
);
|
|
227
|
+
}
|
|
228
|
+
|
|
179
229
|
// ---- Arithmetic precedence, checked by execution -----------------------------------------------
|
|
180
230
|
|
|
181
231
|
fn arb_arith_op() -> impl Strategy<Value = BinaryOp> {
|
|
@@ -309,3 +359,120 @@ proptest! {
|
|
|
309
359
|
prop_assert_eq!(build(), build());
|
|
310
360
|
}
|
|
311
361
|
}
|
|
362
|
+
|
|
363
|
+
// ---- Raw-fragment grouping, checked by execution -----------------------------------------------
|
|
364
|
+
|
|
365
|
+
/// Render one expression on its own, as the text a caller would hand to `raw(...)`.
|
|
366
|
+
fn render_expr(expr: &Expr) -> String {
|
|
367
|
+
let query = Query::empty()
|
|
368
|
+
.select(vec![expr.clone()])
|
|
369
|
+
.expect("select on empty query");
|
|
370
|
+
let sql = to_sql(&query).expect("boolean trees always render");
|
|
371
|
+
sql.strip_prefix("SELECT ")
|
|
372
|
+
.expect("a bare projection renders as `SELECT <expr>`")
|
|
373
|
+
.to_string()
|
|
374
|
+
}
|
|
375
|
+
|
|
376
|
+
fn arb_bool_op() -> impl Strategy<Value = BinaryOp> {
|
|
377
|
+
prop_oneof![
|
|
378
|
+
Just(BinaryOp::And),
|
|
379
|
+
Just(BinaryOp::Or),
|
|
380
|
+
Just(BinaryOp::Eq),
|
|
381
|
+
Just(BinaryOp::NotEq),
|
|
382
|
+
]
|
|
383
|
+
}
|
|
384
|
+
|
|
385
|
+
/// A boolean tree paired with an equivalent tree in which random subtrees were replaced by
|
|
386
|
+
/// `raw(<their rendered SQL>)`. The first is evaluated in Rust; the second is executed.
|
|
387
|
+
fn arb_bool_with_raw() -> impl Strategy<Value = (Expr, Expr)> {
|
|
388
|
+
let leaf = any::<bool>().prop_map(|b| {
|
|
389
|
+
let e = Expr::lit(Scalar::Bool(b));
|
|
390
|
+
(e.clone(), e)
|
|
391
|
+
});
|
|
392
|
+
leaf.prop_recursive(5, 40, 2, |inner| {
|
|
393
|
+
let node = prop_oneof![
|
|
394
|
+
3 => (inner.clone(), arb_bool_op(), inner.clone())
|
|
395
|
+
.prop_map(|((ln, lm), op, (rn, rm))| (ln.binary(op, rn), lm.binary(op, rm))),
|
|
396
|
+
1 => inner.prop_map(|(n, m)| (Expr::unary(UnaryOp::Not, n), Expr::unary(UnaryOp::Not, m))),
|
|
397
|
+
];
|
|
398
|
+
(node, any::<bool>()).prop_map(|((native, mixed), as_raw)| {
|
|
399
|
+
if as_raw {
|
|
400
|
+
let raw = Expr::raw(render_expr(&mixed));
|
|
401
|
+
(native, raw)
|
|
402
|
+
} else {
|
|
403
|
+
(native, mixed)
|
|
404
|
+
}
|
|
405
|
+
})
|
|
406
|
+
})
|
|
407
|
+
}
|
|
408
|
+
|
|
409
|
+
fn eval_bool(expr: &Expr) -> bool {
|
|
410
|
+
match expr {
|
|
411
|
+
Expr::Literal(Scalar::Bool(b)) => *b,
|
|
412
|
+
Expr::Unary {
|
|
413
|
+
op: UnaryOp::Not,
|
|
414
|
+
expr,
|
|
415
|
+
} => !eval_bool(expr),
|
|
416
|
+
Expr::Binary { left, op, right } => {
|
|
417
|
+
let (l, r) = (eval_bool(left), eval_bool(right));
|
|
418
|
+
match op {
|
|
419
|
+
BinaryOp::And => l && r,
|
|
420
|
+
BinaryOp::Or => l || r,
|
|
421
|
+
BinaryOp::Eq => l == r,
|
|
422
|
+
BinaryOp::NotEq => l != r,
|
|
423
|
+
_ => unreachable!("arb_bool_op only emits AND, OR, =, <>"),
|
|
424
|
+
}
|
|
425
|
+
}
|
|
426
|
+
_ => unreachable!("the native tree holds only bool leaves, NOT, and arb_bool_op"),
|
|
427
|
+
}
|
|
428
|
+
}
|
|
429
|
+
|
|
430
|
+
/// Execute `SELECT 1 WHERE <f1> AND <f2> …` built from chained `filter()` calls, and report whether
|
|
431
|
+
/// the row survived.
|
|
432
|
+
fn row_survives(filters: Vec<Expr>) -> Result<bool, String> {
|
|
433
|
+
let mut query = Query::empty()
|
|
434
|
+
.select(vec![Expr::lit(Scalar::Int(1))])
|
|
435
|
+
.expect("select on empty query");
|
|
436
|
+
for f in filters {
|
|
437
|
+
query = query.filter(f).expect("filter on a select");
|
|
438
|
+
}
|
|
439
|
+
let sql = to_sql(&query).map_err(|e| format!("render: {e}"))?;
|
|
440
|
+
runtime().block_on(async move {
|
|
441
|
+
let ctx = SessionContext::new();
|
|
442
|
+
let frame = ctx
|
|
443
|
+
.sql(&sql)
|
|
444
|
+
.await
|
|
445
|
+
.map_err(|e| format!("plan: {e}\n sql: {sql}"))?;
|
|
446
|
+
let batches = frame
|
|
447
|
+
.collect()
|
|
448
|
+
.await
|
|
449
|
+
.map_err(|e| format!("exec: {e}\n sql: {sql}"))?;
|
|
450
|
+
Ok(batches.iter().map(|b| b.num_rows()).sum::<usize>() == 1)
|
|
451
|
+
})
|
|
452
|
+
}
|
|
453
|
+
|
|
454
|
+
proptest! {
|
|
455
|
+
#![proptest_config(ProptestConfig { cases: 96, ..ProptestConfig::default() })]
|
|
456
|
+
|
|
457
|
+
/// Replacing any subtree with `raw(<its SQL>)` must not change the computed value. This is the
|
|
458
|
+
/// guarantee a caller relies on when it passes a user-written predicate through `raw(...)`
|
|
459
|
+
/// without wrapping it in brackets itself.
|
|
460
|
+
#[test]
|
|
461
|
+
fn raw_subtrees_evaluate_like_the_tree_they_replace((native, mixed) in arb_bool_with_raw()) {
|
|
462
|
+
let want = eval_bool(&native);
|
|
463
|
+
let actual = eval_one(mixed).map_err(TestCaseError::fail)?;
|
|
464
|
+
prop_assert_eq!(actual, ScalarValue::Boolean(Some(want)));
|
|
465
|
+
}
|
|
466
|
+
|
|
467
|
+
/// Chained `filter()` calls mean the conjunction of every predicate, whatever each one's own
|
|
468
|
+
/// top-level operator is.
|
|
469
|
+
#[test]
|
|
470
|
+
fn chained_filters_mean_their_conjunction(
|
|
471
|
+
trees in proptest::collection::vec(arb_bool_with_raw(), 1..4),
|
|
472
|
+
) {
|
|
473
|
+
let want = trees.iter().all(|(native, _)| eval_bool(native));
|
|
474
|
+
let mixed = trees.into_iter().map(|(_, mixed)| mixed).collect();
|
|
475
|
+
let got = row_survives(mixed).map_err(TestCaseError::fail)?;
|
|
476
|
+
prop_assert_eq!(got, want);
|
|
477
|
+
}
|
|
478
|
+
}
|
|
@@ -9,7 +9,9 @@ from datafusion_query_builder import (
|
|
|
9
9
|
and_,
|
|
10
10
|
col,
|
|
11
11
|
f,
|
|
12
|
+
lit,
|
|
12
13
|
param,
|
|
14
|
+
query,
|
|
13
15
|
raw,
|
|
14
16
|
table,
|
|
15
17
|
when,
|
|
@@ -112,6 +114,18 @@ def test_cte_cross_join():
|
|
|
112
114
|
assert 'CROSS JOIN short_window' in sql
|
|
113
115
|
|
|
114
116
|
|
|
117
|
+
def test_recursive_cte():
|
|
118
|
+
recursive = query().select(lit(1).alias('n')).union_all(
|
|
119
|
+
table('numbers').select((col('n') + 1).alias('n')).filter(col('n') < 3)
|
|
120
|
+
)
|
|
121
|
+
q = table('numbers').with_recursive_cte('numbers', recursive)
|
|
122
|
+
|
|
123
|
+
assert q.validate() == (
|
|
124
|
+
'WITH RECURSIVE numbers AS (SELECT 1 AS n UNION ALL '
|
|
125
|
+
'SELECT n + 1 AS n FROM numbers WHERE n < 3) SELECT * FROM numbers'
|
|
126
|
+
)
|
|
127
|
+
|
|
128
|
+
|
|
115
129
|
def test_window_function():
|
|
116
130
|
q = table('t').select(
|
|
117
131
|
f.sum(col('v')).over(partition_by=[col('svc')], order_by=[col('t').asc()]).alias('running'),
|
|
@@ -135,6 +149,59 @@ def test_injection_is_escaped():
|
|
|
135
149
|
q.validate() # still a single well-formed statement
|
|
136
150
|
|
|
137
151
|
|
|
152
|
+
def test_json_key_exists_operators():
|
|
153
|
+
# The JSONB key-exists family: `?`, `?|`, `?&`. `has_key` is the cheap presence check the
|
|
154
|
+
# value-extraction fallback (`->> 'k' is not null`) was standing in for.
|
|
155
|
+
q = (
|
|
156
|
+
table('records')
|
|
157
|
+
.select(
|
|
158
|
+
col('attributes').has_key('gen_ai.input.messages').alias('has_msgs'),
|
|
159
|
+
col('attributes').has_any_key(['a', 'b']).alias('has_any'),
|
|
160
|
+
col('attributes').has_all_keys(['a', 'b']).alias('has_all'),
|
|
161
|
+
)
|
|
162
|
+
.filter(col('attributes').has_key('gen_ai.input.messages'))
|
|
163
|
+
)
|
|
164
|
+
assert q.validate() == (
|
|
165
|
+
"SELECT attributes ? 'gen_ai.input.messages' AS has_msgs, "
|
|
166
|
+
"attributes ?| ARRAY['a', 'b'] AS has_any, "
|
|
167
|
+
"attributes ?& ARRAY['a', 'b'] AS has_all "
|
|
168
|
+
"FROM records WHERE attributes ? 'gen_ai.input.messages'"
|
|
169
|
+
)
|
|
170
|
+
|
|
171
|
+
|
|
172
|
+
def test_raw_json_operator_no_longer_silently_truncated():
|
|
173
|
+
# The reported bug: raw() dropped the `?` operator, rendering just `attributes`. It now renders
|
|
174
|
+
# (and re-parses) faithfully.
|
|
175
|
+
q = table('records').select(raw("attributes ? 'gen_ai.input.messages'").alias('x'))
|
|
176
|
+
assert q.validate() == "SELECT attributes ? 'gen_ai.input.messages' AS x FROM records"
|
|
177
|
+
|
|
178
|
+
|
|
179
|
+
def test_raw_predicate_needs_no_hand_written_brackets():
|
|
180
|
+
# A caller used to wrap a user-written predicate as `raw(f'({predicate})')`, because an `OR`
|
|
181
|
+
# inside `raw()` lost its grouping once AND-combined. That text wrapping breaks on a trailing
|
|
182
|
+
# `-- comment` and lets `1=1) OR (1=1` close the bracket. The builder now groups the fragment.
|
|
183
|
+
q = (
|
|
184
|
+
table('records')
|
|
185
|
+
.filter(col('service_name') == 'svc')
|
|
186
|
+
.filter(raw('a = 1 OR b = 2 -- either marker'))
|
|
187
|
+
.select(f.count().filter(raw('c = 1 OR d = 2') & (col('ts') > 0)).alias('bad'))
|
|
188
|
+
)
|
|
189
|
+
assert q.validate() == (
|
|
190
|
+
'SELECT count(*) FILTER (WHERE (c = 1 OR d = 2) AND ts > 0) AS bad '
|
|
191
|
+
"FROM records WHERE service_name = 'svc' AND (a = 1 OR b = 2)"
|
|
192
|
+
)
|
|
193
|
+
|
|
194
|
+
|
|
195
|
+
def test_raw_fragment_with_trailing_tokens_raises():
|
|
196
|
+
# A fragment that parses a leading expression but leaves tokens behind is now a loud error
|
|
197
|
+
# rather than a silent truncation to the prefix.
|
|
198
|
+
try:
|
|
199
|
+
table('t').select(raw('1 + 1 oops')).to_sql()
|
|
200
|
+
except UnparsableSqlError:
|
|
201
|
+
return
|
|
202
|
+
raise AssertionError('expected UnparsableSqlError for trailing tokens')
|
|
203
|
+
|
|
204
|
+
|
|
138
205
|
def test_unparsable_sql_raises_unparsable_sql_error():
|
|
139
206
|
# A bad raw() fragment or cast type is caller-supplied bad SQL -> UnparsableSqlError.
|
|
140
207
|
assert issubclass(UnparsableSqlError, QueryBuilderError)
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|