ide: refactor code actions into their own modules

crates/squawk_ide/src/code_actions/add_explicit_alias.rs

@@ -0,0 +1,176 @@
+use rowan::TextSize;
+use salsa::Database as Db;
+use squawk_syntax::ast::{self, AstNode};
+
+use crate::{
+    column_name::ColumnName, db::File, offsets::token_from_offset, quote::quote_column_alias,
+};
+
+use super::{ActionKind, CodeAction};
+
+pub(super) fn add_explicit_alias(
+    db: &dyn Db,
+    file: File,
+    actions: &mut Vec<CodeAction>,
+    offset: TextSize,
+) -> Option<()> {
+    let token = token_from_offset(db, file, offset)?;
+    let target = token.parent_ancestors().find_map(ast::Target::cast)?;
+
+    if target.as_name().is_some() {
+        return None;
+    }
+
+    if let Some(ast::Expr::FieldExpr(field_expr)) = target.expr()
+        && field_expr.star_token().is_some()
+    {
+        return None;
+    }
+
+    let alias = ColumnName::from_target(target.clone()).and_then(|c| c.0.to_string())?;
+
+    let expr_end = target.expr().map(|e| e.syntax().text_range().end())?;
+
+    let quoted_alias = quote_column_alias(&alias);
+    let replacement = format!(" as {}", quoted_alias);
+
+    actions.push(CodeAction {
+        title: "Add explicit alias".to_owned(),
+        edits: vec![squawk_linter::Edit::insert(replacement, expr_end)],
+        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::add_explicit_alias;
+
+    #[test]
+    fn add_explicit_alias_simple_column() {
+        assert_snapshot!(apply_code_action(
+            add_explicit_alias,
+            "select col_na$0me from t;"),
+            @"select col_name as col_name from t;"
+        );
+    }
+
+    #[test]
+    fn add_explicit_alias_quoted_identifier() {
+        assert_snapshot!(apply_code_action(
+            add_explicit_alias,
+            r#"select "b"$0 from t;"#),
+            @r#"select "b" as b from t;"#
+        );
+    }
+
+    #[test]
+    fn add_explicit_alias_field_expr() {
+        assert_snapshot!(apply_code_action(
+            add_explicit_alias,
+            "select t.col$0umn from t;"),
+            @"select t.column as column from t;"
+        );
+    }
+
+    #[test]
+    fn add_explicit_alias_function_call() {
+        assert_snapshot!(apply_code_action(
+            add_explicit_alias,
+            "select cou$0nt(*) from t;"),
+            @"select count(*) as count from t;"
+        );
+    }
+
+    #[test]
+    fn add_explicit_alias_cast_to_type() {
+        assert_snapshot!(apply_code_action(
+            add_explicit_alias,
+            "select '1'::bigi$0nt from t;"),
+            @"select '1'::bigint as int8 from t;"
+        );
+    }
+
+    #[test]
+    fn add_explicit_alias_cast_column() {
+        assert_snapshot!(apply_code_action(
+            add_explicit_alias,
+            "select col_na$0me::text from t;"),
+            @"select col_name::text as col_name from t;"
+        );
+    }
+
+    #[test]
+    fn add_explicit_alias_case_expr() {
+        assert_snapshot!(apply_code_action(
+            add_explicit_alias,
+            "select ca$0se when true then 'a' end from t;"),
+            @"select case when true then 'a' end as case from t;"
+        );
+    }
+
+    #[test]
+    fn add_explicit_alias_case_with_else() {
+        assert_snapshot!(apply_code_action(
+            add_explicit_alias,
+            "select ca$0se when true then 'a' else now()::text end from t;"),
+            @"select case when true then 'a' else now()::text end as now from t;"
+        );
+    }
+
+    #[test]
+    fn add_explicit_alias_array() {
+        assert_snapshot!(apply_code_action(
+            add_explicit_alias,
+            "select arr$0ay[1, 2, 3] from t;"),
+            @"select array[1, 2, 3] as array from t;"
+        );
+    }
+
+    #[test]
+    fn add_explicit_alias_not_applicable_already_has_alias() {
+        assert!(code_action_not_applicable(
+            add_explicit_alias,
+            "select col_name$0 as foo from t;"
+        ));
+    }
+
+    #[test]
+    fn add_explicit_alias_unknown_column() {
+        assert_snapshot!(apply_code_action(
+            add_explicit_alias,
+            "select 1 $0+ 2 from t;"),
+            @r#"select 1 + 2 as "?column?" from t;"#
+        );
+    }
+
+    #[test]
+    fn add_explicit_alias_not_applicable_star() {
+        assert!(code_action_not_applicable(
+            add_explicit_alias,
+            "select $0* from t;"
+        ));
+    }
+
+    #[test]
+    fn add_explicit_alias_not_applicable_qualified_star() {
+        assert!(code_action_not_applicable(
+            add_explicit_alias,
+            "with t as (select 1 a) select t.*$0 from t;"
+        ));
+    }
+
+    #[test]
+    fn add_explicit_alias_literal() {
+        assert_snapshot!(apply_code_action(
+            add_explicit_alias,
+            "select 'foo$0' from t;"),
+            @r#"select 'foo' as "?column?" from t;"#
+        );
+    }
+}

crates/squawk_ide/src/code_actions/add_schema.rs

@@ -0,0 +1,205 @@
+use rowan::TextSize;
+use salsa::Database as Db;
+use squawk_syntax::ast::{self, AstNode};
+
+use super::{ActionKind, CodeAction};
+use crate::{
+    db::{File, bind},
+    offsets::token_from_offset,
+};
+
+pub(super) fn add_schema(
+    db: &dyn Db,
+    file: File,
+    actions: &mut Vec<CodeAction>,
+    offset: TextSize,
+) -> Option<()> {
+    let token = token_from_offset(db, file, offset)?;
+    let range = token.parent_ancestors().find_map(|node| {
+        if let Some(path) = ast::Path::cast(node.clone()) {
+            if path.qualifier().is_some() {
+                return None;
+            }
+            return Some(path.syntax().text_range());
+        }
+        if let Some(from_item) = ast::FromItem::cast(node.clone()) {
+            let name_ref = from_item.name_ref()?;
+            return Some(name_ref.syntax().text_range());
+        }
+        if let Some(call_expr) = ast::CallExpr::cast(node) {
+            let ast::Expr::NameRef(name_ref) = call_expr.expr()? else {
+                return None;
+            };
+            return Some(name_ref.syntax().text_range());
+        }
+        None
+    })?;
+
+    if !range.contains(offset) {
+        return None;
+    }
+
+    let position = token.text_range().start();
+    let schema = bind(db, file).search_path_at(position).first()?.to_string();
+    let replacement = format!("{}.", schema);
+
+    actions.push(CodeAction {
+        title: "Add schema".to_owned(),
+        edits: vec![squawk_linter::Edit::insert(replacement, position)],
+        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::add_schema;
+
+    #[test]
+    fn add_schema_simple() {
+        assert_snapshot!(apply_code_action(
+            add_schema,
+            "create table t$0(a text, b int);"),
+            @"create table public.t(a text, b int);"
+        );
+    }
+
+    #[test]
+    fn add_schema_create_foreign_table() {
+        assert_snapshot!(apply_code_action(
+            add_schema,
+            "create foreign table t$0(a text, b int) server foo;"),
+            @"create foreign table public.t(a text, b int) server foo;"
+        );
+    }
+
+    #[test]
+    fn add_schema_create_function() {
+        assert_snapshot!(apply_code_action(
+            add_schema,
+            "create function f$0() returns int8\n  as 'select 1'\n  language sql;"),
+            @"create function public.f() returns int8
+  as 'select 1'
+  language sql;"
+        );
+    }
+
+    #[test]
+    fn add_schema_create_type() {
+        assert_snapshot!(apply_code_action(
+            add_schema,
+            "create type t$0 as enum ();"),
+            @"create type public.t as enum ();"
+        );
+    }
+
+    #[test]
+    fn add_schema_table_stmt() {
+        assert_snapshot!(apply_code_action(
+            add_schema,
+            "table t$0;"),
+            @"table public.t;"
+        );
+    }
+
+    #[test]
+    fn add_schema_select_from() {
+        assert_snapshot!(apply_code_action(
+            add_schema,
+            "create table t(a text, b int);
+        select t from t$0;"),
+            @"create table t(a text, b int);
+        select t from public.t;"
+        );
+    }
+
+    #[test]
+    fn add_schema_select_table_value() {
+        // we can't insert the schema here because:
+        // `select public.t from t` isn't valid
+        assert!(code_action_not_applicable(
+            add_schema,
+            "create table t(a text, b int);
+        select t$0 from t;"
+        ));
+    }
+
+    #[test]
+    fn add_schema_select_unqualified_column() {
+        // not applicable since we don't have the table name set
+        // we'll have another quick action to insert table names
+        assert!(code_action_not_applicable(
+            add_schema,
+            "create table t(a text, b int);
+        select a$0 from t;"
+        ));
+    }
+
+    #[test]
+    fn add_schema_select_qualified_column() {
+        // not valid because we haven't specified the schema on the table name
+        // `select public.t.c from t` isn't valid sql
+        assert!(code_action_not_applicable(
+            add_schema,
+            "create table t(c text);
+        select t$0.c from t;"
+        ));
+    }
+
+    #[test]
+    fn add_schema_with_search_path() {
+        assert_snapshot!(
+            apply_code_action(
+                add_schema,
+                "
+set search_path to myschema;
+create table t$0(a text, b int);"
+            ),
+            @"
+set search_path to myschema;
+create table myschema.t(a text, b int);"
+        );
+    }
+
+    #[test]
+    fn add_schema_not_applicable_with_schema() {
+        assert!(code_action_not_applicable(
+            add_schema,
+            "create table myschema.t$0(a text, b int);"
+        ));
+    }
+
+    #[test]
+    fn add_schema_function_call() {
+        assert_snapshot!(apply_code_action(
+            add_schema,
+            "
+create function f() returns int8
+  as 'select 1'
+  language sql;
+
+select f$0();"),
+            @"
+create function f() returns int8
+  as 'select 1'
+  language sql;
+
+select public.f();"
+        );
+    }
+
+    #[test]
+    fn add_schema_function_call_not_applicable_with_schema() {
+        assert!(code_action_not_applicable(
+            add_schema,
+            "
+create function f() returns int8 as 'select 1' language sql;
+select myschema.f$0();"
+        ));
+    }
+}