datafusion-query-builder 0.1.2__tar.gz → 0.3.1__tar.gz

{datafusion_query_builder-0.1.2 → datafusion_query_builder-0.3.1}/Cargo.lock

@@ -1302,7 +1302,7 @@ dependencies = [
 
 [[package]]
 name = "datafusion-query-builder"
-version = "0.1.2"
+version = "0.3.1"
 dependencies = [
  "datafusion",
  "insta",

{datafusion_query_builder-0.1.2 → datafusion_query_builder-0.3.1}/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "datafusion-query-builder"
-version = "0.1.2"
+version = "0.3.1"
 edition = "2024"
 rust-version = "1.85.0"
 description = "Programmatic, injection-safe builder for DataFusion SQL — a typed Rust core with a Python API."

{datafusion_query_builder-0.1.2 → datafusion_query_builder-0.3.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datafusion-query-builder
-Version: 0.1.2
+Version: 0.3.1
 Classifier: Development Status :: 3 - Alpha
 Classifier: Intended Audience :: Developers
 Classifier: License :: OSI Approved :: MIT License
@@ -64,6 +64,17 @@ literals are escaped — values are injection-safe by construction. Reach for `r
 fragment), `f.call("name", ...)` (any function), or `param("name")` (a `${name}` placeholder) when
 you step outside the v1 grammar.
 
+A `raw(...)` fragment must be one complete SQL expression. The builder parses it on its own, drops
+any comments, and adds brackets where its top-level operator binds looser than the surrounding
+one, as it does for native expressions. Pass a user-written predicate as `raw(predicate)`, not
+`raw(f"({predicate})")`. Text wrapping breaks on a trailing `-- comment`, and a predicate such as
+`1=1) OR (1=1` can close the bracket.
+
+The JSONB key-exists operators are first-class: `col("attributes").has_key("gen_ai.input.messages")`
+renders `attributes ? 'gen_ai.input.messages'` — a cheap presence check that never extracts the
+value (contrast `raw("attributes ->> 'k'").is_not_null()`, which reads the whole value out just to
+test it). `has_any_key` and `has_all_keys` render `?|` / `?&`.
+
 ## Architecture
 
 ```

{datafusion_query_builder-0.1.2 → datafusion_query_builder-0.3.1}/README.md

@@ -44,6 +44,17 @@ literals are escaped — values are injection-safe by construction. Reach for `r
 fragment), `f.call("name", ...)` (any function), or `param("name")` (a `${name}` placeholder) when
 you step outside the v1 grammar.
 
+A `raw(...)` fragment must be one complete SQL expression. The builder parses it on its own, drops
+any comments, and adds brackets where its top-level operator binds looser than the surrounding
+one, as it does for native expressions. Pass a user-written predicate as `raw(predicate)`, not
+`raw(f"({predicate})")`. Text wrapping breaks on a trailing `-- comment`, and a predicate such as
+`1=1) OR (1=1` can close the bracket.
+
+The JSONB key-exists operators are first-class: `col("attributes").has_key("gen_ai.input.messages")`
+renders `attributes ? 'gen_ai.input.messages'` — a cheap presence check that never extracts the
+value (contrast `raw("attributes ->> 'k'").is_not_null()`, which reads the whole value out just to
+test it). `has_any_key` and `has_all_keys` render `?|` / `?&`.
+
 ## Architecture
 
 ```

{datafusion_query_builder-0.1.2 → datafusion_query_builder-0.3.1}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "maturin"
 
 [project]
 name = "datafusion-query-builder"
-version = "0.1.2"
+version = "0.3.1"
 description = "Programmatic, injection-safe builder for DataFusion SQL."
 readme = "README.md"
 license = "MIT"

{datafusion_query_builder-0.1.2 → datafusion_query_builder-0.3.1}/python/datafusion_query_builder/__init__.pyi

@@ -18,6 +18,11 @@ class Expr:
     def is_null(self) -> Expr: ...
     def is_not_null(self) -> Expr: ...
     def between(self, low: IntoExpr, high: IntoExpr) -> Expr: ...
+    # JSONB key-exists operators (`?` / `?|` / `?&`). `has_key` is a cheap presence check that,
+    # unlike `col ->> 'k' is not None`, never extracts the value.
+    def has_key(self, key: IntoExpr) -> Expr: ...
+    def has_any_key(self, keys: Sequence[IntoExpr]) -> Expr: ...
+    def has_all_keys(self, keys: Sequence[IntoExpr]) -> Expr: ...
     def asc(self, nulls_first: bool | None = ...) -> SortExpr: ...
     def desc(self, nulls_first: bool | None = ...) -> SortExpr: ...
     # Valid only on a function-call expression:
@@ -79,6 +84,7 @@ class Query:
     ) -> Query: ...
     def cross_join(self, other: str | Query, alias: str | None = ...) -> Query: ...
     def with_cte(self, name: str, query: Query) -> Query: ...
+    def with_recursive_cte(self, name: str, query: Query) -> Query: ...
     def order_by(self, *sorts: IntoSort) -> Query: ...
     def limit(self, n: int) -> Query: ...
     def offset(self, n: int) -> Query: ...

datafusion_query_builder-0.3.1/src/dialect.rs

@@ -0,0 +1,54 @@
+//! The single SQL dialect the builder parses against.
+//!
+//! Two call sites re-parse SQL: [`crate::lower`] parses `raw(...)` fragments and cast-type strings
+//! into AST nodes, and [`crate::render`] re-parses the fully rendered query as a self-test. Both use
+//! this one dialect so a `raw` fragment round-trips through the *same* grammar the final query is
+//! validated against — otherwise a fragment could parse one way going in and fail (or mean something
+//! else) coming out.
+//!
+//! It is `GenericDialect` plus the handful of capabilities DataFusion's own parser accepts that the
+//! generic dialect gates off by default.
+
+use sqlparser::dialect::{Dialect, GenericDialect};
+
+/// Generic SQL plus the DataFusion-flavoured extensions the builder relies on.
+#[derive(Debug, Default)]
+pub(crate) struct BuilderDialect;
+
+impl Dialect for BuilderDialect {
+    fn is_identifier_start(&self, ch: char) -> bool {
+        GenericDialect {}.is_identifier_start(ch)
+    }
+
+    fn is_identifier_part(&self, ch: char) -> bool {
+        GenericDialect {}.is_identifier_part(ch)
+    }
+
+    /// `${var}` dollar-brace placeholders (the pydantic `sqlparser` fork extension the DataFusion
+    /// planner relies on).
+    fn supports_dollar_placeholder(&self) -> bool {
+        true
+    }
+
+    // Capabilities DataFusion's parser accepts that the generic dialect gates off by default.
+    fn supports_filter_during_aggregation(&self) -> bool {
+        true
+    }
+
+    fn supports_group_by_expr(&self) -> bool {
+        true
+    }
+
+    /// Parse `?` / `?|` / `?&` as the Postgres-style JSON key-exists operators (the same ones our
+    /// engine uses), rather than as prepared-statement `?` placeholders.
+    ///
+    /// The method name is an `sqlparser` quirk, not our intent: the fork happens to gate `?`-family
+    /// tokenization behind `supports_geometric_types` — it's the single lever that flips `?` from a
+    /// placeholder to `Token::Question`, and the real Postgres dialect enables the operators through
+    /// the same switch. This has nothing to do with geometry; we're here only for the JSON operators.
+    /// With it off, `attributes ? 'key'` tokenizes `?` as a placeholder, `parse_expr` stops at
+    /// `attributes`, and the operator is silently dropped.
+    fn supports_geometric_types(&self) -> bool {
+        true
+    }
+}

{datafusion_query_builder-0.1.2 → datafusion_query_builder-0.3.1}/src/expr.rs

@@ -36,6 +36,13 @@ pub enum BinaryOp {
     And,
     Or,
     StringConcat,
+    /// JSONB key-exists `?`: does the string on the right exist as a top-level key/element of the
+    /// JSON on the left. DataFusion's JSON extension maps it to `json_contains`.
+    JsonExists,
+    /// JSONB `?|`: does *any* string in the right-hand array exist as a top-level key/element.
+    JsonExistsAny,
+    /// JSONB `?&`: do *all* strings in the right-hand array exist as top-level keys/elements.
+    JsonExistsAll,
 }
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]

{datafusion_query_builder-0.1.2 → datafusion_query_builder-0.3.1}/src/lib.rs

@@ -15,6 +15,7 @@
     clippy::single_match_else
 )]
 
