diff --git a/crates/squawk_ide/src/code_actions/mod.rs b/crates/squawk_ide/src/code_actions/mod.rs index 0b15bace..d321ba20 100644 --- a/crates/squawk_ide/src/code_actions/mod.rs +++ b/crates/squawk_ide/src/code_actions/mod.rs @@ -20,12 +20,15 @@ mod rewrite_from; mod rewrite_integer_radix; mod rewrite_leading_from; mod rewrite_not_equals_operator; +mod rewrite_rows_from_as_unnest; mod rewrite_select_as_table; mod rewrite_select_as_values; mod rewrite_select_into_as_create_table_as; mod rewrite_table_as_select; mod rewrite_timestamp_type; +mod rewrite_unnest_as_rows_from; mod rewrite_values_as_select; +mod unnest; mod unquote_identifier; #[cfg(test)] @@ -47,11 +50,13 @@ use rewrite_from::rewrite_from; use rewrite_integer_radix::rewrite_integer_radix; use rewrite_leading_from::rewrite_leading_from; use rewrite_not_equals_operator::rewrite_not_equals_operator; +use rewrite_rows_from_as_unnest::rewrite_rows_from_as_unnest; use rewrite_select_as_table::rewrite_select_as_table; use rewrite_select_as_values::rewrite_select_as_values; use rewrite_select_into_as_create_table_as::rewrite_select_into_as_create_table_as; use rewrite_table_as_select::rewrite_table_as_select; use rewrite_timestamp_type::rewrite_timestamp_type; +use rewrite_unnest_as_rows_from::rewrite_unnest_as_rows_from; use rewrite_values_as_select::rewrite_values_as_select; use unquote_identifier::unquote_identifier; @@ -93,5 +98,7 @@ pub fn code_actions(db: &dyn Db, position: InFile) -> Option, + actions: &mut Vec, +) -> Option<()> { + let token = token_from_offset(db, position)?; + let rows_from = token.parent_ancestors().find_map(ast::RowsFromItem::cast)?; + + if rows_from + .alias() + .and_then(|alias| alias.columns()) + .is_some_and(|columns| matches!(columns, ast::FromAliasColumns::ColumnDefList(_))) + { + return None; + } + + let mut calls = vec![]; + for rows_from_arg in rows_from.rows_from_args() { + if rows_from_arg.column_def_list().is_some() { + return None; + } + calls.push(unnest_call(&rows_from_arg.call_expr()?)?); + } + let (first, rest) = calls.split_first()?; + + let rows_from_range = TextRange::new( + rows_from.rows_token()?.text_range().start(), + rows_from.r_paren_token()?.text_range().end(), + ); + + if !rest.is_empty() && unnest_shadowed(db, position.file_id, rows_from_range.start()) { + return None; + } + + let args = calls + .iter() + .flat_map(|call| &call.args) + .map(|arg| arg.syntax().text().to_string()) + .collect::>() + .join(", "); + + actions.push(CodeAction { + title: "Rewrite as `unnest`".to_owned(), + edits: vec![Edit::replace( + rows_from_range, + format!("{}({args})", first.name.syntax().text()), + )], + kind: ActionKind::RefactorRewrite, + }); + + Some(()) +} + +#[cfg(test)] +mod test { + use insta::assert_snapshot; + + use crate::code_actions::test_utils::{apply_code_action, code_action_not_applicable}; + + use super::rewrite_rows_from_as_unnest; + + #[test] + fn rewrite_rows_from_as_unnest_simple() { + assert_snapshot!( + apply_code_action( + rewrite_rows_from_as_unnest, + "select * from rows f$0rom (unnest(a), unnest(b));" + ), + @"select * from unnest(a, b);" + ); + } + + #[test] + fn rewrite_rows_from_as_unnest_single_item() { + assert_snapshot!( + apply_code_action( + rewrite_rows_from_as_unnest, + "select * from rows f$0rom (unnest(a));" + ), + @"select * from unnest(a);" + ); + } + + #[test] + fn rewrite_rows_from_as_unnest_multi_arg_item() { + assert_snapshot!( + apply_code_action( + rewrite_rows_from_as_unnest, + "select * from rows f$0rom (unnest(a, b), unnest(c));" + ), + @"select * from unnest(a, b, c);" + ); + } + + #[test] + fn rewrite_rows_from_as_unnest_keeps_alias() { + assert_snapshot!( + apply_code_action( + rewrite_rows_from_as_unnest, + "select * from rows f$0rom (unnest(a), unnest(b)) as z(x, y);" + ), + @"select * from unnest(a, b) as z(x, y);" + ); + } + + #[test] + fn rewrite_rows_from_as_unnest_keeps_with_ordinality() { + assert_snapshot!( + apply_code_action( + rewrite_rows_from_as_unnest, + "select * from rows f$0rom (unnest(a), unnest(b)) with ordinality as z(x, y, n);" + ), + @"select * from unnest(a, b) with ordinality as z(x, y, n);" + ); + } + + #[test] + fn rewrite_rows_from_as_unnest_keeps_lateral() { + assert_snapshot!( + apply_code_action( + rewrite_rows_from_as_unnest, + "select * from t, lateral rows f$0rom (unnest(t.a), unnest(t.b));" + ), + @"select * from t, lateral unnest(t.a, t.b);" + ); + } + + #[test] + fn rewrite_rows_from_as_unnest_array_literals() { + assert_snapshot!( + apply_code_action( + rewrite_rows_from_as_unnest, + "select * from rows from (unnest(array[1,2]$0), unnest(array['a','b']));" + ), + @"select * from unnest(array[1,2], array['a','b']);" + ); + } + + #[test] + fn rewrite_rows_from_as_unnest_preserves_case() { + assert_snapshot!( + apply_code_action( + rewrite_rows_from_as_unnest, + "select * from ROWS F$0ROM (UNNEST(a), UNNEST(b));" + ), + @"select * from UNNEST(a, b);" + ); + } + + #[test] + fn rewrite_rows_from_as_unnest_not_applicable_mixed_functions() { + assert!(code_action_not_applicable( + rewrite_rows_from_as_unnest, + "select * from rows f$0rom (unnest(a), generate_series(1, 10));" + )); + } + + #[test] + fn rewrite_rows_from_as_unnest_not_applicable_qualified() { + assert!(code_action_not_applicable( + rewrite_rows_from_as_unnest, + "select * from rows f$0rom (pg_catalog.unnest(a), unnest(b));" + )); + } + + #[test] + fn rewrite_rows_from_as_unnest_not_applicable_variadic() { + assert!(code_action_not_applicable( + rewrite_rows_from_as_unnest, + "select * from rows f$0rom (unnest(variadic a), unnest(b));" + )); + } + + #[test] + fn rewrite_rows_from_as_unnest_not_applicable_column_def_list() { + assert!(code_action_not_applicable( + rewrite_rows_from_as_unnest, + "select * from rows f$0rom (unnest(a), unnest(b)) as z(x int, y text);" + )); + } + + #[test] + fn rewrite_rows_from_as_unnest_not_applicable_per_item_column_def_list() { + assert!(code_action_not_applicable( + rewrite_rows_from_as_unnest, + "select * from rows f$0rom (unnest(a) as (x int), unnest(b) as (y text));" + )); + } + + // postgres ignores the `all` and merges the call like any other + #[test] + fn rewrite_rows_from_as_unnest_all() { + assert_snapshot!( + apply_code_action( + rewrite_rows_from_as_unnest, + "select * from rows f$0rom (unnest(all a), unnest(b));" + ), + @"select * from unnest(a, b);" + ); + } + + #[test] + fn rewrite_rows_from_as_unnest_not_applicable_distinct() { + assert!(code_action_not_applicable( + rewrite_rows_from_as_unnest, + "select * from rows f$0rom (unnest(distinct a), unnest(b));" + )); + } + + #[test] + fn rewrite_rows_from_as_unnest_not_applicable_star() { + assert!(code_action_not_applicable( + rewrite_rows_from_as_unnest, + "select * from rows f$0rom (unnest(*), unnest(b));" + )); + } + + #[test] + fn rewrite_rows_from_as_unnest_not_applicable_named_args() { + assert!(code_action_not_applicable( + rewrite_rows_from_as_unnest, + "select * from rows f$0rom (unnest(x => a), unnest(b));" + )); + } + + #[test] + fn rewrite_rows_from_as_unnest_not_applicable_order_by() { + assert!(code_action_not_applicable( + rewrite_rows_from_as_unnest, + "select * from rows f$0rom (unnest(a, b order by 1), unnest(c));" + )); + } + + #[test] + fn rewrite_rows_from_as_unnest_not_applicable_over_clause() { + assert!(code_action_not_applicable( + rewrite_rows_from_as_unnest, + "select * from rows f$0rom (unnest(a) over (), unnest(b));" + )); + } + + #[test] + fn rewrite_rows_from_as_unnest_not_applicable_filter_clause() { + assert!(code_action_not_applicable( + rewrite_rows_from_as_unnest, + "select * from rows f$0rom (unnest(a) filter (where x), unnest(b));" + )); + } + + #[test] + fn rewrite_rows_from_as_unnest_not_applicable_per_item_column_def_list_with_alias() { + assert!(code_action_not_applicable( + rewrite_rows_from_as_unnest, + "select * from rows f$0rom (unnest(a) as (x int)) as z(q);" + )); + } + + #[test] + fn rewrite_rows_from_as_unnest_not_applicable_shadowed() { + assert!(code_action_not_applicable( + rewrite_rows_from_as_unnest, + " +create function unnest(x text) returns setof text as $$ select 1 $$ language sql; +select * from rows f$0rom (unnest(a), unnest(b));" + )); + } + + #[test] + fn rewrite_rows_from_as_unnest_other_schema_not_shadowed() { + assert_snapshot!( + apply_code_action( + rewrite_rows_from_as_unnest, + " +create function other.unnest(x text) returns setof text as $$ select 1 $$ language sql; +select * from rows f$0rom (unnest(a), unnest(b));" + ), + @" +create function other.unnest(x text) returns setof text as $$ select 1 $$ language sql; +select * from unnest(a, b);" + ); + } + + #[test] + fn rewrite_rows_from_as_unnest_shadowed_single_arg() { + assert_snapshot!( + apply_code_action( + rewrite_rows_from_as_unnest, + " +create function unnest(x text) returns setof text as $$ select 1 $$ language sql; +select * from rows f$0rom (unnest(a));" + ), + @" +create function unnest(x text) returns setof text as $$ select 1 $$ language sql; +select * from unnest(a);" + ); + } + + #[test] + fn rewrite_rows_from_as_unnest_shadowed_multi_arg_call() { + assert_snapshot!( + apply_code_action( + rewrite_rows_from_as_unnest, + " +create function unnest(x text) returns setof text as $$ select 1 $$ language sql; +select * from rows f$0rom (unnest(a, b));" + ), + @" +create function unnest(x text) returns setof text as $$ select 1 $$ language sql; +select * from unnest(a, b);" + ); + } + + #[test] + fn rewrite_rows_from_as_unnest_not_applicable_plain_function_from_item() { + assert!(code_action_not_applicable( + rewrite_rows_from_as_unnest, + "select * from unn$0est(a, b);" + )); + } +} diff --git a/crates/squawk_ide/src/code_actions/rewrite_unnest_as_rows_from.rs b/crates/squawk_ide/src/code_actions/rewrite_unnest_as_rows_from.rs new file mode 100644 index 00000000..eb2f8dc7 --- /dev/null +++ b/crates/squawk_ide/src/code_actions/rewrite_unnest_as_rows_from.rs @@ -0,0 +1,309 @@ +use rowan::TextSize; +use salsa::Database as Db; +use squawk_linter::Edit; +use squawk_syntax::ast::{self, AstNode}; + +use crate::{file::InFile, offsets::token_from_offset}; + +use super::unnest::{unnest_call, unnest_shadowed}; +use super::{ActionKind, CodeAction}; + +pub(super) fn rewrite_unnest_as_rows_from( + db: &dyn Db, + position: InFile, + actions: &mut Vec, +) -> Option<()> { + let token = token_from_offset(db, position)?; + let from_item = token + .parent_ancestors() + .find_map(ast::FunctionFromItem::cast)?; + + let call_expr = from_item.call_expr()?; + let call = unnest_call(&call_expr)?; + if call.args.len() < 2 { + return None; + } + + if from_item + .alias() + .and_then(|alias| alias.columns()) + .is_some_and(|columns| matches!(columns, ast::FromAliasColumns::ColumnDefList(_))) + { + return None; + } + + let call_range = call_expr.syntax().text_range(); + if unnest_shadowed(db, position.file_id, call_range.start()) { + return None; + } + + let name = call.name.syntax().text(); + let calls = call + .args + .iter() + .map(|arg| format!("{name}({})", arg.syntax().text())) + .collect::>() + .join(", "); + + actions.push(CodeAction { + title: "Rewrite as `rows from`".to_owned(), + edits: vec![Edit::replace(call_range, format!("rows from ({calls})"))], + kind: ActionKind::RefactorRewrite, + }); + + Some(()) +} + +#[cfg(test)] +mod test { + use insta::assert_snapshot; + + use crate::code_actions::test_utils::{apply_code_action, code_action_not_applicable}; + + use super::rewrite_unnest_as_rows_from; + + #[test] + fn rewrite_unnest_as_rows_from_simple() { + assert_snapshot!( + apply_code_action( + rewrite_unnest_as_rows_from, + "select * from unn$0est(a, b);" + ), + @"select * from rows from (unnest(a), unnest(b));" + ); + } + + #[test] + fn rewrite_unnest_as_rows_from_three_args() { + assert_snapshot!( + apply_code_action( + rewrite_unnest_as_rows_from, + "select * from unnest(array[1,2], array['a','b'], c$0);" + ), + @"select * from rows from (unnest(array[1,2]), unnest(array['a','b']), unnest(c));" + ); + } + + #[test] + fn rewrite_unnest_as_rows_from_keeps_alias() { + assert_snapshot!( + apply_code_action( + rewrite_unnest_as_rows_from, + "select * from unn$0est(a, b) as z(x, y);" + ), + @"select * from rows from (unnest(a), unnest(b)) as z(x, y);" + ); + } + + #[test] + fn rewrite_unnest_as_rows_from_keeps_with_ordinality() { + assert_snapshot!( + apply_code_action( + rewrite_unnest_as_rows_from, + "select * from unn$0est(a, b) with ordinality as z(x, y, n);" + ), + @"select * from rows from (unnest(a), unnest(b)) with ordinality as z(x, y, n);" + ); + } + + #[test] + fn rewrite_unnest_as_rows_from_keeps_lateral() { + assert_snapshot!( + apply_code_action( + rewrite_unnest_as_rows_from, + "select * from t, lateral unn$0est(t.a, t.b);" + ), + @"select * from t, lateral rows from (unnest(t.a), unnest(t.b));" + ); + } + + #[test] + fn rewrite_unnest_as_rows_from_preserves_case() { + assert_snapshot!( + apply_code_action( + rewrite_unnest_as_rows_from, + "select * from UNN$0EST(a, b);" + ), + @"select * from rows from (UNNEST(a), UNNEST(b));" + ); + } + + #[test] + fn rewrite_unnest_as_rows_from_in_join() { + assert_snapshot!( + apply_code_action( + rewrite_unnest_as_rows_from, + "select * from t join unn$0est(a, b) on true;" + ), + @"select * from t join rows from (unnest(a), unnest(b)) on true;" + ); + } + + #[test] + fn rewrite_unnest_as_rows_from_not_applicable_single_arg() { + assert!(code_action_not_applicable( + rewrite_unnest_as_rows_from, + "select * from unn$0est(a);" + )); + } + + #[test] + fn rewrite_unnest_as_rows_from_not_applicable_other_function() { + assert!(code_action_not_applicable( + rewrite_unnest_as_rows_from, + "select * from generate_ser$0ies(1, 10);" + )); + } + + #[test] + fn rewrite_unnest_as_rows_from_not_applicable_qualified() { + assert!(code_action_not_applicable( + rewrite_unnest_as_rows_from, + "select * from pg_catalog.unn$0est(a, b);" + )); + } + + #[test] + fn rewrite_unnest_as_rows_from_not_applicable_variadic() { + assert!(code_action_not_applicable( + rewrite_unnest_as_rows_from, + "select * from unn$0est(variadic a, b);" + )); + } + + #[test] + fn rewrite_unnest_as_rows_from_all() { + assert_snapshot!( + apply_code_action( + rewrite_unnest_as_rows_from, + "select * from unn$0est(all a, b);" + ), + @"select * from rows from (unnest(a), unnest(b));" + ); + } + + #[test] + fn rewrite_unnest_as_rows_from_not_applicable_distinct() { + assert!(code_action_not_applicable( + rewrite_unnest_as_rows_from, + "select * from unn$0est(distinct a, b);" + )); + } + + #[test] + fn rewrite_unnest_as_rows_from_not_applicable_star() { + assert!(code_action_not_applicable( + rewrite_unnest_as_rows_from, + "select * from unn$0est(*);" + )); + } + + #[test] + fn rewrite_unnest_as_rows_from_not_applicable_named_args() { + assert!(code_action_not_applicable( + rewrite_unnest_as_rows_from, + "select * from unn$0est(x => a, y => b);" + )); + } + + #[test] + fn rewrite_unnest_as_rows_from_not_applicable_order_by() { + assert!(code_action_not_applicable( + rewrite_unnest_as_rows_from, + "select * from unn$0est(a, b order by 1);" + )); + } + + // the clause is part of the call expr we replace, so without the check it + // would be dropped + #[test] + fn rewrite_unnest_as_rows_from_not_applicable_over_clause() { + assert!(code_action_not_applicable( + rewrite_unnest_as_rows_from, + "select * from unn$0est(a, b) over ();" + )); + } + + #[test] + fn rewrite_unnest_as_rows_from_not_applicable_filter_clause() { + assert!(code_action_not_applicable( + rewrite_unnest_as_rows_from, + "select * from unn$0est(a, b) filter (where x);" + )); + } + + #[test] + fn rewrite_unnest_as_rows_from_not_applicable_column_def_list() { + assert!(code_action_not_applicable( + rewrite_unnest_as_rows_from, + "select * from unn$0est(a, b) as z(x int, y text);" + )); + } + + #[test] + fn rewrite_unnest_as_rows_from_not_applicable_shadowed() { + assert!(code_action_not_applicable( + rewrite_unnest_as_rows_from, + " +create function unnest(x text) returns setof text as $$ select 1 $$ language sql; +select * from unn$0est(a, b);" + )); + } + + #[test] + fn rewrite_unnest_as_rows_from_other_schema_not_shadowed() { + assert_snapshot!( + apply_code_action( + rewrite_unnest_as_rows_from, + " +create function other.unnest(x text) returns setof text as $$ select 1 $$ language sql; +select * from unn$0est(a, b);" + ), + @" +create function other.unnest(x text) returns setof text as $$ select 1 $$ language sql; +select * from rows from (unnest(a), unnest(b));" + ); + } + + #[test] + fn rewrite_unnest_as_rows_from_not_applicable_shadowed_via_search_path() { + assert!(code_action_not_applicable( + rewrite_unnest_as_rows_from, + " +create function other.unnest(x text) returns setof text as $$ select 1 $$ language sql; +set search_path to other, public, pg_catalog; +select * from unn$0est(a, b);" + )); + } + + #[test] + fn rewrite_unnest_as_rows_from_unrelated_function_definition() { + assert_snapshot!( + apply_code_action( + rewrite_unnest_as_rows_from, + " +create function f(x text) returns setof text as $$ select 1 $$ language sql; +select * from unn$0est(a, b);" + ), + @" +create function f(x text) returns setof text as $$ select 1 $$ language sql; +select * from rows from (unnest(a), unnest(b));" + ); + } + + #[test] + fn rewrite_unnest_as_rows_from_not_applicable_in_target_list() { + assert!(code_action_not_applicable( + rewrite_unnest_as_rows_from, + "select unn$0est(a, b);" + )); + } + + #[test] + fn rewrite_unnest_as_rows_from_not_applicable_already_rows_from() { + assert!(code_action_not_applicable( + rewrite_unnest_as_rows_from, + "select * from rows from (unn$0est(a), unnest(b));" + )); + } +} diff --git a/crates/squawk_ide/src/code_actions/unnest.rs b/crates/squawk_ide/src/code_actions/unnest.rs new file mode 100644 index 00000000..7a828715 --- /dev/null +++ b/crates/squawk_ide/src/code_actions/unnest.rs @@ -0,0 +1,64 @@ +use rowan::TextSize; +use salsa::Database as Db; +use squawk_syntax::ast; + +use crate::{ + db::{File, bind}, + name::Name, + symbols::SymbolKind, +}; + +pub(super) struct UnnestCall { + pub(super) name: ast::NameRef, + pub(super) args: Vec, +} + +pub(super) fn unnest_call(call_expr: &ast::CallExpr) -> Option { + let ast::Expr::NameRef(name_ref) = call_expr.expr()? else { + return None; + }; + if Name::from_node(&name_ref) != "unnest" { + return None; + } + if call_expr.over_clause().is_some() + || call_expr.filter_clause().is_some() + || call_expr.within_clause().is_some() + || call_expr.null_treatment().is_some() + { + return None; + } + let arg_list = call_expr.arg_list()?; + if arg_list.star_token().is_some() + || matches!( + arg_list.all_or_distinct(), + Some(ast::AllOrDistinct::Distinct(_)) + ) + { + return None; + } + let mut args = vec![]; + for arg in arg_list.args() { + if arg.variadic_token().is_some() + || arg.named_arg().is_some() + || arg.order_by_clause().is_some() + { + return None; + } + args.push(arg); + } + if args.is_empty() { + return None; + } + Some(UnnestCall { + name: name_ref, + args, + }) +} + +pub(super) fn unnest_shadowed(db: &dyn Db, file: File, position: TextSize) -> bool { + let binder = bind(db, file); + let schemas = binder.resolved_schemas(position, None); + binder + .lookup_with("unnest", SymbolKind::Function, &schemas) + .is_some() +}