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 45f2b2f069..fa028d95b3 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 @@ -154,7 +154,9 @@ const REMOVED_CONTEXT_KEYS: [&str; 12] = [ /// pass /// ``` fn check_function_parameters(checker: &mut Checker, function_def: &StmtFunctionDef) { - if !is_airflow_task(function_def, checker.semantic()) { + if !is_airflow_task(function_def, checker.semantic()) + && !is_execute_method_inherits_from_airflow_operator(function_def, checker.semantic()) + { return; } @@ -1076,3 +1078,39 @@ fn is_airflow_task(function_def: &StmtFunctionDef, semantic: &SemanticModel) -> }) }) } + +/// Check it's "execute" method inherits from Airflow base operator +/// +/// For example: +/// +/// ```python +/// from airflow.models.baseoperator import BaseOperator +/// +/// class CustomOperator(BaseOperator): +/// def execute(self): +/// pass +/// ``` +fn is_execute_method_inherits_from_airflow_operator( + function_def: &StmtFunctionDef, + semantic: &SemanticModel, +) -> bool { + if function_def.name.as_str() != "execute" { + return false; + } + + let ScopeKind::Class(class_def) = semantic.current_scope().kind else { + return false; + }; + + if !class_def.bases().iter().any(|class_base| { + semantic + .resolve_qualified_name(class_base) + .is_some_and(|qualified_name| { + matches!(qualified_name.segments(), ["airflow", .., "BaseOperator"]) + }) + }) { + return false; + }; + + true +} 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 index dfcc106812..84332e02b0 100644 --- 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 @@ -327,3 +327,12 @@ AIR302_context.py:115:13: AIR302 `airflow.operators.dummy.DummyOperator` is remo 117 | params={ | = help: Use `airflow.operators.empty.EmptyOperator` instead + +AIR302_context.py:135:23: AIR302 `next_ds` is removed in Airflow 3.0 + | +134 | class CustomOperator(BaseOperator): +135 | def execute(self, next_ds, context): + | ^^^^^^^ AIR302 +136 | execution_date = context["execution_date"] +137 | next_ds = context["next_ds"] + |