From 89ed2ce014811ab089d7a21f1bb58ffbbc20bcde Mon Sep 17 00:00:00 2001 From: Ankit Chaurasia <8670962+sunank200@users.noreply.github.com> Date: Mon, 23 Dec 2024 12:02:30 +0545 Subject: [PATCH] add lint rule to show error for removed context variables in airflow --- .../test/fixtures/airflow/AIR302_context.py | 19 +++ .../src/checkers/ast/analyze/expression.rs | 4 +- crates/ruff_linter/src/rules/airflow/mod.rs | 1 + .../src/rules/airflow/rules/removal_in_3.rs | 67 ++++++++++- ...flow__tests__AIR302_AIR302_context.py.snap | 109 ++++++++++++++++++ 5 files changed, 195 insertions(+), 5 deletions(-) create mode 100644 crates/ruff_linter/resources/test/fixtures/airflow/AIR302_context.py create mode 100644 crates/ruff_linter/src/rules/airflow/snapshots/ruff_linter__rules__airflow__tests__AIR302_AIR302_context.py.snap diff --git a/crates/ruff_linter/resources/test/fixtures/airflow/AIR302_context.py b/crates/ruff_linter/resources/test/fixtures/airflow/AIR302_context.py new file mode 100644 index 00000000000000..afbe57e8301060 --- /dev/null +++ b/crates/ruff_linter/resources/test/fixtures/airflow/AIR302_context.py @@ -0,0 +1,19 @@ +from airflow.decorators import task + +@task +def print_config(**context): + # This should not throw an error as logical_date is part of airflow context. + logical_date = context["logical_date"] + + # Removed usage - should trigger violations + execution_date = context["execution_date"] + next_ds = context["next_ds"] + next_ds_nodash = context["next_ds_nodash"] + next_execution_date = context["next_execution_date"] + prev_ds = context["prev_ds"] + prev_ds_nodash = context["prev_ds_nodash"] + prev_execution_date = context["prev_execution_date"] + prev_execution_date_success = context["prev_execution_date_success"] + tomorrow_ds = context["tomorrow_ds"] + yesterday_ds = context["yesterday_ds"] + yesterday_ds_nodash = context["yesterday_ds_nodash"] diff --git a/crates/ruff_linter/src/checkers/ast/analyze/expression.rs b/crates/ruff_linter/src/checkers/ast/analyze/expression.rs index a3d753750f4e33..4a1ddd412bea7c 100644 --- a/crates/ruff_linter/src/checkers/ast/analyze/expression.rs +++ b/crates/ruff_linter/src/checkers/ast/analyze/expression.rs @@ -171,7 +171,9 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) { if checker.enabled(Rule::NonPEP646Unpack) { pyupgrade::rules::use_pep646_unpack(checker, subscript); } - + if checker.enabled(Rule::Airflow3Removal) { + airflow::rules::removed_in_3(checker, expr); + } pandas_vet::rules::subscript(checker, value, expr); } Expr::Tuple(ast::ExprTuple { diff --git a/crates/ruff_linter/src/rules/airflow/mod.rs b/crates/ruff_linter/src/rules/airflow/mod.rs index 4aa5d618c43295..e0b475367dbe13 100644 --- a/crates/ruff_linter/src/rules/airflow/mod.rs +++ b/crates/ruff_linter/src/rules/airflow/mod.rs @@ -16,6 +16,7 @@ mod tests { #[test_case(Rule::AirflowDagNoScheduleArgument, Path::new("AIR301.py"))] #[test_case(Rule::Airflow3Removal, Path::new("AIR302_args.py"))] #[test_case(Rule::Airflow3Removal, Path::new("AIR302_names.py"))] + #[test_case(Rule::Airflow3Removal, Path::new("AIR302_context.py"))] #[test_case(Rule::Airflow3MovedToProvider, Path::new("AIR303.py"))] fn rules(rule_code: Rule, path: &Path) -> Result<()> { let snapshot = format!("{}_{}", rule_code.noqa_code(), path.to_string_lossy()); diff --git a/crates/ruff_linter/src/rules/airflow/rules/removal_in_3.rs b/crates/ruff_linter/src/rules/airflow/rules/removal_in_3.rs index 1374d0cb65da10..d7ae86dca6bc2c 100644 --- a/crates/ruff_linter/src/rules/airflow/rules/removal_in_3.rs +++ b/crates/ruff_linter/src/rules/airflow/rules/removal_in_3.rs @@ -1,6 +1,6 @@ use ruff_diagnostics::{Diagnostic, Edit, Fix, FixAvailability, Violation}; use ruff_macros::{derive_message_formats, ViolationMetadata}; -use ruff_python_ast::{name::QualifiedName, Arguments, Expr, ExprAttribute, ExprCall}; +use ruff_python_ast::{name::QualifiedName, Arguments, Expr, ExprAttribute, ExprCall,ExprName, ExprStringLiteral, ExprSubscript}; use ruff_python_semantic::analyze::typing; use ruff_python_semantic::Modules; use ruff_text_size::Ranged; @@ -41,6 +41,7 @@ enum Replacement { pub(crate) struct Airflow3Removal { deprecated: String, replacement: Replacement, + is_context_variable: bool, } impl Violation for Airflow3Removal { @@ -51,6 +52,7 @@ impl Violation for Airflow3Removal { let Airflow3Removal { deprecated, replacement, + is_context_variable, } = self; match replacement { Replacement::None => format!("`{deprecated}` is removed in Airflow 3.0"), @@ -58,7 +60,11 @@ impl Violation for Airflow3Removal { format!("`{deprecated}` is removed in Airflow 3.0") } Replacement::Message(message) => { - format!("`{deprecated}` is removed in Airflow 3.0; {message}") + if *is_context_variable { + format!("`{deprecated}` {message}") + } else { + format!("`{deprecated}` is removed in Airflow 3.0; {message}") + } } } } @@ -86,6 +92,7 @@ fn diagnostic_for_argument( Some(name) => Replacement::Name(name), None => Replacement::None, }, + is_context_variable: false, }, keyword .arg @@ -203,6 +210,7 @@ fn removed_method(checker: &mut Checker, expr: &Expr) { Airflow3Removal { deprecated: attr.to_string(), replacement, + is_context_variable: false, }, attr.range(), )); @@ -640,12 +648,58 @@ fn removed_name(checker: &mut Checker, expr: &Expr, ranged: impl Ranged) { Airflow3Removal { deprecated, replacement, + is_context_variable: false, }, ranged.range(), )); } } +fn extract_name_from_slice(slice: &Expr) -> Option { + if let Expr::StringLiteral(ExprStringLiteral { value, .. }) = slice { + Some(value.to_string()) + } else { + None + } +} + + +pub(crate) fn removed_context_variable(checker: &mut Checker, expr: &Expr) { + if let Expr::Subscript(ExprSubscript { value, slice, .. }) = expr { + if let Expr::Name(ExprName { id, .. }) = &**value { + if id.as_str() == "context" { + if let Some(key) = extract_name_from_slice(slice) { + const REMOVED_CONTEXT_KEYS: [&str; 11] = [ + "execution_date", + "next_ds", + "next_ds_nodash", + "next_execution_date", + "prev_ds", + "prev_ds_nodash", + "prev_execution_date", + "prev_execution_date_success", + "tomorrow_ds", + "yesterday_ds", + "yesterday_ds_nodash", + ]; + if REMOVED_CONTEXT_KEYS.contains(&key.as_str()) { + checker.diagnostics.push(Diagnostic::new( + Airflow3Removal { + deprecated: key.clone(), + replacement: Replacement::Message( + "is removed in the Airflow context in Airflow 3.0", + ), + is_context_variable: true, + }, + slice.range(), + )); + } + } + } + } + } +} + /// AIR302 pub(crate) fn removed_in_3(checker: &mut Checker, expr: &Expr) { if !checker.semantic().seen_module(Modules::AIRFLOW) { @@ -662,8 +716,13 @@ pub(crate) fn removed_in_3(checker: &mut Checker, expr: &Expr) { removed_method(checker, expr); } - Expr::Attribute(ExprAttribute { attr: ranged, .. }) => removed_name(checker, expr, ranged), + Expr::Attribute(ExprAttribute { attr, .. }) => { + removed_name(checker, expr, attr); + } ranged @ Expr::Name(_) => removed_name(checker, expr, ranged), - _ => {} + Expr::Subscript(_) => { + removed_context_variable(checker, expr) + }, + _ => {}, } } diff --git a/crates/ruff_linter/src/rules/airflow/snapshots/ruff_linter__rules__airflow__tests__AIR302_AIR302_context.py.snap b/crates/ruff_linter/src/rules/airflow/snapshots/ruff_linter__rules__airflow__tests__AIR302_AIR302_context.py.snap new file mode 100644 index 00000000000000..257237a9d27963 --- /dev/null +++ b/crates/ruff_linter/src/rules/airflow/snapshots/ruff_linter__rules__airflow__tests__AIR302_AIR302_context.py.snap @@ -0,0 +1,109 @@ +--- +source: crates/ruff_linter/src/rules/airflow/mod.rs +snapshot_kind: text +--- +AIR302_context.py:9:30: AIR302 `execution_date` is removed in the Airflow context in Airflow 3.0 + | + 8 | # Removed usage - should trigger violations + 9 | execution_date = context["execution_date"] + | ^^^^^^^^^^^^^^^^ AIR302 +10 | next_ds = context["next_ds"] +11 | next_ds_nodash = context["next_ds_nodash"] + | + +AIR302_context.py:10:23: AIR302 `next_ds` is removed in the Airflow context in Airflow 3.0 + | + 8 | # Removed usage - should trigger violations + 9 | execution_date = context["execution_date"] +10 | next_ds = context["next_ds"] + | ^^^^^^^^^ AIR302 +11 | next_ds_nodash = context["next_ds_nodash"] +12 | next_execution_date = context["next_execution_date"] + | + +AIR302_context.py:11:30: AIR302 `next_ds_nodash` is removed in the Airflow context in Airflow 3.0 + | + 9 | execution_date = context["execution_date"] +10 | next_ds = context["next_ds"] +11 | next_ds_nodash = context["next_ds_nodash"] + | ^^^^^^^^^^^^^^^^ AIR302 +12 | next_execution_date = context["next_execution_date"] +13 | prev_ds = context["prev_ds"] + | + +AIR302_context.py:12:35: AIR302 `next_execution_date` is removed in the Airflow context in Airflow 3.0 + | +10 | next_ds = context["next_ds"] +11 | next_ds_nodash = context["next_ds_nodash"] +12 | next_execution_date = context["next_execution_date"] + | ^^^^^^^^^^^^^^^^^^^^^ AIR302 +13 | prev_ds = context["prev_ds"] +14 | prev_ds_nodash = context["prev_ds_nodash"] + | + +AIR302_context.py:13:23: AIR302 `prev_ds` is removed in the Airflow context in Airflow 3.0 + | +11 | next_ds_nodash = context["next_ds_nodash"] +12 | next_execution_date = context["next_execution_date"] +13 | prev_ds = context["prev_ds"] + | ^^^^^^^^^ AIR302 +14 | prev_ds_nodash = context["prev_ds_nodash"] +15 | prev_execution_date = context["prev_execution_date"] + | + +AIR302_context.py:14:30: AIR302 `prev_ds_nodash` is removed in the Airflow context in Airflow 3.0 + | +12 | next_execution_date = context["next_execution_date"] +13 | prev_ds = context["prev_ds"] +14 | prev_ds_nodash = context["prev_ds_nodash"] + | ^^^^^^^^^^^^^^^^ AIR302 +15 | prev_execution_date = context["prev_execution_date"] +16 | prev_execution_date_success = context["prev_execution_date_success"] + | + +AIR302_context.py:15:35: AIR302 `prev_execution_date` is removed in the Airflow context in Airflow 3.0 + | +13 | prev_ds = context["prev_ds"] +14 | prev_ds_nodash = context["prev_ds_nodash"] +15 | prev_execution_date = context["prev_execution_date"] + | ^^^^^^^^^^^^^^^^^^^^^ AIR302 +16 | prev_execution_date_success = context["prev_execution_date_success"] +17 | tomorrow_ds = context["tomorrow_ds"] + | + +AIR302_context.py:16:43: AIR302 `prev_execution_date_success` is removed in the Airflow context in Airflow 3.0 + | +14 | prev_ds_nodash = context["prev_ds_nodash"] +15 | prev_execution_date = context["prev_execution_date"] +16 | prev_execution_date_success = context["prev_execution_date_success"] + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ AIR302 +17 | tomorrow_ds = context["tomorrow_ds"] +18 | yesterday_ds = context["yesterday_ds"] + | + +AIR302_context.py:17:27: AIR302 `tomorrow_ds` is removed in the Airflow context in Airflow 3.0 + | +15 | prev_execution_date = context["prev_execution_date"] +16 | prev_execution_date_success = context["prev_execution_date_success"] +17 | tomorrow_ds = context["tomorrow_ds"] + | ^^^^^^^^^^^^^ AIR302 +18 | yesterday_ds = context["yesterday_ds"] +19 | yesterday_ds_nodash = context["yesterday_ds_nodash"] + | + +AIR302_context.py:18:28: AIR302 `yesterday_ds` is removed in the Airflow context in Airflow 3.0 + | +16 | prev_execution_date_success = context["prev_execution_date_success"] +17 | tomorrow_ds = context["tomorrow_ds"] +18 | yesterday_ds = context["yesterday_ds"] + | ^^^^^^^^^^^^^^ AIR302 +19 | yesterday_ds_nodash = context["yesterday_ds_nodash"] + | + +AIR302_context.py:19:35: AIR302 `yesterday_ds_nodash` is removed in the Airflow context in Airflow 3.0 + | +17 | tomorrow_ds = context["tomorrow_ds"] +18 | yesterday_ds = context["yesterday_ds"] +19 | yesterday_ds_nodash = context["yesterday_ds_nodash"] + | ^^^^^^^^^^^^^^^^^^^^^ AIR302 + |