diff --git a/crates/squawk_ide/src/goto_definition.rs b/crates/squawk_ide/src/goto_definition.rs index 883e115b..15cd9330 100644 --- a/crates/squawk_ide/src/goto_definition.rs +++ b/crates/squawk_ide/src/goto_definition.rs @@ -15302,4 +15302,67 @@ select ordinality$0 from rows from (unnest(array[1,2])) with ordinality; ╰╴ ─ 1. source ────────── 2. destination "); } + + #[test] + fn goto_rows_from_alias_column_with_per_item_column_def_list() { + assert_snapshot!(goto(" +select z.q$0 from rows from (unnest(array[1,2]) as (x int)) as z(q); +"), @" + ╭▸ + 2 │ select z.q from rows from (unnest(array[1,2]) as (x int)) as z(q); + ╰╴ ─ 1. source ─ 2. destination + "); + } + + #[test] + fn goto_rows_from_column_def_after_partial_item_alias() { + assert_snapshot!(goto(" +create function f_rec() returns record as $$ select 1 $$ language sql; +select x$0 from rows from (unnest(array[1,2]), f_rec() as (x int)) as z(c1); +"), @" + ╭▸ + 3 │ select x from rows from (unnest(array[1,2]), f_rec() as (x int)) as z(c1); + ╰╴ ─ 1. source ─ 2. destination + "); + } + + #[test] + fn goto_rows_from_applies_outer_alias_skip_across_items() { + assert_snapshot!(goto(" +create function f() returns record as $$ select 1 $$ language sql; +create function g() returns record as $$ select 1 $$ language sql; +select y$0 from rows from (f() as (x int), g() as (y int)) as z(q); +"), @" + ╭▸ + 4 │ select y from rows from (f() as (x int), g() as (y int)) as z(q); + ╰╴ ─ 1. source ─ 2. destination + "); + } + + #[test] + fn goto_rows_from_per_item_column_def_list_column() { + assert_snapshot!(goto(" +select x$0 from rows from (unnest(array[1,2]) as (x int)); +"), @" + ╭▸ + 2 │ select x from rows from (unnest(array[1,2]) as (x int)); + ╰╴ ─ 1. source ─ 2. destination + "); + } + + #[test] + fn goto_rows_from_column_after_outer_alias_for_composite_return() { + assert_snapshot!(goto(" +create type pair as (a int, b int); +create function f() returns setof pair as $$ select 1, 2 $$ language sql; +create function g() returns table(y int) as $$ select 3 $$ language sql; +select y$0 from rows from (f(), g()) as z(q, r); +"), @" + ╭▸ + 4 │ create function g() returns table(y int) as $$ select 3 $$ language sql; + │ ─ 2. destination + 5 │ select y from rows from (f(), g()) as z(q, r); + ╰╴ ─ 1. source + "); + } } diff --git a/crates/squawk_ide/src/resolve.rs b/crates/squawk_ide/src/resolve.rs index 8ecf892a..8d8e6156 100644 --- a/crates/squawk_ide/src/resolve.rs +++ b/crates/squawk_ide/src/resolve.rs @@ -2445,15 +2445,43 @@ fn resolve_from_item_column_by_name_after_index( } if let ast::FromItem::RowsFromItem(rows_from) = from_item { - for call_expr in rows_from.call_exprs() { - if let Some(ptr) = resolve_column_from_call_expr_return_table( - db, - InFile::new(file, &call_expr), - scope_name_ref, - column_name, - skip_column_count, - ) { - return Some(ptr); + let mut remaining_skip = skip_column_count; + for arg in rows_from.rows_from_args() { + if let Some(column_def_list) = arg + .column_def_list() + .filter(|alias| alias.column_list().is_some()) + { + let (column_count, column) = resolve_column_list_column( + file, + alias_column_names(Some(column_def_list)), + column_name, + remaining_skip, + ); + if let Some(column) = column { + return Some(column); + } + remaining_skip = remaining_skip.saturating_sub(column_count); + continue; + } + if let Some(call_expr) = arg.call_expr() { + if let Some(ptr) = resolve_column_from_call_expr_return_table( + db, + InFile::new(file, &call_expr), + scope_name_ref, + column_name, + remaining_skip, + ) { + return Some(ptr); + } + if remaining_skip > 0 { + let column_count = count_columns_for_call_expr_return_table( + db, + InFile::new(file, &call_expr), + scope_name_ref, + ) + .unwrap_or(1); + remaining_skip = remaining_skip.saturating_sub(column_count); + } } } } @@ -4892,6 +4920,60 @@ fn resolve_json_table_column( None } +fn count_columns_for_call_expr_return_table( + db: &dyn Db, + call_expr: InFile<&ast::CallExpr>, + name_ref: &impl ast::NameLike, +) -> Option { + let file = call_expr.file_id; + let call_expr = call_expr.value; + let position = name_ref.syntax().text_range().start(); + let (schema, function_name) = name::schema_and_func_name(call_expr)?; + let schemas = bind(db, file).resolved_schemas(position, schema.as_ref()); + let function_locs = resolve_function(db, &function_name, &schemas, None, file)?; + let function_loc = function_locs.first()?; + let function_node = function_loc.to_node(db)?; + let create_function = function_node + .ancestors() + .find_map(ast::CreateFunction::cast)?; + + if let Some(table_arg_list) = create_function.ret_type().and_then(|r| r.table_arg_list()) { + return Some( + table_arg_list + .args() + .filter(|arg| matches!(arg, ast::TableArg::Column(_))) + .count(), + ); + } + + if let Some(param_list) = create_function.param_list() { + let output_count = param_list + .params() + .filter(|param| { + matches!( + param.mode(), + Some(ast::ParamMode::ParamInOut(_) | ast::ParamMode::ParamOut(_)) + ) + }) + .count(); + if output_count > 0 { + return Some(output_count); + } + } + + if let Some(ast::Type::PathType(path_type)) = create_function.ret_type().and_then(|r| r.ty()) + && let Some(path) = path_type.path_ref() + && let Some(column_count) = + count_columns_for_path(db, InFile::new(function_loc.file, &path)).or_else(|| { + count_columns_for_composite_type_path(db, InFile::new(function_loc.file, &path)) + }) + { + return Some(column_count); + } + + create_function.ret_type().map(|_| 1) +} + fn resolve_column_from_call_expr_return_table( db: &dyn Db, call_expr: InFile<&ast::CallExpr>, @@ -5142,6 +5224,30 @@ fn resolve_composite_type_field_ptr( composite_type_field_location(file, &type_node, &field_name) } +fn count_columns_for_composite_type_path( + db: &dyn Db, + path: InFile<&ast::PathRef>, +) -> Option { + let file = path.file_id; + let path = path.value; + let (schema, type_name) = name::schema_and_name_path(path)?; + let position = path.syntax().text_range().start(); + let schemas = bind(db, file).resolved_schemas(position, schema.as_ref()); + let type_name_ptr = resolve_type_name_ptr(db, &type_name, &schemas, file)?; + let tree = parse(db, file).tree(); + let type_node = type_name_ptr.to_node(tree.syntax()); + let create_type = type_node.ancestors().find_map(ast::CreateType::cast)?; + let ast::CreateTypeKind::CompositeType(composite) = create_type.kind()? else { + return None; + }; + Some( + composite + .composite_field_list()? + .composite_field_defs() + .count(), + ) +} + fn resolve_composite_type_field_for_path( db: &dyn Db, path: InFile<&ast::PathRef>, diff --git a/crates/squawk_parser/src/generated/syntax_kind.rs b/crates/squawk_parser/src/generated/syntax_kind.rs index 0fbdd59d..0923e893 100644 --- a/crates/squawk_parser/src/generated/syntax_kind.rs +++ b/crates/squawk_parser/src/generated/syntax_kind.rs @@ -1608,6 +1608,7 @@ pub enum SyntaxKind { ROUTINE_SIG, ROUTINE_SIG_LIST, ROW, + ROWS_FROM_ARG, ROWS_FROM_ITEM, ROWS_FUNC_OPTION, ROW_EXCLUSIVE, diff --git a/crates/squawk_parser/src/grammar.rs b/crates/squawk_parser/src/grammar.rs index 883dca6b..78da919c 100644 --- a/crates/squawk_parser/src/grammar.rs +++ b/crates/squawk_parser/src/grammar.rs @@ -3801,8 +3801,10 @@ fn opt_row_from_expr(p: &mut Parser<'_>) -> bool { if !p.at_ts(EXPR_FIRST) { return false; } + let m = p.start(); call_expr(p); opt_from_alias(p); + m.complete(p, ROWS_FROM_ARG); true } diff --git a/crates/squawk_parser/tests/data/ok/select.sql b/crates/squawk_parser/tests/data/ok/select.sql index 7e0aacd1..3142f56a 100644 --- a/crates/squawk_parser/tests/data/ok/select.sql +++ b/crates/squawk_parser/tests/data/ok/select.sql @@ -233,6 +233,9 @@ select * from rows from(f()) with ordinality; select * from rows from(f()) with ordinality as t; select * from rows from(f()) as t(a, b, c); select * from rows from(f()) t(a, b, c); +select * from rows from(f() as (x int), g() as (y int)); +select * from rows from(f() as (x int)) as t(a); +select * from rows from(f() as (x int)) with ordinality as t(a, n); -- select_with_where_clause -- simple diff --git a/crates/squawk_parser/tests/snapshots/tests__merge_ok.snap b/crates/squawk_parser/tests/snapshots/tests__merge_ok.snap index 1cb11fee..0799f4c6 100644 --- a/crates/squawk_parser/tests/snapshots/tests__merge_ok.snap +++ b/crates/squawk_parser/tests/snapshots/tests__merge_ok.snap @@ -147,20 +147,21 @@ SOURCE_FILE FROM_KW "from" WHITESPACE " " L_PAREN "(" - CALL_EXPR - NAME_REF - IDENT "f" - ARG_LIST - L_PAREN "(" - ARG - LITERAL - INT_NUMBER "1" - COMMA "," - WHITESPACE " " - ARG - LITERAL - INT_NUMBER "2" - R_PAREN ")" + ROWS_FROM_ARG + CALL_EXPR + NAME_REF + IDENT "f" + ARG_LIST + L_PAREN "(" + ARG + LITERAL + INT_NUMBER "1" + COMMA "," + WHITESPACE " " + ARG + LITERAL + INT_NUMBER "2" + R_PAREN ")" R_PAREN ")" WHITESPACE "\n " ON_CLAUSE diff --git a/crates/squawk_parser/tests/snapshots/tests__select_ok.snap b/crates/squawk_parser/tests/snapshots/tests__select_ok.snap index e5aed5ba..a8499d05 100644 --- a/crates/squawk_parser/tests/snapshots/tests__select_ok.snap +++ b/crates/squawk_parser/tests/snapshots/tests__select_ok.snap @@ -3342,20 +3342,21 @@ SOURCE_FILE WHITESPACE " " FROM_KW "from" L_PAREN "(" - CALL_EXPR - NAME_REF - IDENT "f" - ARG_LIST - L_PAREN "(" - ARG - NAME_REF - IDENT "a" - COMMA "," - WHITESPACE " " - ARG - NAME_REF - IDENT "b" - R_PAREN ")" + ROWS_FROM_ARG + CALL_EXPR + NAME_REF + IDENT "f" + ARG_LIST + L_PAREN "(" + ARG + NAME_REF + IDENT "a" + COMMA "," + WHITESPACE " " + ARG + NAME_REF + IDENT "b" + R_PAREN ")" R_PAREN ")" WHITESPACE " " FROM_ALIAS @@ -3410,20 +3411,21 @@ SOURCE_FILE WHITESPACE " " FROM_KW "from" L_PAREN "(" - CALL_EXPR - NAME_REF - IDENT "f" - ARG_LIST - L_PAREN "(" - ARG - NAME_REF - IDENT "a" - COMMA "," - WHITESPACE " " - ARG - NAME_REF - IDENT "b" - R_PAREN ")" + ROWS_FROM_ARG + CALL_EXPR + NAME_REF + IDENT "f" + ARG_LIST + L_PAREN "(" + ARG + NAME_REF + IDENT "a" + COMMA "," + WHITESPACE " " + ARG + NAME_REF + IDENT "b" + R_PAREN ")" R_PAREN ")" SEMICOLON ";" WHITESPACE "\n" @@ -3443,12 +3445,13 @@ SOURCE_FILE WHITESPACE " " FROM_KW "from" L_PAREN "(" - CALL_EXPR - NAME_REF - IDENT "f" - ARG_LIST - L_PAREN "(" - R_PAREN ")" + ROWS_FROM_ARG + CALL_EXPR + NAME_REF + IDENT "f" + ARG_LIST + L_PAREN "(" + R_PAREN ")" R_PAREN ")" SEMICOLON ";" WHITESPACE "\n" @@ -3468,12 +3471,13 @@ SOURCE_FILE WHITESPACE " " FROM_KW "from" L_PAREN "(" - CALL_EXPR - NAME_REF - IDENT "f" - ARG_LIST - L_PAREN "(" - R_PAREN ")" + ROWS_FROM_ARG + CALL_EXPR + NAME_REF + IDENT "f" + ARG_LIST + L_PAREN "(" + R_PAREN ")" R_PAREN ")" WHITESPACE " " WITH_ORDINALITY @@ -3498,12 +3502,13 @@ SOURCE_FILE WHITESPACE " " FROM_KW "from" L_PAREN "(" - CALL_EXPR - NAME_REF - IDENT "f" - ARG_LIST - L_PAREN "(" - R_PAREN ")" + ROWS_FROM_ARG + CALL_EXPR + NAME_REF + IDENT "f" + ARG_LIST + L_PAREN "(" + R_PAREN ")" R_PAREN ")" WHITESPACE " " WITH_ORDINALITY @@ -3534,12 +3539,13 @@ SOURCE_FILE WHITESPACE " " FROM_KW "from" L_PAREN "(" - CALL_EXPR - NAME_REF - IDENT "f" - ARG_LIST - L_PAREN "(" - R_PAREN ")" + ROWS_FROM_ARG + CALL_EXPR + NAME_REF + IDENT "f" + ARG_LIST + L_PAREN "(" + R_PAREN ")" R_PAREN ")" WHITESPACE " " FROM_ALIAS @@ -3581,12 +3587,13 @@ SOURCE_FILE WHITESPACE " " FROM_KW "from" L_PAREN "(" - CALL_EXPR - NAME_REF - IDENT "f" - ARG_LIST - L_PAREN "(" - R_PAREN ")" + ROWS_FROM_ARG + CALL_EXPR + NAME_REF + IDENT "f" + ARG_LIST + L_PAREN "(" + R_PAREN ")" R_PAREN ")" WHITESPACE " " FROM_ALIAS @@ -3609,6 +3616,187 @@ SOURCE_FILE IDENT "c" R_PAREN ")" SEMICOLON ";" + WHITESPACE "\n" + SELECT + SELECT_CLAUSE + SELECT_KW "select" + WHITESPACE " " + TARGET_LIST + TARGET + STAR "*" + WHITESPACE " " + FROM_CLAUSE + FROM_KW "from" + WHITESPACE " " + ROWS_FROM_ITEM + ROWS_KW "rows" + WHITESPACE " " + FROM_KW "from" + L_PAREN "(" + ROWS_FROM_ARG + CALL_EXPR + NAME_REF + IDENT "f" + ARG_LIST + L_PAREN "(" + R_PAREN ")" + WHITESPACE " " + FROM_ALIAS + AS_KW "as" + WHITESPACE " " + ALIAS_COLUMN_LIST + L_PAREN "(" + ALIAS_COLUMN + COLUMN_NAME + IDENT "x" + WHITESPACE " " + PATH_TYPE + PATH_REF + PATH_SEGMENT_REF + INT_KW "int" + R_PAREN ")" + COMMA "," + WHITESPACE " " + ROWS_FROM_ARG + CALL_EXPR + NAME_REF + IDENT "g" + ARG_LIST + L_PAREN "(" + R_PAREN ")" + WHITESPACE " " + FROM_ALIAS + AS_KW "as" + WHITESPACE " " + ALIAS_COLUMN_LIST + L_PAREN "(" + ALIAS_COLUMN + COLUMN_NAME + IDENT "y" + WHITESPACE " " + PATH_TYPE + PATH_REF + PATH_SEGMENT_REF + INT_KW "int" + R_PAREN ")" + R_PAREN ")" + SEMICOLON ";" + WHITESPACE "\n" + SELECT + SELECT_CLAUSE + SELECT_KW "select" + WHITESPACE " " + TARGET_LIST + TARGET + STAR "*" + WHITESPACE " " + FROM_CLAUSE + FROM_KW "from" + WHITESPACE " " + ROWS_FROM_ITEM + ROWS_KW "rows" + WHITESPACE " " + FROM_KW "from" + L_PAREN "(" + ROWS_FROM_ARG + CALL_EXPR + NAME_REF + IDENT "f" + ARG_LIST + L_PAREN "(" + R_PAREN ")" + WHITESPACE " " + FROM_ALIAS + AS_KW "as" + WHITESPACE " " + ALIAS_COLUMN_LIST + L_PAREN "(" + ALIAS_COLUMN + COLUMN_NAME + IDENT "x" + WHITESPACE " " + PATH_TYPE + PATH_REF + PATH_SEGMENT_REF + INT_KW "int" + R_PAREN ")" + R_PAREN ")" + WHITESPACE " " + FROM_ALIAS + AS_KW "as" + WHITESPACE " " + TABLE_ALIAS + IDENT "t" + ALIAS_COLUMN_LIST + L_PAREN "(" + ALIAS_COLUMN + COLUMN_NAME + IDENT "a" + R_PAREN ")" + SEMICOLON ";" + WHITESPACE "\n" + SELECT + SELECT_CLAUSE + SELECT_KW "select" + WHITESPACE " " + TARGET_LIST + TARGET + STAR "*" + WHITESPACE " " + FROM_CLAUSE + FROM_KW "from" + WHITESPACE " " + ROWS_FROM_ITEM + ROWS_KW "rows" + WHITESPACE " " + FROM_KW "from" + L_PAREN "(" + ROWS_FROM_ARG + CALL_EXPR + NAME_REF + IDENT "f" + ARG_LIST + L_PAREN "(" + R_PAREN ")" + WHITESPACE " " + FROM_ALIAS + AS_KW "as" + WHITESPACE " " + ALIAS_COLUMN_LIST + L_PAREN "(" + ALIAS_COLUMN + COLUMN_NAME + IDENT "x" + WHITESPACE " " + PATH_TYPE + PATH_REF + PATH_SEGMENT_REF + INT_KW "int" + R_PAREN ")" + R_PAREN ")" + WHITESPACE " " + WITH_ORDINALITY + WITH_KW "with" + WHITESPACE " " + ORDINALITY_KW "ordinality" + WHITESPACE " " + FROM_ALIAS + AS_KW "as" + WHITESPACE " " + TABLE_ALIAS + IDENT "t" + ALIAS_COLUMN_LIST + L_PAREN "(" + ALIAS_COLUMN + COLUMN_NAME + IDENT "a" + COMMA "," + WHITESPACE " " + ALIAS_COLUMN + COLUMN_NAME + IDENT "n" + R_PAREN ")" + SEMICOLON ";" WHITESPACE "\n\n" COMMENT "-- select_with_where_clause" WHITESPACE "\n" diff --git a/crates/squawk_syntax/src/ast/generated/nodes.rs b/crates/squawk_syntax/src/ast/generated/nodes.rs index 4796c3c7..f136e44f 100644 --- a/crates/squawk_syntax/src/ast/generated/nodes.rs +++ b/crates/squawk_syntax/src/ast/generated/nodes.rs @@ -22275,6 +22275,21 @@ impl RowShare { } } +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct RowsFromArg { + pub(crate) syntax: SyntaxNode, +} +impl RowsFromArg { + #[inline] + pub fn call_expr(&self) -> Option { + support::child(&self.syntax) + } + #[inline] + pub fn column_def_list(&self) -> Option { + support::child(&self.syntax) + } +} + #[derive(Debug, Clone, PartialEq, Eq, Hash)] pub struct RowsFromItem { pub(crate) syntax: SyntaxNode, @@ -22285,7 +22300,7 @@ impl RowsFromItem { support::child(&self.syntax) } #[inline] - pub fn call_exprs(&self) -> AstChildren { + pub fn rows_from_args(&self) -> AstChildren { support::children(&self.syntax) } #[inline] @@ -48386,6 +48401,24 @@ impl AstNode for RowShare { &self.syntax } } +impl AstNode for RowsFromArg { + #[inline] + fn can_cast(kind: SyntaxKind) -> bool { + kind == SyntaxKind::ROWS_FROM_ARG + } + #[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 RowsFromItem { #[inline] fn can_cast(kind: SyntaxKind) -> bool { diff --git a/crates/squawk_syntax/src/postgresql.ungram b/crates/squawk_syntax/src/postgresql.ungram index 875d33d6..d7069032 100644 --- a/crates/squawk_syntax/src/postgresql.ungram +++ b/crates/squawk_syntax/src/postgresql.ungram @@ -2911,7 +2911,10 @@ ExprFromItem = 'lateral'? (CastExpr | CollationForFn) alias:FromAlias? RowsFromItem = - 'lateral'? 'rows' 'from' '(' (CallExpr (',' CallExpr)*) ')' WithOrdinality? alias:FromAlias? + 'lateral'? 'rows' 'from' '(' (RowsFromArg (',' RowsFromArg)*) ')' WithOrdinality? alias:FromAlias? + +RowsFromArg = + CallExpr column_def_list:FromAlias? FromClause = 'from'