Support predicates in function arguments

src/parser/builder.rs

@@ -1936,6 +1936,147 @@ mod tests {
         assert!(result.is_ok());
     }
 
+    #[test]
+    fn test_boolean_operators_in_function_args() {
+        let query = r#"
+            SELECT IFF(x = 'a' OR x = 'b', 1, 0) AS rate
+            FROM t
+            VISUALISE rate AS x
+            DRAW point
+        "#;
+
+        let result = parse_test_query(query);
+        assert!(result.is_ok());
+    }
+
+    #[test]
+    fn test_and_operator_in_function_args() {
+        let query = r#"
+            SELECT IFF(x > 0 AND x < 10, 1, 0) AS flag
+            FROM t
+            VISUALISE flag AS x
+            DRAW point
+        "#;
+
+        let result = parse_test_query(query);
+        assert!(result.is_ok());
+    }
+
+    #[test]
+    fn test_between_in_function_args() {
+        let query = r#"
+            SELECT IFF(x BETWEEN 1 AND 10, 1, 0) AS flag
+            FROM t
+            VISUALISE flag AS x
+            DRAW point
+        "#;
+
+        let result = parse_test_query(query);
+        assert!(result.is_ok());
+    }
+
+    #[test]
+    fn test_not_operator_in_function_args() {
+        let query = r#"
+            SELECT IFF(NOT x = 'a', 1, 0) AS flag
+            FROM t
+            VISUALISE flag AS x
+            DRAW point
+        "#;
+
+        let result = parse_test_query(query);
+        assert!(result.is_ok());
+    }
+
+    #[test]
+    fn test_in_predicate_in_function_args() {
+        let query = r#"
+            SELECT IFF(x IN ('a', 'b'), 1, 0) AS flag
+            FROM t
+            VISUALISE flag AS x
+            DRAW point
+        "#;
+
+        let result = parse_test_query(query);
+        assert!(result.is_ok());
+    }
+
+    #[test]
+    fn test_not_in_predicate_in_function_args() {
+        let query = r#"
+            SELECT IFF(x NOT IN ('a', 'b'), 1, 0) AS flag
+            FROM t
+            VISUALISE flag AS x
+            DRAW point
+        "#;
+
+        let result = parse_test_query(query);
+        assert!(result.is_ok());
+    }
+
+    #[test]
+    fn test_is_null_in_function_args() {
+        let query = r#"
+            SELECT IFF(x IS NULL, 1, 0) AS flag
+            FROM t
+            VISUALISE flag AS x
+            DRAW point
+        "#;
+
+        let result = parse_test_query(query);
+        assert!(result.is_ok());
+    }
+
+    #[test]
+    fn test_like_predicate_in_function_args() {
+        let query = r#"
+            SELECT IFF(x LIKE 'a%', 1, 0) AS flag
+            FROM t
+            VISUALISE flag AS x
+            DRAW point
+        "#;
+
+        let result = parse_test_query(query);
+        assert!(result.is_ok());
+    }
+
+    #[test]
+    fn test_ilike_predicate_in_function_args() {
+        let query = r#"
+            SELECT IFF(x ILIKE 'a%', 1, 0) AS flag
+            FROM t
+            VISUALISE flag AS x
+            DRAW point
+        "#;
+
+        let result = parse_test_query(query);
+        assert!(result.is_ok());
+    }
+
+    #[test]
+    fn test_extract_named_arg_predicate_value_spans_full_expression() {
+        let query = "SELECT FOO(flag => x = 'a' OR x = 'b') FROM t VISUALISE x AS x DRAW point";
+        let source = make_source(query);
+        let root = source.root();
+        let named_arg = source.find_node(&root, "(named_arg) @arg").unwrap();
+
+        let (_, value_node) = extract_name_value_nodes(&named_arg, "named_arg").unwrap();
+        assert_eq!(source.get_text(&value_node), "x = 'a' OR x = 'b'");
+    }
+
+    #[test]
+    fn test_named_arg_predicate_in_function_args() {
+        let query = r#"
+            SELECT FOO(flag => x = 'a' OR x = 'b')
+            FROM t
+            VISUALISE x AS x
+            DRAW point
+        "#;
+
+        let result = parse_test_query(query);
+        assert!(result.is_ok());
+    }
+
     // ========================================
     // Negative Test Cases - Should Error
     // ========================================

tree-sitter-ggsql/grammar.js

@@ -389,13 +389,80 @@ module.exports = grammar({
     // Function argument: position or named
     function_arg: $ => choice(
       $.named_arg,
-      $.position_arg
+      $.function_arg_expression
     ),
 
     named_arg: $ => seq(
       field('name', $.identifier),
       choice(':=', '=>'),
-      field('value', $.position_arg)
+      field('value', $.function_arg_expression)
+    ),
+
+    // Function arguments support a slightly richer predicate grammar than the
+    // generic position_arg rule.
+    function_arg_expression: $ => prec.left(seq(
+      $._function_arg_and_expression,
+      repeat(seq(caseInsensitive('OR'), $._function_arg_and_expression))
+    )),
+
+    _function_arg_and_expression: $ => prec.left(seq(
+      $._function_arg_not_expression,
+      repeat(seq(caseInsensitive('AND'), $._function_arg_not_expression))
+    )),
+
+    _function_arg_not_expression: $ => choice(
+      $.not_predicate,
+      $._function_arg_predicate
+    ),
+
+    not_predicate: $ => prec.right(seq(
+      caseInsensitive('NOT'),
+      $._function_arg_not_expression
+    )),
+
+    _function_arg_predicate: $ => seq(
+      $.position_arg,
+      optional($._predicate_suffix)
+    ),
+
+    _predicate_suffix: $ => choice(
+      $.between_suffix,
+      $.in_suffix,
+      $.is_suffix,
+      $.like_suffix,
+    ),
+
+    between_suffix: $ => seq(
+      optional(caseInsensitive('NOT')),
+      caseInsensitive('BETWEEN'),
+      $.position_arg,
+      caseInsensitive('AND'),
+      $.position_arg
+    ),
+
+    in_suffix: $ => seq(
+      optional(caseInsensitive('NOT')),
+      caseInsensitive('IN'),
+      choice($.in_value_list, $.scalar_subquery)
+    ),
+
+    in_value_list: $ => seq(
+      '(',
+      $.position_arg,
+      repeat(seq(',', $.position_arg)),
+      ')'
+    ),
+
+    is_suffix: $ => seq(
+      caseInsensitive('IS'),
+      optional(caseInsensitive('NOT')),
+      caseInsensitive('NULL')
+    ),
+
+    like_suffix: $ => seq(
+      optional(caseInsensitive('NOT')),
+      choice(caseInsensitive('LIKE'), caseInsensitive('ILIKE')),
+      $.position_arg
     ),
 
     // Position argument: supports complex expressions including:
@@ -420,7 +487,11 @@ module.exports = grammar({
       // Scalar subquery: (SELECT ...) or (WITH ... SELECT ...)
       $.scalar_subquery,
       // Arithmetic/comparison expression (binary operators)
-      seq($.position_arg, choice('+', '-', '*', '/', '%', '||', '::', '<', '>', '<=', '>=', '=', '!=', '<>'), $.position_arg),
+      seq(
+        $.position_arg,
+        choice('+', '-', '*', '/', '%', '||', '::', '<', '>', '<=', '>=', '=', '!=', '<>'),
+        $.position_arg
+      ),
       // Parenthesized expression
       seq('(', $.position_arg, ')')
     )),