+mod dialect;
 pub mod expr;
 pub mod functions;
 pub mod lower;

{datafusion_query_builder-0.1.2 → datafusion_query_builder-0.3.1}/src/lower.rs

@@ -6,10 +6,10 @@
 //! high-churn nodes are isolated into the small `lower_*` helpers below for exactly that reason.
 
 use sqlparser::ast;
-use sqlparser::dialect::GenericDialect;
 use sqlparser::parser::Parser;
 use sqlparser::tokenizer::Span;
 
+use crate::dialect::BuilderDialect;
 use crate::expr::{BinaryOp, Call, Expr, Scalar, SortExpr, UnaryOp, Window};
 use crate::query::{Body, BuildError, Cte, Join, JoinKind, Query, Result, Select, SetOp, TableRef};
 
@@ -78,7 +78,7 @@ fn string_literal(s: &str) -> ast::Expr {
 }
 
 fn parse_data_type(text: &str) -> Result<ast::DataType> {
-    let dialect = GenericDialect {};
+    let dialect = BuilderDialect;
     let mut parser = Parser::new(&dialect)
         .try_with_sql(text)
         .map_err(|e| BuildError::UnparsableSql(format!("invalid cast type {text:?}: {e}")))?;
@@ -95,37 +95,131 @@ fn parse_data_type(text: &str) -> Result<ast::DataType> {
     Ok(data_type)
 }
 
-/// Binding tightness of an operator (higher binds tighter), used to decide where to insert
-/// parentheses. `sqlparser`'s `Display` is purely structural — it never adds precedence parens —
-/// so the builder must wrap operands itself or a tree like `a / (b - c)` would render as the
-/// semantically different `a / b - c`.
+// Binding tightness tiers (higher binds tighter), used to decide where to insert parentheses.
+// `sqlparser`'s `Display` is purely structural — it never adds precedence parens — so the builder
+// must wrap operands itself or a tree like `a / (b - c)` would render as the semantically different
+// `a / b - c`.
+const OR_PREC: u8 = 10;
+const AND_PREC: u8 = 20;
+const NOT_PREC: u8 = 30;
+/// `IS [NOT] NULL`, `IN`, `BETWEEN`, `LIKE` and friends. Dialects disagree on where these sit
+/// relative to comparisons (Postgres binds `LIKE` tighter than `=`, `GenericDialect` looser), so
+/// they rank just below comparisons: always wrapped inside a comparison or arithmetic operand,
+/// never wrapped under `AND` / `OR` / `NOT`.
+const PREDICATE_PREC: u8 = 35;
+const COMPARISON_PREC: u8 = 40;
+/// Operators outside the builder's own set (`->>`, `@>`, bitwise, custom). Postgres binds these
+/// tighter than comparisons and looser than arithmetic.
+const OTHER_OP_PREC: u8 = 45;
+const ADDITIVE_PREC: u8 = 50;
+const MULTIPLICATIVE_PREC: u8 = 60;
+const AT_TIME_ZONE_PREC: u8 = 65;
+const PREFIX_PREC: u8 = 70;
+const ATOM_PREC: u8 = u8::MAX;
+
 fn binary_op_prec(op: BinaryOp) -> u8 {
     match op {
-        BinaryOp::Or => 1,
-        BinaryOp::And => 2,
+        BinaryOp::Or => OR_PREC,
+        BinaryOp::And => AND_PREC,
         BinaryOp::Eq
         | BinaryOp::NotEq
         | BinaryOp::Lt
         | BinaryOp::LtEq
         | BinaryOp::Gt
-        | BinaryOp::GtEq => 4,
-        BinaryOp::Plus | BinaryOp::Minus | BinaryOp::StringConcat => 5,
-        BinaryOp::Multiply | BinaryOp::Divide | BinaryOp::Modulo => 6,
+        | BinaryOp::GtEq => COMPARISON_PREC,
+        // JSONB key-exists operators return a boolean and are typically combined with AND/OR/NOT;
+        // one comparison-level tier keeps `(a ? 'x') AND (b ? 'y')` and `NOT (a ? 'x')` grouped
+        // correctly, which is all realistic usage needs.
+        BinaryOp::JsonExists | BinaryOp::JsonExistsAny | BinaryOp::JsonExistsAll => COMPARISON_PREC,
+        BinaryOp::Plus | BinaryOp::Minus | BinaryOp::StringConcat => ADDITIVE_PREC,
+        BinaryOp::Multiply | BinaryOp::Divide | BinaryOp::Modulo => MULTIPLICATIVE_PREC,
     }
 }
 
-/// Binding tightness of an expression; atoms (columns, literals, function calls, casts, …) bind
-/// tightest and never need wrapping.
-fn expr_prec(expr: &Expr) -> u8 {
+/// Binding tightness of an already-lowered expression.
+///
+/// This reads the `sqlparser` AST rather than the façade, because a `raw(...)` fragment is opaque
+/// until it is parsed: `raw("a = 1 OR b = 2")` is an `OR` and must be wrapped like one, or
+/// `raw("a = 1 OR b = 2") & c` renders as `a = 1 OR b = 2 AND c`. Anything not recognized as an
+/// atom ranks lowest, so an unfamiliar node gets redundant parentheses rather than a wrong grouping.
+fn ast_prec(expr: &ast::Expr) -> u8 {
+    use ast::BinaryOperator as B;
+    use ast::Expr as E;
     match expr {
-        Expr::Binary { op, .. } => binary_op_prec(*op),
-        Expr::Unary {
-            op: UnaryOp::Not, ..
-        } => 3,
-        Expr::Unary {
-            op: UnaryOp::Neg, ..
-        } => 7,
-        _ => u8::MAX,
+        E::BinaryOp { op, .. } => match op {
+            B::Or => OR_PREC,
+            B::And => AND_PREC,
+            B::Eq
+            | B::NotEq
+            | B::Lt
+            | B::LtEq
+            | B::Gt
+            | B::GtEq
+            | B::Spaceship
+            | B::Question
+            | B::QuestionAnd
+            | B::QuestionPipe => COMPARISON_PREC,
+            B::Plus | B::Minus | B::StringConcat => ADDITIVE_PREC,
+            B::Multiply | B::Divide | B::Modulo => MULTIPLICATIVE_PREC,
+            _ => OTHER_OP_PREC,
+        },
+        E::UnaryOp {
+            op: ast::UnaryOperator::Not,
+            ..
+        }
+        | E::Exists { negated: true, .. } => NOT_PREC,
+        E::UnaryOp { .. } => PREFIX_PREC,
+        E::IsFalse(_)
+        | E::IsNotFalse(_)
+        | E::IsTrue(_)
+        | E::IsNotTrue(_)
+        | E::IsNull(_)
+        | E::IsNotNull(_)
+        | E::IsUnknown(_)
+        | E::IsNotUnknown(_)
+        | E::IsDistinctFrom(..)
+        | E::IsNotDistinctFrom(..)
+        | E::IsNormalized { .. }
+        | E::InList { .. }
+        | E::InSubquery { .. }
+        | E::InUnnest { .. }
+        | E::Between { .. }
+        | E::Like { .. }
+        | E::ILike { .. }
+        | E::SimilarTo { .. }
+        | E::RLike { .. }
+        | E::AnyOp { .. }
+        | E::AllOp { .. }
+        | E::MemberOf(_) => PREDICATE_PREC,
+        E::AtTimeZone { .. } => AT_TIME_ZONE_PREC,
+        E::Identifier(_)
+        | E::CompoundIdentifier(_)
+        | E::CompoundFieldAccess { .. }
+        | E::JsonAccess { .. }
+        | E::Value(_)
+        | E::TypedString { .. }
+        | E::Function(_)
+        | E::Cast { .. }
+        | E::Convert { .. }
+        | E::Case { .. }
+        | E::Nested(_)
+        | E::Tuple(_)
+        | E::Array(_)
+        | E::Subquery(_)
+        | E::Exists { negated: false, .. }
+        | E::Interval(_)
+        | E::Extract { .. }
+        | E::Ceil { .. }
+        | E::Floor { .. }
+        | E::Position { .. }
+        | E::Substring { .. }
+        | E::Trim { .. }
+        | E::Overlay { .. }
+        | E::Collate { .. }
+        | E::Struct { .. }
+        | E::Map(_)
+        | E::Dictionary(_) => ATOM_PREC,
+        _ => 0,
     }
 }
 
@@ -154,6 +248,9 @@ fn lower_binary_op(op: BinaryOp) -> ast::BinaryOperator {
         BinaryOp::And => B::And,
         BinaryOp::Or => B::Or,
         BinaryOp::StringConcat => B::StringConcat,
+        BinaryOp::JsonExists => B::Question,
+        BinaryOp::JsonExistsAny => B::QuestionPipe,
+        BinaryOp::JsonExistsAll => B::QuestionAnd,
     }
 }
 
@@ -229,13 +326,20 @@ fn lower_call(call: &Call) -> Result<ast::Expr> {
     }))
 }
 
