diff --git a/crates/squawk_ide/src/classify.rs b/crates/squawk_ide/src/classify.rs index 8a208d9e..2b9c91e6 100644 --- a/crates/squawk_ide/src/classify.rs +++ b/crates/squawk_ide/src/classify.rs @@ -28,6 +28,7 @@ pub(crate) enum NameRefClass { Index, InsertColumn, InsertQualifiedColumnTable, + InsertTable, JoinUsingColumn, LikeTable, MergeColumn, @@ -706,7 +707,7 @@ pub(crate) fn classify_name_ref(node: &SyntaxNode) -> Option { { return Some(NameRefClass::InsertColumn); } - return Some(NameRefClass::Table); + return Some(NameRefClass::InsertTable); } if ast::JoinUsingClause::can_cast(ancestor.kind()) && in_column_list { return Some(NameRefClass::JoinUsingColumn); diff --git a/crates/squawk_ide/src/collect.rs b/crates/squawk_ide/src/collect.rs index d4213483..d37b5f3b 100644 --- a/crates/squawk_ide/src/collect.rs +++ b/crates/squawk_ide/src/collect.rs @@ -1,14 +1,14 @@ use crate::ast_nav; use crate::builtins::builtins_file; use crate::column_name::ColumnName; -use crate::db::{File, bind, parse}; +use crate::db::{File, parse}; use crate::infer::{Type, infer_type_from_expr, infer_type_from_ty}; +use crate::location::LocationKind; use crate::name::{self, Name}; use crate::resolve::{ - ResolvedTableName, find_from_item_in_from_clause, qualified_star_table_name, resolve_cte_table, - resolve_table_name, table_ptr_from_from_item, + ResolvedTableName, find_from_item_in_from_clause, qualified_star_table_name, + resolve_table_like, resolve_table_name, table_ptr_from_from_item, }; -use crate::symbols::SymbolKind; use salsa::Database as Db; use squawk_syntax::{ SyntaxNodePtr, @@ -227,11 +227,7 @@ pub(crate) fn view_like_columns_with_types( results } -pub(crate) fn with_table_column_names( - db: &dyn Db, - file: File, - with_table: ast::WithTable, -) -> Vec { +fn with_table_column_names(db: &dyn Db, file: File, with_table: ast::WithTable) -> Vec { with_table_columns_with_types(db, file, with_table) .into_iter() .map(|(name, _)| name) @@ -435,9 +431,7 @@ fn columns_for_star_from_table_ptr( } } -pub(crate) fn target_list_columns_with_types( - target_list: &ast::TargetList, -) -> Vec<(Name, Option)> { +fn target_list_columns_with_types(target_list: &ast::TargetList) -> Vec<(Name, Option)> { let mut columns = vec![]; for target in target_list.targets() { if let Some((col_name, _node)) = ColumnName::from_target(target.clone()) @@ -466,7 +460,6 @@ fn select_variant_columns_with_types( file: File, select_variant: &ast::SelectVariant, ) -> Vec<(Name, Option)> { - let binder = bind(db, file); let tree = parse(db, file).tree(); let root = tree.syntax(); match select_variant { @@ -517,35 +510,34 @@ fn select_variant_columns_with_types( let Some((schema, table_name)) = name::schema_and_name_path(&path) else { return vec![]; }; - - if schema.is_none() - && let Some(name_ref) = path.segment().and_then(|segment| segment.name_ref()) - && let Some(with_table_ptr) = resolve_cte_table(&name_ref, &table_name) - && let Some(with_table) = with_table_ptr - .to_node(root) - .ancestors() - .find_map(ast::WithTable::cast) - { - return with_table_column_names(db, file, with_table) - .into_iter() - .map(|name| (name, None)) - .collect(); - } - + let name_ref = path.segment().and_then(|segment| segment.name_ref()); let position = table.syntax().text_range().start(); - let Some(table_ptr) = - binder.lookup_with(&table_name, SymbolKind::Table, position, &schema) + let Some((ptr, kind)) = + resolve_table_like(db, file, name_ref.as_ref(), &table_name, &schema, position) else { return vec![]; }; - let Some(create_table) = table_ptr - .to_node(root) - .ancestors() - .find_map(ast::CreateTableLike::cast) - else { - return vec![]; - }; - table_columns(db, file, &create_table) + let node = ptr.to_node(root); + match kind { + LocationKind::View => node + .ancestors() + .find_map(ast::CreateViewLike::cast) + .map(|v| view_like_columns_with_types(&v)) + .unwrap_or_default(), + LocationKind::Table => { + if let Some(with_table) = node.ancestors().find_map(ast::WithTable::cast) { + return with_table_column_names(db, file, with_table) + .into_iter() + .map(|name| (name, None)) + .collect(); + } + node.ancestors() + .find_map(ast::CreateTableLike::cast) + .map(|t| table_columns(db, file, &t)) + .unwrap_or_default() + } + _ => vec![], + } } } } diff --git a/crates/squawk_ide/src/goto_definition.rs b/crates/squawk_ide/src/goto_definition.rs index 7b3f53e6..b598f6df 100644 --- a/crates/squawk_ide/src/goto_definition.rs +++ b/crates/squawk_ide/src/goto_definition.rs @@ -533,6 +533,24 @@ select t$0.* from t; "); } + #[test] + fn goto_cte_shadowed_table_star_column_count_not_found() { + goto_not_found( + " +create table t(a int, b int); +with + t as ( + select 1 + ), + -- yy overrides y since there's only 1 column in the * + u(x, yy) as ( + select *, 2 y, 3 z from t + ) +select y$0 from u; +", + ); + } + #[test] fn goto_cross_join_func_column() { assert_snapshot!(goto(r#" @@ -4192,6 +4210,34 @@ select f2$0 from dup(42) as u(x, y); ); } + #[test] + fn goto_fn_call_column_from_cte() { + assert_snapshot!(goto(" +with cte as (select 1 as a) +select a$0(cte) from cte; +"), @" + ╭▸ + 2 │ with cte as (select 1 as a) + │ ─ 2. destination + 3 │ select a(cte) from cte; + ╰╴ ─ 1. source + "); + } + + #[test] + fn goto_fn_call_column_from_view() { + assert_snapshot!(goto(" +create view v as select 1 as a; +select a$0(v) from v; +"), @" + ╭▸ + 2 │ create view v as select 1 as a; + │ ─ 2. destination + 3 │ select a(v) from v; + ╰╴ ─ 1. source + "); + } + #[test] fn goto_select_aggregate_call() { assert_snapshot!(goto(" @@ -5341,6 +5387,26 @@ insert into public.users$0(id, email) values (1, 'test@example.com'); "); } + #[test] + fn goto_insert_view() { + assert_snapshot!(goto(" +create table users as select 1 id, 'joe' name, 'joe@example.com' email, 'active' status; +create view active_users as + select id, name, email + from users + where status = 'active'; +insert into active_users$0 (name, email) +values ('Alice', 'alice@example.com'); +"), @" + ╭▸ + 3 │ create view active_users as + │ ──────────── 2. destination + ‡ + 7 │ insert into active_users (name, email) + ╰╴ ─ 1. source + "); + } + #[test] fn goto_insert_column() { assert_snapshot!(goto(" diff --git a/crates/squawk_ide/src/hover.rs b/crates/squawk_ide/src/hover.rs index 3e392d00..dcdbeeff 100644 --- a/crates/squawk_ide/src/hover.rs +++ b/crates/squawk_ide/src/hover.rs @@ -5,12 +5,14 @@ use crate::column_name::ColumnName; use crate::comments::preceding_comment; use crate::db::{File, bind, parse}; use crate::location::{Location, LocationKind}; +use crate::name; use crate::offsets::token_from_offset; use crate::symbols::{Name, Schema}; use crate::{goto_definition, resolve}; use rowan::TextSize; use salsa::Database as Db; use squawk_syntax::SyntaxNode; +use squawk_syntax::SyntaxNodePtr; use squawk_syntax::{ SyntaxKind, ast::{self, AstNode}, @@ -467,7 +469,7 @@ fn format_alias_with_column_list(db: &dyn Db, file: File, alias: ast::Alias) -> } fn hover_qualified_star(db: &dyn Db, file: File, field_expr: ast::FieldExpr) -> Option { - let table_ptr = resolve::resolve_qualified_star_table_ptr(db, file, field_expr)?; + let table_ptr = qualified_star_table_ptr(db, file, field_expr)?; hover_qualified_star_columns(db, file, &table_ptr) } @@ -484,7 +486,7 @@ fn hover_unqualified_star(db: &dyn Db, file: File, target: ast::Target) -> Optio fn hover_unqualified_star_with_binder(db: &dyn Db, file: File, target: &ast::Target) -> Vec { let mut results = vec![]; - if let Some(table_ptrs) = resolve::resolve_unqualified_star_table_ptrs(db, file, target) { + if let Some(table_ptrs) = unqualified_star_table_ptrs(db, file, target) { for table_ptr in table_ptrs { if let Some(columns) = hover_qualified_star_columns(db, file, &table_ptr) { results.push(columns); @@ -517,7 +519,7 @@ fn hover_unqualified_star_in_arg_list( file: File, arg_list: ast::ArgList, ) -> Option { - let table_ptrs = resolve::resolve_unqualified_star_in_arg_list_ptrs(db, file, &arg_list)?; + let table_ptrs = unqualified_star_in_arg_list_ptrs(db, file, &arg_list)?; let mut results = vec![]; for table_ptr in table_ptrs { if let Some(columns) = hover_qualified_star_columns(db, file, &table_ptr) { @@ -736,7 +738,7 @@ fn hover_qualified_star_columns_from_subquery( for target in target_list.targets() { if target.star_token().is_some() { - let table_ptrs = resolve::resolve_unqualified_star_table_ptrs(db, file, &target)?; + let table_ptrs = unqualified_star_table_ptrs(db, file, &target)?; for table_ptr in table_ptrs { if let Some(columns) = hover_qualified_star_columns(db, file, &table_ptr) { results.push(columns) @@ -1394,6 +1396,123 @@ fn hover_routine(db: &dyn Db, def: Location) -> Option { None } +fn qualified_star_table_ptr( + db: &dyn Db, + file: File, + field_expr: ast::FieldExpr, +) -> Option { + let table_name = resolve::qualified_star_table_name(&field_expr)?; + let position = field_expr.syntax().text_range().start(); + let target = field_expr + .syntax() + .ancestors() + .find_map(ast::Target::cast)?; + + let path = match ast_nav::target_parent_query(target)? { + ast_nav::ParentQuery::Select(select) => { + let from_clause = select.from_clause()?; + let from_item = resolve::find_from_item_in_from_clause(&from_clause, &table_name)?; + + if let Some(alias) = from_item.alias() + && alias.column_list().is_some() + { + return Some(SyntaxNodePtr::new(alias.syntax())); + } + + let (schema, table_name) = name::schema_and_table_from_from_item(&from_item)?; + + let name_ref = from_item.name_ref(); + if let Some((table_like_ptr, _kind)) = resolve::resolve_table_like( + db, + file, + name_ref.as_ref(), + &table_name, + &schema, + position, + ) { + return Some(table_like_ptr); + } + + return None; + } + ast_nav::ParentQuery::Update(update) => update.relation_name()?.path()?, + ast_nav::ParentQuery::Delete(delete) => delete.relation_name()?.path()?, + ast_nav::ParentQuery::Insert(insert) => insert.path()?, + ast_nav::ParentQuery::Merge(merge) => merge.relation_name()?.path()?, + }; + + table_or_view_or_cte_ptrs(db, file, position, &path)? + .into_iter() + .next() +} + +fn table_or_view_or_cte_ptrs( + db: &dyn Db, + file: File, + position: TextSize, + path: &ast::Path, +) -> Option> { + let (schema, table_name) = name::schema_and_name_path(path)?; + let mut results = vec![]; + let name_ref = path.segment().and_then(|x| x.name_ref()); + + if let Some((table_like_ptr, _kind)) = + resolve::resolve_table_like(db, file, name_ref.as_ref(), &table_name, &schema, position) + { + results.push(table_like_ptr); + } + + if results.is_empty() { + return None; + } + Some(results) +} + +fn unqualified_star_table_ptrs( + db: &dyn Db, + file: File, + target: &ast::Target, +) -> Option> { + target.star_token()?; + + let path = match ast_nav::target_parent_query(target.clone())? { + ast_nav::ParentQuery::Select(select) => { + let from_clause = select.from_clause()?; + let results = resolve::table_ptrs_from_clause(db, file, &from_clause); + if results.is_empty() { + return None; + } + return Some(results); + } + ast_nav::ParentQuery::Update(update) => update.relation_name()?.path(), + ast_nav::ParentQuery::Insert(insert) => insert.path(), + ast_nav::ParentQuery::Delete(delete) => delete.relation_name()?.path(), + ast_nav::ParentQuery::Merge(merge) => merge.relation_name()?.path(), + }?; + + let position = target.syntax().text_range().start(); + table_or_view_or_cte_ptrs(db, file, position, &path) +} + +fn unqualified_star_in_arg_list_ptrs( + db: &dyn Db, + file: File, + arg_list: &ast::ArgList, +) -> Option> { + let from_clause = arg_list + .syntax() + .ancestors() + .find_map(ast::Select::cast)? + .from_clause()?; + let results = resolve::table_ptrs_from_clause(db, file, &from_clause); + + if results.is_empty() { + return None; + } + + Some(results) +} + #[cfg(test)] mod test { use crate::db::{Database, File}; diff --git a/crates/squawk_ide/src/resolve.rs b/crates/squawk_ide/src/resolve.rs index 16dd5bd1..330010a9 100644 --- a/crates/squawk_ide/src/resolve.rs +++ b/crates/squawk_ide/src/resolve.rs @@ -57,6 +57,25 @@ pub(crate) fn resolve_name_ref( LocationKind::Table )]) } + NameRefClass::InsertTable => { + let (schema, relation_name) = name::schema_and_table_name(name_ref)?; + let position = name_ref.syntax().text_range().start(); + + if let Some(ptr) = resolve_table_name_ptr(db, file, &relation_name, &schema, position) { + return Some(smallvec![Location::new( + file, + ptr.text_range(), + LocationKind::Table + )]); + } + + let ptr = resolve_view_name_ptr(db, file, &relation_name, &schema, position)?; + Some(smallvec![Location::new( + file, + ptr.text_range(), + LocationKind::View + )]) + } NameRefClass::NamedArgParameter => { let (schema, function_name) = name_ref.syntax().ancestors().find_map(|a| { if let Some(call_expr) = ast::CallExpr::cast(a.clone()) { @@ -118,34 +137,10 @@ pub(crate) fn resolve_name_ref( } NameRefClass::FromTable => { let (schema, table_name) = name::schema_and_name(name_ref); - if schema.is_none() - && let Some(cte_ptr) = resolve_cte_table(name_ref, &table_name) - { - return Some(smallvec![Location::new( - file, - cte_ptr.text_range(), - LocationKind::Table - )]); - } - let position = name_ref.syntax().text_range().start(); - - if let Some(table_name_ptr) = - resolve_table_name_ptr(db, file, &table_name, &schema, position) - { - return Some(smallvec![Location::new( - file, - table_name_ptr.text_range(), - LocationKind::Table - )]); - } - - let ptr = resolve_view_name_ptr(db, file, &table_name, &schema, position)?; - Some(smallvec![Location::new( - file, - ptr.text_range(), - LocationKind::View - )]) + let (ptr, kind) = + resolve_table_like(db, file, Some(name_ref), &table_name, &schema, position)?; + Some(smallvec![Location::new(file, ptr.text_range(), kind)]) } NameRefClass::Index => { let position = name_ref.syntax().text_range().start(); @@ -786,6 +781,32 @@ fn resolve_view_name_ptr( binder.lookup_with(view_name, SymbolKind::View, position, schema) } +pub(crate) fn resolve_table_like( + db: &dyn Db, + file: File, + name_ref: Option<&ast::NameRef>, + table_name: &Name, + schema: &Option, + position: TextSize, +) -> Option<(SyntaxNodePtr, LocationKind)> { + if let Some(view_name_ptr) = resolve_view_name_ptr(db, file, table_name, schema, position) { + return Some((view_name_ptr, LocationKind::View)); + } + + if let Some(table_name_ptr) = resolve_table_name_ptr(db, file, table_name, schema, position) { + return Some((table_name_ptr, LocationKind::Table)); + } + + if schema.is_none() + && let Some(name_ref) = name_ref + && let Some(cte_ptr) = resolve_cte_table(name_ref, table_name) + { + return Some((cte_ptr, LocationKind::Table)); + } + + None +} + fn resolve_for_kind_with_params( db: &dyn Db, file: File, @@ -1011,40 +1032,21 @@ fn resolve_select_qualified_column_table_name_ptr( return resolve_function(db, file, &function_name, &function_schema, None, position); } - if let Some(schema) = explicit_schema { - let position = table_name_ref.syntax().text_range().start(); - let ptr = resolve_table_name_ptr(db, file, &table_name, &Some(schema), position)?; - return Some(smallvec![Location::new( - file, - ptr.text_range(), - LocationKind::Table - )]); - } - - let (schema, table_name) = name::schema_and_table_from_from_item(&from_item)?; - if schema.is_none() - && let Some(cte_ptr) = resolve_cte_table(table_name_ref, &table_name) - { - return Some(smallvec![Location::new( - file, - cte_ptr.text_range(), - LocationKind::Table - )]); - } + let (schema, table_name) = if let Some(schema) = explicit_schema { + (Some(schema), table_name) + } else { + name::schema_and_table_from_from_item(&from_item)? + }; let position = table_name_ref.syntax().text_range().start(); - if let Some(ptr) = resolve_table_name_ptr(db, file, &table_name, &schema, position) { - return Some(smallvec![Location::new( - file, - ptr.text_range(), - LocationKind::Table - )]); - } - let ptr = resolve_view_name_ptr(db, file, &table_name, &schema, position)?; - Some(smallvec![Location::new( + let (ptr, kind) = resolve_table_like( + db, file, - ptr.text_range(), - LocationKind::View - )]) + Some(table_name_ref), + &table_name, + &schema, + position, + )?; + Some(smallvec![Location::new(file, ptr.text_range(), kind)]) } enum ReturningClauseMatch { @@ -1297,6 +1299,8 @@ pub(crate) enum ResolvedTableName { SelectInto(ast::SelectInto), View(ast::CreateViewLike), } +// TODO: basically goto def for a table name, we should refactor into an ast_nav +// + goto def setup pub(crate) fn resolve_table_name( db: &dyn Db, file: File, @@ -1304,38 +1308,31 @@ pub(crate) fn resolve_table_name( schema: &Option, position: TextSize, ) -> Option { + use ResolvedTableName::*; + let (ptr, kind) = resolve_table_like(db, file, None, table_name, schema, position)?; let tree = parse(db, file).tree(); let root = tree.syntax(); - use ResolvedTableName::*; - if let Some(table_name_ptr) = resolve_table_name_ptr(db, file, table_name, schema, position) { - let table_name_node = table_name_ptr.to_node(root); - if let Some(create_table) = table_name_node - .ancestors() - .find_map(ast::CreateTableLike::cast) - { - return Some(Table(create_table)); - } - if let Some(create_table_as) = table_name_node - .ancestors() - .find_map(ast::CreateTableAs::cast) - { - return Some(TableAs(create_table_as)); - } - if let Some(select_into) = table_name_node.ancestors().find_map(ast::SelectInto::cast) { - return Some(SelectInto(select_into)); - } - } + let node = ptr.to_node(root); - if let Some(view_name_ptr) = resolve_view_name_ptr(db, file, table_name, schema, position) { - let view_name_node = view_name_ptr.to_node(root); - if let Some(create_view) = view_name_node + match kind { + LocationKind::Table => { + if let Some(create_table) = node.ancestors().find_map(ast::CreateTableLike::cast) { + return Some(Table(create_table)); + } + if let Some(create_table_as) = node.ancestors().find_map(ast::CreateTableAs::cast) { + return Some(TableAs(create_table_as)); + } + if let Some(select_into) = node.ancestors().find_map(ast::SelectInto::cast) { + return Some(SelectInto(select_into)); + } + None + } + LocationKind::View => node .ancestors() .find_map(ast::CreateViewLike::cast) - { - return Some(View(create_view)); - } + .map(View), + _ => None, } - None } fn resolve_merge_alias(name_ref: &ast::NameRef, table_name: &Name) -> Option { @@ -1448,34 +1445,7 @@ fn resolve_from_item_column_ptr( let (schema, table_name) = name::schema_and_table_from_from_item(from_item)?; - if schema.is_none() - && let Some(cte_ptr) = resolve_cte_table(column_name_ref, &table_name) - { - if let Some(cte_column_ptr) = - resolve_cte_column(db, file, column_name_ref, &table_name, &column_name) - { - return Some(cte_column_ptr); - } - if column_name == table_name { - return Some(smallvec![Location::new( - file, - cte_ptr.text_range(), - LocationKind::Table - )]); - } - if let Some(alias_name) = from_item.alias().and_then(|x| x.name()) - && Name::from_node(&alias_name) == column_name - { - return Some(smallvec![Location::new( - file, - alias_name.syntax().text_range(), - LocationKind::Table - )]); - } - return None; - } - - if let Some(ptr) = resolve_column_from_table_or_view( + if let Some(ptr) = resolve_column_from_table_or_view_or_cte( db, file, column_name_ref, @@ -1499,7 +1469,7 @@ fn resolve_from_item_column_ptr( None } -fn resolve_column_from_table_or_view( +fn resolve_column_from_table_or_view_or_cte( db: &dyn Db, file: File, name_ref: &ast::NameRef, @@ -1507,10 +1477,18 @@ fn resolve_column_from_table_or_view( schema: &Option, column_name: &Name, ) -> Option> { - resolve_column_from_table_or_view_impl(db, file, name_ref, table_name, schema, column_name, 0) + resolve_column_from_table_or_view_or_cte_impl( + db, + file, + name_ref, + table_name, + schema, + column_name, + 0, + ) } -fn resolve_column_from_table_or_view_impl( +fn resolve_column_from_table_or_view_or_cte_impl( db: &dyn Db, file: File, name_ref: &ast::NameRef, @@ -1519,114 +1497,120 @@ fn resolve_column_from_table_or_view_impl( column_name: &Name, depth: u32, ) -> Option> { - let tree = parse(db, file).tree(); - let root = tree.syntax(); if depth > 40 { log::info!("max resolve depth reached, probably in a cycle"); return None; } let position = name_ref.syntax().text_range().start(); + let (table_like_ptr, kind) = + resolve_table_like(db, file, Some(name_ref), table_name, schema, position)?; - if let Some(table_name_ptr) = resolve_table_name_ptr(db, file, table_name, schema, position) { - let table_name_node = table_name_ptr.to_node(root); - - if let Some(create_table) = table_name_node - .ancestors() - .find_map(ast::CreateTableLike::cast) - { - // 1. try to find a matching column - if let Some(ptr) = find_column_in_create_table(db, file, &create_table, column_name) { - return Some(ptr); + match kind { + LocationKind::Table => { + if schema.is_none() && resolve_cte_table(name_ref, table_name).is_some() { + return resolve_cte_column(db, file, name_ref, table_name, column_name); } - // 2. No column found, 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()) - { - let (parent_schema, parent_table_name) = name::schema_and_name_path(&parent_path)?; - return resolve_column_from_table_or_view_impl( - db, - file, - name_ref, - &parent_table_name, - &parent_schema, - column_name, - depth + 1, - ); - } + let tree = parse(db, file).tree(); + let root = tree.syntax(); + let node = table_like_ptr.to_node(root); - // 3. No column found, check if the name matches the table name. - // For example, in: - // ```sql - // create table t(a int); - // select t from t; - // ``` - if column_name == table_name { - return Some(smallvec![Location::new( - file, - table_name_ptr.text_range(), - LocationKind::Table - )]); - } - } + if let Some(create_table) = node.ancestors().find_map(ast::CreateTableLike::cast) { + if let Some(cols) = + find_column_in_create_table(db, file, &create_table, column_name) + { + return Some(cols); + } - if let Some(create_table_as) = table_name_node - .ancestors() - .find_map(ast::CreateTableAs::cast) - { - if let Some(ptr) = find_column_in_create_table_as(file, &create_table_as, column_name) { - return Some(ptr); - } + // 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()) + { + let (parent_schema, parent_table_name) = + name::schema_and_name_path(&parent_path)?; + return resolve_column_from_table_or_view_or_cte_impl( + db, + file, + name_ref, + &parent_table_name, + &parent_schema, + column_name, + depth + 1, + ); + } - if column_name == table_name { - return Some(smallvec![Location::new( - file, - table_name_ptr.text_range(), - LocationKind::Table - )]); + // For example, in: + // ```sql + // create table t(a int); + // select t from t; + // ``` + if column_name == table_name { + return Some(smallvec![Location::new( + file, + table_like_ptr.text_range(), + LocationKind::Table + )]); + } } - } - if let Some(select_into) = table_name_node.ancestors().find_map(ast::SelectInto::cast) { - if let Some(ptr) = find_column_in_select_into(file, &select_into, column_name) { - return Some(ptr); + if let Some(create_table_as) = node.ancestors().find_map(ast::CreateTableAs::cast) { + if let Some(cols) = + find_column_in_create_table_as(file, &create_table_as, column_name) + { + return Some(cols); + } + + if column_name == table_name { + return Some(smallvec![Location::new( + file, + table_like_ptr.text_range(), + LocationKind::Table + )]); + } } - if column_name == table_name { - return Some(smallvec![Location::new( - file, - table_name_ptr.text_range(), - LocationKind::Table - )]); + if let Some(select_into) = node.ancestors().find_map(ast::SelectInto::cast) { + if let Some(cols) = find_column_in_select_into(file, &select_into, column_name) { + return Some(cols); + } + + if column_name == table_name { + return Some(smallvec![Location::new( + file, + table_like_ptr.text_range(), + LocationKind::Table + )]); + } } + + None } - } + LocationKind::View => { + let tree = parse(db, file).tree(); + let root = tree.syntax(); + let node = table_like_ptr.to_node(root); - // ditto as above but with view - if let Some(view_name_ptr) = resolve_view_name_ptr(db, file, table_name, schema, position) { - let view_name_node = view_name_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) + { + return Some(cols); + } - if let Some(create_view) = view_name_node - .ancestors() - .find_map(ast::CreateViewLike::cast) - { - if let Some(ptr) = find_column_in_create_view_like(file, &create_view, column_name) { - return Some(ptr); + if column_name == table_name { + return Some(smallvec![Location::new( + file, + table_like_ptr.text_range(), + LocationKind::View + )]); + } } - if column_name == table_name { - return Some(smallvec![Location::new( - file, - view_name_ptr.text_range(), - LocationKind::View - )]); - } + None } + _ => None, } - - None } fn resolve_from_item_for_cte_star( @@ -1640,14 +1624,8 @@ fn resolve_from_item_for_cte_star( if let Some((schema, table_name)) = name::schema_and_table_from_from_item(from_item) && table_name == *cte_name { - return resolve_column_from_table_or_view( - db, - file, - name_ref, - &table_name, - &schema, - column_name, - ); + let position = name_ref.syntax().text_range().start(); + return resolve_column_for_table(db, file, &table_name, &schema, column_name, position); } resolve_from_item_column_ptr(db, file, from_item, name_ref) @@ -1696,8 +1674,15 @@ fn resolve_fn_call_column( let from_clause = select.from_clause()?; for from_item in ast_nav::iter_from_clause(&from_clause) { - if let Some(result) = - resolve_from_item_for_fn_call_column(db, file, &from_item, &column_name, name_ref) + if let Some((schema, table_name)) = name::schema_and_table_from_from_item(&from_item) + && let Some(result) = resolve_column_from_table_or_view_or_cte( + db, + file, + name_ref, + &table_name, + &schema, + &column_name, + ) { return Some(result); } @@ -1706,28 +1691,6 @@ fn resolve_fn_call_column( None } -fn resolve_from_item_for_fn_call_column( - db: &dyn Db, - file: File, - from_item: &ast::FromItem, - column_name: &Name, - name_ref: &ast::NameRef, -) -> Option> { - let tree = parse(db, file).tree(); - let root = tree.syntax(); - let (schema, table_name) = name::schema_and_table_from_from_item(from_item)?; - - let position = name_ref.syntax().text_range().start(); - let table_ptr = resolve_table_name_ptr(db, file, &table_name, &schema, position)?; - - let table_name_node = table_ptr.to_node(root); - let create_table = table_name_node - .ancestors() - .find_map(ast::CreateTableLike::cast)?; - - find_column_in_create_table(db, file, &create_table, column_name) -} - fn is_from_item_match(from_item: &ast::FromItem, qualifier: &Name) -> bool { if let Some(alias_name) = from_item.alias().and_then(|a| a.name()) && Name::from_node(&alias_name) == *qualifier @@ -1983,7 +1946,7 @@ fn find_column_in_select_into( None } -pub(crate) fn resolve_cte_table(name_ref: &ast::NameRef, cte_name: &Name) -> Option { +fn resolve_cte_table(name_ref: &ast::NameRef, cte_name: &Name) -> Option { let with_table = ast_nav::find_cte_with_table(name_ref, cte_name)?; Some(SyntaxNodePtr::new(with_table.name()?.syntax())) } @@ -1992,7 +1955,7 @@ fn count_columns_for_path(db: &dyn Db, file: File, path: &ast::Path) -> Option, position: TextSize, + name_ref: Option<&ast::NameRef>, ) -> Option { let tree = parse(db, file).tree(); let root = tree.syntax(); - if let Some(table_name_ptr) = resolve_table_name_ptr(db, file, table_name, schema, position) { - let table_name_node = table_name_ptr.to_node(root); + let (table_like_ptr, kind) = + resolve_table_like(db, file, name_ref, table_name, schema, position)?; - if let Some(create_table) = table_name_node - .ancestors() - .find_map(ast::CreateTableLike::cast) - { - let mut count: usize = 0; - if let Some(args) = create_table.table_arg_list() { - for arg in args.args() { - if matches!(arg, ast::TableArg::Column(_)) { - count = count.saturating_add(1); + match kind { + LocationKind::Table => { + if schema.is_none() + && let Some(name_ref) = name_ref + && let Some(with_table) = ast_nav::find_cte_with_table(name_ref, table_name) + { + return count_columns_for_with_table(with_table); + } + + let table_like_node = table_like_ptr.to_node(root); + if let Some(create_table) = table_like_node + .ancestors() + .find_map(ast::CreateTableLike::cast) + { + let mut count: usize = 0; + if let Some(args) = create_table.table_arg_list() { + for arg in args.args() { + if matches!(arg, ast::TableArg::Column(_)) { + count = count.saturating_add(1); + } } } + return Some(count); } - return Some(count); - } - - if let Some(create_table_as) = table_name_node - .ancestors() - .find_map(ast::CreateTableAs::cast) - { - let select = ast_nav::select_from_variant(create_table_as.query()?)?; - if let Some(target_list) = select.select_clause().and_then(|c| c.target_list()) { - return Some(target_list.targets().count()); + if let Some(create_table_as) = table_like_node + .ancestors() + .find_map(ast::CreateTableAs::cast) + { + let select = ast_nav::select_from_variant(create_table_as.query()?)?; + if let Some(target_list) = select.select_clause().and_then(|c| c.target_list()) { + return Some(target_list.targets().count()); + } } - } - if let Some(select_into) = table_name_node.ancestors().find_map(ast::SelectInto::cast) { - if let Some(target_list) = select_into.select_clause().and_then(|c| c.target_list()) { - return Some(target_list.targets().count()); + if let Some(select_into) = table_like_node.ancestors().find_map(ast::SelectInto::cast) { + if let Some(target_list) = select_into.select_clause().and_then(|c| c.target_list()) + { + return Some(target_list.targets().count()); + } } - } - } - - if let Some(view_name_ptr) = resolve_view_name_ptr(db, file, table_name, schema, position) { - let view_name_node = view_name_ptr.to_node(root); - if let Some(create_view) = view_name_node - .ancestors() - .find_map(ast::CreateViewLike::cast) - { + None + } + LocationKind::View => { + let table_like_node = table_like_ptr.to_node(root); + let create_view = table_like_node + .ancestors() + .find_map(ast::CreateViewLike::cast)?; if let Some(column_list) = create_view.column_list() { return Some(column_list.columns().count()); } let select = ast_nav::select_from_variant(create_view.query()?)?; - if let Some(target_list) = select.select_clause().and_then(|c| c.target_list()) { // This is not quite right if there's a `*` in the view definition. // It becomes recursive. // For now, let's assume simple views. return Some(target_list.targets().count()); } + + None } + _ => None, } - None +} + +fn count_columns_for_with_table(with_table: ast::WithTable) -> Option { + if let Some(column_list) = with_table.column_list() { + return Some(column_list.columns().count()); + } + + let query = with_table.query()?; + if let ast::WithQuery::Values(values) = query { + return values + .row_list() + .and_then(|x| x.rows().next()) + .map(|row| row.exprs().count()); + } + + let select = ast_nav::select_from_with_query(query)?; + select + .select_clause() + .and_then(|c| c.target_list()) + .map(|t| t.targets().count()) } fn resolve_cte_column( @@ -2116,14 +2110,7 @@ fn resolve_cte_column( let path = table.relation_name()?.path()?; let (schema, table_name) = name::schema_and_name_path(&path)?; - if schema.is_none() - && let Some(nested_cte_column) = - resolve_cte_column(db, file, name_ref, &table_name, column_name) - { - return Some(nested_cte_column); - } - - return resolve_column_from_table_or_view( + return resolve_column_from_table_or_view_or_cte( db, file, name_ref, @@ -2201,6 +2188,16 @@ fn resolve_cte_column( column_index = column_list_end; } + if column_name == cte_name + && let Some(name_node) = with_table.name() + { + return Some(smallvec![Location::new( + file, + name_node.syntax().text_range(), + LocationKind::Table + )]); + } + None } @@ -2305,14 +2302,7 @@ fn resolve_subquery_column_ptr( let path = table.relation_name()?.path()?; let (schema, table_name) = name::schema_and_name_path(&path)?; - if schema.is_none() - && let Some(cte_column_ptr) = - resolve_cte_column(db, file, name_ref, &table_name, column_name) - { - return Some(cte_column_ptr); - } - - return resolve_column_from_table_or_view( + return resolve_column_from_table_or_view_or_cte( db, file, name_ref, @@ -2415,139 +2405,6 @@ pub(crate) fn qualified_star_table_name(field_expr: &ast::FieldExpr) -> Option Option { - let table_name = qualified_star_table_name(&field_expr)?; - let position = field_expr.syntax().text_range().start(); - let target = field_expr - .syntax() - .ancestors() - .find_map(ast::Target::cast)?; - - let res = ast_nav::target_parent_query(target)?; - - let path = match res { - ast_nav::ParentQuery::Select(select) => { - let from_clause = select.from_clause()?; - let from_item = find_from_item_in_from_clause(&from_clause, &table_name)?; - - if let Some(alias) = from_item.alias() - && alias.column_list().is_some() - { - return Some(SyntaxNodePtr::new(alias.syntax())); - } - - let (schema, table_name) = name::schema_and_table_from_from_item(&from_item)?; - - if let Some(table_name_ptr) = - resolve_table_name_ptr(db, file, &table_name, &schema, position) - { - return Some(table_name_ptr); - } - - if let Some(view_name_ptr) = - resolve_view_name_ptr(db, file, &table_name, &schema, position) - { - return Some(view_name_ptr); - } - - if schema.is_none() - && let Some(name_ref) = from_item.name_ref() - { - return resolve_cte_table(&name_ref, &table_name); - } - - return None; - } - ast_nav::ParentQuery::Update(update) => update.relation_name()?.path()?, - ast_nav::ParentQuery::Delete(delete) => delete.relation_name()?.path()?, - ast_nav::ParentQuery::Insert(insert) => insert.path()?, - ast_nav::ParentQuery::Merge(merge) => merge.relation_name()?.path()?, - }; - - return resolve_table_or_view_or_cte_ptrs(db, file, position, &path)? - .into_iter() - .next(); -} - -fn resolve_table_or_view_or_cte_ptrs( - db: &dyn Db, - file: File, - position: TextSize, - path: &ast::Path, -) -> Option> { - let (schema, table_name) = name::schema_and_name_path(path)?; - - let mut results = vec![]; - - if let Some(table_name_ptr) = resolve_table_name_ptr(db, file, &table_name, &schema, position) { - results.push(table_name_ptr); - } - - if let Some(view_name_ptr) = resolve_view_name_ptr(db, file, &table_name, &schema, position) { - results.push(view_name_ptr); - } - - if schema.is_none() - && let Some(name_ref) = path.segment().and_then(|x| x.name_ref()) - && let Some(cte_ptr) = resolve_cte_table(&name_ref, &table_name) - { - results.push(cte_ptr); - } - - if results.is_empty() { - return None; - } - - Some(results) -} - -pub(crate) fn resolve_unqualified_star_table_ptrs( - db: &dyn Db, - file: File, - target: &ast::Target, -) -> Option> { - target.star_token()?; - - let position = target.syntax().text_range().start(); - - let path = match ast_nav::target_parent_query(target.clone())? { - ast_nav::ParentQuery::Select(select) => { - let from_clause = select.from_clause()?; - let results = table_ptrs_from_clause(db, file, &from_clause); - if results.is_empty() { - return None; - } - return Some(results); - } - ast_nav::ParentQuery::Update(update) => update.relation_name()?.path(), - ast_nav::ParentQuery::Insert(insert) => insert.path(), - ast_nav::ParentQuery::Delete(delete) => delete.relation_name()?.path(), - ast_nav::ParentQuery::Merge(merge) => merge.relation_name()?.path(), - }?; - - resolve_table_or_view_or_cte_ptrs(db, file, position, &path) -} - -pub(crate) fn resolve_unqualified_star_in_arg_list_ptrs( - db: &dyn Db, - file: File, - arg_list: &ast::ArgList, -) -> Option> { - let select = arg_list.syntax().ancestors().find_map(ast::Select::cast)?; - let from_clause = select.from_clause()?; - let results = table_ptrs_from_clause(db, file, &from_clause); - - if results.is_empty() { - return None; - } - - Some(results) -} - // TODO: I don't think we should be passing in file here and using it, we want // to use the content from the updated file with the completion marker pub(crate) fn table_ptrs_from_clause( @@ -2574,25 +2431,12 @@ pub(crate) fn table_ptrs_from_clause( continue; }; + let name_ref = from_item.name_ref(); let position = from_item.syntax().text_range().start(); - if let Some(table_name_ptr) = - resolve_table_name_ptr(db, file, &table_name, &schema, position) + if let Some((table_like_ptr, _kind)) = + resolve_table_like(db, file, name_ref.as_ref(), &table_name, &schema, position) { - results.push(table_name_ptr); - continue; - } - - if let Some(view_name_ptr) = resolve_view_name_ptr(db, file, &table_name, &schema, position) - { - results.push(view_name_ptr); - continue; - } - - if schema.is_none() - && let Some(name_ref) = from_item.name_ref() - && let Some(cte_ptr) = resolve_cte_table(&name_ref, &table_name) - { - results.push(cte_ptr); + results.push(table_like_ptr); } } @@ -2613,24 +2457,11 @@ pub(crate) fn table_ptr_from_from_item( } let (schema, table_name) = name::schema_and_table_from_from_item(from_item)?; + let name_ref = from_item.name_ref(); let position = from_item.syntax().text_range().start(); - if let Some(table_ptr) = resolve_table_name_ptr(db, file, &table_name, &schema, position) { - return Some(table_ptr); - } - - if let Some(view_name_ptr) = resolve_view_name_ptr(db, file, &table_name, &schema, position) { - return Some(view_name_ptr); - } - - if schema.is_none() - && let Some(name_ref) = from_item.name_ref() - && let Some(cte_ptr) = resolve_cte_table(&name_ref, &table_name) - { - return Some(cte_ptr); - } - - None + resolve_table_like(db, file, name_ref.as_ref(), &table_name, &schema, position) + .map(|(table_like_ptr, _kind)| table_like_ptr) } fn table_ptr_from_paren_expr( @@ -2695,43 +2526,7 @@ fn count_columns_for_from_item( ) -> Option { let (schema, table_name) = name::schema_and_table_from_from_item(from_item)?; let position = name_ref.syntax().text_range().start(); - - if let Some(count) = count_columns_for_table_name(db, file, &table_name, &schema, position) { - return Some(count); - } - - if schema.is_none() - && let Some(cte_column_count) = count_columns_for_cte(name_ref, &table_name) - { - return Some(cte_column_count); - } - - None -} - -fn count_columns_for_cte(name_ref: &ast::NameRef, cte_name: &Name) -> Option { - let with_table = ast_nav::find_cte_with_table(name_ref, cte_name)?; - - if let Some(column_list) = with_table.column_list() { - return Some(column_list.columns().count()); - } - - let query = with_table.query()?; - - if let ast::WithQuery::Values(values) = query { - if let Some(first_row) = values.row_list().and_then(|x| x.rows().next()) { - return Some(first_row.exprs().count()); - } - return None; - } - - let select = ast_nav::select_from_with_query(query)?; - - if let Some(target_list) = select.select_clause().and_then(|c| c.target_list()) { - return Some(target_list.targets().count()); - } - - None + count_columns_for_table_name(db, file, &table_name, &schema, position, Some(name_ref)) } fn resolve_from_clause_for_cte_star( @@ -3115,20 +2910,16 @@ fn resolve_merge_table_name_ptr( if let Some(item_name_ref) = from_item.name_ref() { let item_name = Name::from_node(&item_name_ref); if item_name == table_name { - if let Some(cte_ptr) = resolve_cte_table(table_name_ref, &table_name) { - return Some(smallvec![Location::new( - file, - cte_ptr.text_range(), - LocationKind::Table - )]); - } let position = table_name_ref.syntax().text_range().start(); - let ptr = resolve_table_name_ptr(db, file, &item_name, &None, position)?; - return Some(smallvec![Location::new( + let (ptr, kind) = resolve_table_like( + db, file, - ptr.text_range(), - LocationKind::Table - )]); + Some(table_name_ref), + &item_name, + &None, + position, + )?; + return Some(smallvec![Location::new(file, ptr.text_range(), kind)]); } } // Check for aliased source tables