From f3d90889d53b2653d0e5c81f5dfd0fb0084ed40d Mon Sep 17 00:00:00 2001 From: Steve Dignam Date: Sat, 21 Feb 2026 16:43:35 -0500 Subject: [PATCH 1/2] ide: fix col names for collation for, at time zone, overlaps --- crates/squawk_ide/src/column_name.rs | 38 +++++- .../src/generated/syntax_kind.rs | 1 + crates/squawk_parser/src/grammar.rs | 25 +++- crates/squawk_parser/tests/data/ok/select.sql | 1 + .../tests/snapshots/tests__select_ok.snap | 38 +++++- crates/squawk_syntax/src/ast.rs | 2 +- .../squawk_syntax/src/ast/generated/nodes.rs | 49 ++++++++ crates/squawk_syntax/src/ast/node_ext.rs | 114 ++++++++++++------ crates/squawk_syntax/src/postgresql.ungram | 4 + 9 files changed, 226 insertions(+), 46 deletions(-) diff --git a/crates/squawk_ide/src/column_name.rs b/crates/squawk_ide/src/column_name.rs index dc3d4d48..0f812320 100644 --- a/crates/squawk_ide/src/column_name.rs +++ b/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"); diff --git a/crates/squawk_parser/src/generated/syntax_kind.rs b/crates/squawk_parser/src/generated/syntax_kind.rs index 9896bbd0..2b2ea529 100644 --- a/crates/squawk_parser/src/generated/syntax_kind.rs +++ b/crates/squawk_parser/src/generated/syntax_kind.rs @@ -638,6 +638,7 @@ pub enum SyntaxKind { CLUSTER, CLUSTER_ON, COLLATE, + COLLATION_FOR_FN, COLON_COLON, COLON_EQ, COLUMN, diff --git a/crates/squawk_parser/src/grammar.rs b/crates/squawk_parser/src/grammar.rs index c65b41ed..edc99c15 100644 --- a/crates/squawk_parser/src/grammar.rs +++ b/crates/squawk_parser/src/grammar.rs @@ -809,6 +809,7 @@ fn atom_expr(p: &mut Parser<'_>) -> Option { (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 { } 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"); diff --git a/crates/squawk_parser/tests/data/ok/select.sql b/crates/squawk_parser/tests/data/ok/select.sql index b9483744..45f6fd62 100644 --- a/crates/squawk_parser/tests/data/ok/select.sql +++ b/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; diff --git a/crates/squawk_parser/tests/snapshots/tests__select_ok.snap b/crates/squawk_parser/tests/snapshots/tests__select_ok.snap index ecc023f5..7c92a795 100644 --- a/crates/squawk_parser/tests/snapshots/tests__select_ok.snap +++ b/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" diff --git a/crates/squawk_syntax/src/ast.rs b/crates/squawk_syntax/src/ast.rs index caa21d63..53dd25bf 100644 --- a/crates/squawk_syntax/src/ast.rs +++ b/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}, }; diff --git a/crates/squawk_syntax/src/ast/generated/nodes.rs b/crates/squawk_syntax/src/ast/generated/nodes.rs index fb73eca6..3fb79184 100644 --- a/crates/squawk_syntax/src/ast/generated/nodes.rs +++ b/crates/squawk_syntax/src/ast/generated/nodes.rs @@ -2185,6 +2185,10 @@ impl CallExpr { support::child(&self.syntax) } #[inline] + pub fn collation_for_fn(&self) -> Option { + support::child(&self.syntax) + } + #[inline] pub fn exists_fn(&self) -> Option { support::child(&self.syntax) } @@ -2539,6 +2543,33 @@ impl Collate { } } +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct CollationForFn { + pub(crate) syntax: SyntaxNode, +} +impl CollationForFn { + #[inline] + pub fn expr(&self) -> Option { + support::child(&self.syntax) + } + #[inline] + pub fn l_paren_token(&self) -> Option { + support::token(&self.syntax, SyntaxKind::L_PAREN) + } + #[inline] + pub fn r_paren_token(&self) -> Option { + support::token(&self.syntax, SyntaxKind::R_PAREN) + } + #[inline] + pub fn collation_token(&self) -> Option { + support::token(&self.syntax, SyntaxKind::COLLATION_KW) + } + #[inline] + pub fn for_token(&self) -> Option { + support::token(&self.syntax, SyntaxKind::FOR_KW) + } +} + #[derive(Debug, Clone, PartialEq, Eq, Hash)] pub struct ColonColon { pub(crate) syntax: SyntaxNode, @@ -19042,6 +19073,24 @@ impl AstNode for Collate { &self.syntax } } +impl AstNode for CollationForFn { + #[inline] + fn can_cast(kind: SyntaxKind) -> bool { + kind == SyntaxKind::COLLATION_FOR_FN + } + #[inline] + fn cast(syntax: SyntaxNode) -> Option { + if Self::can_cast(syntax.kind()) { + Some(Self { syntax }) + } else { + None + } + } + #[inline] + fn syntax(&self) -> &SyntaxNode { + &self.syntax + } +} impl AstNode for ColonColon { #[inline] fn can_cast(kind: SyntaxKind) -> bool { diff --git a/crates/squawk_syntax/src/ast/node_ext.rs b/crates/squawk_syntax/src/ast/node_ext.rs index 7940847a..4c58ee9b 100644 --- a/crates/squawk_syntax/src/ast/node_ext.rs +++ b/crates/squawk_syntax/src/ast/node_ext.rs @@ -97,18 +97,8 @@ pub enum BinOp { In(SyntaxToken), Is(SyntaxToken), IsDistinctFrom(ast::IsDistinctFrom), - IsJson(ast::IsJson), - IsJsonArray(ast::IsJsonArray), - IsJsonObject(ast::IsJsonObject), - IsJsonScalar(ast::IsJsonScalar), - IsJsonValue(ast::IsJsonValue), IsNot(ast::IsNot), IsNotDistinctFrom(ast::IsNotDistinctFrom), - IsNotJson(ast::IsNotJson), - IsNotJsonArray(ast::IsNotJsonArray), - IsNotJsonObject(ast::IsNotJsonObject), - IsNotJsonScalar(ast::IsNotJsonScalar), - IsNotJsonValue(ast::IsNotJsonValue), LAngle(SyntaxToken), Like(SyntaxToken), Lteq(ast::Lteq), @@ -130,6 +120,25 @@ pub enum BinOp { Star(SyntaxToken), } +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum PostfixOp { + AtLocal(SyntaxToken), + IsJson(ast::IsJson), + IsJsonArray(ast::IsJsonArray), + IsJsonObject(ast::IsJsonObject), + IsJsonScalar(ast::IsJsonScalar), + IsJsonValue(ast::IsJsonValue), + IsNormalized(ast::IsNormalized), + IsNotJson(ast::IsNotJson), + IsNotJsonArray(ast::IsNotJsonArray), + IsNotJsonObject(ast::IsNotJsonObject), + IsNotJsonScalar(ast::IsNotJsonScalar), + IsNotJsonValue(ast::IsNotJsonValue), + IsNotNormalized(ast::IsNotNormalized), + IsNull(SyntaxToken), + NotNull(SyntaxToken), +} + impl ast::BinExpr { #[inline] pub fn lhs(&self) -> Option { @@ -183,57 +192,92 @@ impl ast::BinExpr { SyntaxKind::IS_DISTINCT_FROM => { BinOp::IsDistinctFrom(ast::IsDistinctFrom { syntax: node }) } - SyntaxKind::IS_JSON => BinOp::IsJson(ast::IsJson { syntax: node }), + SyntaxKind::IS_NOT => BinOp::IsNot(ast::IsNot { syntax: node }), + SyntaxKind::IS_NOT_DISTINCT_FROM => { + BinOp::IsNotDistinctFrom(ast::IsNotDistinctFrom { syntax: node }) + } + SyntaxKind::LTEQ => BinOp::Lteq(ast::Lteq { syntax: node }), + SyntaxKind::NEQ => BinOp::Neq(ast::Neq { syntax: node }), + SyntaxKind::NEQB => BinOp::Neqb(ast::Neqb { syntax: node }), + SyntaxKind::NOT_ILIKE => BinOp::NotIlike(ast::NotIlike { syntax: node }), + SyntaxKind::NOT_IN => BinOp::NotIn(ast::NotIn { syntax: node }), + SyntaxKind::NOT_LIKE => BinOp::NotLike(ast::NotLike { syntax: node }), + SyntaxKind::NOT_SIMILAR_TO => { + BinOp::NotSimilarTo(ast::NotSimilarTo { syntax: node }) + } + SyntaxKind::OPERATOR_CALL => { + BinOp::OperatorCall(ast::OperatorCall { syntax: node }) + } + SyntaxKind::SIMILAR_TO => BinOp::SimilarTo(ast::SimilarTo { syntax: node }), + _ => continue, + }; + return Some(op); + } + } + } + None + } +} + +impl ast::PostfixExpr { + pub fn op(&self) -> Option { + let lhs = self.expr()?; + + let siblings = lhs.syntax().siblings_with_tokens(Direction::Next).skip(1); + for child in siblings { + match child { + NodeOrToken::Token(token) => { + let op = match token.kind() { + SyntaxKind::AT_KW => PostfixOp::AtLocal(token), + SyntaxKind::ISNULL_KW => PostfixOp::IsNull(token), + SyntaxKind::NOTNULL_KW => PostfixOp::NotNull(token), + _ => continue, + }; + return Some(op); + } + NodeOrToken::Node(node) => { + let op = match node.kind() { + SyntaxKind::IS_JSON => PostfixOp::IsJson(ast::IsJson { syntax: node }), SyntaxKind::IS_JSON_ARRAY => { - BinOp::IsJsonArray(ast::IsJsonArray { syntax: node }) + PostfixOp::IsJsonArray(ast::IsJsonArray { syntax: node }) } SyntaxKind::IS_JSON_OBJECT => { - BinOp::IsJsonObject(ast::IsJsonObject { syntax: node }) + PostfixOp::IsJsonObject(ast::IsJsonObject { syntax: node }) } SyntaxKind::IS_JSON_SCALAR => { - BinOp::IsJsonScalar(ast::IsJsonScalar { syntax: node }) + PostfixOp::IsJsonScalar(ast::IsJsonScalar { syntax: node }) } SyntaxKind::IS_JSON_VALUE => { - BinOp::IsJsonValue(ast::IsJsonValue { syntax: node }) + PostfixOp::IsJsonValue(ast::IsJsonValue { syntax: node }) } - SyntaxKind::IS_NOT => BinOp::IsNot(ast::IsNot { syntax: node }), - SyntaxKind::IS_NOT_DISTINCT_FROM => { - BinOp::IsNotDistinctFrom(ast::IsNotDistinctFrom { syntax: node }) + SyntaxKind::IS_NORMALIZED => { + PostfixOp::IsNormalized(ast::IsNormalized { syntax: node }) } SyntaxKind::IS_NOT_JSON => { - BinOp::IsNotJson(ast::IsNotJson { syntax: node }) + PostfixOp::IsNotJson(ast::IsNotJson { syntax: node }) } SyntaxKind::IS_NOT_JSON_ARRAY => { - BinOp::IsNotJsonArray(ast::IsNotJsonArray { syntax: node }) + PostfixOp::IsNotJsonArray(ast::IsNotJsonArray { syntax: node }) } SyntaxKind::IS_NOT_JSON_OBJECT => { - BinOp::IsNotJsonObject(ast::IsNotJsonObject { syntax: node }) + PostfixOp::IsNotJsonObject(ast::IsNotJsonObject { syntax: node }) } SyntaxKind::IS_NOT_JSON_SCALAR => { - BinOp::IsNotJsonScalar(ast::IsNotJsonScalar { syntax: node }) + PostfixOp::IsNotJsonScalar(ast::IsNotJsonScalar { syntax: node }) } SyntaxKind::IS_NOT_JSON_VALUE => { - BinOp::IsNotJsonValue(ast::IsNotJsonValue { syntax: node }) + PostfixOp::IsNotJsonValue(ast::IsNotJsonValue { syntax: node }) } - SyntaxKind::LTEQ => BinOp::Lteq(ast::Lteq { syntax: node }), - SyntaxKind::NEQ => BinOp::Neq(ast::Neq { syntax: node }), - SyntaxKind::NEQB => BinOp::Neqb(ast::Neqb { syntax: node }), - SyntaxKind::NOT_ILIKE => BinOp::NotIlike(ast::NotIlike { syntax: node }), - SyntaxKind::NOT_IN => BinOp::NotIn(ast::NotIn { syntax: node }), - SyntaxKind::NOT_LIKE => BinOp::NotLike(ast::NotLike { syntax: node }), - SyntaxKind::NOT_SIMILAR_TO => { - BinOp::NotSimilarTo(ast::NotSimilarTo { syntax: node }) + SyntaxKind::IS_NOT_NORMALIZED => { + PostfixOp::IsNotNormalized(ast::IsNotNormalized { syntax: node }) } - SyntaxKind::OPERATOR_CALL => { - BinOp::OperatorCall(ast::OperatorCall { syntax: node }) - } - SyntaxKind::SIMILAR_TO => BinOp::SimilarTo(ast::SimilarTo { syntax: node }), _ => continue, }; return Some(op); } } } + None } } diff --git a/crates/squawk_syntax/src/postgresql.ungram b/crates/squawk_syntax/src/postgresql.ungram index 98b204a1..3353096b 100644 --- a/crates/squawk_syntax/src/postgresql.ungram +++ b/crates/squawk_syntax/src/postgresql.ungram @@ -89,6 +89,7 @@ CallExpr = | SomeFn | AnyFn | AllFn +| CollationForFn | ExistsFn JsonArrayFn = @@ -330,6 +331,9 @@ XmlPiFn = (',' Expr)? ')' +CollationForFn = + 'collation' 'for' '(' Expr ')' + ExistsFn = 'exists' '(' SelectVariant ')' From ace7818761875197369a6021eb3e7bd2de0c6344 Mon Sep 17 00:00:00 2001 From: Steve Dignam Date: Sat, 21 Feb 2026 16:46:11 -0500 Subject: [PATCH 2/2] fix --- crates/squawk_ide/src/classify.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/crates/squawk_ide/src/classify.rs b/crates/squawk_ide/src/classify.rs index 03381ee6..107412b1 100644 --- a/crates/squawk_ide/src/classify.rs +++ b/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