diff --git a/crates/squawk_ide/src/classify.rs b/crates/squawk_ide/src/classify.rs index 6b3cdcae..8f1f52a8 100644 --- a/crates/squawk_ide/src/classify.rs +++ b/crates/squawk_ide/src/classify.rs @@ -724,6 +724,13 @@ pub(crate) fn classify_name_ref(node: &SyntaxNode) -> Option { // (anything in SELECT except FROM clause) return Some(NameRefClass::SelectColumn); } + if ast::CompoundSelect::can_cast(ancestor.kind()) + && in_order_by_clause + && let Some(parent) = node.parent() + && ast::SortBy::can_cast(parent.kind()) + { + return Some(NameRefClass::SelectOrderByAliasOrColumn); + } if ast::ColumnList::can_cast(ancestor.kind()) { in_column_list = true; } diff --git a/crates/squawk_ide/src/goto_definition.rs b/crates/squawk_ide/src/goto_definition.rs index b7f60dc1..713f80d7 100644 --- a/crates/squawk_ide/src/goto_definition.rs +++ b/crates/squawk_ide/src/goto_definition.rs @@ -2388,6 +2388,49 @@ select a$0 from t; "); } + #[test] + fn goto_view_select_star_column_gap() { + assert_snapshot!(goto(" +create table t(a int, b int); +create view v as select * from t; +select a$0 from v; +"), @" + ╭▸ + 2 │ create table t(a int, b int); + │ ─ 2. destination + 3 │ create view v as select * from t; + 4 │ select a from v; + ╰╴ ─ 1. source + "); + } + + #[test] + fn goto_view_table_query_column_gap() { + assert_snapshot!(goto(" +create table t(a int); +create view v as table t; +select a$0 from v; +"), @" + ╭▸ + 2 │ create table t(a int); + │ ─ 2. destination + 3 │ create view v as table t; + 4 │ select a from v; + ╰╴ ─ 1. source + "); + } + + #[test] + fn goto_values_partial_alias_remaining_column_gap() { + assert_snapshot!(goto(" +select column2$0 from (values (1, 2)) v(a); +"), @" + ╭▸ + 2 │ select column2 from (values (1, 2)) v(a); + ╰╴ ─ 1. source ─ 2. destination + "); + } + #[test] fn goto_create_table_inherits() { assert_snapshot!(goto(" @@ -5504,6 +5547,25 @@ select a$0 from inserted; "); } + #[test] + fn goto_cte_returning_qualified_star_column_gap() { + assert_snapshot!(goto(" +create table t(a int, b int); +with changed as ( + insert into t values (1, 2) + returning new.* +) +select a$0 from changed; +"), @" + ╭▸ + 2 │ create table t(a int, b int); + │ ─ 2. destination + ‡ + 7 │ select a from changed; + ╰╴ ─ 1. source + "); + } + #[test] fn goto_cte_delete_returning_star_column() { assert_snapshot!(goto(" @@ -5592,6 +5654,16 @@ select a$0 from (select 1 a); "); } + #[test] + fn goto_subquery_star_partial_alias_masks_original_column_gap() { + goto_not_found( + " +create table t(a int, b int); +select a$0 from (select * from t) u(x); +", + ); + } + #[test] fn goto_subquery_column_with_as() { assert_snapshot!(goto(" @@ -5614,6 +5686,18 @@ select c$0 from (select 1 c union select 2 c); "); } + #[test] + fn goto_subquery_compound_select_column_order_by() { + assert_snapshot!(goto(" +with t as (select 1 a) +select 2 a from t union select 1 order by a$0; +"), @" + ╭▸ + 3 │ select 2 a from t union select 1 order by a; + ╰╴ ─ 2. destination ─ 1. source + "); + } + #[test] fn goto_subquery_compound_select_column_with_nested_parens() { assert_snapshot!(goto(" diff --git a/crates/squawk_ide/src/resolve.rs b/crates/squawk_ide/src/resolve.rs index 536a808c..53f32275 100644 --- a/crates/squawk_ide/src/resolve.rs +++ b/crates/squawk_ide/src/resolve.rs @@ -1048,8 +1048,8 @@ enum ReturningClauseMatch { fn match_table_in_returning_clause( table_name: &Name, stmt_table_name: &Name, - alias: Option, - returning_clause: Option, + alias: Option<&ast::Alias>, + returning_clause: Option<&ast::ReturningClause>, ) -> Option { // Check `returning with (old as alias, new as alias)` if let Some(option_list) = returning_clause.and_then(|x| x.returning_option_list()) { @@ -1253,7 +1253,8 @@ fn resolve_column_for_table( let file = resolved.file_id; match resolved.value { ResolvedTableName::View(create_view) => { - if let Some(ptr) = find_column_in_create_view_like(file, &create_view, column_name) { + if let Some(ptr) = find_column_in_create_view_like(db, file, &create_view, column_name) + { return Some(ptr); } return resolve_function(db, column_name, schemas, None, file); @@ -1367,21 +1368,35 @@ fn resolve_from_item_column_by_name( from_item: InFile<&ast::FromItem>, scope_name_ref: &ast::NameRef, column_name: &Name, +) -> Option> { + resolve_from_item_column_by_name_after_index(db, from_item, scope_name_ref, column_name, 0) +} + +fn resolve_from_item_column_by_name_after_index( + db: &dyn Db, + from_item: InFile<&ast::FromItem>, + scope_name_ref: &ast::NameRef, + column_name: &Name, + skip_column_count: usize, ) -> Option> { let file = from_item.file_id; let from_item = from_item.value; + let original_skip = skip_column_count; + if let Some(paren_select) = from_item.paren_select() { let alias = from_item.alias(); - if let Some(ptr) = resolve_subquery_column_ptr( + if let Some(ptr) = resolve_subquery_column_ptr_with_skip( db, InFile::new(file, &paren_select), scope_name_ref, column_name, alias.as_ref(), + skip_column_count, ) { return Some(ptr); } - if let Some(alias_name) = alias.and_then(|x| x.name()) + if original_skip == 0 + && let Some(alias_name) = alias.and_then(|x| x.name()) && Name::from_node(&alias_name) == *column_name { return Some(smallvec![Location::new( @@ -1394,29 +1409,28 @@ fn resolve_from_item_column_by_name( } if let Some(paren_expr) = from_item.paren_expr() { - if let Some(column_list) = from_item.alias().and_then(|x| x.column_list()) { - for col in column_list.columns() { - if let Some(col_name) = col.name() - && Name::from_node(&col_name) == *column_name - { - return Some(smallvec![Location::new( - file, - col_name.syntax().text_range(), - LocationKind::Column - )]); - } - } + let (alias_len, alias_column) = resolve_column_list_column( + file, + from_item.alias().and_then(|x| x.column_list()), + column_name, + skip_column_count, + ); + if let Some(alias_column) = alias_column { + return Some(alias_column); } + let skip_column_count = skip_column_count.max(alias_len); - if let Some(ptr) = resolve_column_from_paren_expr( + if let Some(ptr) = resolve_column_from_paren_expr_with_skip( db, InFile::new(file, &paren_expr), scope_name_ref, column_name, + skip_column_count, ) { return Some(ptr); } - if let Some(alias_name) = from_item.alias().and_then(|x| x.name()) + if original_skip == 0 + && let Some(alias_name) = from_item.alias().and_then(|x| x.name()) && Name::from_node(&alias_name) == *column_name { return Some(smallvec![Location::new( @@ -1428,22 +1442,31 @@ fn resolve_from_item_column_by_name( return None; } - let alias_len = if let Some(column_list) = from_item.alias().and_then(|x| x.column_list()) { - for col in column_list.columns() { - if let Some(col_name) = col.name() - && Name::from_node(&col_name) == *column_name - { - return Some(smallvec![Location::new( - file, - col_name.syntax().text_range(), - LocationKind::Column - )]); - } - } - column_list.columns().count() - } else { - 0 - }; + if let Some(select_variant) = from_item + .syntax() + .children() + .find_map(ast::SelectVariant::cast) + { + let select = ast_nav::select_from_variant(select_variant)?; + return resolve_column_from_select_targets( + db, + InFile::new(file, &select), + scope_name_ref, + column_name, + skip_column_count, + ); + } + + let (alias_len, alias_column) = resolve_column_list_column( + file, + from_item.alias().and_then(|x| x.column_list()), + column_name, + skip_column_count, + ); + if let Some(alias_column) = alias_column { + return Some(alias_column); + } + let skip_column_count = skip_column_count.max(alias_len); if let Some(call_expr) = from_item.call_expr() && let Some(ptr) = resolve_column_from_call_expr_return_table( @@ -1451,7 +1474,7 @@ fn resolve_from_item_column_by_name( InFile::new(file, &call_expr), scope_name_ref, column_name, - alias_len, + skip_column_count, ) { return Some(ptr); @@ -1460,17 +1483,20 @@ fn resolve_from_item_column_by_name( let (schema, table_name) = name::schema_and_table_from_from_item(from_item)?; let scope_name_ref = relation_name_ref_from_from_item(from_item)?; - if let Some(ptr) = resolve_column_from_table_or_view_or_cte( + if let Some(ptr) = resolve_column_from_table_or_view_or_cte_impl( db, InFile::new(file, &scope_name_ref), &table_name, schema.as_ref(), column_name, + 0, + skip_column_count, ) { return Some(ptr); } - if let Some(alias_name) = from_item.alias().and_then(|x| x.name()) + if original_skip == 0 + && let Some(alias_name) = from_item.alias().and_then(|x| x.name()) && Name::from_node(&alias_name) == *column_name { return Some(smallvec![Location::new( @@ -1490,7 +1516,15 @@ fn resolve_column_from_table_or_view_or_cte( schema: Option<&Schema>, column_name: &Name, ) -> Option> { - resolve_column_from_table_or_view_or_cte_impl(db, name_ref, table_name, schema, column_name, 0) + resolve_column_from_table_or_view_or_cte_impl( + db, + name_ref, + table_name, + schema, + column_name, + 0, + 0, + ) } fn resolve_column_from_table_or_view_or_cte_impl( @@ -1500,6 +1534,7 @@ fn resolve_column_from_table_or_view_or_cte_impl( schema: Option<&Schema>, column_name: &Name, depth: u32, + skip_column_count: usize, ) -> Option> { let file = name_ref.file_id; let name_ref = name_ref.value; @@ -1516,11 +1551,12 @@ fn resolve_column_from_table_or_view_or_cte_impl( match kind { LocationKind::Table => { if schemas.unqualified() && resolve_cte_table(name_ref, table_name).is_some() { - return resolve_cte_column( + return resolve_cte_column_with_skip( db, InFile::new(file, name_ref), table_name, column_name, + skip_column_count, ); } @@ -1529,16 +1565,21 @@ fn resolve_column_from_table_or_view_or_cte_impl( let node = table_like_ptr.to_node(root); if let Some(create_table) = node.ancestors().find_map(ast::CreateTableLike::cast) { - if let Some(cols) = - find_column_in_create_table(db, InFile::new(file, &create_table), column_name) - { + if let Some(cols) = find_column_in_create_table_impl( + db, + InFile::new(file, &create_table), + column_name, + 0, + skip_column_count, + ) { return Some(cols); } // check if this is a partitioned table - if let Some(parent_path) = ast::CreateTable::cast(create_table.syntax().clone()) - .and_then(|x| x.partition_of()) - .and_then(|x| x.path()) + if skip_column_count == 0 + && let Some(parent_path) = ast::CreateTable::cast(create_table.syntax().clone()) + .and_then(|x| x.partition_of()) + .and_then(|x| x.path()) { let (parent_schema, parent_table_name) = name::schema_and_name_path(&parent_path)?; @@ -1549,6 +1590,7 @@ fn resolve_column_from_table_or_view_or_cte_impl( parent_schema.as_ref(), column_name, depth + 1, + skip_column_count, ); } @@ -1557,7 +1599,7 @@ fn resolve_column_from_table_or_view_or_cte_impl( // create table t(a int); // select t from t; // ``` - if column_name == table_name { + if skip_column_count == 0 && column_name == table_name { return Some(smallvec![Location::new( file, table_like_ptr.text_range(), @@ -1567,15 +1609,16 @@ fn resolve_column_from_table_or_view_or_cte_impl( } if let Some(create_table_as) = node.ancestors().find_map(ast::CreateTableAs::cast) { - if let Some(cols) = find_column_in_create_table_as( + if let Some(cols) = find_column_in_create_table_as_with_skip( db, InFile::new(file, &create_table_as), column_name, + skip_column_count, ) { return Some(cols); } - if column_name == table_name { + if skip_column_count == 0 && column_name == table_name { return Some(smallvec![Location::new( file, table_like_ptr.text_range(), @@ -1585,12 +1628,17 @@ fn resolve_column_from_table_or_view_or_cte_impl( } if let Some(select_into) = node.ancestors().find_map(ast::SelectInto::cast) { - if let Some(cols) = find_column_in_select_into(db, file, &select_into, column_name) - { + if let Some(cols) = find_column_in_select_into_with_skip( + db, + file, + &select_into, + column_name, + skip_column_count, + ) { return Some(cols); } - if column_name == table_name { + if skip_column_count == 0 && column_name == table_name { return Some(smallvec![Location::new( file, table_like_ptr.text_range(), @@ -1607,12 +1655,17 @@ fn resolve_column_from_table_or_view_or_cte_impl( let node = table_like_ptr.to_node(root); if let Some(create_view) = node.ancestors().find_map(ast::CreateViewLike::cast) { - if let Some(cols) = find_column_in_create_view_like(file, &create_view, column_name) - { + if let Some(cols) = find_column_in_create_view_like_with_skip( + db, + file, + &create_view, + column_name, + skip_column_count, + ) { return Some(cols); } - if column_name == table_name { + if skip_column_count == 0 && column_name == table_name { return Some(smallvec![Location::new( file, table_like_ptr.text_range(), @@ -1633,18 +1686,31 @@ fn resolve_from_item_for_cte_star( name_ref: &ast::NameRef, cte_name: &Name, column_name: &Name, + skip_column_count: usize, ) -> Option> { let file = from_item.file_id; if let Some((schema, table_name)) = name::schema_and_table_from_from_item(from_item.value) && table_name == *cte_name { let scope_name_ref = relation_name_ref_from_from_item(from_item.value)?; - let position = scope_name_ref.syntax().text_range().start(); - let schemas = bind(db, file).resolved_schemas(position, schema.as_ref()); - return resolve_column_for_table(db, &table_name, &schemas, column_name, file); + return resolve_column_from_table_or_view_or_cte_impl( + db, + InFile::new(file, &scope_name_ref), + &table_name, + schema.as_ref(), + column_name, + 0, + skip_column_count, + ); } - resolve_from_item_column_by_name(db, from_item, name_ref, column_name) + resolve_from_item_column_by_name_after_index( + db, + from_item, + name_ref, + column_name, + skip_column_count, + ) } fn resolve_select_column_ptr( @@ -1684,12 +1750,45 @@ fn resolve_select_order_by_alias_or_column_ptr( db: &dyn Db, column_name_ref: InFile<&ast::NameRef>, ) -> Option> { + if let Some(compound_select) = compound_select_for_order_by_name_ref(column_name_ref.value) { + return resolve_compound_select_order_by_column_ptr(db, column_name_ref, &compound_select); + } + if let Some(ptr) = resolve_select_target_alias_ptr(column_name_ref) { return Some(ptr); } resolve_select_column_ptr(db, column_name_ref) } +fn compound_select_for_order_by_name_ref(name_ref: &ast::NameRef) -> Option { + let order_by_clause = name_ref + .syntax() + .ancestors() + .find_map(ast::OrderByClause::cast)?; + order_by_clause + .syntax() + .parent() + .and_then(ast::CompoundSelect::cast) +} + +fn resolve_compound_select_order_by_column_ptr( + db: &dyn Db, + column_name_ref: InFile<&ast::NameRef>, + compound_select: &ast::CompoundSelect, +) -> Option> { + let file = column_name_ref.file_id; + let column_name_ref = column_name_ref.value; + let column_name = Name::from_node(column_name_ref); + let select = ast_nav::select_from_variant(compound_select.lhs()?)?; + resolve_column_from_select_targets( + db, + InFile::new(file, &select), + column_name_ref, + &column_name, + 0, + ) +} + fn resolve_select_target_alias_ptr( column_name_ref: InFile<&ast::NameRef>, ) -> Option> { @@ -1827,7 +1926,7 @@ pub(crate) fn find_column_in_create_table( create_table: InFile<&impl ast::HasCreateTable>, column_name: &Name, ) -> Option> { - find_column_in_create_table_impl(db, create_table, column_name, 0) + find_column_in_create_table_impl(db, create_table, column_name, 0, 0) } fn find_column_in_create_table_impl( @@ -1835,6 +1934,7 @@ fn find_column_in_create_table_impl( create_table: InFile<&impl ast::HasCreateTable>, column_name: &Name, depth: usize, + skip_column_count: usize, ) -> Option> { let file = create_table.file_id; let create_table = create_table.value; @@ -1843,20 +1943,25 @@ fn find_column_in_create_table_impl( return None; } + let mut column_index = 0usize; for arg in ast_nav::create_table_args(create_table) { match arg { ast_nav::CreateTableArg::Inherits(path) => { - let (schema, table_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()); - if let Some(resolved) = resolve_table_name(db, &table_name, &schemas, file) - && let Some(ptr) = find_column_in_resolved(db, resolved, column_name, depth + 1) - { - return Some(ptr); + if skip_column_count == 0 { + let (schema, table_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()); + if let Some(resolved) = resolve_table_name(db, &table_name, &schemas, file) + && let Some(ptr) = + find_column_in_resolved(db, resolved, column_name, depth + 1) + { + return Some(ptr); + } } } ast_nav::CreateTableArg::Column(column) => { - if let Some(name) = column.name() + if column_index >= skip_column_count + && let Some(name) = column.name() && Name::from_node(&name) == *column_name { return Some(smallvec![Location::new( @@ -1865,16 +1970,20 @@ fn find_column_in_create_table_impl( LocationKind::Column )]); } + column_index = column_index.saturating_add(1); } ast_nav::CreateTableArg::LikeClause(like_clause) => { - let path = like_clause.path()?; - let (schema, table_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()); - if let Some(resolved) = resolve_table_name(db, &table_name, &schemas, file) - && let Some(ptr) = find_column_in_resolved(db, resolved, column_name, depth + 1) - { - return Some(ptr); + if skip_column_count == 0 { + let path = like_clause.path()?; + let (schema, table_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()); + if let Some(resolved) = resolve_table_name(db, &table_name, &schemas, file) + && let Some(ptr) = + find_column_in_resolved(db, resolved, column_name, depth + 1) + { + return Some(ptr); + } } } ast_nav::CreateTableArg::TableConstraint(_) => (), @@ -1897,6 +2006,7 @@ fn find_column_in_resolved( InFile::new(file, &parent_table), column_name, depth, + 0, ), ResolvedTableName::TableAs(create_table_as) => { find_column_in_create_table_as(db, InFile::new(file, &create_table_as), column_name) @@ -1905,43 +2015,97 @@ fn find_column_in_resolved( find_column_in_select_into(db, file, &select_into, column_name) } ResolvedTableName::View(create_view) => { - find_column_in_create_view_like(file, &create_view, column_name) + find_column_in_create_view_like(db, file, &create_view, column_name) } } } -// TODO: this is similar to the CTE funcs, maybe we can simplify -fn find_column_in_create_view_like( +fn resolve_column_list_column( file: File, - create_view: &ast::CreateViewLike, + column_list: Option, column_name: &Name, -) -> Option> { - let column_list_len = if let Some(column_list) = create_view.column_list() { - for column in column_list.columns() { - if let Some(col_name) = column.name() - && Name::from_node(&col_name) == *column_name - { - return Some(smallvec![Location::new( + skip_column_count: usize, +) -> (usize, Option>) { + let Some(column_list) = column_list else { + return (0, None); + }; + + let mut column_list_len = 0usize; + for (idx, column) in column_list.columns().enumerate() { + column_list_len = idx + 1; + if idx >= skip_column_count + && let Some(col_name) = column.name() + && Name::from_node(&col_name) == *column_name + { + return ( + column_list_len, + Some(smallvec![Location::new( file, col_name.syntax().text_range(), LocationKind::Column - )]); - } + )]), + ); } - column_list.columns().count() - } else { - 0 - }; + } + + (column_list_len, None) +} + +fn resolve_values_column_after_index( + file: File, + values: &ast::Values, + column_name: &Name, + skip_column_count: usize, +) -> Option> { + for (idx, (col_name, expr)) in ast_nav::iter_values_columns(values).enumerate() { + if idx >= skip_column_count && col_name == *column_name { + return Some(smallvec![Location::new( + file, + expr.syntax().text_range(), + LocationKind::Column + )]); + } + } + None +} - let select = ast_nav::select_from_variant(create_view.query()?)?; +fn resolve_column_from_targets( + db: &dyn Db, + file: File, + target_list: &ast::TargetList, + from_clause: Option<&ast::FromClause>, + column_name: &Name, + skip_column_count: usize, + mut resolve_star: impl FnMut( + InFile<&ast::FromClause>, + &Name, + usize, + ) -> Option>, + mut resolve_qualified_star: impl FnMut( + InFile<&ast::FromClause>, + &Name, + &Name, + usize, + ) -> Option>, +) -> Option> { + let mut column_index = 0usize; - for (idx, target) in select.select_clause()?.target_list()?.targets().enumerate() { - if idx < column_list_len { + for target in target_list.targets() { + let target_column_count = from_clause + .and_then(|from_clause| { + count_columns_for_target(db, InFile::new(file, &target), from_clause) + }) + .unwrap_or(1); + let column_list_end = column_index.saturating_add(target_column_count); + if column_list_end <= skip_column_count { + column_index = column_list_end; continue; } + let target_skip = skip_column_count.saturating_sub(column_index); if let Some((col_name, node)) = ColumnName::from_target(target.clone()) { if let Some(col_name_str) = col_name.to_string() + && column_index >= skip_column_count && Name::from_string(col_name_str) == *column_name { return Some(smallvec![Location::new( @@ -1950,96 +2114,204 @@ fn find_column_in_create_view_like( LocationKind::Column )]); } + + if matches!(col_name, ColumnName::Star) + && let Some(from_clause) = from_clause + && let Some(result) = + resolve_star(InFile::new(file, from_clause), column_name, target_skip) + { + return Some(result); + } + } + + if let Some(ast::Expr::FieldExpr(field_expr)) = target.expr() + && let Some(table_name) = qualified_star_table_name(&field_expr) + && let Some(from_clause) = from_clause + && let Some(result) = resolve_qualified_star( + InFile::new(file, from_clause), + column_name, + &table_name, + target_skip, + ) + { + return Some(result); } + + column_index = column_list_end; } None } +fn find_column_in_target_list_with_skip( + db: &dyn Db, + file: File, + target_list: &ast::TargetList, + from_clause: Option<&ast::FromClause>, + column_name: &Name, + skip_column_count: usize, +) -> Option> { + resolve_column_from_targets( + db, + file, + target_list, + from_clause, + column_name, + skip_column_count, + |from_clause, column_name, skip_column_count| { + find_column_in_from_clause_with_skip(db, from_clause, column_name, skip_column_count) + }, + |from_clause, column_name, table_name, skip_column_count| { + find_column_in_qualified_from_clause_with_skip( + db, + from_clause, + column_name, + table_name, + skip_column_count, + ) + }, + ) +} + +// TODO: this is similar to the CTE funcs, maybe we can simplify +fn find_column_in_create_view_like( + db: &dyn Db, + file: File, + create_view: &ast::CreateViewLike, + column_name: &Name, +) -> Option> { + find_column_in_create_view_like_with_skip(db, file, create_view, column_name, 0) +} + +fn find_column_in_create_view_like_with_skip( + db: &dyn Db, + file: File, + create_view: &ast::CreateViewLike, + column_name: &Name, + skip_column_count: usize, +) -> Option> { + let (column_list_len, alias_column) = resolve_column_list_column( + file, + create_view.column_list(), + column_name, + skip_column_count, + ); + if let Some(alias_column) = alias_column { + return Some(alias_column); + } + let skip_column_count = skip_column_count.max(column_list_len); + + let query = create_view.query()?; + if let ast::SelectVariant::Table(table) = &query { + let path = table.relation_name()?.path()?; + let (schema, table_name) = name::schema_and_name_path(&path)?; + let table_name_ref = relation_name_ref_from_table(table)?; + return resolve_column_from_table_or_view_or_cte_impl( + db, + InFile::new(file, &table_name_ref), + &table_name, + schema.as_ref(), + column_name, + 0, + skip_column_count, + ); + } + + let select = ast_nav::select_from_variant(query)?; + let target_list = select.select_clause()?.target_list()?; + let from_clause = select.from_clause(); + find_column_in_target_list_with_skip( + db, + file, + &target_list, + from_clause.as_ref(), + column_name, + skip_column_count, + ) +} + fn find_column_in_create_table_as( db: &dyn Db, create_table_as: InFile<&ast::CreateTableAs>, column_name: &Name, +) -> Option> { + find_column_in_create_table_as_with_skip(db, create_table_as, column_name, 0) +} + +fn find_column_in_create_table_as_with_skip( + db: &dyn Db, + create_table_as: InFile<&ast::CreateTableAs>, + column_name: &Name, + skip_column_count: usize, ) -> Option> { let file = create_table_as.file_id; let create_table_as = create_table_as.value; let query = create_table_as.query()?; if let ast::SelectVariant::Values(values) = &query { - for (col_name, expr) in ast_nav::iter_values_columns(values) { - if col_name == *column_name { - return Some(smallvec![Location::new( - file, - expr.syntax().text_range(), - LocationKind::Column - )]); - } - } - return None; + return resolve_values_column_after_index(file, values, column_name, skip_column_count); } if let ast::SelectVariant::Table(table) = &query { let path = table.relation_name()?.path()?; let (schema, table_name) = name::schema_and_name_path(&path)?; let table_name_ref = relation_name_ref_from_table(table)?; - return resolve_column_from_table_or_view_or_cte( + return resolve_column_from_table_or_view_or_cte_impl( db, InFile::new(file, &table_name_ref), &table_name, schema.as_ref(), column_name, + 0, + skip_column_count, ); } let select = ast_nav::select_from_variant(query)?; let target_list = select.select_clause()?.target_list()?; let from_clause = select.from_clause(); - - for target in target_list.targets() { - if target.star_token().is_some() { - if let Some(from_clause) = &from_clause - && let Some(result) = - find_column_in_from_clause(db, InFile::new(file, from_clause), column_name) - { - return Some(result); - } - continue; - } - if let Some((col_name, node)) = ColumnName::from_target(target.clone()) - && let Some(col_name_str) = col_name.to_string() - && Name::from_string(col_name_str) == *column_name - { - return Some(smallvec![Location::new( - file, - node.text_range(), - LocationKind::Column - )]); - } - } - - None + find_column_in_target_list_with_skip( + db, + file, + &target_list, + from_clause.as_ref(), + column_name, + skip_column_count, + ) } -fn find_column_in_from_clause( +fn find_column_in_from_clause_with_skip( db: &dyn Db, from_clause: InFile<&ast::FromClause>, column_name: &Name, + skip_column_count: usize, ) -> Option> { let file = from_clause.file_id; let from_clause = from_clause.value; + let mut column_index = 0usize; for from_item in ast_nav::iter_from_clause(from_clause) { - let Some((schema, table_name)) = name::schema_and_table_from_from_item(&from_item) else { - continue; - }; - let Some(relation_name_ref) = relation_name_ref_from_from_item(&from_item) else { + let item_skip = skip_column_count.saturating_sub(column_index); + if let Some(count) = count_columns_for_from_item(db, InFile::new(file, &from_item)) { + if item_skip >= count { + column_index = column_index.saturating_add(count); + continue; + } + } + + let Some((schema, table_name)) = name::schema_and_table_from_from_item(&from_item) else { continue; }; - if let Some(result) = resolve_column_from_table_or_view_or_cte( + let Some(relation_name_ref) = relation_name_ref_from_from_item(&from_item) else { + continue; + }; + if let Some(result) = resolve_column_from_table_or_view_or_cte_impl( db, InFile::new(file, &relation_name_ref), &table_name, schema.as_ref(), column_name, + 0, + item_skip, ) { return Some(result); } @@ -2047,36 +2319,54 @@ fn find_column_in_from_clause( None } +fn find_column_in_qualified_from_clause_with_skip( + db: &dyn Db, + from_clause: InFile<&ast::FromClause>, + column_name: &Name, + table_name: &Name, + skip_column_count: usize, +) -> Option> { + let file = from_clause.file_id; + let from_item = find_from_item_in_from_clause(from_clause.value, table_name)?; + let (schema, table_name) = name::schema_and_table_from_from_item(&from_item)?; + let relation_name_ref = relation_name_ref_from_from_item(&from_item)?; + resolve_column_from_table_or_view_or_cte_impl( + db, + InFile::new(file, &relation_name_ref), + &table_name, + schema.as_ref(), + column_name, + 0, + skip_column_count, + ) +} + fn find_column_in_select_into( db: &dyn Db, file: File, select_into: &ast::SelectInto, column_name: &Name, ) -> Option> { - let from_clause = select_into.from_clause(); - for target in select_into.select_clause()?.target_list()?.targets() { - if target.star_token().is_some() { - if let Some(from_clause) = &from_clause - && let Some(result) = - find_column_in_from_clause(db, InFile::new(file, from_clause), column_name) - { - return Some(result); - } - continue; - } - if let Some((col_name, node)) = ColumnName::from_target(target.clone()) - && let Some(col_name_str) = col_name.to_string() - && Name::from_string(col_name_str) == *column_name - { - return Some(smallvec![Location::new( - file, - node.text_range(), - LocationKind::Column - )]); - } - } + find_column_in_select_into_with_skip(db, file, select_into, column_name, 0) +} - None +fn find_column_in_select_into_with_skip( + db: &dyn Db, + file: File, + select_into: &ast::SelectInto, + column_name: &Name, + skip_column_count: usize, +) -> Option> { + let from_clause = select_into.from_clause(); + let target_list = select_into.select_clause()?.target_list()?; + find_column_in_target_list_with_skip( + db, + file, + &target_list, + from_clause.as_ref(), + column_name, + skip_column_count, + ) } fn relation_name_ref_from_table(table: &ast::Table) -> Option { @@ -2208,48 +2498,36 @@ fn resolve_cte_column( name_ref: InFile<&ast::NameRef>, cte_name: &Name, column_name: &Name, +) -> Option> { + resolve_cte_column_with_skip(db, name_ref, cte_name, column_name, 0) +} + +fn resolve_cte_column_with_skip( + db: &dyn Db, + name_ref: InFile<&ast::NameRef>, + cte_name: &Name, + column_name: &Name, + skip_column_count: usize, ) -> Option> { let file = name_ref.file_id; let name_ref = name_ref.value; let with_table = ast_nav::find_cte_with_table(name_ref, cte_name)?; - let column_list_len = if let Some(column_list) = with_table.column_list() { - for column in column_list.columns() { - if let Some(col_name) = column.name() - && Name::from_node(&col_name) == *column_name - { - return Some(smallvec![Location::new( - file, - col_name.syntax().text_range(), - LocationKind::Column - )]); - } - } - column_list.columns().count() - } else { - 0 - }; + let (column_list_len, alias_column) = resolve_column_list_column( + file, + with_table.column_list(), + column_name, + skip_column_count, + ); + if let Some(alias_column) = alias_column { + return Some(alias_column); + } + let skip_column_count = skip_column_count.max(column_list_len); let query = with_table.query()?; if let ast::WithQuery::Values(values) = query { - if column_list_len > 0 { - return None; - } - if let Some(num_str) = column_name.0.strip_prefix("column") - && let Ok(col_num) = num_str.parse::() - && col_num > 0 - && let Some(expr) = values - .row_list() - .and_then(|x| x.rows().next()?.exprs().nth(col_num - 1)) - { - return Some(smallvec![Location::new( - file, - expr.syntax().text_range(), - LocationKind::Column - )]); - } - return None; + return resolve_values_column_after_index(file, &values, column_name, skip_column_count); } if let ast::WithQuery::Table(table) = query { @@ -2257,83 +2535,65 @@ fn resolve_cte_column( let (schema, table_name) = name::schema_and_name_path(&path)?; let table_name_ref = relation_name_ref_from_table(&table)?; - return resolve_column_from_table_or_view_or_cte( + return resolve_column_from_table_or_view_or_cte_impl( db, InFile::new(file, &table_name_ref), &table_name, schema.as_ref(), column_name, + 0, + skip_column_count, ); } - if let Some(column) = - column_in_with_query(InFile::new(file, &query), db, column_name, column_list_len) - { + if let Some(column) = column_in_with_query( + InFile::new(file, &query), + db, + column_name, + skip_column_count, + ) { return Some(column); } let cte_select = ast_nav::select_from_with_query(query)?; - let target_list = cte_select.select_clause()?.target_list()?; let from_clause = cte_select.from_clause(); - let mut column_index: usize = 0; - - for target in target_list.targets() { - let target_column_count = from_clause - .as_ref() - .and_then(|from_clause| { - count_columns_for_target(db, InFile::new(file, &target), from_clause) - }) - .unwrap_or(1); - let column_list_end = column_index.saturating_add(target_column_count); - if column_list_end <= column_list_len { - column_index = column_list_end; - continue; - } - - if let Some((col_name, node)) = ColumnName::from_target(target.clone()) { - if let Some(col_name_str) = col_name.to_string() - && Name::from_string(col_name_str) == *column_name - { - return Some(smallvec![Location::new( - file, - node.text_range(), - LocationKind::Column - )]); - } - if matches!(col_name, ColumnName::Star) - && let Some(from_clause) = &from_clause - && let Some(result) = resolve_from_clause_for_cte_star( - db, - InFile::new(file, from_clause), - name_ref, - cte_name, - column_name, - ) - { - return Some(result); - } - } - if let Some(ast::Expr::FieldExpr(field_expr)) = target.expr() - && let Some(table_name) = qualified_star_table_name(&field_expr) - && let Some(from_clause) = &from_clause - && let Some(result) = resolve_qualified_star_in_from_clause( + let result = resolve_column_from_targets( + db, + file, + &target_list, + from_clause.as_ref(), + column_name, + skip_column_count, + |from_clause, column_name, skip_column_count| { + resolve_from_clause_for_cte_star( db, - InFile::new(file, from_clause), + from_clause, name_ref, cte_name, column_name, - &table_name, + skip_column_count, ) - { - return Some(result); - } - - column_index = column_list_end; + }, + |from_clause, column_name, table_name, skip_column_count| { + resolve_qualified_star_in_from_clause( + db, + from_clause, + name_ref, + cte_name, + column_name, + table_name, + skip_column_count, + ) + }, + ); + if result.is_some() { + return result; } - if column_name == cte_name + if skip_column_count == 0 + && column_name == cte_name && let Some(name_node) = with_table.name() { return Some(smallvec![Location::new( @@ -2354,28 +2614,43 @@ fn column_in_with_query( ) -> Option> { let file = query.file_id; let query = query.value; - let (returning_clause, path) = match query { - ast::WithQuery::Delete(delete) => { - (delete.returning_clause(), delete.relation_name()?.path()?) - } - ast::WithQuery::Insert(insert) => (insert.returning_clause(), insert.path()?), - ast::WithQuery::Merge(merge) => (merge.returning_clause(), merge.relation_name()?.path()?), - ast::WithQuery::Update(update) => { - (update.returning_clause(), update.relation_name()?.path()?) + let (returning_clause, alias, path) = match query { + ast::WithQuery::Delete(delete) => ( + delete.returning_clause(), + delete.alias(), + delete.relation_name()?.path()?, + ), + ast::WithQuery::Insert(insert) => { + (insert.returning_clause(), insert.alias(), insert.path()?) } + ast::WithQuery::Merge(merge) => ( + merge.returning_clause(), + merge.alias(), + merge.relation_name()?.path()?, + ), + ast::WithQuery::Update(update) => ( + update.returning_clause(), + update.alias(), + update.relation_name()?.path()?, + ), ast::WithQuery::Select(_) | ast::WithQuery::CompoundSelect(_) | ast::WithQuery::Table(_) | ast::WithQuery::Values(_) | ast::WithQuery::ParenSelect(_) => return None, }; + let returning_clause = returning_clause?; + let (_, stmt_table_name) = name::schema_and_name_path(&path)?; let mut column_index: usize = 0; - for target in returning_clause?.target_list()?.targets() { - let target_column_count = if target.star_token().is_some() { - count_columns_for_path(db, InFile::new(file, &path)).unwrap_or(1) - } else { - 1 + for target in returning_clause.target_list()?.targets() { + let target_kind = + returning_target_kind(&target, &stmt_table_name, alias.as_ref(), &returning_clause); + let target_column_count = match target_kind { + ReturningTargetKind::Column => 1, + ReturningTargetKind::QualifiedStar | ReturningTargetKind::UnqualifiedStar => { + count_columns_for_path(db, InFile::new(file, &path)).unwrap_or(1) + } }; let column_list_end = column_index.saturating_add(target_column_count); @@ -2401,42 +2676,83 @@ fn column_in_with_query( return Some(ptr); } } + if matches!(target_kind, ReturningTargetKind::QualifiedStar) + && let Some(ptr) = + resolve_column_for_path(db, InFile::new(file, &path), column_name.clone()) + { + return Some(ptr); + } column_index = column_list_end; } None } +#[derive(Clone, Copy)] +enum ReturningTargetKind { + Column, + QualifiedStar, + UnqualifiedStar, +} + +fn returning_target_kind( + target: &ast::Target, + stmt_table_name: &Name, + alias: Option<&ast::Alias>, + returning_clause: &ast::ReturningClause, +) -> ReturningTargetKind { + if target.star_token().is_some() { + return ReturningTargetKind::UnqualifiedStar; + } + + if let Some(ast::Expr::FieldExpr(field_expr)) = target.expr() + && let Some(table_name) = qualified_star_table_name(&field_expr) + && match_table_in_returning_clause( + &table_name, + stmt_table_name, + alias, + Some(returning_clause), + ) + .is_some() + { + return ReturningTargetKind::QualifiedStar; + } + + ReturningTargetKind::Column +} + fn resolve_subquery_column_ptr( db: &dyn Db, paren_select: InFile<&ast::ParenSelect>, name_ref: &ast::NameRef, column_name: &Name, alias: Option<&ast::Alias>, +) -> Option> { + resolve_subquery_column_ptr_with_skip(db, paren_select, name_ref, column_name, alias, 0) +} + +fn resolve_subquery_column_ptr_with_skip( + db: &dyn Db, + paren_select: InFile<&ast::ParenSelect>, + name_ref: &ast::NameRef, + column_name: &Name, + alias: Option<&ast::Alias>, + skip_column_count: usize, ) -> Option> { let file = paren_select.file_id; let paren_select = paren_select.value; let select_variant = paren_select.select()?; - let column_list_len = if let Some(column_list) = alias.and_then(|x| x.column_list()) { - for col in column_list.columns() { - if let Some(col_name) = col.name() - && Name::from_node(&col_name) == *column_name - { - return Some(smallvec![Location::new( - file, - col_name.syntax().text_range(), - LocationKind::Column - )]); - } - } - if matches!(select_variant, ast::SelectVariant::Values(_)) { - return None; - } - column_list.columns().count() - } else { - 0 - }; + let (column_list_len, alias_column) = resolve_column_list_column( + file, + alias.and_then(|x| x.column_list()), + column_name, + skip_column_count, + ); + if let Some(alias_column) = alias_column { + return Some(alias_column); + } + let skip_column_count = skip_column_count.max(column_list_len); // TODO: this should just be a match stmt if let ast::SelectVariant::Table(table) = select_variant { @@ -2444,30 +2760,19 @@ fn resolve_subquery_column_ptr( let (schema, table_name) = name::schema_and_name_path(&path)?; let table_name_ref = relation_name_ref_from_table(&table)?; - return resolve_column_from_table_or_view_or_cte( + return resolve_column_from_table_or_view_or_cte_impl( db, InFile::new(file, &table_name_ref), &table_name, schema.as_ref(), column_name, + 0, + skip_column_count, ); } if let ast::SelectVariant::Values(values) = select_variant { - if let Some(num_str) = column_name.0.strip_prefix("column") - && let Ok(col_num) = num_str.parse::() - && col_num > 0 - && let Some(expr) = values - .row_list() - .and_then(|x| x.rows().next()?.exprs().nth(col_num - 1)) - { - return Some(smallvec![Location::new( - file, - expr.syntax().text_range(), - LocationKind::Column - )]); - } - return None; + return resolve_values_column_after_index(file, &values, column_name, skip_column_count); } let subquery_select = ast_nav::select_from_variant(select_variant)?; @@ -2476,7 +2781,7 @@ fn resolve_subquery_column_ptr( InFile::new(file, &subquery_select), name_ref, column_name, - column_list_len, + skip_column_count, ) } @@ -2485,58 +2790,41 @@ fn resolve_column_from_select_targets( select: InFile<&ast::Select>, name_ref: &ast::NameRef, column_name: &Name, - skip_target_count: usize, + skip_column_count: usize, ) -> Option> { let file = select.file_id; let select = select.value; + let target_list = select.select_clause()?.target_list()?; let from_clause = select.from_clause(); - for (i, target) in select.select_clause()?.target_list()?.targets().enumerate() { - if let Some((col_name, node)) = ColumnName::from_target(target.clone()) { - if i < skip_target_count && !matches!(col_name, ColumnName::Star) { - continue; - } - if let Some(col_name_str) = col_name.to_string() - && Name::from_string(col_name_str) == *column_name - { - return Some(smallvec![Location::new( - file, - node.text_range(), - LocationKind::Column - )]); - } - if matches!(col_name, ColumnName::Star) - && let Some(from_clause) = &from_clause - { - for from_item in ast_nav::iter_from_clause(from_clause) { - if let Some(result) = resolve_from_item_column_by_name( - db, - InFile::new(file, &from_item), - name_ref, - column_name, - ) { - return Some(result); - } - } - } - } - - if let Some(ast::Expr::FieldExpr(field_expr)) = target.expr() - && let Some(table_name) = qualified_star_table_name(&field_expr) - && let Some(from_clause) = &from_clause - && let Some(from_item) = find_from_item_in_from_clause(from_clause, &table_name) - && let Some(result) = resolve_from_item_column_by_name( + resolve_column_from_targets( + db, + file, + &target_list, + from_clause.as_ref(), + column_name, + skip_column_count, + |from_clause, column_name, skip_column_count| { + resolve_from_clause_column_after_index( + db, + from_clause, + name_ref, + column_name, + skip_column_count, + ) + }, + |from_clause, column_name, table_name, skip_column_count| { + let file = from_clause.file_id; + let from_item = find_from_item_in_from_clause(from_clause.value, table_name)?; + resolve_from_item_column_by_name_after_index( db, InFile::new(file, &from_item), name_ref, column_name, + skip_column_count, ) - { - return Some(result); - } - } - - None + }, + ) } pub(crate) fn qualified_star_table_name(field_expr: &ast::FieldExpr) -> Option { @@ -2680,21 +2968,64 @@ fn count_columns_for_from_item(db: &dyn Db, from_item: InFile<&ast::FromItem>) - count_columns_for_table_name(db, &table_name, &schemas, file, Some(&scope_name_ref)) } +fn resolve_from_clause_column_after_index( + db: &dyn Db, + from_clause: InFile<&ast::FromClause>, + name_ref: &ast::NameRef, + column_name: &Name, + skip_column_count: usize, +) -> Option> { + let file = from_clause.file_id; + let mut column_index = 0usize; + for from_item in ast_nav::iter_from_clause(from_clause.value) { + let item_skip = skip_column_count.saturating_sub(column_index); + if let Some(count) = count_columns_for_from_item(db, InFile::new(file, &from_item)) { + if item_skip >= count { + column_index = column_index.saturating_add(count); + continue; + } + } + + if let Some(result) = resolve_from_item_column_by_name_after_index( + db, + InFile::new(file, &from_item), + name_ref, + column_name, + item_skip, + ) { + return Some(result); + } + } + + None +} + fn resolve_from_clause_for_cte_star( db: &dyn Db, from_clause: InFile<&ast::FromClause>, name_ref: &ast::NameRef, cte_name: &Name, column_name: &Name, + skip_column_count: usize, ) -> Option> { let file = from_clause.file_id; + let mut column_index = 0usize; for from_item in ast_nav::iter_from_clause(from_clause.value) { + let item_skip = skip_column_count.saturating_sub(column_index); + if let Some(count) = count_columns_for_from_item(db, InFile::new(file, &from_item)) { + if item_skip >= count { + column_index = column_index.saturating_add(count); + continue; + } + } + if let Some(result) = resolve_from_item_for_cte_star( db, InFile::new(file, &from_item), name_ref, cte_name, column_name, + item_skip, ) { return Some(result); } @@ -2710,6 +3041,7 @@ fn resolve_qualified_star_in_from_clause( cte_name: &Name, column_name: &Name, table_name: &Name, + skip_column_count: usize, ) -> Option> { let file = from_clause.file_id; let from_item = find_from_item_in_from_clause(from_clause.value, table_name)?; @@ -2719,6 +3051,7 @@ fn resolve_qualified_star_in_from_clause( name_ref, cte_name, column_name, + skip_column_count, ) } @@ -2727,6 +3060,16 @@ fn resolve_column_from_paren_expr( paren_expr: InFile<&ast::ParenExpr>, name_ref: &ast::NameRef, column_name: &Name, +) -> Option> { + resolve_column_from_paren_expr_with_skip(db, paren_expr, name_ref, column_name, 0) +} + +fn resolve_column_from_paren_expr_with_skip( + db: &dyn Db, + paren_expr: InFile<&ast::ParenExpr>, + name_ref: &ast::NameRef, + column_name: &Name, + skip_column_count: usize, ) -> Option> { let file = paren_expr.file_id; let paren_expr = paren_expr.value; @@ -2736,7 +3079,7 @@ fn resolve_column_from_paren_expr( InFile::new(file, &select), name_ref, column_name, - 0, + skip_column_count, ); } @@ -2746,56 +3089,30 @@ fn resolve_column_from_paren_expr( InFile::new(file, &call_expr), name_ref, column_name, - 0, + skip_column_count, ) { return Some(result); } if let Some(ast::Expr::ParenExpr(paren_expr)) = paren_expr.expr() { - return resolve_column_from_paren_expr( + return resolve_column_from_paren_expr_with_skip( db, InFile::new(file, &paren_expr), name_ref, column_name, + skip_column_count, ); } if let Some(from_item) = paren_expr.from_item() { - if let Some(paren_select) = from_item.paren_select() { - let alias = from_item.alias(); - return resolve_subquery_column_ptr( - db, - InFile::new(file, &paren_select), - name_ref, - column_name, - alias.as_ref(), - ); - } - - if let Some(select_variant) = from_item - .syntax() - .children() - .find_map(ast::SelectVariant::cast) - { - let select = ast_nav::select_from_variant(select_variant)?; - return resolve_column_from_select_targets( - db, - InFile::new(file, &select), - name_ref, - column_name, - 0, - ); - } - - if let Some(nested_paren_expr) = from_item.paren_expr() { - return resolve_column_from_paren_expr( - db, - InFile::new(file, &nested_paren_expr), - name_ref, - column_name, - ); - } + return resolve_from_item_column_by_name_after_index( + db, + InFile::new(file, &from_item), + name_ref, + column_name, + skip_column_count, + ); } None @@ -3030,8 +3347,12 @@ fn resolve_table_in_returning_clause( let table_name = Name::from_node(table_name_ref); let (schema, stmt_table_name) = name::schema_and_name_path(path)?; - let matched = - match_table_in_returning_clause(&table_name, &stmt_table_name, alias, returning_clause)?; + let matched = match_table_in_returning_clause( + &table_name, + &stmt_table_name, + alias.as_ref(), + returning_clause.as_ref(), + )?; let position = table_name_ref.syntax().text_range().start(); let schemas = bind(db, file).resolved_schemas(position, schema.as_ref());