[airflow] Add ruff rules to catch deprecated Airflow imports for Airflow 3.1 (AIR321)

crates/ruff_linter/resources/test/fixtures/airflow/AIR301_names.py

@@ -14,7 +14,6 @@
 from airflow.datasets import DatasetAliasEvent, DatasetEvent
 from airflow.operators.postgres_operator import Mapping
 from airflow.operators.subdag import SubDagOperator
-from airflow.secrets.cache import SecretCache
 from airflow.secrets.local_filesystem import LocalFilesystemBackend
 from airflow.triggers.external_task import TaskStateTrigger
 from airflow.utils import dates
@@ -61,9 +60,6 @@
 # get_connection
 LocalFilesystemBackend()
 
-# airflow.secrets.cache
-SecretCache()
-
 
 # airflow.triggers.external_task
 TaskStateTrigger()

crates/ruff_linter/resources/test/fixtures/airflow/AIR321_names.py

@@ -0,0 +1,102 @@
+from __future__ import annotations
+
+# airflow.utils.setup_teardown
+from airflow.utils.setup_teardown import (
+    BaseSetupTeardownContext,
+    SetupTeardownContext,
+)
+
+BaseSetupTeardownContext()
+SetupTeardownContext()
+
+# airflow.utils.xcom
+from airflow.utils.xcom import XCOM_RETURN_KEY
+
+XCOM_RETURN_KEY
+
+# airflow.utils.task_group
+from airflow.utils.task_group import TaskGroup
+
+TaskGroup()
+
+# airflow.utils.timezone
+from datetime import datetime as dt
+from airflow.utils.timezone import (
+    coerce_datetime,
+    convert_to_utc,
+    datetime,
+    make_naive,
+    parse,
+    utc,
+    utcnow,
+)
+
+current_time = dt.now()
+
+coerce_datetime(v=current_time, tz=None)
+convert_to_utc(current_time)
+datetime(2026, 1, 1, 21, 30, 0)
+make_naive(current_time, tz=None)
+parse("2026-01-01 21:30:00")
+utc
+utcnow()
+
+# airflow.utils.decorators
+from airflow.utils.decorators import (
+    remove_task_decorator,
+    fixup_decorator_warning_stack,
+)
+
+remove_task_decorator()
+fixup_decorator_warning_stack()
+
+# airflow.models.abstractoperator
+from airflow.models.abstractoperator import (
+    AbstractOperator,
+    NotMapped,
+    TaskStateChangeCallback,
+)
+
+AbstractOperator()
+NotMapped()
+TaskStateChangeCallback()
+
+# airflow.models.baseoperator
+from airflow.models.baseoperator import BaseOperator
+
+BaseOperator()
+
+# airflow.macros
+from airflow.macros import (
+    ds_add,
+    ds_format,
+    datetime_diff_for_humans,
+    ds_format_locale,
+)
+
+ds_add("2026-01-01", 5)
+ds_format("2026-01-01", "%Y-%m-%d", "%m-%d-%y")
+datetime_diff_for_humans(
+    dt(2026, 1, 1, 21, 30, 0),
+    since=dt(2026, 1, 2, 21, 30, 0)
+)
+ds_format_locale("01/01/2026", "%d/%m/%Y", "EEEE dd MMMM yyyy", "en_US")
+
+# airflow.io
+from airflow.io import (
+    get_fs,
+    has_fs,
+    Properties,
+)
+
+get_fs()
+has_fs()
+Properties()
+
+# airflow.secrets.cache
+from airflow.secrets.cache import SecretCache
+SecretCache()
+
+# airflow.hooks
+from airflow.hooks.base import BaseHook
+BaseHook()

crates/ruff_linter/src/checkers/ast/analyze/expression.rs

