[airflow] Extract AIR311 from AIR301 rules (AIR301, AIR311) (#17310)

<!--
Thank you for contributing to Ruff! To help us out with reviewing,
please consider the following:

- Does this pull request include a summary of the change? (See below.)
- Does this pull request include a descriptive title?
- Does this pull request include references to any relevant issues?
-->

## Summary

<!-- What's the purpose of the change? What does it do, and why? -->

As discussed in
#14626 (comment),
we're to separate suggested changes from required changes.

The following symbols have been moved to AIR311 from AIR301. They still
work in Airflow 3.0, but they're suggested to be changed as they're
expected to be removed in a future version.

* arguments
    * `airflow..DAG | dag`
        * `sla_miss_callback`
    * operators
        * `sla`
* name
* `airflow.Dataset] | [airflow.datasets.Dataset` → `airflow.sdk.Asset`
    * `airflow.datasets, rest @ ..`
        * `DatasetAlias` → `airflow.sdk.AssetAlias`
        * `DatasetAll` → `airflow.sdk.AssetAll`
        * `DatasetAny` → `airflow.sdk.AssetAny`
* `expand_alias_to_datasets` → `airflow.sdk.expand_alias_to_assets`
        * `metadata.Metadata` → `airflow.sdk.Metadata`
    <!--airflow.models.baseoperator-->
    * `airflow.models.baseoperator.chain` → `airflow.sdk.chain`
* `airflow.models.baseoperator.chain_linear` →
`airflow.sdk.chain_linear`
* `airflow.models.baseoperator.cross_downstream` →
`airflow.sdk.cross_downstream`
* `airflow.models.baseoperatorlink.BaseOperatorLink` →
`airflow.sdk.definitions.baseoperatorlink.BaseOperatorLink`
    * `airflow.timetables, rest @ ..`
* `datasets.DatasetOrTimeSchedule` → *
`airflow.timetables.assets.AssetOrTimeSchedule`
    * `airflow.utils, rest @ ..`
        <!--airflow.utils.dag_parsing_context-->
* `dag_parsing_context.get_parsing_context` →
`airflow.sdk.get_parsing_context`

## Test Plan

<!-- How was it tested? -->

The test fixture has been updated acccordingly

Author: Lee-W

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

@@ -22,12 +22,6 @@
 DAG(dag_id="class_timetable", timetable=NullTimetable())
 
 
-def sla_callback(*arg, **kwargs):
-    pass
-
-
-DAG(dag_id="class_sla_callback", sla_miss_callback=sla_callback)
-
 DAG(dag_id="class_fail_stop", fail_stop=True)
 
 DAG(dag_id="class_default_view", default_view="dag_default_view")
@@ -53,11 +47,6 @@ def decorator_timetable():
     pass
 
 
-@dag(sla_miss_callback=sla_callback)
-def decorator_sla_callback():
-    pass
-
-
 @dag()
 def decorator_deprecated_operator_args():
     trigger_dagrun_op = trigger_dagrun.TriggerDagRunOperator(

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

@@ -9,9 +9,6 @@
     PY311,
     PY312,
 )
-from airflow import (
-    Dataset as DatasetFromRoot,
-)
 from airflow.api_connexion.security import requires_access, requires_access_dataset
 from airflow.auth.managers.base_auth_manager import is_authorized_dataset
 from airflow.auth.managers.models.resource_details import DatasetDetails
@@ -26,25 +23,16 @@
     set,
 )
 from airflow.contrib.aws_athena_hook import AWSAthenaHook
-from airflow.datasets import (
-    Dataset,
-    DatasetAlias,
-    DatasetAliasEvent,
-    DatasetAll,
-    DatasetAny,
-    expand_alias_to_datasets,
-)
+from airflow.datasets import DatasetAliasEvent
 from airflow.datasets.manager import (
     DatasetManager,
     dataset_manager,
     resolve_dataset_manager,
 )
-from airflow.datasets.metadata import Metadata
 from airflow.hooks.base_hook import BaseHook
 from airflow.lineage.hook import DatasetLineageInfo
 from airflow.listeners.spec.dataset import on_dataset_changed, on_dataset_created
 from airflow.metrics.validators import AllowListValidator, BlockListValidator
-from airflow.models.baseoperator import chain, chain_linear, cross_downstream
 from airflow.operators.subdag import SubDagOperator
 from airflow.providers.amazon.aws.auth_manager.avp.entities import AvpEntities
 from airflow.providers.amazon.aws.datasets import s3
@@ -61,12 +49,10 @@
 from airflow.secrets.local_filesystem import LocalFilesystemBackend, load_connections
 from airflow.security.permissions import RESOURCE_DATASET
 from airflow.sensors.base_sensor_operator import BaseSensorOperator
