From 1e07990e79e646fd7ff36a030c64bc730489d022 Mon Sep 17 00:00:00 2001 From: Wei Lee Date: Sat, 4 Jan 2025 00:27:02 +0900 Subject: [PATCH] feat(AIR302): check whether removed context key "conf" is access through context.get(...) --- .../test/fixtures/airflow/AIR302_context.py | 35 ++++++++ crates/ruff_linter/src/rules/airflow/mod.rs | 1 + .../src/rules/airflow/rules/removal_in_3.rs | 89 ++++++++++++++++++- ...airflow__tests__AIR302_AIR302_args.py.snap | 1 + ...flow__tests__AIR302_AIR302_context.py.snap | 21 +++++ 5 files changed, 143 insertions(+), 4 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 0000000000000..6d8fc79dbabc6 --- /dev/null +++ b/crates/ruff_linter/resources/test/fixtures/airflow/AIR302_context.py @@ -0,0 +1,35 @@ +import pendulum + +from airflow.decorators import dag, task +from airflow.providers.standard.operators.python import PythonOperator + + +def access_invalid_key_in_context(**context): + print("access invalid key", context["conf"]) + + +@task +def access_invalid_key_task_out_of_dag(**context): + print("access invalid key", context.get("conf")) + + +@dag( + schedule=None, + start_date=pendulum.datetime(2021, 1, 1, tz="UTC"), + catchup=False, + tags=[""], +) +def invalid_dag(): + @task() + def access_invalid_key_task(**context): + print("access invalid key", context.get("conf")) + + task1 = PythonOperator( + task_id="task1", + python_callable=access_invalid_key_in_context, + ) + access_invalid_key_task() >> task1 + access_invalid_key_task_out_of_dag() + + +invalid_dag() diff --git a/crates/ruff_linter/src/rules/airflow/mod.rs b/crates/ruff_linter/src/rules/airflow/mod.rs index 2f2e7dded7bb0..d7020097863ac 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 f5639de80926f..bf9dc299ce16a 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,10 @@ +use crate::checkers::ast::Checker; use ruff_diagnostics::{Diagnostic, Edit, Fix, FixAvailability, Violation}; use ruff_macros::{derive_message_formats, ViolationMetadata}; +use ruff_python_ast::helpers::map_callable; use ruff_python_ast::{ - name::QualifiedName, Arguments, Expr, ExprAttribute, ExprCall, ExprContext, ExprName, - StmtClassDef, + name::QualifiedName, Arguments, Expr, ExprAttribute, ExprCall, ExprContext, ExprName, Stmt, + StmtClassDef, StmtFunctionDef, }; use ruff_python_semantic::analyze::typing; use ruff_python_semantic::Modules; @@ -10,8 +12,6 @@ use ruff_python_semantic::ScopeKind; use ruff_text_size::Ranged; use ruff_text_size::TextRange; -use crate::checkers::ast::Checker; - /// ## What it does /// Checks for uses of deprecated Airflow functions and values. /// @@ -87,6 +87,7 @@ pub(crate) fn removed_in_3(checker: &mut Checker, expr: &Expr) { check_call_arguments(checker, &qualname, arguments); }; check_method(checker, call_expr); + check_context_get(checker, call_expr); } Expr::Attribute(attribute_expr @ ExprAttribute { attr, .. }) => { check_name(checker, expr, attr.range()); @@ -203,6 +204,54 @@ fn check_call_arguments(checker: &mut Checker, qualname: &QualifiedName, argumen }; } +/// Check whether a removed context key is access through context.get("key"). +/// +/// ```python +/// from airflow.decorators import task +/// +/// +/// @task +/// def access_invalid_key_task_out_of_dag(**context): +/// print("access invalid key", context.get("conf")) +/// ``` +fn check_context_get(checker: &mut Checker, call_expr: &ExprCall) { + const REMOVED_CONTEXT_KEYS: [&str; 1] = ["conf"]; + + if !is_taskflow(checker) { + return; + } + + let Expr::Attribute(ExprAttribute { value, attr, .. }) = &*call_expr.func else { + return; + }; + + // Ensure the method called is `context` + if !value + .as_name_expr() + .is_some_and(|name| matches!(name.id.as_str(), "context")) + { + return; + } + + // Ensure the method called is `get` + if attr.as_str() != "get" { + return; + } + + for removed_key in REMOVED_CONTEXT_KEYS { + if let Some(argument) = call_expr.arguments.find_argument_value(removed_key, 0) { + checker.diagnostics.push(Diagnostic::new( + Airflow3Removal { + deprecated: removed_key.to_string(), + replacement: Replacement::None, + }, + argument.range(), + )); + return; + } + } +} + /// Check whether a removed Airflow class attribute (include property) is called. /// /// For example: @@ -849,3 +898,35 @@ fn is_airflow_builtin_or_provider(segments: &[&str], module: &str, symbol_suffix _ => false, } } + +/// Check whether the function is decorated by @task +/// +/// +/// Examples for the above patterns: +/// ```python +/// from airflow.decorators import task +/// +/// +/// @task +/// def access_invalid_key_task_out_of_dag(**context): +/// print("access invalid key", context.get("conf")) +/// ``` +fn is_taskflow(checker: &mut Checker) -> bool { + let mut parents = checker.semantic().current_statements(); + if let Some(Stmt::FunctionDef(StmtFunctionDef { decorator_list, .. })) = + parents.find(|stmt| stmt.is_function_def_stmt()) + { + for decorator in decorator_list { + if checker + .semantic() + .resolve_qualified_name(map_callable(&decorator.expression)) + .is_some_and(|qualified_name| { + matches!(qualified_name.segments(), ["airflow", "decorators", "task"]) + }) + { + return true; + } + } + } + false +} diff --git a/crates/ruff_linter/src/rules/airflow/snapshots/ruff_linter__rules__airflow__tests__AIR302_AIR302_args.py.snap b/crates/ruff_linter/src/rules/airflow/snapshots/ruff_linter__rules__airflow__tests__AIR302_AIR302_args.py.snap index 1f9f4a94f2ad7..d1dd8fbba7277 100644 --- a/crates/ruff_linter/src/rules/airflow/snapshots/ruff_linter__rules__airflow__tests__AIR302_AIR302_args.py.snap +++ b/crates/ruff_linter/src/rules/airflow/snapshots/ruff_linter__rules__airflow__tests__AIR302_AIR302_args.py.snap @@ -1,5 +1,6 @@ --- source: crates/ruff_linter/src/rules/airflow/mod.rs +snapshot_kind: text --- AIR302_args.py:18:39: AIR302 [*] `schedule_interval` is removed in Airflow 3.0 | 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 0000000000000..d6db0872a4d88 --- /dev/null +++ b/crates/ruff_linter/src/rules/airflow/snapshots/ruff_linter__rules__airflow__tests__AIR302_AIR302_context.py.snap @@ -0,0 +1,21 @@ +--- +source: crates/ruff_linter/src/rules/airflow/mod.rs +snapshot_kind: text +--- +AIR302_context.py:13:45: AIR302 `conf` is removed in Airflow 3.0 + | +11 | @task +12 | def access_invalid_key_task_out_of_dag(**context): +13 | print("access invalid key", context.get("conf")) + | ^^^^^^ AIR302 + | + +AIR302_context.py:25:49: AIR302 `conf` is removed in Airflow 3.0 + | +23 | @task() +24 | def access_invalid_key_task(**context): +25 | print("access invalid key", context.get("conf")) + | ^^^^^^ AIR302 +26 | +27 | task1 = PythonOperator( + |