ide: fix col names for collation for, at time zone, overlaps

crates/squawk_ide/src/classify.rs

@@ -62,6 +62,7 @@ fn is_special_fn(kind: SyntaxKind) -> bool {
     matches!(
         kind,
         SyntaxKind::EXTRACT_FN
+            | SyntaxKind::COLLATION_FOR_FN
             | SyntaxKind::JSON_EXISTS_FN
             | SyntaxKind::JSON_ARRAY_FN
             | SyntaxKind::JSON_OBJECT_FN

crates/squawk_ide/src/column_name.rs

@@ -216,9 +216,18 @@ fn name_from_expr(expr: ast::Expr, in_type: bool) -> Option<(ColumnName, SyntaxN
         ast::Expr::ArrayExpr(_) => {
             return Some((ColumnName::Column("array".to_string()), node));
         }
-        ast::Expr::BetweenExpr(_) | ast::Expr::BinExpr(_) => {
+        ast::Expr::BetweenExpr(_) => {
             return Some((ColumnName::UnknownColumn(None), node));
         }
+        ast::Expr::BinExpr(bin_expr) => match bin_expr.op() {
+            Some(ast::BinOp::AtTimeZone(_)) => {
+                return Some((ColumnName::Column("timezone".to_string()), node));
+            }
+            Some(ast::BinOp::Overlaps(_)) => {
+                return Some((ColumnName::Column("overlaps".to_string()), node));
+            }
+            _ => return Some((ColumnName::UnknownColumn(None), node)),
+        },
         ast::Expr::CallExpr(call_expr) => {
             if let Some(exists_fn) = call_expr.exists_fn() {
                 return Some((
@@ -365,6 +374,12 @@ fn name_from_expr(expr: ast::Expr, in_type: bool) -> Option<(ColumnName, SyntaxN
                     xml_pi_fn.syntax().clone(),
                 ));
             }
+            if let Some(collation_for_fn) = call_expr.collation_for_fn() {
+                return Some((
+                    ColumnName::Column("pg_collation_for".to_string()),
+                    collation_for_fn.syntax().clone(),
+                ));
+            }
             if let Some(func_name) = call_expr.expr() {
                 match func_name {
                     ast::Expr::ArrayExpr(_)
@@ -432,9 +447,18 @@ fn name_from_expr(expr: ast::Expr, in_type: bool) -> Option<(ColumnName, SyntaxN
                 return name_from_expr(base, in_type);
             }
         }
-        ast::Expr::Literal(_) | ast::Expr::PrefixExpr(_) | ast::Expr::PostfixExpr(_) => {
+        ast::Expr::Literal(_) | ast::Expr::PrefixExpr(_) => {
             return Some((ColumnName::UnknownColumn(None), node));
         }
+        ast::Expr::PostfixExpr(postfix_expr) => match postfix_expr.op() {
+            Some(ast::PostfixOp::AtLocal(_)) => {
+                return Some((ColumnName::Column("timezone".to_string()), node));
+            }
+            Some(ast::PostfixOp::IsNormalized(_)) => {
+                return Some((ColumnName::Column("is_normalized".to_string()), node));
+            }
+            _ => return Some((ColumnName::UnknownColumn(None), node)),
+        },
         ast::Expr::NameRef(name_ref) => {
             return name_from_name_ref(name_ref, in_type);
         }
@@ -477,6 +501,15 @@ fn examples() {
     // postfix
     assert_snapshot!(name("x is null"), @"?column?");
     assert_snapshot!(name("x is not null"), @"?column?");
+    assert_snapshot!(name("'foo' is normalized"), @"is_normalized");
+    assert_snapshot!(name("'foo' is not normalized"), @"?column?");
+    assert_snapshot!(name("now() at local"), @"timezone");
+    // bin expr
+    assert_snapshot!(name("now() at time zone 'America/Chicago'"), @"timezone");
+    assert_snapshot!(
+        name("(DATE '2001-02-16', DATE '2001-12-21') OVERLAPS (DATE '2001-10-30', DATE '2002-10-30')"),
+        @"overlaps"
+    );
     // paren expr
     assert_snapshot!(name("(1 * 2)"), @"?column?");
     assert_snapshot!(name("(select 1 as a)"), @"a");
@@ -486,6 +519,7 @@ fn examples() {
     assert_snapshot!(name("schema.func_name(1)"), @"func_name");
 
     // special funcs
+    assert_snapshot!(name("collation for ('bar')"), @"pg_collation_for");
     assert_snapshot!(name("extract(year from now())"), @"extract");
     assert_snapshot!(name("exists(select 1)"), @"exists");
     assert_snapshot!(name(r#"json_exists('{"a":1}', '$.a')"#), @"json_exists");

crates/squawk_parser/src/grammar.rs

@@ -809,6 +809,7 @@ fn atom_expr(p: &mut Parser<'_>) -> Option<CompletedMarker> {
         (XMLPI_KW, L_PAREN) => xmlpi_fn(p),
         (SOME_KW | ALL_KW | ANY_KW, L_PAREN) => some_any_all_fn(p),
         (EXISTS_KW, L_PAREN) => exists_fn(p),
+        (COLLATION_KW, FOR_KW) => collation_for_fn(p),
         _ if p.at_ts(NAME_REF_FIRST) => name_ref_(p)?,
         (L_PAREN, _) => tuple_expr(p),
         (ARRAY_KW, L_BRACK | L_PAREN) => {
@@ -863,6 +864,21 @@ fn exists_fn(p: &mut Parser<'_>) -> CompletedMarker {
     m.complete(p, CALL_EXPR)
 }
 
+fn collation_for_fn(p: &mut Parser<'_>) -> CompletedMarker {
+    assert!(p.at(COLLATION_KW) && p.nth_at(1, FOR_KW));
+    let m = p.start();
+    p.bump(COLLATION_KW);
+    p.bump(FOR_KW);
+    p.expect(L_PAREN);
+    if expr(p).is_none() {
+        p.error("expected expression");
+    }
+    p.expect(R_PAREN);
+    let m = m.complete(p, COLLATION_FOR_FN).precede(p);
+    opt_agg_clauses(p);
+    m.complete(p, CALL_EXPR)
+}
+
 // XMLPI '(' NAME_P ColLabel ',' a_expr ')'
 // XMLPI '(' NAME_P ColLabel ')'
 fn xmlpi_fn(p: &mut Parser<'_>) -> CompletedMarker {
@@ -2042,11 +2058,6 @@ fn name_ref_(p: &mut Parser<'_>) -> Option<CompletedMarker> {
     }
     let m = p.start();
     let kind = match p.current() {
-        COLLATION_KW => {
-            p.bump(COLLATION_KW);
-            p.expect(FOR_KW);
-            NAME_REF
-        }
         TIMESTAMP_KW | TIME_KW => {
             p.bump_any();
             if p.eat(L_PAREN) {
@@ -3180,6 +3191,10 @@ fn data_source(p: &mut Parser<'_>) {
             }
             opt_alias(p);
         }
+        COLLATION_KW if p.nth_at(1, FOR_KW) => {
+            collation_for_fn(p);
+            opt_alias(p);
+        }
         _ if p.at_ts(FROM_ITEM_KEYWORDS_FIRST) => from_item_name(p),
         _ => {
             p.error("expected table reference");

crates/squawk_parser/tests/data/ok/select.sql

@@ -472,6 +472,7 @@ select c null;
 -- select_special_funcs
 -- collation
 select collation for ( b + c );
+select * from collation for ('x'::text);
 
 -- current_role
 select current_role;

crates/squawk_parser/tests/snapshots/tests__select_ok.snap

@@ -5918,12 +5918,11 @@ SOURCE_FILE
       TARGET_LIST
         TARGET
           CALL_EXPR
-            NAME_REF
+            COLLATION_FOR_FN
               COLLATION_KW "collation"
               WHITESPACE " "
               FOR_KW "for"
-            WHITESPACE " "
-            ARG_LIST
+              WHITESPACE " "
               L_PAREN "("
               WHITESPACE " "
               BIN_EXPR
@@ -5937,6 +5936,39 @@ SOURCE_FILE
               WHITESPACE " "
               R_PAREN ")"
   SEMICOLON ";"
+  WHITESPACE "\n"
+  SELECT
+    SELECT_CLAUSE
+      SELECT_KW "select"
+      WHITESPACE " "
+      TARGET_LIST
+        TARGET
+          STAR "*"
+    WHITESPACE " "
+    FROM_CLAUSE
+      FROM_KW "from"
+      WHITESPACE " "
+      FROM_ITEM
+        CALL_EXPR
+          COLLATION_FOR_FN
+            COLLATION_KW "collation"
+            WHITESPACE " "
+            FOR_KW "for"
+            WHITESPACE " "
+            L_PAREN "("
+            CAST_EXPR
+              LITERAL
+                STRING "'x'"
+              COLON_COLON
+                COLON ":"
+                COLON ":"
+              PATH_TYPE
+                PATH
+                  PATH_SEGMENT
+                    NAME_REF
+                      TEXT_KW "text"
+            R_PAREN ")"
+  SEMICOLON ";"
   WHITESPACE "\n\n"
   COMMENT "-- current_role"
   WHITESPACE "\n"

crates/squawk_syntax/src/ast.rs

@@ -37,7 +37,7 @@ use squawk_parser::SyntaxKind;
 
 pub use self::{
     generated::tokens::*,
-    node_ext::{BinOp, LitKind},
+    node_ext::{BinOp, LitKind, PostfixOp},
     nodes::*,
     traits::{HasCreateTable, HasParamList, HasWithClause, NameLike},
 };