ide: code actions unnest <-> rows from

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<TextSize>) -> Option<Vec<CodeA
     rewrite_between_as_binary_expression(db, position, &mut actions);
     rewrite_not_equals_operator(db, position, &mut actions);
     rewrite_timestamp_type(db, position, &mut actions);
+    rewrite_unnest_as_rows_from(db, position, &mut actions);
+    rewrite_rows_from_as_unnest(db, position, &mut actions);
     Some(actions)
 }

crates/squawk_ide/src/code_actions/rewrite_rows_from_as_unnest.rs

@@ -0,0 +1,330 @@
+use rowan::{TextRange, 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_rows_from_as_unnest(
+    db: &dyn Db,
+    position: InFile<TextSize>,
+    actions: &mut Vec<CodeAction>,
+) -> 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::<Vec<_>>()
+        .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);"
+        ));
+    }
+}