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};

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

use super::{ActionKind, CodeAction};

pub(super) fn rewrite_at_time_zone_as_timezone(
    db: &dyn Db,
    position: InFile<TextSize>,
    actions: &mut Vec<CodeAction>,
) -> Option<()> {
    let token = token_from_offset(db, position)?;
    let bin_expr = token
        .parent_ancestors()
        .filter_map(ast::BinExpr::cast)
        .find(|expr| matches!(expr.op(), Some(ast::BinOp::AtTimeZone(_))))?;

    let value = bin_expr.lhs()?;
    let zone = bin_expr.rhs()?;
    let replacement = format!(
        "timezone({}, {})",
        zone.syntax().text(),
        value.syntax().text()
    );

    actions.push(CodeAction {
        title: "Rewrite as timezone function `timezone()`".to_owned(),
        edits: vec![Edit::replace(bin_expr.syntax().text_range(), replacement)],
        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_at_time_zone_as_timezone;

    #[test]
    fn rewrites_at_time_zone_as_timezone() {
        assert_snapshot!(
            apply_code_action(
                rewrite_at_time_zone_as_timezone,
                "select TIMESTAMP '2026-08-17 09:30' AT TIME $0ZONE 'America/New_York';",
            ),
            @"select timezone('America/New_York', TIMESTAMP '2026-08-17 09:30');"
        );
    }

    #[test]
    fn applies_from_the_value_and_preserves_surrounding_expression() {
        assert_snapshot!(
            apply_code_action(
                rewrite_at_time_zone_as_timezone,
                "select 1 + created$0_at AT TIME ZONE zone_name;",
            ),
            @"select 1 + timezone(zone_name, created_at);"
        );
    }

    #[test]
    fn rewrites_expression_zone() {
        assert_snapshot!(
            apply_code_action(
                rewrite_at_time_zone_as_timezone,
                "select ts AT TIME ZONE coale$0sce(zone_name, 'UTC');",
            ),
            @"select timezone(coalesce(zone_name, 'UTC'), ts);"
        );
    }

    #[test]
    fn rewrites_nearest_nested_at_time_zone() {
        assert_snapshot!(
            apply_code_action(
                rewrite_at_time_zone_as_timezone,
                "select (ts AT TIME ZONE z$01) AT TIME ZONE z2;",
            ),
            @"select (timezone(z1, ts)) AT TIME ZONE z2;"
        );
    }

    #[test]
    fn not_applicable_outside_at_time_zone() {
        assert!(code_action_not_applicable(
            rewrite_at_time_zone_as_timezone,
            "select created_at +$0 interval '1 hour';"
        ));
    }
}