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..31aad3834427e8 --- /dev/null +++ b/crates/ruff_linter/resources/test/fixtures/airflow/AIR302_context.py @@ -0,0 +1,76 @@ +from airflow.models import DAG +from airflow.operators.dummy import DummyOperator +from datetime import datetime +from airflow.plugins_manager import AirflowPlugin +from airflow.decorators import task, get_current_context +from airflow.models.baseoperator import BaseOperator + +@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"] + +with DAG( + dag_id="example_dag", + schedule_interval="@daily", + start_date=datetime(2023, 1, 1), + template_searchpath=["/templates"], +) as dag: + task1 = DummyOperator( + task_id="task1", + params={ + # Removed variables in template + "execution_date": "{{ execution_date }}", + "next_ds": "{{ next_ds }}", + "prev_ds": "{{ prev_ds }}" + }, + ) + +class CustomMacrosPlugin(AirflowPlugin): + name = "custom_macros" + macros = { + "execution_date_macro": lambda context: context["execution_date"], + "next_ds_macro": lambda context: context["next_ds"] + } + +@task +def print_config(): + context = get_current_context() + 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"] + +class CustomOperator(BaseOperator): + def execute(self, context): + 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 d3e55f12a74758..bdd8c1a7715fa9 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 2f2e7dded7bb08..d7020097863acf 100644 --- a/crates/ruff_linter/src/rules/airflow/mod.rs +++ b/crates/ruff_linter/src/rules/airflow/mod.rs @@ -18,6 +18,7 @@ mod tests { #[test_case(Rule::Airflow3Removal, Path::new("AIR302_names.py"))] #[test_case(Rule::Airflow3Removal, Path::new("AIR302_class_attribute.py"))] #[test_case(Rule::Airflow3Removal, Path::new("AIR302_airflow_plugin.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 3c35a64d869980..8a0c1e442d7492 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,8 +1,8 @@ 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, ExprContext, ExprName, - StmtClassDef, + name::QualifiedName, Arguments, Expr, ExprAttribute, ExprCall, ExprContext, ExprName, ExprStringLiteral, + StmtClassDef, ExprSubscript, }; use ruff_python_semantic::analyze::typing; use ruff_python_semantic::Modules; @@ -71,6 +71,63 @@ impl Violation for Airflow3Removal { } } +fn extract_name_from_slice(slice: &Expr) -> Option { + match slice { + Expr::StringLiteral(ExprStringLiteral { value, .. }) => Some(value.to_string()), + _ => None, + } +} + +pub(crate) fn removed_context_variable(checker: &mut Checker, expr: &Expr) { + const REMOVED_CONTEXT_KEYS: [&str; 12] = [ + "conf", + "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 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) { + if REMOVED_CONTEXT_KEYS.contains(&key.as_str()) { + checker.diagnostics.push(Diagnostic::new( + Airflow3Removal { + deprecated: key, + replacement: Replacement::None, + }, + slice.range(), + )); + } + } + } + } + } + + if let Expr::StringLiteral(ExprStringLiteral { value, .. }) = expr { + let value_str = value.to_string(); + for key in REMOVED_CONTEXT_KEYS { + if value_str.contains(&format!("{{{{ {key} }}}}")) { + checker.diagnostics.push(Diagnostic::new( + Airflow3Removal { + deprecated: key.to_string(), + replacement: Replacement::None, + }, + expr.range(), + )); + } + } + } +} + /// AIR302 pub(crate) fn removed_in_3(checker: &mut Checker, expr: &Expr) { if !checker.semantic().seen_module(Modules::AIRFLOW) { @@ -100,6 +157,9 @@ pub(crate) fn removed_in_3(checker: &mut Checker, expr: &Expr) { } } } + 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..1403c8e025ef57 --- /dev/null +++ b/crates/ruff_linter/src/rules/airflow/snapshots/ruff_linter__rules__airflow__tests__AIR302_AIR302_context.py.snap @@ -0,0 +1,378 @@ +--- +source: crates/ruff_linter/src/rules/airflow/mod.rs +snapshot_kind: text +--- +AIR302_context.py:14:30: AIR302 `execution_date` is removed in Airflow 3.0 + | +13 | # Removed usage - should trigger violations +14 | execution_date = context["execution_date"] + | ^^^^^^^^^^^^^^^^ AIR302 +15 | next_ds = context["next_ds"] +16 | next_ds_nodash = context["next_ds_nodash"] + | + +AIR302_context.py:15:23: AIR302 `next_ds` is removed in Airflow 3.0 + | +13 | # Removed usage - should trigger violations +14 | execution_date = context["execution_date"] +15 | next_ds = context["next_ds"] + | ^^^^^^^^^ AIR302 +16 | next_ds_nodash = context["next_ds_nodash"] +17 | next_execution_date = context["next_execution_date"] + | + +AIR302_context.py:16:30: AIR302 `next_ds_nodash` is removed in Airflow 3.0 + | +14 | execution_date = context["execution_date"] +15 | next_ds = context["next_ds"] +16 | next_ds_nodash = context["next_ds_nodash"] + | ^^^^^^^^^^^^^^^^ AIR302 +17 | next_execution_date = context["next_execution_date"] +18 | prev_ds = context["prev_ds"] + | + +AIR302_context.py:17:35: AIR302 `next_execution_date` is removed in Airflow 3.0 + | +15 | next_ds = context["next_ds"] +16 | next_ds_nodash = context["next_ds_nodash"] +17 | next_execution_date = context["next_execution_date"] + | ^^^^^^^^^^^^^^^^^^^^^ AIR302 +18 | prev_ds = context["prev_ds"] +19 | prev_ds_nodash = context["prev_ds_nodash"] + | + +AIR302_context.py:18:23: AIR302 `prev_ds` is removed in Airflow 3.0 + | +16 | next_ds_nodash = context["next_ds_nodash"] +17 | next_execution_date = context["next_execution_date"] +18 | prev_ds = context["prev_ds"] + | ^^^^^^^^^ AIR302 +19 | prev_ds_nodash = context["prev_ds_nodash"] +20 | prev_execution_date = context["prev_execution_date"] + | + +AIR302_context.py:19:30: AIR302 `prev_ds_nodash` is removed in Airflow 3.0 + | +17 | next_execution_date = context["next_execution_date"] +18 | prev_ds = context["prev_ds"] +19 | prev_ds_nodash = context["prev_ds_nodash"] + | ^^^^^^^^^^^^^^^^ AIR302 +20 | prev_execution_date = context["prev_execution_date"] +21 | prev_execution_date_success = context["prev_execution_date_success"] + | + +AIR302_context.py:20:35: AIR302 `prev_execution_date` is removed in Airflow 3.0 + | +18 | prev_ds = context["prev_ds"] +19 | prev_ds_nodash = context["prev_ds_nodash"] +20 | prev_execution_date = context["prev_execution_date"] + | ^^^^^^^^^^^^^^^^^^^^^ AIR302 +21 | prev_execution_date_success = context["prev_execution_date_success"] +22 | tomorrow_ds = context["tomorrow_ds"] + | + +AIR302_context.py:21:43: AIR302 `prev_execution_date_success` is removed in Airflow 3.0 + | +19 | prev_ds_nodash = context["prev_ds_nodash"] +20 | prev_execution_date = context["prev_execution_date"] +21 | prev_execution_date_success = context["prev_execution_date_success"] + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ AIR302 +22 | tomorrow_ds = context["tomorrow_ds"] +23 | yesterday_ds = context["yesterday_ds"] + | + +AIR302_context.py:22:27: AIR302 `tomorrow_ds` is removed in Airflow 3.0 + | +20 | prev_execution_date = context["prev_execution_date"] +21 | prev_execution_date_success = context["prev_execution_date_success"] +22 | tomorrow_ds = context["tomorrow_ds"] + | ^^^^^^^^^^^^^ AIR302 +23 | yesterday_ds = context["yesterday_ds"] +24 | yesterday_ds_nodash = context["yesterday_ds_nodash"] + | + +AIR302_context.py:23:28: AIR302 `yesterday_ds` is removed in Airflow 3.0 + | +21 | prev_execution_date_success = context["prev_execution_date_success"] +22 | tomorrow_ds = context["tomorrow_ds"] +23 | yesterday_ds = context["yesterday_ds"] + | ^^^^^^^^^^^^^^ AIR302 +24 | yesterday_ds_nodash = context["yesterday_ds_nodash"] + | + +AIR302_context.py:24:35: AIR302 `yesterday_ds_nodash` is removed in Airflow 3.0 + | +22 | tomorrow_ds = context["tomorrow_ds"] +23 | yesterday_ds = context["yesterday_ds"] +24 | yesterday_ds_nodash = context["yesterday_ds_nodash"] + | ^^^^^^^^^^^^^^^^^^^^^ AIR302 +25 | +26 | with DAG( + | + +AIR302_context.py:28:5: AIR302 [*] `schedule_interval` is removed in Airflow 3.0 + | +26 | with DAG( +27 | dag_id="example_dag", +28 | schedule_interval="@daily", + | ^^^^^^^^^^^^^^^^^ AIR302 +29 | start_date=datetime(2023, 1, 1), +30 | template_searchpath=["/templates"], + | + = help: Use `schedule` instead + +ℹ Safe fix +25 25 | +26 26 | with DAG( +27 27 | dag_id="example_dag", +28 |- schedule_interval="@daily", + 28 |+ schedule="@daily", +29 29 | start_date=datetime(2023, 1, 1), +30 30 | template_searchpath=["/templates"], +31 31 | ) as dag: + +AIR302_context.py:32:13: AIR302 `airflow.operators.dummy.DummyOperator` is removed in Airflow 3.0 + | +30 | template_searchpath=["/templates"], +31 | ) as dag: +32 | task1 = DummyOperator( + | ^^^^^^^^^^^^^ AIR302 +33 | task_id="task1", +34 | params={ + | + = help: Use `airflow.operators.empty.EmptyOperator` instead + +AIR302_context.py:45:57: AIR302 `execution_date` is removed in Airflow 3.0 + | +43 | name = "custom_macros" +44 | macros = { +45 | "execution_date_macro": lambda context: context["execution_date"], + | ^^^^^^^^^^^^^^^^ AIR302 +46 | "next_ds_macro": lambda context: context["next_ds"] +47 | } + | + +AIR302_context.py:46:50: AIR302 `next_ds` is removed in Airflow 3.0 + | +44 | macros = { +45 | "execution_date_macro": lambda context: context["execution_date"], +46 | "next_ds_macro": lambda context: context["next_ds"] + | ^^^^^^^^^ AIR302 +47 | } + | + +AIR302_context.py:52:30: AIR302 `execution_date` is removed in Airflow 3.0 + | +50 | def print_config(): +51 | context = get_current_context() +52 | execution_date = context["execution_date"] + | ^^^^^^^^^^^^^^^^ AIR302 +53 | next_ds = context["next_ds"] +54 | next_ds_nodash = context["next_ds_nodash"] + | + +AIR302_context.py:53:23: AIR302 `next_ds` is removed in Airflow 3.0 + | +51 | context = get_current_context() +52 | execution_date = context["execution_date"] +53 | next_ds = context["next_ds"] + | ^^^^^^^^^ AIR302 +54 | next_ds_nodash = context["next_ds_nodash"] +55 | next_execution_date = context["next_execution_date"] + | + +AIR302_context.py:54:30: AIR302 `next_ds_nodash` is removed in Airflow 3.0 + | +52 | execution_date = context["execution_date"] +53 | next_ds = context["next_ds"] +54 | next_ds_nodash = context["next_ds_nodash"] + | ^^^^^^^^^^^^^^^^ AIR302 +55 | next_execution_date = context["next_execution_date"] +56 | prev_ds = context["prev_ds"] + | + +AIR302_context.py:55:35: AIR302 `next_execution_date` is removed in Airflow 3.0 + | +53 | next_ds = context["next_ds"] +54 | next_ds_nodash = context["next_ds_nodash"] +55 | next_execution_date = context["next_execution_date"] + | ^^^^^^^^^^^^^^^^^^^^^ AIR302 +56 | prev_ds = context["prev_ds"] +57 | prev_ds_nodash = context["prev_ds_nodash"] + | + +AIR302_context.py:56:23: AIR302 `prev_ds` is removed in Airflow 3.0 + | +54 | next_ds_nodash = context["next_ds_nodash"] +55 | next_execution_date = context["next_execution_date"] +56 | prev_ds = context["prev_ds"] + | ^^^^^^^^^ AIR302 +57 | prev_ds_nodash = context["prev_ds_nodash"] +58 | prev_execution_date = context["prev_execution_date"] + | + +AIR302_context.py:57:30: AIR302 `prev_ds_nodash` is removed in Airflow 3.0 + | +55 | next_execution_date = context["next_execution_date"] +56 | prev_ds = context["prev_ds"] +57 | prev_ds_nodash = context["prev_ds_nodash"] + | ^^^^^^^^^^^^^^^^ AIR302 +58 | prev_execution_date = context["prev_execution_date"] +59 | prev_execution_date_success = context["prev_execution_date_success"] + | + +AIR302_context.py:58:35: AIR302 `prev_execution_date` is removed in Airflow 3.0 + | +56 | prev_ds = context["prev_ds"] +57 | prev_ds_nodash = context["prev_ds_nodash"] +58 | prev_execution_date = context["prev_execution_date"] + | ^^^^^^^^^^^^^^^^^^^^^ AIR302 +59 | prev_execution_date_success = context["prev_execution_date_success"] +60 | tomorrow_ds = context["tomorrow_ds"] + | + +AIR302_context.py:59:43: AIR302 `prev_execution_date_success` is removed in Airflow 3.0 + | +57 | prev_ds_nodash = context["prev_ds_nodash"] +58 | prev_execution_date = context["prev_execution_date"] +59 | prev_execution_date_success = context["prev_execution_date_success"] + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ AIR302 +60 | tomorrow_ds = context["tomorrow_ds"] +61 | yesterday_ds = context["yesterday_ds"] + | + +AIR302_context.py:60:27: AIR302 `tomorrow_ds` is removed in Airflow 3.0 + | +58 | prev_execution_date = context["prev_execution_date"] +59 | prev_execution_date_success = context["prev_execution_date_success"] +60 | tomorrow_ds = context["tomorrow_ds"] + | ^^^^^^^^^^^^^ AIR302 +61 | yesterday_ds = context["yesterday_ds"] +62 | yesterday_ds_nodash = context["yesterday_ds_nodash"] + | + +AIR302_context.py:61:28: AIR302 `yesterday_ds` is removed in Airflow 3.0 + | +59 | prev_execution_date_success = context["prev_execution_date_success"] +60 | tomorrow_ds = context["tomorrow_ds"] +61 | yesterday_ds = context["yesterday_ds"] + | ^^^^^^^^^^^^^^ AIR302 +62 | yesterday_ds_nodash = context["yesterday_ds_nodash"] + | + +AIR302_context.py:62:35: AIR302 `yesterday_ds_nodash` is removed in Airflow 3.0 + | +60 | tomorrow_ds = context["tomorrow_ds"] +61 | yesterday_ds = context["yesterday_ds"] +62 | yesterday_ds_nodash = context["yesterday_ds_nodash"] + | ^^^^^^^^^^^^^^^^^^^^^ AIR302 +63 | +64 | class CustomOperator(BaseOperator): + | + +AIR302_context.py:66:34: AIR302 `execution_date` is removed in Airflow 3.0 + | +64 | class CustomOperator(BaseOperator): +65 | def execute(self, context): +66 | execution_date = context["execution_date"] + | ^^^^^^^^^^^^^^^^ AIR302 +67 | next_ds = context["next_ds"] +68 | next_ds_nodash = context["next_ds_nodash"] + | + +AIR302_context.py:67:27: AIR302 `next_ds` is removed in Airflow 3.0 + | +65 | def execute(self, context): +66 | execution_date = context["execution_date"] +67 | next_ds = context["next_ds"] + | ^^^^^^^^^ AIR302 +68 | next_ds_nodash = context["next_ds_nodash"] +69 | next_execution_date = context["next_execution_date"] + | + +AIR302_context.py:68:34: AIR302 `next_ds_nodash` is removed in Airflow 3.0 + | +66 | execution_date = context["execution_date"] +67 | next_ds = context["next_ds"] +68 | next_ds_nodash = context["next_ds_nodash"] + | ^^^^^^^^^^^^^^^^ AIR302 +69 | next_execution_date = context["next_execution_date"] +70 | prev_ds = context["prev_ds"] + | + +AIR302_context.py:69:39: AIR302 `next_execution_date` is removed in Airflow 3.0 + | +67 | next_ds = context["next_ds"] +68 | next_ds_nodash = context["next_ds_nodash"] +69 | next_execution_date = context["next_execution_date"] + | ^^^^^^^^^^^^^^^^^^^^^ AIR302 +70 | prev_ds = context["prev_ds"] +71 | prev_ds_nodash = context["prev_ds_nodash"] + | + +AIR302_context.py:70:27: AIR302 `prev_ds` is removed in Airflow 3.0 + | +68 | next_ds_nodash = context["next_ds_nodash"] +69 | next_execution_date = context["next_execution_date"] +70 | prev_ds = context["prev_ds"] + | ^^^^^^^^^ AIR302 +71 | prev_ds_nodash = context["prev_ds_nodash"] +72 | prev_execution_date = context["prev_execution_date"] + | + +AIR302_context.py:71:34: AIR302 `prev_ds_nodash` is removed in Airflow 3.0 + | +69 | next_execution_date = context["next_execution_date"] +70 | prev_ds = context["prev_ds"] +71 | prev_ds_nodash = context["prev_ds_nodash"] + | ^^^^^^^^^^^^^^^^ AIR302 +72 | prev_execution_date = context["prev_execution_date"] +73 | prev_execution_date_success = context["prev_execution_date_success"] + | + +AIR302_context.py:72:39: AIR302 `prev_execution_date` is removed in Airflow 3.0 + | +70 | prev_ds = context["prev_ds"] +71 | prev_ds_nodash = context["prev_ds_nodash"] +72 | prev_execution_date = context["prev_execution_date"] + | ^^^^^^^^^^^^^^^^^^^^^ AIR302 +73 | prev_execution_date_success = context["prev_execution_date_success"] +74 | tomorrow_ds = context["tomorrow_ds"] + | + +AIR302_context.py:73:47: AIR302 `prev_execution_date_success` is removed in Airflow 3.0 + | +71 | prev_ds_nodash = context["prev_ds_nodash"] +72 | prev_execution_date = context["prev_execution_date"] +73 | prev_execution_date_success = context["prev_execution_date_success"] + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ AIR302 +74 | tomorrow_ds = context["tomorrow_ds"] +75 | yesterday_ds = context["yesterday_ds"] + | + +AIR302_context.py:74:31: AIR302 `tomorrow_ds` is removed in Airflow 3.0 + | +72 | prev_execution_date = context["prev_execution_date"] +73 | prev_execution_date_success = context["prev_execution_date_success"] +74 | tomorrow_ds = context["tomorrow_ds"] + | ^^^^^^^^^^^^^ AIR302 +75 | yesterday_ds = context["yesterday_ds"] +76 | yesterday_ds_nodash = context["yesterday_ds_nodash"] + | + +AIR302_context.py:75:32: AIR302 `yesterday_ds` is removed in Airflow 3.0 + | +73 | prev_execution_date_success = context["prev_execution_date_success"] +74 | tomorrow_ds = context["tomorrow_ds"] +75 | yesterday_ds = context["yesterday_ds"] + | ^^^^^^^^^^^^^^ AIR302 +76 | yesterday_ds_nodash = context["yesterday_ds_nodash"] + | + +AIR302_context.py:76:39: AIR302 `yesterday_ds_nodash` is removed in Airflow 3.0 + | +74 | tomorrow_ds = context["tomorrow_ds"] +75 | yesterday_ds = context["yesterday_ds"] +76 | yesterday_ds_nodash = context["yesterday_ds_nodash"] + | ^^^^^^^^^^^^^^^^^^^^^ AIR302 + |