[airflow] Third positional parameter not named ti_key should be flagged for BaseOperatorLink.get_link (AIR303) (#22828)

sjyangkevin · web-flow · commit c2ffad5a6648 · 2026-02-06T09:10:00.000-05:00
## Summary Context: apache/airflow#41641 apache/airflow#46415 >In the old signature, `dttm` is used by Airflow as a positional argument, so the argument name in the function declaration may not be `dttm`. In the new signature, however, the argument name must be `ti_key` (although it may be non-keyword-only; self, operator, ti_key should be considered new signature). So the deprecated signature detection rule should be exactly three positional arguments, the third one not named `ti_key`. Old Signature ```python class MyOperatorLink(BaseOperatorLink): def get_link(self, operator, dttm): ... ``` New Signature ```python class MyOperatorLink(BaseOperatorLink): def get_link(self, operator, *, ti_key): ... ``` This PR extends AIR303 to detect deprecated `get_link` method signatures in `BaseOperatorLink` subclasses, which is a method definition. `BaseOperatorLink` is a base class can be extended by the user, and the `get_link` method is called by Airflow. As the way Airflow internally call this method change, we need to flag the user's implementation if the method definition is not match with the new signature (@Lee-W , please correct me if I understand it wrong, whenever you have time, thanks). The rule detects deprecated signatures where: 1. The class inherits from `BaseOperatorLink` (from `airflow.models` or `airflow.sdk`) 2. The method is named `get_link` 3. There are exactly 3 positional parameters 4. The 3rd parameter is not named as `ti_key` <img width="584" height="342" alt="Screenshot from 2026-01-23 22-25-34" src="https://github.com/user-attachments/assets/3628b15f-0b87-4de9-bc24-9c84ea88d583" /> ## Test Plan The test cases are added to `AIR303.py` and the test snapshot is updated.
diff --git a/crates/ruff_linter/resources/test/fixtures/airflow/AIR303.py b/crates/ruff_linter/resources/test/fixtures/airflow/AIR303.py
@@ -72,3 +72,113 @@
 Asset("asset1")
 Asset("asset1", extra={"key": "value"})
 Asset(uri="asset1", extra={"key": "value"})
+
+
+from airflow.models.baseoperatorlink import BaseOperatorLink
+
+# 3 positional args, 3rd NOT named 'ti_key'
+class DeprecatedOperatorLink1(BaseOperatorLink):
+    def get_link(self, operator, dttm):
+        pass
+
+# keyword-only ti_key
+class ValidOperatorLink1(BaseOperatorLink):
+    def get_link(self, operator, *, ti_key):
+        pass
+
+# positional ti_key with correct name
+class ValidOperatorLink2(BaseOperatorLink):
+    def get_link(self, operator, ti_key):
+        pass
+
+# Multiple keyword-only args with ti_key
+class ValidOperatorLink3(BaseOperatorLink):
+    def get_link(self, operator, *, ti_key, extra=None):
+        pass
+
+# *args with ti_key keyword-only
+class ValidOperatorLink4(BaseOperatorLink):
+    def get_link(self, operator, *args, ti_key):
+        pass
+
+# **kwargs alongside ti_key
+class ValidOperatorLink5(BaseOperatorLink):
+    def get_link(self, operator, *, ti_key, **kwargs):
+        pass
+
+# not exactly 3 positional args
+class InvalidOperatorLink1(BaseOperatorLink):
+    def get_link(self, operator):
+        pass
+
+class InvalidOperatorLink2(BaseOperatorLink):
+    def get_link(self, operator, ti_key, extra):
+        pass
+
+class InvalidOperatorLink3(BaseOperatorLink):
+    def get_link(self, operator, *, something):
+        pass
+
+# Not a subclass of BaseOperatorLink
+class NotOperatorLinkSubclass:
+    def get_link(self, operator, dttm):
+        pass
+
+# Not a class method
+def get_link(operator, dttm):
+    pass
+
+
+from airflow.sdk import BaseOperatorLink
+
+# 3 positional args, 3rd NOT named 'ti_key'
+class DeprecatedOperatorLink2(BaseOperatorLink):
+    def get_link(self, operator, dttm):
+        pass
+
+# keyword-only ti_key
+class ValidOperatorLink6(BaseOperatorLink):
+    def get_link(self, operator, *, ti_key):
+        pass
+
+# positional ti_key with correct name
+class ValidOperatorLink7(BaseOperatorLink):
+    def get_link(self, operator, ti_key):
+        pass
+
+# Multiple keyword-only args with ti_key
+class ValidOperatorLink8(BaseOperatorLink):
+    def get_link(self, operator, *, ti_key, extra=None):
+        pass
+
+# *args with ti_key keyword-only
+class ValidOperatorLink9(BaseOperatorLink):
+    def get_link(self, operator, *args, ti_key):
+        pass
+
+# **kwargs alongside ti_key
+class ValidOperatorLink10(BaseOperatorLink):
+    def get_link(self, operator, *, ti_key, **kwargs):
+        pass
+
+# not exactly 3 positional args
+class InvalidOperatorLink4(BaseOperatorLink):
+    def get_link(self, operator):
+        pass
+
+class InvalidOperatorLink5(BaseOperatorLink):
+    def get_link(self, operator, ti_key, extra):
+        pass
+
+class InvalidOperatorLink6(BaseOperatorLink):
+    def get_link(self, operator, *, something):
+        pass
+
+# Not a subclass of BaseOperatorLink
+class NotOperatorLinkSubclass2:
+    def get_link(self, operator, dttm):
+        pass
+
+# Not a class method
+def get_link(operator, dttm):
+    pass
diff --git a/crates/ruff_linter/src/checkers/ast/analyze/statement.rs b/crates/ruff_linter/src/checkers/ast/analyze/statement.rs
@@ -342,6 +342,9 @@ pub(crate) fn statement(stmt: &Stmt, checker: &mut Checker) {
             if checker.is_rule_enabled(Rule::Airflow3Removal) {
                 airflow::rules::airflow_3_removal_function_def(checker, function_def);
             }
+            if checker.is_rule_enabled(Rule::Airflow3IncompatibleFunctionSignature) {
+                airflow::rules::airflow_3_incompatible_method_signature_def(checker, function_def);
+            }
             if checker.is_rule_enabled(Rule::NonPEP695GenericFunction) {
                 pyupgrade::rules::non_pep695_generic_function(checker, function_def);
             }
diff --git a/crates/ruff_linter/src/rules/airflow/helpers.rs b/crates/ruff_linter/src/rules/airflow/helpers.rs
@@ -3,12 +3,14 @@ use crate::fix::edits::remove_unused_imports;
 use crate::importer::ImportRequest;
 use crate::rules::numpy::helpers::{AttributeSearcher, ImportSearcher};
 use ruff_diagnostics::{Edit, Fix};
-use ruff_python_ast::name::QualifiedNameBuilder;
+use ruff_python_ast::name::{QualifiedName, QualifiedNameBuilder};
 use ruff_python_ast::statement_visitor::StatementVisitor;
 use ruff_python_ast::visitor::Visitor;
-use ruff_python_ast::{Expr, ExprAttribute, ExprName, StmtTry};
+use ruff_python_ast::{Expr, ExprAttribute, ExprName, StmtFunctionDef, StmtTry};
 use ruff_python_semantic::Exceptions;
+use ruff_python_semantic::ScopeKind;
 use ruff_python_semantic::SemanticModel;
+use ruff_python_semantic::analyze::class::any_qualified_base_class;
 use ruff_python_semantic::{MemberNameImport, NameImport};
 use ruff_text_size::Ranged;
 use ruff_text_size::TextRange;
@@ -68,6 +70,12 @@ pub(crate) enum ProviderReplacement {
     },
 }
 
+#[derive(Clone, Debug, Eq, PartialEq)]
+pub(crate) enum FunctionSignatureChange {
+    /// Carries a message describing the function signature change.
+    Message(&'static str),
+}
+
 pub(crate) fn is_guarded_by_try_except(
     expr: &Expr,
     module: &str,
@@ -260,3 +268,25 @@ pub(crate) fn generate_remove_and_runtime_import_edit(
 
     Some(Fix::unsafe_edits(remove_edit, [import_edit]))
 }
+
+/// This is a helper function to check if the given function definition is a method
+/// that inherits from a base class.
+pub(crate) fn is_method_in_subclass<F>(
+    function_def: &StmtFunctionDef,
+    semantic: &SemanticModel,
+    method_name: &str,
+    is_base_class: F,
+) -> bool
+where
+    F: Fn(QualifiedName) -> bool,
+{
+    if function_def.name.as_str() != method_name {
+        return false;
+    }
+
+    let ScopeKind::Class(class_def) = semantic.current_scope().kind else {
+        return false;
+    };
+
+    any_qualified_base_class(class_def, semantic, &is_base_class)
+}
diff --git a/crates/ruff_linter/src/rules/airflow/rules/function_signature_change_in_3.rs b/crates/ruff_linter/src/rules/airflow/rules/function_signature_change_in_3.rs
@@ -1,8 +1,9 @@
 use crate::checkers::ast::Checker;
+use crate::rules::airflow::helpers::{FunctionSignatureChange, is_method_in_subclass};
 use crate::{FixAvailability, Violation};
 use ruff_macros::{ViolationMetadata, derive_message_formats};
 use ruff_python_ast::name::QualifiedName;
-use ruff_python_ast::{Arguments, Expr, ExprAttribute, ExprCall, Identifier};
+use ruff_python_ast::{Arguments, Expr, ExprAttribute, ExprCall, Identifier, StmtFunctionDef};
 use ruff_python_semantic::Modules;
 use ruff_python_semantic::analyze::typing;
 use ruff_text_size::Ranged;
@@ -38,34 +39,22 @@ use ruff_text_size::Ranged;
 #[violation_metadata(preview_since = "0.14.11")]
 pub(crate) struct Airflow3IncompatibleFunctionSignature {
     function_name: String,
-    change_type: FunctionSignatureChangeType,
+    change: FunctionSignatureChange,
 }
 
 impl Violation for Airflow3IncompatibleFunctionSignature {
     const FIX_AVAILABILITY: FixAvailability = FixAvailability::None;
 
     #[derive_message_formats]
     fn message(&self) -> String {
-        let Airflow3IncompatibleFunctionSignature {
-            function_name,
-            change_type,
-        } = self;
-        match change_type {
-            FunctionSignatureChangeType::KeywordOnly { .. }
-            | FunctionSignatureChangeType::PositionalArgumentChange { .. } => {
-                format!("`{function_name}` signature is changed in Airflow 3.0")
-            }
-        }
+        let Airflow3IncompatibleFunctionSignature { function_name, .. } = self;
+        format!("`{function_name}` signature is changed in Airflow 3.0")
     }
 
     fn fix_title(&self) -> Option<String> {
-        let Airflow3IncompatibleFunctionSignature { change_type, .. } = self;
-        match change_type {
-            FunctionSignatureChangeType::KeywordOnly { message }
-            | FunctionSignatureChangeType::PositionalArgumentChange { message } => {
-                Some(message.to_string())
-            }
-        }
+        let Airflow3IncompatibleFunctionSignature { change, .. } = self;
+        let FunctionSignatureChange::Message(message) = change;
+        Some((*message).to_string())
     }
 }
 
@@ -106,6 +95,60 @@ pub(crate) fn airflow_3_incompatible_function_signature(checker: &Checker, expr:
     check_constructor_arguments(checker, &qualified_name, arguments, func);
 }
 
+/// AIR303
+pub(crate) fn airflow_3_incompatible_method_signature_def(
+    checker: &Checker,
+    function_def: &StmtFunctionDef,
+) {
+    if !checker.semantic().seen_module(Modules::AIRFLOW) {
+        return;
+    }
+
+    // Check for deprecated get_link signature in BaseOperatorLink subclasses
+    if is_method_in_subclass(
+        function_def,
+        checker.semantic(),
+        "get_link",
+        |qualified_name| {
+            matches!(
+                qualified_name.segments(),
+                ["airflow", "models" | "sdk", .., "BaseOperatorLink"]
+            )
+        },
+    ) {
+        let parameters = &function_def.parameters;
+        let positional_count = parameters.posonlyargs.len() + parameters.args.len();
+
+        let is_valid_signature = match positional_count {
+            // check valid signature `def get_link(self, operator, *, ti_key)`
+            2 => parameters
+                .kwonlyargs
+                .iter()
+                .any(|p| p.name().as_str() == "ti_key"),
+            // check valid signature `def get_link(self, operator, ti_key)`
+            3 => parameters
+                .posonlyargs
+                .iter()
+                .chain(parameters.args.iter())
+                .nth(2)
+                .is_some_and(|p| p.name().as_str() == "ti_key"),
+            _ => false,
+        };
+
+        if !is_valid_signature {
+            checker.report_diagnostic(
+                Airflow3IncompatibleFunctionSignature {
+                    function_name: "get_link".to_string(),
+                    change: FunctionSignatureChange::Message(
+                        "Use `def get_link(self, operator, *, ti_key)` or `def get_link(self, operator, ti_key)` as the method signature.",
+                    ),
+                },
+                function_def.name.range(),
+            );
+        }
+    }
+}
+
 fn check_method_arguments(
     checker: &Checker,
     qualified_name: &QualifiedName,
@@ -120,9 +163,9 @@ fn check_method_arguments(
             checker.report_diagnostic(
                 Airflow3IncompatibleFunctionSignature {
                     function_name: attr.to_string(),
-                    change_type: FunctionSignatureChangeType::KeywordOnly {
-                        message: "Pass positional arguments as keyword arguments (e.g., `create_asset(uri=...)`)",
-                    },
+                    change: FunctionSignatureChange::Message(
+                        "Pass positional arguments as keyword arguments (e.g., `create_asset(uri=...)`)",
+                    ),
                 },
                 attr.range(),
             );
@@ -146,9 +189,9 @@ fn check_constructor_arguments(
                 checker.report_diagnostic(
                     Airflow3IncompatibleFunctionSignature {
                         function_name,
-                        change_type: FunctionSignatureChangeType::PositionalArgumentChange {
-                            message: "Use keyword argument `extra` instead of passing a dict as the second positional argument (e.g., `Asset(name=..., uri=..., extra=...)`)",
-                        },
+                        change: FunctionSignatureChange::Message(
+                            "Use keyword argument `extra` instead of passing a dict as the second positional argument (e.g., `Asset(name=..., uri=..., extra=...)`)",
+                        ),
                     },
                     func.range(),
                 );
@@ -174,11 +217,3 @@ fn is_dict_expression(checker: &Checker, expr: &Expr) -> bool {
         _ => false,
     }
 }
-
-#[derive(Clone, Debug, Eq, PartialEq)]
-pub(crate) enum FunctionSignatureChangeType {
-    /// Function signature changed to only accept keyword arguments.
-    KeywordOnly { message: &'static str },
-    /// Function signature changed to not accept certain positional arguments.
-    PositionalArgumentChange { message: &'static str },
-}
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
@@ -1,7 +1,7 @@
 use crate::checkers::ast::Checker;
 use crate::rules::airflow::helpers::{
     Replacement, generate_import_edit, generate_remove_and_runtime_import_edit,
-    is_airflow_builtin_or_provider, is_guarded_by_try_except,
+    is_airflow_builtin_or_provider, is_guarded_by_try_except, is_method_in_subclass,
 };
 use crate::{Edit, Fix, FixAvailability, Violation};
 use ruff_macros::{ViolationMetadata, derive_message_formats};
@@ -1199,19 +1199,7 @@ 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;
-    };
-
-    class_def.bases().iter().any(|class_base| {
-        semantic
-            .resolve_qualified_name(class_base)
-            .is_some_and(|qualified_name| {
-                matches!(qualified_name.segments(), ["airflow", .., "BaseOperator"])
-            })
+    is_method_in_subclass(function_def, semantic, "execute", |qualified_name| {
+        matches!(qualified_name.segments(), ["airflow", .., "BaseOperator"])
     })
 }
diff --git a/crates/ruff_linter/src/rules/airflow/snapshots/ruff_linter__rules__airflow__tests__AIR303_AIR303.py.snap b/crates/ruff_linter/src/rules/airflow/snapshots/ruff_linter__rules__airflow__tests__AIR303_AIR303.py.snap

Original file line number	Diff line number	Diff line change
`@@ -342,6 +342,9 @@ pub(crate) fn statement(stmt: &Stmt, checker: &mut Checker) {`
`342`	`342`	`if checker.is_rule_enabled(Rule::Airflow3Removal) {`
`343`	`343`	`airflow::rules::airflow_3_removal_function_def(checker, function_def);`
`344`	`344`	`}`
	`345`	`+ if checker.is_rule_enabled(Rule::Airflow3IncompatibleFunctionSignature) {`
	`346`	`+ airflow::rules::airflow_3_incompatible_method_signature_def(checker, function_def);`
	`347`	`+ }`
`345`	`348`	`if checker.is_rule_enabled(Rule::NonPEP695GenericFunction) {`
`346`	`349`	`pyupgrade::rules::non_pep695_generic_function(checker, function_def);`
`347`	`350`	`}`