@@ -262,6 +262,9 @@ pub(crate) fn expression(expr: &Expr, checker: &Checker) {
                     if checker.is_rule_enabled(Rule::Airflow3Removal) {
                         airflow::rules::airflow_3_removal_expr(checker, expr);
                     }
+                    if checker.is_rule_enabled(Rule::Airflow31Moved) {
+                        airflow::rules::airflow_3_1_moved_expr(checker, expr);
+                    }
                     if checker.is_rule_enabled(Rule::Airflow3SuggestedUpdate) {
                         airflow::rules::airflow_3_0_suggested_update_expr(checker, expr);
                     }
@@ -487,6 +490,9 @@ pub(crate) fn expression(expr: &Expr, checker: &Checker) {
             if checker.is_rule_enabled(Rule::Airflow3Removal) {
                 airflow::rules::airflow_3_removal_expr(checker, expr);
             }
+            if checker.is_rule_enabled(Rule::Airflow31Moved) {
+                airflow::rules::airflow_3_1_moved_expr(checker, expr);
+            }
             if checker.is_rule_enabled(Rule::Airflow3SuggestedUpdate) {
                 airflow::rules::airflow_3_0_suggested_update_expr(checker, expr);
             }

crates/ruff_linter/src/codes.rs

@@ -1128,6 +1128,7 @@ pub fn code_to_rule(linter: Linter, code: &str) -> Option<(RuleGroup, Rule)> {
         (Airflow, "303") => rules::airflow::rules::Airflow3IncompatibleFunctionSignature,
         (Airflow, "311") => rules::airflow::rules::Airflow3SuggestedUpdate,
         (Airflow, "312") => rules::airflow::rules::Airflow3SuggestedToMoveToProvider,
+        (Airflow, "321") => rules::airflow::rules::Airflow31Moved,
 
         // perflint
         (Perflint, "101") => rules::perflint::rules::UnnecessaryListCast,

crates/ruff_linter/src/rules/airflow/helpers.rs

@@ -13,6 +13,9 @@ use ruff_python_semantic::{MemberNameImport, NameImport};
 use ruff_text_size::Ranged;
 use ruff_text_size::TextRange;
 
+/// Warning message for internal modules that are not part of the public API.
+pub(crate) const INTERNAL_MODULE_WARNING: &str = "This is an internal module which is not suggested to be used and is subject to change without notice.";
+
 #[derive(Clone, Debug, Eq, PartialEq)]
 pub(crate) enum Replacement {
     // There's no replacement or suggestion other than removal
@@ -33,6 +36,20 @@ pub(crate) enum Replacement {
         module: &'static str,
         name: String,
     },
+    // Symbols moved to Task SDK in Airflow 3. Used when we want to match multiple names.
+    // e.g., `airflow.io.get_fs | has_fs | Properties` to `airflow.sdk.io.get_fs | has_fs | Properties`
+    SourceModuleMovedToSDK {
+        module: &'static str,
+        name: String,
+        version: &'static str,
+    },
+    // Symbols updated in Airflow 3 with only module changed. Used when we want to include custom message and optionally report diagnostics.
+    SourceModuleMovedWithMessage {
+        module: &'static str,
+        name: String,
+        message: &'static str,
+        suggest_fix: bool,
+    },
 }
 
 #[derive(Clone, Debug, Eq, PartialEq)]

crates/ruff_linter/src/rules/airflow/mod.rs

@@ -23,6 +23,7 @@ mod tests {
     #[test_case(Rule::Airflow3Removal, Path::new("AIR301_class_attribute.py"))]
     #[test_case(Rule::Airflow3Removal, Path::new("AIR301_airflow_plugin.py"))]
     #[test_case(Rule::Airflow3Removal, Path::new("AIR301_context.py"))]
+    #[test_case(Rule::Airflow31Moved, Path::new("AIR321_names.py"))]
     #[test_case(Rule::Airflow3MovedToProvider, Path::new("AIR302_amazon.py"))]
     #[test_case(Rule::Airflow3MovedToProvider, Path::new("AIR302_celery.py"))]
     #[test_case(Rule::Airflow3MovedToProvider, Path::new("AIR302_common_sql.py"))]

crates/ruff_linter/src/rules/airflow/rules/mod.rs

@@ -1,5 +1,6 @@
 pub(crate) use dag_schedule_argument::*;
 pub(crate) use function_signature_change_in_3::*;
+pub(crate) use moved_in_3_1::*;
 pub(crate) use moved_to_provider_in_3::*;
 pub(crate) use removal_in_3::*;
 pub(crate) use suggested_to_move_to_provider_in_3::*;
@@ -8,6 +9,7 @@ pub(crate) use task_variable_name::*;
 
 mod dag_schedule_argument;
 mod function_signature_change_in_3;
+mod moved_in_3_1;
 mod moved_to_provider_in_3;
 mod removal_in_3;
 mod suggested_to_move_to_provider_in_3;