+/// Lower an operand, and parenthesize it when its precedence is `<= wrap_at_or_below`.
+fn lower_operand(expr: &Expr, wrap_at_or_below: u8) -> Result<ast::Expr> {
+    let lowered = lower_expr(expr)?;
+    let wrap = ast_prec(&lowered) <= wrap_at_or_below;
+    Ok(maybe_paren(lowered, wrap))
+}
+
 fn lower_binary(left: &Expr, op: BinaryOp, right: &Expr) -> Result<ast::Expr> {
     let parent = binary_op_prec(op);
     // Left keeps its operator unwrapped while it binds at least as tight (left-associative); the
     // right operand is wrapped even at equal precedence, so `a - (b - c)` and `a / (b - c)` keep
     // their grouping.
-    let left_ast = maybe_paren(lower_expr(left)?, expr_prec(left) < parent);
-    let right_ast = maybe_paren(lower_expr(right)?, expr_prec(right) <= parent);
+    let left_ast = lower_operand(left, parent - 1)?;
+    let right_ast = lower_operand(right, parent)?;
     Ok(ast::Expr::BinaryOp {
         left: Box::new(left_ast),
         op: lower_binary_op(op),
@@ -243,31 +347,20 @@ fn lower_binary(left: &Expr, op: BinaryOp, right: &Expr) -> Result<ast::Expr> {
     })
 }
 
-/// Whether an expression's `Display` begins with a `-` token (a negative numeric literal or a
-/// nested unary negation). A unary minus placed directly in front of one emits `--`, which
-/// DataFusion lexes as a line comment — silently swallowing the rest of the query — so such an
-/// operand must be parenthesized.
-fn renders_with_leading_minus(expr: &Expr) -> bool {
-    match expr {
-        Expr::Literal(Scalar::Int(i)) => *i < 0,
-        Expr::Literal(Scalar::Float(f)) => f.is_sign_negative(),
-        Expr::Unary {
-            op: UnaryOp::Neg, ..
-        } => true,
-        _ => false,
-    }
-}
-
 fn lower_unary(op: UnaryOp, operand: &Expr) -> Result<ast::Expr> {
     let (parent, ast_op) = match op {
-        UnaryOp::Not => (3, ast::UnaryOperator::Not),
-        UnaryOp::Neg => (7, ast::UnaryOperator::Minus),
+        UnaryOp::Not => (NOT_PREC, ast::UnaryOperator::Not),
+        UnaryOp::Neg => (PREFIX_PREC, ast::UnaryOperator::Minus),
     };
-    let wrap = expr_prec(operand) < parent
-        || (matches!(op, UnaryOp::Neg) && renders_with_leading_minus(operand));
+    let lowered = lower_expr(operand)?;
+    // A unary minus placed directly in front of an operand that renders with a leading `-` (a
+    // negative literal, a nested negation, a `raw("-x")`) emits `--`, which DataFusion lexes as a
+    // line comment that swallows the rest of the query.
+    let wrap = ast_prec(&lowered) < parent
+        || (matches!(op, UnaryOp::Neg) && lowered.to_string().starts_with('-'));
     Ok(ast::Expr::UnaryOp {
         op: ast_op,
-        expr: Box::new(maybe_paren(lower_expr(operand)?, wrap)),
+        expr: Box::new(maybe_paren(lowered, wrap)),
     })
 }
 
@@ -288,8 +381,10 @@ pub fn lower_expr(expr: &Expr) -> Result<ast::Expr> {
         Expr::Raw(sql) => parse_raw_expr(sql)?,
         Expr::Binary { left, op, right } => lower_binary(left, *op, right)?,
         Expr::Unary { op, expr } => lower_unary(*op, expr)?,
+        // The operand of a predicate form is wrapped when it is itself a comparison or looser, so
+        // `(a = b) IS NULL` and `(a OR b) IN (…)` keep their grouping in every dialect.
         Expr::IsNull { expr, negated } => {
-            let inner = Box::new(lower_expr(expr)?);
+            let inner = Box::new(lower_operand(expr, COMPARISON_PREC)?);
             if *negated {
                 ast::Expr::IsNotNull(inner)
             } else {
@@ -301,7 +396,7 @@ pub fn lower_expr(expr: &Expr) -> Result<ast::Expr> {
             list,
             negated,
         } => ast::Expr::InList {
-            expr: Box::new(lower_expr(expr)?),
+            expr: Box::new(lower_operand(expr, COMPARISON_PREC)?),
             list: list.iter().map(lower_expr).collect::<Result<Vec<_>>>()?,
             negated: *negated,
         },
@@ -310,7 +405,7 @@ pub fn lower_expr(expr: &Expr) -> Result<ast::Expr> {
             subquery,
             negated,
         } => ast::Expr::InSubquery {
-            expr: Box::new(lower_expr(expr)?),
+            expr: Box::new(lower_operand(expr, COMPARISON_PREC)?),
             subquery: Box::new(lower_query(subquery)?),
             negated: *negated,
         },
@@ -320,10 +415,10 @@ pub fn lower_expr(expr: &Expr) -> Result<ast::Expr> {
             high,
             negated,
         } => ast::Expr::Between {
-            expr: Box::new(lower_expr(expr)?),
+            expr: Box::new(lower_operand(expr, COMPARISON_PREC)?),
             negated: *negated,
-            low: Box::new(lower_expr(low)?),
-            high: Box::new(lower_expr(high)?),
+            low: Box::new(lower_operand(low, COMPARISON_PREC)?),
+            high: Box::new(lower_operand(high, COMPARISON_PREC)?),
         },
         Expr::Case {
             when_then,
@@ -361,14 +456,26 @@ pub fn lower_expr(expr: &Expr) -> Result<ast::Expr> {
 }
 
 /// Parse a raw fragment into an expression so it composes with the rest of the AST.
+///
+/// `parse_expr` stops at the first token it can't continue on and returns the partial expression it
+/// has so far — so a fragment like `attributes ? 'k'` parsed under a dialect that doesn't know `?`
+/// would yield just `attributes`, silently dropping the rest. We parse with [`BuilderDialect`] (which
+/// *does* understand the JSONB operators) and then assert the parser reached EOF, turning any
+/// leftover tokens into a loud error instead of a truncated query. Mirrors [`parse_data_type`].
 fn parse_raw_expr(sql: &str) -> Result<ast::Expr> {
-    let dialect = GenericDialect {};
+    let dialect = BuilderDialect;
     let mut parser = Parser::new(&dialect)
         .try_with_sql(sql)
         .map_err(|e| BuildError::UnparsableSql(format!("invalid raw expression {sql:?}: {e}")))?;
-    parser
+    let expr = parser
         .parse_expr()
-        .map_err(|e| BuildError::UnparsableSql(format!("invalid raw expression {sql:?}: {e}")))
+        .map_err(|e| BuildError::UnparsableSql(format!("invalid raw expression {sql:?}: {e}")))?;
+    if parser.peek_token().token != sqlparser::tokenizer::Token::EOF {
+        return Err(BuildError::UnparsableSql(format!(
+            "invalid raw expression {sql:?}: unexpected trailing tokens"
+        )));
+    }
+    Ok(expr)
 }
 
 fn lower_select_item(expr: &Expr) -> Result<ast::SelectItem> {
@@ -428,12 +535,17 @@ fn and_combine(filters: &[Expr]) -> Result<Option<ast::Expr>> {
     let Some(first) = iter.next() else {
         return Ok(None);
     };
-    let mut acc = lower_expr(first)?;
+    // A lone filter needs no parentheses. Once filters are AND-combined, each one that binds looser
+    // than `AND` is wrapped, or `.filter(a | b).filter(c)` renders as `a OR b AND c`.
+    if filters.len() == 1 {
+        return Ok(Some(lower_expr(first)?));
+    }
+    let mut acc = lower_operand(first, AND_PREC - 1)?;
     for f in iter {
         acc = ast::Expr::BinaryOp {
             left: Box::new(acc),
             op: ast::BinaryOperator::And,
-            right: Box::new(lower_expr(f)?),
+            right: Box::new(lower_operand(f, AND_PREC - 1)?),
         };
     }
     Ok(Some(acc))
@@ -621,7 +733,7 @@ pub fn lower_query(query: &Query) -> Result<ast::Query> {
     } else {
         Some(ast::With {
             with_token: ast::helpers::attached_token::AttachedToken::empty(),
-            recursive: false,
+            recursive: query.recursive,
             cte_tables: query
                 .ctes
                 .iter()

{datafusion_query_builder-0.1.2 → datafusion_query_builder-0.3.1}/src/python.rs

@@ -172,7 +172,7 @@ impl PyExpr {
                 self.inner.clone().in_subquery(q.inner.clone(), false),
             ));
         }
-        let list = coerce_list(values)?;
+        let list = coerce_list(values, "expected a list/tuple of values or a subquery")?;
         Ok(PyExpr::new(self.inner.clone().in_list(list, false)))
     }
 
@@ -183,7 +183,7 @@ impl PyExpr {
                 self.inner.clone().in_subquery(q.inner.clone(), true),
             ));
         }
-        let list = coerce_list(values)?;
+        let list = coerce_list(values, "expected a list/tuple of values or a subquery")?;
         Ok(PyExpr::new(self.inner.clone().in_list(list, true)))
     }
 
@@ -199,6 +199,28 @@ impl PyExpr {
         PyExpr::new(self.inner.clone().between(low.0, high.0, false))
     }
 
+    /// JSONB key-exists `?`: does `key` exist as a top-level key/element of this JSON value.
+    /// A cheap presence check that (unlike `->> key IS NOT NULL`) never extracts the value.
+    fn has_key(&self, key: ExprArg) -> PyExpr {
+        self.bin(BinaryOp::JsonExists, key)
+    }
+
+    /// JSONB `?|`: does *any* of `keys` exist as a top-level key/element.
+    fn has_any_key(&self, keys: &Bound<'_, PyAny>) -> PyResult<PyExpr> {
+        let array = Expr::Array(coerce_list(keys, "expected a list/tuple of keys")?);
+        Ok(PyExpr::new(
+            self.inner.clone().binary(BinaryOp::JsonExistsAny, array),
+        ))
+    }
+
+    /// JSONB `?&`: do *all* of `keys` exist as top-level keys/elements.
+    fn has_all_keys(&self, keys: &Bound<'_, PyAny>) -> PyResult<PyExpr> {
+        let array = Expr::Array(coerce_list(keys, "expected a list/tuple of keys")?);
+        Ok(PyExpr::new(
+            self.inner.clone().binary(BinaryOp::JsonExistsAll, array),
+        ))
+    }
+
     #[pyo3(signature = (nulls_first=None))]
     fn asc(&self, nulls_first: Option<bool>) -> PySortExpr {
         PySortExpr {
@@ -334,10 +356,10 @@ impl PyExpr {
     }
 }
 
-fn coerce_list(obj: &Bound<'_, PyAny>) -> PyResult<Vec<Expr>> {
+fn coerce_list(obj: &Bound<'_, PyAny>, expected: &'static str) -> PyResult<Vec<Expr>> {
     let items: Vec<ExprArg> = obj
         .extract()
-        .map_err(|_| QueryBuilderError::new_err("expected a list/tuple of values or a subquery"))?;
+        .map_err(|_| QueryBuilderError::new_err(expected))?;
     Ok(items.into_iter().map(|a| a.0).collect())
 }
 
@@ -519,6 +541,14 @@ impl PyQuery {
         PyQuery::new(self.inner.clone().with_cte(name, query.inner.clone()))
     }
 
+    fn with_recursive_cte(&self, name: &str, query: PyRef<'_, PyQuery>) -> PyQuery {
+        PyQuery::new(
+            self.inner
+                .clone()
+                .with_recursive_cte(name, query.inner.clone()),
+        )
+    }
+
     #[pyo3(signature = (*sorts))]
     fn order_by(&self, sorts: Vec<SortArg>) -> PyQuery {
         let sorts = sorts.into_iter().map(|s| s.0).collect();

{datafusion_query_builder-0.1.2 → datafusion_query_builder-0.3.1}/src/query.rs

@@ -110,6 +110,7 @@ pub struct Cte {
 #[derive(Debug, Clone, PartialEq)]
 pub struct Query {
     pub ctes: Vec<Cte>,
+    pub recursive: bool,
     pub body: Body,
     pub order_by: Vec<SortExpr>,
     pub limit: Option<i64>,
@@ -138,6 +139,7 @@ impl Query {
     fn from_select(select: Select) -> Self {
         Query {
             ctes: vec![],
+            recursive: false,
             body: Body::Select(Box::new(select)),
             order_by: vec![],
             limit: None,
@@ -232,6 +234,17 @@ impl Query {
         self
     }
 
+    /// Add a recursive CTE; reference it later by name in a `FROM` slot.
+    #[must_use]
+    pub fn with_recursive_cte(mut self, name: impl Into<String>, query: Query) -> Self {
+        self.recursive = true;
+        self.ctes.push(Cte {
+            name: name.into(),
+            query,
+        });
+        self
+    }
+
     /// Combine with another query via a set operation. The new outer query starts with no
     /// CTEs/ORDER BY/LIMIT (add them afterward to apply to the whole result); each operand keeps
     /// its own, and any operand carrying them is wrapped in a derived subquery during lowering.
@@ -239,6 +252,7 @@ impl Query {
     pub fn set_op(self, op: SetOp, all: bool, other: Query) -> Self {
         Query {
             ctes: vec![],
+            recursive: false,
             body: Body::SetOp {
                 op,
                 all,

{datafusion_query_builder-0.1.2 → datafusion_query_builder-0.3.1}/src/render.rs

@@ -1,8 +1,8 @@
 //! Rendering façade queries to SQL text, plus a cheap round-trip parse self-test.
 
-use sqlparser::dialect::{Dialect, GenericDialect};
 use sqlparser::parser::Parser;
 
+use crate::dialect::BuilderDialect;
 use crate::lower::lower_query;
 use crate::query::{BuildError, Query, Result};
 
@@ -11,42 +11,14 @@ pub fn to_sql(query: &Query) -> Result<String> {
     Ok(lower_query(query)?.to_string())
 }
 
-/// A parsing dialect that mirrors what DataFusion accepts closely enough to prove a generated
-/// query is well-formed: generic SQL plus `${var}` dollar-brace placeholders (the pydantic
-/// `sqlparser` fork extension).
-#[derive(Debug, Default)]
-struct ValidationDialect;
-
-impl Dialect for ValidationDialect {
-    fn is_identifier_start(&self, ch: char) -> bool {
-        GenericDialect {}.is_identifier_start(ch)
-    }
-
-    fn is_identifier_part(&self, ch: char) -> bool {
-        GenericDialect {}.is_identifier_part(ch)
-    }
-
-    fn supports_dollar_placeholder(&self) -> bool {
-        true
-    }
-
-    // Capabilities DataFusion's parser accepts that the generic dialect gates off by default.
-    fn supports_filter_during_aggregation(&self) -> bool {
-        true
-    }
-
-    fn supports_group_by_expr(&self) -> bool {
-        true
-    }
-}
-
 /// Render and re-parse the query, returning the SQL on success. Proves *parseability* (catches a
 /// malformed `raw(...)` fragment or a structural bug); it does not prove the query will *plan*
 /// (unknown columns/functions/types are caught by DataFusion at plan time). Returns the SQL so callers can
-/// validate-and-use in one step.
+/// validate-and-use in one step. Re-parses with [`BuilderDialect`] — the same grammar `raw(...)`
+/// fragments are parsed against — so a fragment that lowered cleanly always re-parses.
 pub fn validate(query: &Query) -> Result<String> {
     let sql = to_sql(query)?;
-    let dialect = ValidationDialect;
+    let dialect = BuilderDialect;
     Parser::parse_sql(&dialect, &sql).map_err(|e| {
         // The builder produced SQL that doesn't parse — a bug in the builder, not caller input.
         BuildError::Misuse(format!(

{datafusion_query_builder-0.1.2 → datafusion_query_builder-0.3.1}/tests/core.rs

@@ -6,7 +6,7 @@ use insta::assert_snapshot;
 
 use datafusion_query_builder::expr::{BinaryOp, Expr, Scalar, UnaryOp};
 use datafusion_query_builder::functions::{call, count_star};
-use datafusion_query_builder::query::{JoinKind, Query, SetOp, TableRef};
+use datafusion_query_builder::query::{BuildError, JoinKind, Query, SetOp, TableRef};
 use datafusion_query_builder::{to_sql, validate};
 
 fn lit_str(s: &str) -> Expr {
@@ -142,6 +142,26 @@ fn slo_burn_rate_ctes_cross_join() {
     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");
 }
 
+#[test]
+fn recursive_cte() {
+    let seed = Query::empty()
+        .select(vec![Expr::lit(Scalar::Int(1)).alias("n")])
+        .unwrap();
+    let step = Query::table("numbers")
+        .select(vec![
+            Expr::column("n")
+                .binary(BinaryOp::Plus, Expr::lit(Scalar::Int(1)))
+                .alias("n"),
+        ])
+        .unwrap()
+        .filter(Expr::column("n").binary(BinaryOp::Lt, Expr::lit(Scalar::Int(3))))
+        .unwrap();
+    let recursive = seed.set_op(SetOp::Union, true, step);
+    let q = Query::table("numbers").with_recursive_cte("numbers", recursive);
+
+    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");
+}
+
 #[test]
 fn union_all_distinct_keys() {
     let attr = Query::table("metrics")
@@ -224,6 +244,84 @@ fn invalid_raw_fragment_is_rejected() {
     assert!(to_sql(&q).is_err());
 }
 
+#[test]
+fn raw_json_key_exists_no_longer_truncates() {
+    // Regression: `raw("attributes ? '...'")` used to render as just `attributes` — the JSONB `?`
+    // key-exists operator was silently dropped because the raw fragment was parsed with a dialect
+    // that tokenized `?` as a prepared-statement placeholder and `parse_expr` stopped early. It now
+    // parses (and re-parses) faithfully.
+    let q = Query::table("records")
+        .select(vec![
+            Expr::raw("attributes ? 'gen_ai.input.messages'").alias("x"),
+        ])
+        .unwrap();
+    assert_snapshot!(
+        validate(&q).unwrap(),
+        @"SELECT attributes ? 'gen_ai.input.messages' AS x FROM records"
+    );
+}
+
+#[test]
+fn raw_fragment_with_trailing_tokens_is_rejected() {
+    // The other half of the truncation bug: a fragment that parses a valid leading expression but
+    // leaves tokens behind must error, not silently keep only the prefix. Pre-fix `1 + 1 oops`
+    // rendered as `1 + 1`.
+    let q = Query::table("t")
+        .select(vec![Expr::raw("1 + 1 oops").alias("x")])
+        .unwrap();
+    let err = to_sql(&q).unwrap_err();
+    assert!(
+        matches!(&err, BuildError::UnparsableSql(m) if m.contains("trailing")),
+        "unexpected error: {err:?}"
+    );
+}
+
+#[test]
+fn json_key_exists_operators_render_natively() {
+    // The three JSONB key-exists operators as first-class `BinaryOp`s, so callers don't need `raw()`.
+    let exists = Expr::column("attributes")
+        .binary(BinaryOp::JsonExists, lit_str("gen_ai.input.messages"))
+        .alias("has_msgs");
+    let any = Expr::column("attributes")
+        .binary(
+            BinaryOp::JsonExistsAny,
+            Expr::Array(vec![lit_str("a"), lit_str("b")]),
+        )
+        .alias("has_any");
+    let all = Expr::column("attributes")
+        .binary(
+            BinaryOp::JsonExistsAll,
+            Expr::Array(vec![lit_str("a"), lit_str("b")]),
+        )
+        .alias("has_all");
+    let q = Query::table("records")
+        .select(vec![exists, any, all])
+        .unwrap();
+    assert_snapshot!(
+        validate(&q).unwrap(),
+        @"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"
+    );
+}
+
+#[test]
+fn json_key_exists_combines_with_boolean_ops_without_losing_grouping() {
+    // Combined with AND / NOT: the key-exists operands must stay grouped so the meaning survives a
+    // round-trip through the parser.
+    let pred = Expr::unary(
+        UnaryOp::Not,
+        Expr::column("attributes").binary(BinaryOp::JsonExists, lit_str("a")),
+    )
+    .binary(
+        BinaryOp::And,
+        Expr::column("resource").binary(BinaryOp::JsonExists, lit_str("b")),
+    );
+    let q = Query::table("records").filter(pred).unwrap();
+    assert_snapshot!(
+        validate(&q).unwrap(),
+        @"SELECT * FROM records WHERE NOT attributes ? 'a' AND resource ? 'b'"
+    );
+}
+
 #[test]
 fn string_literal_with_backslash_quote_round_trips() {
     // 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() {
         @"SELECT * FROM records WHERE service_name = 'O''Brien' AND path = 'C:\\logs'"
     );
 }
+
+#[test]
+fn raw_fragments_keep_their_grouping_when_combined() {
+    // A `raw(...)` fragment is opaque until parsed, so it used to be treated as an atom and never
+    // wrapped: `raw("a = 1 OR b = 2") & c` rendered `a = 1 OR b = 2 AND c`, which means
+    // `a = 1 OR (b = 2 AND c)`. Callers had to wrap the text themselves as `({fragment})`, which
+    // breaks on a trailing `-- comment` and lets a fragment like `1=1) OR (1=1` close the bracket.
+    let or_fragment = || Expr::raw("a = 1 OR b = 2");
+    let c = || Expr::column("c");
+    let q = Query::table("t")
+        .select(vec![
+            or_fragment().binary(BinaryOp::And, c()).alias("and_right"),
+            c().binary(BinaryOp::And, or_fragment()).alias("and_left"),
+            Expr::unary(UnaryOp::Not, or_fragment()).alias("not_or"),
+            Expr::lit(Scalar::Bool(true))
+                .binary(BinaryOp::Eq, Expr::raw("x = 1"))
+                .alias("cmp_in_cmp"),
+            Expr::unary(UnaryOp::Neg, Expr::raw("x + 1")).alias("neg_sum"),
+            Expr::unary(UnaryOp::Neg, Expr::raw("-x")).alias("neg_neg"),
+            Expr::raw("x IS NULL")
+                .binary(BinaryOp::Eq, Expr::lit(Scalar::Bool(false)))
+                .alias("is_null_in_cmp"),
+        ])
+        .unwrap();
+    assert_snapshot!(
+        validate(&q).unwrap(),
+        @"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"
+    );
+}
+
+#[test]
+fn raw_fragments_that_bind_tightly_stay_unwrapped() {
+    // Output compatibility: a fragment that already binds tighter than its parent renders exactly
+    // as before, so existing callers see no change in the common shapes.
+    let q = Query::table("t")
+        .filter(Expr::raw("attributes ->> 'k'").binary(BinaryOp::Eq, lit_str("v")))
+        .unwrap()
+        .filter(Expr::raw("x IS NULL"))
+        .unwrap()
+        .filter(Expr::raw("(a OR b)"))
+        .unwrap()
+        .filter(Expr::column("ts").binary(BinaryOp::Gt, Expr::raw("now() - INTERVAL '1 hour'")))
+        .unwrap()
+        .select(vec![
+            Expr::raw("attributes ->> 'n'")
+                .binary(BinaryOp::Plus, Expr::lit(Scalar::Int(1)))
+                .alias("n"),
+        ])
+        .unwrap();
+    assert_snapshot!(
+        validate(&q).unwrap(),
+        @"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'"
+    );
+}
+
+#[test]
+fn raw_fragment_with_trailing_line_comment_combines_safely() {
+    // The fragment is parsed on its own, so a trailing `-- comment` ends at the end of the
+    // fragment and is dropped. It can no longer swallow a bracket or operator that follows it.
+    let q = Query::table("t")
+        .filter(Expr::column("service_name").binary(BinaryOp::Eq, lit_str("svc")))
+        .unwrap()
+        .filter(Expr::raw("a = 1 OR b = 2 -- either marker"))
+        .unwrap()
+        .select(vec![count_star().alias("n")])
+        .unwrap();
+    assert_snapshot!(
+        validate(&q).unwrap(),
+        @"SELECT count(*) AS n FROM t WHERE service_name = 'svc' AND (a = 1 OR b = 2)"
+    );
+}
+
+#[test]
+fn chained_filters_keep_disjunctions_grouped() {
+    // `filter()` calls are AND-combined. A disjunction passed to one of them used to render bare:
+    // `.filter(a | b).filter(c)` became `a OR b AND c`.
+    let a_or_b = || eq("a", "1").binary(BinaryOp::Or, eq("b", "1"));
+    let q = Query::table("t")
+        .filter(a_or_b())
+        .unwrap()
+        .filter(eq("c", "1"))
+        .unwrap()
+        .filter(a_or_b())
+        .unwrap()
+        .select(vec![Expr::column("x")])
+        .unwrap();
+    assert_snapshot!(
+        validate(&q).unwrap(),
+        @"SELECT x FROM t WHERE (a = '1' OR b = '1') AND c = '1' AND (a = '1' OR b = '1')"
+    );
+
+    // A single filter is the whole `WHERE` clause and needs no brackets.
+    let single = Query::table("t")
+        .filter(a_or_b())
+        .unwrap()
+        .select(vec![Expr::column("x")])
+        .unwrap();
+    assert_snapshot!(
+        validate(&single).unwrap(),
+        @"SELECT x FROM t WHERE a = '1' OR b = '1'"
+    );
+}
+
+#[test]
+fn predicate_operands_keep_their_grouping() {
+    // `IS NULL`, `IN` and `BETWEEN` bind tighter than `AND` / `OR` and, in some dialects, looser
+    // than comparisons, so a boolean operand must be wrapped.
+    let a_or_b = || Expr::column("a").binary(BinaryOp::Or, Expr::column("b"));
+    let q = Query::table("t")
+        .select(vec![
+            a_or_b().is_null(false).alias("n"),
+            Expr::raw("x = 1")
+                .in_list(vec![Expr::lit(Scalar::Bool(true))], false)
+                .alias("i"),
+            Expr::column("x")
+                .between(
+                    Expr::raw("lo OR hi"),
+                    Expr::column("y").binary(BinaryOp::Plus, Expr::lit(Scalar::Int(1))),
+                    false,
+                )
+                .alias("b"),
+            Expr::raw("attributes ->> 'k'")
+                .is_null(true)
+                .alias("unchanged"),
+        ])
+        .unwrap();
+    assert_snapshot!(
+        validate(&q).unwrap(),
+        @"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"
+    );
+}
+
+#[test]
+fn raw_fragment_cannot_close_a_bracket_it_did_not_open() {
+    // Without text wrapping there is no caller bracket to close, so an unbalanced fragment is a
+    // parse error instead of an expression that escapes its scope.
+    let q = Query::table("t")
+        .filter(eq("service_name", "svc"))
+        .unwrap()
+        .filter(Expr::raw("1=1) OR (1=1"))
+        .unwrap();
+    assert!(matches!(to_sql(&q), Err(BuildError::UnparsableSql(_))));
+}

{datafusion_query_builder-0.1.2 → datafusion_query_builder-0.3.1}/tests/properties.proptest-regressions

@@ -5,3 +5,4 @@
 # It is recommended to check this file in to source control so that
 # everyone who runs the test benefits from these saved cases.
 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)) } } }
+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)) } })

{datafusion_query_builder-0.1.2 → datafusion_query_builder-0.3.1}/tests/properties.rs

@@ -176,6 +176,56 @@ fn known_nasty_scalars_round_trip() {
     }
 }
 
+#[test]
+fn recursive_cte_executes_in_datafusion() {
+    let seed = Query::empty()
+        .select(vec![Expr::lit(Scalar::Int(1)).alias("n")])
+        .expect("select seed");
+    let step = Query::table("numbers")
+        .select(vec![
+            Expr::column("n")
+                .binary(BinaryOp::Plus, Expr::lit(Scalar::Int(1)))
+                .alias("n"),
+        ])
+        .expect("select recursive step")
+        .filter(Expr::column("n").binary(BinaryOp::Lt, Expr::lit(Scalar::Int(3))))
+        .expect("filter recursive step");
+    let query = Query::table("numbers").with_recursive_cte(
+        "numbers",
+        seed.set_op(datafusion_query_builder::query::SetOp::Union, true, step),
+    );
+    let sql = to_sql(&query).expect("render recursive query");
+
+    let values = runtime().block_on(async move {
+        let ctx = SessionContext::new();
+        let batches = ctx
+            .sql(&sql)
+            .await
+            .expect("plan recursive query")
+            .collect()
+            .await
+            .expect("execute recursive query");
+        batches
+            .iter()
+            .flat_map(|batch| {
+                (0..batch.num_rows()).map(|row| {
+                    ScalarValue::try_from_array(batch.column(0), row)
+                        .expect("extract recursive value")
+                })
+            })
+            .collect::<Vec<_>>()
+    });
+
+    assert_eq!(
+        values,
+        vec![
+            ScalarValue::Int64(Some(1)),
+            ScalarValue::Int64(Some(2)),
+            ScalarValue::Int64(Some(3)),
+        ]
+    );
+}
+
 // ---- Arithmetic precedence, checked by execution -----------------------------------------------
 
 fn arb_arith_op() -> impl Strategy<Value = BinaryOp> {
@@ -309,3 +359,120 @@ proptest! {
         prop_assert_eq!(build(), build());
     }
 }
+
+// ---- Raw-fragment grouping, checked by execution -----------------------------------------------
+
+/// Render one expression on its own, as the text a caller would hand to `raw(...)`.
+fn render_expr(expr: &Expr) -> String {
+    let query = Query::empty()
+        .select(vec![expr.clone()])
+        .expect("select on empty query");
+    let sql = to_sql(&query).expect("boolean trees always render");
+    sql.strip_prefix("SELECT ")
+        .expect("a bare projection renders as `SELECT <expr>`")
+        .to_string()
+}
+
+fn arb_bool_op() -> impl Strategy<Value = BinaryOp> {
+    prop_oneof![
+        Just(BinaryOp::And),
+        Just(BinaryOp::Or),
+        Just(BinaryOp::Eq),
+        Just(BinaryOp::NotEq),
+    ]
+}
+
+/// A boolean tree paired with an equivalent tree in which random subtrees were replaced by
+/// `raw(<their rendered SQL>)`. The first is evaluated in Rust; the second is executed.
+fn arb_bool_with_raw() -> impl Strategy<Value = (Expr, Expr)> {
+    let leaf = any::<bool>().prop_map(|b| {
+        let e = Expr::lit(Scalar::Bool(b));
+        (e.clone(), e)
+    });
+    leaf.prop_recursive(5, 40, 2, |inner| {
+        let node = prop_oneof![
+            3 => (inner.clone(), arb_bool_op(), inner.clone())
+                .prop_map(|((ln, lm), op, (rn, rm))| (ln.binary(op, rn), lm.binary(op, rm))),
+            1 => inner.prop_map(|(n, m)| (Expr::unary(UnaryOp::Not, n), Expr::unary(UnaryOp::Not, m))),
+        ];
+        (node, any::<bool>()).prop_map(|((native, mixed), as_raw)| {
+            if as_raw {
+                let raw = Expr::raw(render_expr(&mixed));
+                (native, raw)
+            } else {
+                (native, mixed)
+            }
+        })
+    })
+}
+
+fn eval_bool(expr: &Expr) -> bool {
+    match expr {
+        Expr::Literal(Scalar::Bool(b)) => *b,
+        Expr::Unary {
+            op: UnaryOp::Not,
+            expr,
+        } => !eval_bool(expr),
+        Expr::Binary { left, op, right } => {
+            let (l, r) = (eval_bool(left), eval_bool(right));
+            match op {
+                BinaryOp::And => l && r,
+                BinaryOp::Or => l || r,
+                BinaryOp::Eq => l == r,
+                BinaryOp::NotEq => l != r,
+                _ => unreachable!("arb_bool_op only emits AND, OR, =, <>"),
+            }
+        }
+        _ => unreachable!("the native tree holds only bool leaves, NOT, and arb_bool_op"),
+    }
+}
+
+/// Execute `SELECT 1 WHERE <f1> AND <f2> …` built from chained `filter()` calls, and report whether
+/// the row survived.
+fn row_survives(filters: Vec<Expr>) -> Result<bool, String> {
+    let mut query = Query::empty()
+        .select(vec![Expr::lit(Scalar::Int(1))])
+        .expect("select on empty query");
+    for f in filters {
+        query = query.filter(f).expect("filter on a select");
+    }
+    let sql = to_sql(&query).map_err(|e| format!("render: {e}"))?;
+    runtime().block_on(async move {
+        let ctx = SessionContext::new();
+        let frame = ctx
+            .sql(&sql)
+            .await
+            .map_err(|e| format!("plan: {e}\n  sql: {sql}"))?;
+        let batches = frame
+            .collect()
+            .await
+            .map_err(|e| format!("exec: {e}\n  sql: {sql}"))?;
+        Ok(batches.iter().map(|b| b.num_rows()).sum::<usize>() == 1)
+    })
+}
+
+proptest! {
+    #![proptest_config(ProptestConfig { cases: 96, ..ProptestConfig::default() })]
+
+    /// Replacing any subtree with `raw(<its SQL>)` must not change the computed value. This is the
+    /// guarantee a caller relies on when it passes a user-written predicate through `raw(...)`
+    /// without wrapping it in brackets itself.
+    #[test]
+    fn raw_subtrees_evaluate_like_the_tree_they_replace((native, mixed) in arb_bool_with_raw()) {
+        let want = eval_bool(&native);
+        let actual = eval_one(mixed).map_err(TestCaseError::fail)?;
+        prop_assert_eq!(actual, ScalarValue::Boolean(Some(want)));
+    }
+
+    /// Chained `filter()` calls mean the conjunction of every predicate, whatever each one's own
+    /// top-level operator is.
+    #[test]
+    fn chained_filters_mean_their_conjunction(
+        trees in proptest::collection::vec(arb_bool_with_raw(), 1..4),
+    ) {
+        let want = trees.iter().all(|(native, _)| eval_bool(native));
+        let mixed = trees.into_iter().map(|(_, mixed)| mixed).collect();
+        let got = row_survives(mixed).map_err(TestCaseError::fail)?;
+        prop_assert_eq!(got, want);
+    }
+}

{datafusion_query_builder-0.1.2 → datafusion_query_builder-0.3.1}/tests/test_python.py

@@ -9,7 +9,9 @@ from datafusion_query_builder import (
     and_,
     col,
     f,
+    lit,
     param,
+    query,
     raw,
     table,
     when,
@@ -112,6 +114,18 @@ def test_cte_cross_join():
     assert 'CROSS JOIN short_window' in sql
 
 
+def test_recursive_cte():
+    recursive = query().select(lit(1).alias('n')).union_all(
+        table('numbers').select((col('n') + 1).alias('n')).filter(col('n') < 3)
+    )
+    q = table('numbers').with_recursive_cte('numbers', recursive)
+
+    assert q.validate() == (
+        'WITH RECURSIVE numbers AS (SELECT 1 AS n UNION ALL '
+        'SELECT n + 1 AS n FROM numbers WHERE n < 3) SELECT * FROM numbers'
+    )
+
+
 def test_window_function():
     q = table('t').select(
         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():
     q.validate()  # still a single well-formed statement
 
 
+def test_json_key_exists_operators():
+    # The JSONB key-exists family: `?`, `?|`, `?&`. `has_key` is the cheap presence check the
+    # value-extraction fallback (`->> 'k' is not null`) was standing in for.
+    q = (
+        table('records')
+        .select(
+            col('attributes').has_key('gen_ai.input.messages').alias('has_msgs'),
+            col('attributes').has_any_key(['a', 'b']).alias('has_any'),
+            col('attributes').has_all_keys(['a', 'b']).alias('has_all'),
+        )
+        .filter(col('attributes').has_key('gen_ai.input.messages'))
+    )
+    assert q.validate() == (
+        "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 WHERE attributes ? 'gen_ai.input.messages'"
+    )
+
+
+def test_raw_json_operator_no_longer_silently_truncated():
+    # The reported bug: raw() dropped the `?` operator, rendering just `attributes`. It now renders
+    # (and re-parses) faithfully.
+    q = table('records').select(raw("attributes ? 'gen_ai.input.messages'").alias('x'))
+    assert q.validate() == "SELECT attributes ? 'gen_ai.input.messages' AS x FROM records"
+
+
+def test_raw_predicate_needs_no_hand_written_brackets():
+    # A caller used to wrap a user-written predicate as `raw(f'({predicate})')`, because an `OR`
+    # inside `raw()` lost its grouping once AND-combined. That text wrapping breaks on a trailing
+    # `-- comment` and lets `1=1) OR (1=1` close the bracket. The builder now groups the fragment.
+    q = (
+        table('records')
+        .filter(col('service_name') == 'svc')
+        .filter(raw('a = 1 OR b = 2 -- either marker'))
+        .select(f.count().filter(raw('c = 1 OR d = 2') & (col('ts') > 0)).alias('bad'))
+    )
+    assert q.validate() == (
+        'SELECT count(*) FILTER (WHERE (c = 1 OR d = 2) AND ts > 0) AS bad '
+        "FROM records WHERE service_name = 'svc' AND (a = 1 OR b = 2)"
+    )
+
+
+def test_raw_fragment_with_trailing_tokens_raises():
+    # A fragment that parses a leading expression but leaves tokens behind is now a loud error
+    # rather than a silent truncation to the prefix.
+    try:
+        table('t').select(raw('1 + 1 oops')).to_sql()
+    except UnparsableSqlError:
+        return
+    raise AssertionError('expected UnparsableSqlError for trailing tokens')
+
+
 def test_unparsable_sql_raises_unparsable_sql_error():
     # A bad raw() fragment or cast type is caller-supplied bad SQL -> UnparsableSqlError.
     assert issubclass(UnparsableSqlError, QueryBuilderError)