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.
Files changed (23) hide show
  1. {datafusion_query_builder-0.1.2 → datafusion_query_builder-0.3.1}/Cargo.lock +1 -1
  2. {datafusion_query_builder-0.1.2 → datafusion_query_builder-0.3.1}/Cargo.toml +1 -1
  3. {datafusion_query_builder-0.1.2 → datafusion_query_builder-0.3.1}/PKG-INFO +12 -1
  4. {datafusion_query_builder-0.1.2 → datafusion_query_builder-0.3.1}/README.md +11 -0
  5. {datafusion_query_builder-0.1.2 → datafusion_query_builder-0.3.1}/pyproject.toml +1 -1
  6. {datafusion_query_builder-0.1.2 → datafusion_query_builder-0.3.1}/python/datafusion_query_builder/__init__.pyi +6 -0
  7. datafusion_query_builder-0.3.1/src/dialect.rs +54 -0
  8. {datafusion_query_builder-0.1.2 → datafusion_query_builder-0.3.1}/src/expr.rs +7 -0
  9. {datafusion_query_builder-0.1.2 → datafusion_query_builder-0.3.1}/src/lib.rs +1 -0
  10. {datafusion_query_builder-0.1.2 → datafusion_query_builder-0.3.1}/src/lower.rs +168 -56
  11. {datafusion_query_builder-0.1.2 → datafusion_query_builder-0.3.1}/src/python.rs +34 -4
  12. {datafusion_query_builder-0.1.2 → datafusion_query_builder-0.3.1}/src/query.rs +14 -0
  13. {datafusion_query_builder-0.1.2 → datafusion_query_builder-0.3.1}/src/render.rs +4 -32
  14. {datafusion_query_builder-0.1.2 → datafusion_query_builder-0.3.1}/tests/core.rs +242 -1
  15. {datafusion_query_builder-0.1.2 → datafusion_query_builder-0.3.1}/tests/properties.proptest-regressions +1 -0
  16. {datafusion_query_builder-0.1.2 → datafusion_query_builder-0.3.1}/tests/properties.rs +167 -0
  17. {datafusion_query_builder-0.1.2 → datafusion_query_builder-0.3.1}/tests/test_python.py +67 -0
  18. {datafusion_query_builder-0.1.2 → datafusion_query_builder-0.3.1}/.github/workflows/ci.yml +0 -0
  19. {datafusion_query_builder-0.1.2 → datafusion_query_builder-0.3.1}/.gitignore +0 -0
  20. {datafusion_query_builder-0.1.2 → datafusion_query_builder-0.3.1}/LICENSE +0 -0
  21. {datafusion_query_builder-0.1.2 → datafusion_query_builder-0.3.1}/python/datafusion_query_builder/__init__.py +0 -0
  22. {datafusion_query_builder-0.1.2 → datafusion_query_builder-0.3.1}/python/datafusion_query_builder/py.typed +0 -0
  23. {datafusion_query_builder-0.1.2 → datafusion_query_builder-0.3.1}/src/functions.rs +0 -0
@@ -1302,7 +1302,7 @@ dependencies = [
1302
1302
 
1303
1303
  [[package]]
1304
1304
  name = "datafusion-query-builder"
1305
- version = "0.1.2"
1305
+ version = "0.3.1"
1306
1306
  dependencies = [
1307
1307
  "datafusion",
1308
1308
  "insta",
@@ -1,6 +1,6 @@
1
1
  [package]
2
2
  name = "datafusion-query-builder"
3
- version = "0.1.2"
3
+ version = "0.3.1"
4
4
  edition = "2024"
5
5
  rust-version = "1.85.0"
6
6
  description = "Programmatic, injection-safe builder for DataFusion SQL — a typed Rust core with a Python API."
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: datafusion-query-builder
3
- Version: 0.1.2
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
  ```
@@ -4,7 +4,7 @@ build-backend = "maturin"
4
4
 
5
5
  [project]
6
6
  name = "datafusion-query-builder"
7
- version = "0.1.2"
7
+ version = "0.3.1"
8
8
  description = "Programmatic, injection-safe builder for DataFusion SQL."
9
9
  readme = "README.md"
10
10
  license = "MIT"
@@ -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)]
@@ -15,6 +15,7 @@
15
15
  clippy::single_match_else
16
16
  )]
17
17
 
18
+ mod dialect;
18
19
  pub mod expr;
19
20
  pub mod functions;
20
21
  pub mod lower;
@@ -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 = GenericDialect {};
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
- /// Binding tightness of an operator (higher binds tighter), used to decide where to insert
99
- /// parentheses. `sqlparser`'s `Display` is purely structural — it never adds precedence parens —
100
- /// so the builder must wrap operands itself or a tree like `a / (b - c)` would render as the
101
- /// semantically different `a / b - c`.
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 => 1,
105
- BinaryOp::And => 2,
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 => 4,
112
- BinaryOp::Plus | BinaryOp::Minus | BinaryOp::StringConcat => 5,
113
- BinaryOp::Multiply | BinaryOp::Divide | BinaryOp::Modulo => 6,
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; atoms (columns, literals, function calls, casts, …) bind
118
- /// tightest and never need wrapping.
119
- fn expr_prec(expr: &Expr) -> u8 {
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
- Expr::Binary { op, .. } => binary_op_prec(*op),
122
- Expr::Unary {
123
- op: UnaryOp::Not, ..
124
- } => 3,
125
- Expr::Unary {
126
- op: UnaryOp::Neg, ..
127
- } => 7,
128
- _ => u8::MAX,
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 = maybe_paren(lower_expr(left)?, expr_prec(left) < parent);
238
- let right_ast = maybe_paren(lower_expr(right)?, expr_prec(right) <= parent);
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 => (3, ast::UnaryOperator::Not),
264
- UnaryOp::Neg => (7, ast::UnaryOperator::Minus),
352
+ UnaryOp::Not => (NOT_PREC, ast::UnaryOperator::Not),
353
+ UnaryOp::Neg => (PREFIX_PREC, ast::UnaryOperator::Minus),
265
354
  };
266
- let wrap = expr_prec(operand) < parent
267
- || (matches!(op, UnaryOp::Neg) && renders_with_leading_minus(operand));
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(lower_expr(operand)?, wrap)),
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(lower_expr(expr)?);
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(lower_expr(expr)?),
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(lower_expr(expr)?),
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(lower_expr(expr)?),
418
+ expr: Box::new(lower_operand(expr, COMPARISON_PREC)?),
324
419
  negated: *negated,
325
- low: Box::new(lower_expr(low)?),
326
- high: Box::new(lower_expr(high)?),
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 = GenericDialect {};
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
- let mut acc = lower_expr(first)?;
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(lower_expr(f)?),
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: false,
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>) -> PyResult<Vec<Expr>> {
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("expected a list/tuple of values or a subquery"))?;
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 = ValidationDialect;
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)