-from airflow.timetables.datasets import DatasetOrTimeSchedule
 from airflow.timetables.simple import DatasetTriggeredTimetable
 from airflow.triggers.external_task import TaskStateTrigger
 from airflow.utils import dates
 from airflow.utils.dag_cycle_tester import test_cycle
-from airflow.utils.dag_parsing_context import get_parsing_context
 from airflow.utils.dates import (
     date_range,
     datetime_to_nano,
@@ -105,14 +91,9 @@
 # airflow.contrib.*
 AWSAthenaHook()
 
+
 # airflow.datasets
-Dataset()
-DatasetAlias()
 DatasetAliasEvent()
-DatasetAll()
-DatasetAny()
-Metadata()
-expand_alias_to_datasets
 
 # airflow.datasets.manager
 DatasetManager()
@@ -134,10 +115,6 @@
 BlockListValidator()
 
 
-# airflow.models.baseoperator
-chain, chain_linear, cross_downstream
-
-
 # airflow.operators.branch_operator
 BaseBranchOperator()
 
@@ -207,7 +184,6 @@
 
 
 # airflow.timetables
-DatasetOrTimeSchedule()
 DatasetTriggeredTimetable()
 
 # airflow.triggers.external_task
@@ -231,8 +207,6 @@
 # airflow.utils.dag_cycle_tester
 test_cycle
 
-# airflow.utils.dag_parsing_context
-get_parsing_context
 
 # airflow.utils.db
 create_session

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

@@ -0,0 +1,28 @@
+from __future__ import annotations
+
+from datetime import timedelta
+
+from airflow import DAG, dag
+from airflow.operators.datetime import BranchDateTimeOperator
+
+
+def sla_callback(*arg, **kwargs):
+    pass
+
+
+DAG(dag_id="class_sla_callback", sla_miss_callback=sla_callback)
+
+
+@dag(sla_miss_callback=sla_callback)
+def decorator_sla_callback():
+    pass
+
+
+@dag()
+def decorator_deprecated_operator_args():
+    branch_dt_op2 = BranchDateTimeOperator(
+        task_id="branch_dt_op2",
+        sla=timedelta(seconds=10),
+    )
+
+    branch_dt_op2

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

@@ -0,0 +1,39 @@
+from __future__ import annotations
+
+from airflow import Dataset as DatasetFromRoot
+from airflow.datasets import (
+    Dataset,
+    DatasetAlias,
+    DatasetAll,
+    DatasetAny,
+    expand_alias_to_datasets,
+)
+from airflow.datasets.metadata import Metadata
+from airflow.models.baseoperator import chain, chain_linear, cross_downstream
+from airflow.models.baseoperatorlink import BaseOperatorLink
+from airflow.timetables.datasets import DatasetOrTimeSchedule
+from airflow.utils.dag_parsing_context import get_parsing_context
+
+DatasetFromRoot()
+
+# airflow.datasets
+Dataset()
+DatasetAlias()
+DatasetAll()
+DatasetAny()
+Metadata()
+expand_alias_to_datasets()
+
+# airflow.models.baseoperator
+chain()
+chain_linear()
+cross_downstream()
+
+# airflow.models.baseoperatolinker
+BaseOperatorLink()
+
+# airflow.timetables.datasets
+DatasetOrTimeSchedule()
+
+# airflow.utils.dag_parsing_context
+get_parsing_context()

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

@@ -229,6 +229,9 @@ pub(crate) fn expression(expr: &Expr, checker: &Checker) {
                     if checker.enabled(Rule::Airflow3Removal) {
                         airflow::rules::airflow_3_removal_expr(checker, expr);
                     }
+                    if checker.enabled(Rule::Airflow3SuggestedUpdate) {
+                        airflow::rules::airflow_3_0_suggested_update_expr(checker, expr);
+                    }
                     if checker.enabled(Rule::Airflow3MovedToProvider) {
                         airflow::rules::moved_to_provider_in_3(checker, expr);
                     }
@@ -451,6 +454,9 @@ pub(crate) fn expression(expr: &Expr, checker: &Checker) {
             if checker.enabled(Rule::Airflow3Removal) {
                 airflow::rules::airflow_3_removal_expr(checker, expr);
             }
+            if checker.enabled(Rule::Airflow3SuggestedUpdate) {
+                airflow::rules::airflow_3_0_suggested_update_expr(checker, expr);
+            }
             if checker.enabled(Rule::Airflow3MovedToProvider) {
                 airflow::rules::moved_to_provider_in_3(checker, expr);
             }
@@ -1152,6 +1158,9 @@ pub(crate) fn expression(expr: &Expr, checker: &Checker) {
             if checker.enabled(Rule::Airflow3Removal) {
                 airflow::rules::airflow_3_removal_expr(checker, expr);
             }
+            if checker.enabled(Rule::Airflow3SuggestedUpdate) {
+                airflow::rules::airflow_3_0_suggested_update_expr(checker, expr);
+            }
             if checker.enabled(Rule::UnnecessaryCastToInt) {
                 ruff::rules::unnecessary_cast_to_int(checker, call);
             }

crates/ruff_linter/src/codes.rs

@@ -1072,6 +1072,7 @@ pub fn code_to_rule(linter: Linter, code: &str) -> Option<(RuleGroup, Rule)> {
         (Airflow, "002") => (RuleGroup::Preview, rules::airflow::rules::AirflowDagNoScheduleArgument),
         (Airflow, "301") => (RuleGroup::Preview, rules::airflow::rules::Airflow3Removal),
         (Airflow, "302") => (RuleGroup::Preview, rules::airflow::rules::Airflow3MovedToProvider),
+        (Airflow, "311") => (RuleGroup::Preview, rules::airflow::rules::Airflow3SuggestedUpdate),
         (Airflow, "312") => (RuleGroup::Preview, rules::airflow::rules::Airflow3SuggestedToMoveToProvider),
 
         // perflint

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

@@ -81,3 +81,45 @@ fn try_block_contains_undeprecated_import(try_node: &StmtTry, replacement: &Repl
     import_searcher.visit_body(&try_node.body);
     import_searcher.found_import
 }
+
+/// Check whether the segments corresponding to the fully qualified name points to a symbol that's
+/// either a builtin or coming from one of the providers in Airflow.
+///
+/// The pattern it looks for are:
+/// - `airflow.providers.**.<module>.**.*<symbol_suffix>` for providers
+/// - `airflow.<module>.**.*<symbol_suffix>` for builtins
+///
+/// where `**` is one or more segments separated by a dot, and `*` is one or more characters.
+///
+/// Examples for the above patterns:
+/// - `airflow.providers.google.cloud.secrets.secret_manager.CloudSecretManagerBackend` (provider)
+/// - `airflow.secrets.base_secrets.BaseSecretsBackend` (builtin)
+pub(crate) fn is_airflow_builtin_or_provider(
+    segments: &[&str],
+    module: &str,
+    symbol_suffix: &str,
+) -> bool {
+    match segments {
+        ["airflow", "providers", rest @ ..] => {
+            if let (Some(pos), Some(last_element)) =
+                (rest.iter().position(|&s| s == module), rest.last())
+            {
+                // Check that the module is not the last element i.e., there's a symbol that's
+                // being used from the `module` that ends with `symbol_suffix`.
+                pos + 1 < rest.len() && last_element.ends_with(symbol_suffix)
+            } else {
+                false
+            }
+        }
+
+        ["airflow", first, rest @ ..] => {
+            if let Some(last) = rest.last() {
+                *first == module && last.ends_with(symbol_suffix)
+            } else {
+                false
+            }
+        }
+
+        _ => false,
+    }
+}

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

@@ -22,6 +22,8 @@ mod tests {
     #[test_case(Rule::Airflow3Removal, Path::new("AIR301_airflow_plugin.py"))]
     #[test_case(Rule::Airflow3Removal, Path::new("AIR301_context.py"))]
     #[test_case(Rule::Airflow3MovedToProvider, Path::new("AIR302.py"))]
+    #[test_case(Rule::Airflow3SuggestedUpdate, Path::new("AIR311_args.py"))]
+    #[test_case(Rule::Airflow3SuggestedUpdate, Path::new("AIR311_names.py"))]
     #[test_case(Rule::Airflow3SuggestedToMoveToProvider, Path::new("AIR312.py"))]
     fn rules(rule_code: Rule, path: &Path) -> Result<()> {
         let snapshot = format!("{}_{}", rule_code.noqa_code(), path.to_string_lossy());

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

@@ -2,10 +2,12 @@ pub(crate) use dag_schedule_argument::*;
 pub(crate) use moved_to_provider_in_3::*;
 pub(crate) use removal_in_3::*;
 pub(crate) use suggested_to_move_to_provider_in_3::*;
+pub(crate) use suggested_to_update_3_0::*;
 pub(crate) use task_variable_name::*;
 
 mod dag_schedule_argument;
 mod moved_to_provider_in_3;
 mod removal_in_3;
 mod suggested_to_move_to_provider_in_3;
+mod suggested_to_update_3_0;
 mod task_variable_name;