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

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

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

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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datafusion-query-builder
-Version: 0.1.2
+Version: 0.2.0
 Classifier: Development Status :: 3 - Alpha
 Classifier: Intended Audience :: Developers
 Classifier: License :: OSI Approved :: MIT License
@@ -64,6 +64,11 @@ 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.
 
+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.2.0}/README.md

@@ -44,6 +44,11 @@ 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.
 
+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.2.0}/pyproject.toml

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

{datafusion_query_builder-0.1.2 → datafusion_query_builder-0.2.0}/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:

datafusion_query_builder-0.2.0/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.2.0}/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.2.0}/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.2.0}/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}")))?;
@@ -109,6 +109,10 @@ fn binary_op_prec(op: BinaryOp) -> u8 {
         | BinaryOp::LtEq
         | BinaryOp::Gt
         | BinaryOp::GtEq => 4,
+        // 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,
     }
@@ -154,6 +158,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,
     }
 }
 
@@ -361,14 +368,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> {

{datafusion_query_builder-0.1.2 → datafusion_query_builder-0.2.0}/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())
 }
 

{datafusion_query_builder-0.1.2 → datafusion_query_builder-0.2.0}/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.2.0}/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 {
@@ -224,6 +224,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

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

@@ -135,6 +135,43 @@ 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_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)