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

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

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

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

@@ -1,6 +1,6 @@
 [package]
 name = "datafusion-query-builder"
-version = "0.2.0"
+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.2.0 → datafusion_query_builder-0.3.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datafusion-query-builder
-Version: 0.2.0
+Version: 0.3.1
 Classifier: Development Status :: 3 - Alpha
 Classifier: Intended Audience :: Developers
 Classifier: License :: OSI Approved :: MIT License
@@ -64,6 +64,12 @@ 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

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

@@ -44,6 +44,12 @@ 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

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

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

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

@@ -84,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.2.0 → datafusion_query_builder-0.3.1}/src/lower.rs

@@ -95,41 +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::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 => 4,
-        BinaryOp::Plus | BinaryOp::Minus | BinaryOp::StringConcat => 5,
-        BinaryOp::Multiply | BinaryOp::Divide | BinaryOp::Modulo => 6,
+        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,
     }
 }
 
@@ -236,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),
@@ -250,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)),
     })
 }
 
@@ -295,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 {
@@ -308,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,
         },
@@ -317,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,
         },
@@ -327,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,
@@ -447,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))
@@ -640,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.2.0 → datafusion_query_builder-0.3.1}/src/python.rs

@@ -541,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.2.0 → 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.2.0 → datafusion_query_builder-0.3.1}/tests/core.rs

@@ -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")
@@ -329,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.2.0 → 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.2.0 → 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.2.0 → 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'),
@@ -162,6 +176,22 @@ def test_raw_json_operator_no_longer_silently_truncated():
     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.