squawk-ide 2.63.0

Linter for Postgres migrations & SQL
Documentation
use rowan::TextSize;
use salsa::Database as Db;
use squawk_linter::Edit;
use squawk_syntax::{
    ast::{self, AstNode, LitKind},
    quote::quote_string_literal,
};

use crate::{file::InFile, offsets::token_from_offset};

use super::{ActionKind, CodeAction};

pub(super) fn rewrite_extract_as_function_call(
    db: &dyn Db,
    position: InFile<TextSize>,
    actions: &mut Vec<CodeAction>,
) -> Option<()> {
    let token = token_from_offset(db, position)?;
    let extract = token.parent_ancestors().find_map(ast::ExtractFn::cast)?;
    let field = extract.extract_field()?;
    let expr = extract.expr()?;

    let field = extract_field_argument(&field)?;
    let replacement = format!("pg_catalog.extract({field}, {})", expr.syntax().text());

    actions.push(CodeAction {
        title: "Rewrite as function call `pg_catalog.extract()`".to_owned(),
        edits: vec![Edit::replace(extract.syntax().text_range(), replacement)],
        kind: ActionKind::RefactorRewrite,
    });

    Some(())
}

fn extract_field_argument(field: &ast::ExtractField) -> Option<String> {
    match field {
        ast::ExtractField::ExtractFieldLiteral(field)
            if matches!(
                field.literal()?.kind()?,
                LitKind::String(_)
                    | LitKind::EscString(_)
                    | LitKind::NationalString(_)
                    | LitKind::UnicodeEscString(_)
                    | LitKind::DollarQuotedString(_)
            ) =>
        {
            Some(field.syntax().text().to_string())
        }
        ast::ExtractField::ExtractFieldName(name) => Some(quote_string_literal(&name.text())),
        ast::ExtractField::ExtractFieldLiteral(_) => None,
    }
}

#[cfg(test)]
mod test {
    use insta::assert_snapshot;

    use crate::code_actions::test_utils::{apply_code_action, code_action_not_applicable};

    use super::rewrite_extract_as_function_call;

    #[test]
    fn rewrites_extract_as_function_call() {
        assert_snapshot!(
            apply_code_action(
                rewrite_extract_as_function_call,
                "select EXTR$0ACT(YEAR FROM DATE '2026-08-17');",
            ),
            @"select pg_catalog.extract('year', DATE '2026-08-17');"
        );
    }

    #[test]
    fn applies_with_cursor_in_source_expression() {
        assert_snapshot!(
            apply_code_action(
                rewrite_extract_as_function_call,
                "select extract(epoch from current_time$0stamp);",
            ),
            @"select pg_catalog.extract('epoch', current_timestamp);"
        );
    }

    #[test]
    fn preserves_string_field_and_aggregate_clause() {
        assert_snapshot!(
            apply_code_action(
                rewrite_extract_as_function_call,
                "select extract('EPOCH' from ts$0) filter (where ok);",
            ),
            @"select pg_catalog.extract('EPOCH', ts) filter (where ok);"
        );
    }

    #[test]
    fn converts_quoted_field_to_string() {
        assert_snapshot!(
            apply_code_action(
                rewrite_extract_as_function_call,
                r#"select extract("time""zone" from ts$0);"#,
            ),
            @"select pg_catalog.extract('time\"zone', ts);"
        );
    }

    #[test]
    fn not_applicable_outside_extract() {
        assert!(code_action_not_applicable(
            rewrite_extract_as_function_call,
            "select date_part('year', ts$0);"
        ));
    }

    #[test]
    fn converts_unicode_identifier_with_custom_escape() {
        assert_snapshot!(
            apply_code_action(
                rewrite_extract_as_function_call,
                r#"select extract(U&"@0079ear" UESCAPE '@' from ts$0);"#,
            ),
            @"select pg_catalog.extract('year', ts);"
        );
    }

    #[test]
    fn does_not_reinterpret_bit_string_field_as_identifier() {
        assert!(code_action_not_applicable(
            rewrite_extract_as_function_call,
            "select extract(B'0101' from ts$0);"
        ));
    }
}