[airflow] Extend AIR311 rules (#17422)

<!--
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? -->

* Extend the following AIR311 rules
* `airflow.io.path.ObjectStoragePath` → `airflow.sdk.ObjectStoragePath`
    * `airflow.io.storage.attach` → `airflow.sdk.io.attach`
    * `airflow.models.dag.DAG` → `airflow.sdk.DAG`
    * `airflow.models.DAG` → `airflow.sdk.DAG`
    * `airflow.decorators.dag` → `airflow.sdk.dag`
    * `airflow.decorators.task` → `airflow.sdk.task`
    * `airflow.decorators.task_group` → `airflow.sdk.task_group`
    * `airflow.decorators.setup` → `airflow.sdk.setup`
    * `airflow.decorators.teardown` → `airflow.sdk.teardown`

## Test Plan

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

The test case has been added to the button of the existing test
fixtures, confirmed to be correct and later reorgnaized

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

@@ -9,11 +9,17 @@ from airflow.datasets import (
     expand_alias_to_datasets,
 )
 from airflow.datasets.metadata import Metadata
+from airflow.decorators import dag, setup, task, task_group, teardown
+from airflow.io.path import ObjectStoragePath
+from airflow.io.storage import attach
+from airflow.models import DAG as DAGFromModel
 from airflow.models.baseoperator import chain, chain_linear, cross_downstream
 from airflow.models.baseoperatorlink import BaseOperatorLink
+from airflow.models.dag import DAG as DAGFromDag
 from airflow.timetables.datasets import DatasetOrTimeSchedule
 from airflow.utils.dag_parsing_context import get_parsing_context
 
+# airflow
 DatasetFromRoot()
 
 # airflow.datasets
@@ -24,6 +30,20 @@ DatasetAny()
 Metadata()
 expand_alias_to_datasets()
 
+# airflow.decorators
+dag()
+task()
+task_group()
+setup()
+teardown()
+
+# airflow.io
+ObjectStoragePath()
+attach()
+
+# airflow.models
+DAGFromModel()
+
 # airflow.models.baseoperator
 chain()
 chain_linear()
@@ -32,6 +52,8 @@ cross_downstream()
 # airflow.models.baseoperatolinker
 BaseOperatorLink()
 
+# airflow.models.dag
+DAGFromDag()
 # airflow.timetables.datasets
 DatasetOrTimeSchedule()
 

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

@@ -210,6 +210,25 @@ fn check_name(checker: &Checker, expr: &Expr, range: TextRange) {
             "expand_alias_to_datasets" => Replacement::Name("airflow.sdk.expand_alias_to_assets"),
             _ => return,
         },
+
+        // airflow.decorators
+        ["airflow", "decorators", rest @ ("dag" | "task" | "task_group" | "setup" | "teardown")] => {
+            Replacement::SourceModuleMoved {
+                module: "airflow.sdk",
+                name: (*rest).to_string(),
+            }
+        }
+
+        // airflow.io
+        ["airflow", "io", "path", "ObjectStoragePath"] => Replacement::SourceModuleMoved {
+            module: "airflow.sdk",
+            name: "ObjectStoragePath".to_string(),
+        },
+        ["airflow", "io", "storage", "attach"] => Replacement::SourceModuleMoved {
+            module: "airflow.sdk.io",
+            name: "attach".to_string(),
+        },
+
         // airflow.models.baseoperator
         ["airflow", "models", "baseoperator", rest] => match *rest {
             "chain" | "chain_linear" | "cross_downstream" => Replacement::SourceModuleMoved {
@@ -221,6 +240,11 @@ fn check_name(checker: &Checker, expr: &Expr, range: TextRange) {
             }
             _ => return,
         },
+        // airflow.model..DAG
+        ["airflow", "models", .., "DAG"] => Replacement::SourceModuleMoved {
+            module: "airflow.sdk",
+            name: "DAG".to_string(),
+        },
         // airflow.timetables
         ["airflow", "timetables", "datasets", "DatasetOrTimeSchedule"] => {
             Replacement::Name("airflow.timetables.assets.AssetOrTimeSchedule")

crates/ruff_linter/src/rules/airflow/snapshots/ruff_linter__rules__airflow__tests__AIR311_AIR311_names.py.snap

@@ -1,153 +1,248 @@
 ---
 source: crates/ruff_linter/src/rules/airflow/mod.rs
 ---
-AIR311_names.py:17:1: AIR311 [*] `airflow.Dataset` is removed in Airflow 3.0; It still works in Airflow 3.0 but is expected to be removed in a future version.
+AIR311_names.py:23:1: AIR311 [*] `airflow.Dataset` is removed in Airflow 3.0; It still works in Airflow 3.0 but is expected to be removed in a future version.
    |
-15 | from airflow.utils.dag_parsing_context import get_parsing_context
+22 | # airflow
-16 |
+23 | DatasetFromRoot()
-17 | DatasetFromRoot()
    | ^^^^^^^^^^^^^^^ AIR311
-18 |
+24 |
-19 | # airflow.datasets
+25 | # airflow.datasets
    |
    = help: Use `airflow.sdk.Asset` instead
 
 ℹ Safe fix
-13 13 | from airflow.models.baseoperatorlink import BaseOperatorLink
+18 18 | from airflow.models.dag import DAG as DAGFromDag
-14 14 | from airflow.timetables.datasets import DatasetOrTimeSchedule
+19 19 | from airflow.timetables.datasets import DatasetOrTimeSchedule
-15 15 | from airflow.utils.dag_parsing_context import get_parsing_context
+20 20 | from airflow.utils.dag_parsing_context import get_parsing_context
-   16 |+from airflow.sdk import Asset
+   21 |+from airflow.sdk import Asset
-16 17 | 
+21 22 | 
-17    |-DatasetFromRoot()
+22 23 | # airflow
-   18 |+Asset()
+23    |-DatasetFromRoot()
-18 19 | 
+   24 |+Asset()
-19 20 | # airflow.datasets
+24 25 | 
-20 21 | Dataset()
+25 26 | # airflow.datasets
+26 27 | Dataset()
 
-AIR311_names.py:20:1: AIR311 [*] `airflow.datasets.Dataset` is removed in Airflow 3.0; It still works in Airflow 3.0 but is expected to be removed in a future version.
+AIR311_names.py:26:1: AIR311 [*] `airflow.datasets.Dataset` is removed in Airflow 3.0; It still works in Airflow 3.0 but is expected to be removed in a future version.
    |
-19 | # airflow.datasets
+25 | # airflow.datasets
-20 | Dataset()
+26 | Dataset()
    | ^^^^^^^ AIR311
-21 | DatasetAlias()
+27 | DatasetAlias()
-22 | DatasetAll()
+28 | DatasetAll()
    |
    = help: Use `airflow.sdk.Asset` instead
 
 ℹ Safe fix
-13 13 | from airflow.models.baseoperatorlink import BaseOperatorLink
+18 18 | from airflow.models.dag import DAG as DAGFromDag
-14 14 | from airflow.timetables.datasets import DatasetOrTimeSchedule
+19 19 | from airflow.timetables.datasets import DatasetOrTimeSchedule
-15 15 | from airflow.utils.dag_parsing_context import get_parsing_context
+20 20 | from airflow.utils.dag_parsing_context import get_parsing_context
-   16 |+from airflow.sdk import Asset
+   21 |+from airflow.sdk import Asset
-16 17 | 
+21 22 | 
-17 18 | DatasetFromRoot()
+22 23 | # airflow
-18 19 | 
+23 24 | DatasetFromRoot()
-19 20 | # airflow.datasets
+24 25 | 
-20    |-Dataset()
+25 26 | # airflow.datasets
-   21 |+Asset()
+26    |-Dataset()
-21 22 | DatasetAlias()
+   27 |+Asset()
-22 23 | DatasetAll()
+27 28 | DatasetAlias()
-23 24 | DatasetAny()
+28 29 | DatasetAll()
+29 30 | DatasetAny()
 
-AIR311_names.py:21:1: AIR311 `airflow.datasets.DatasetAlias` is removed in Airflow 3.0; It still works in Airflow 3.0 but is expected to be removed in a future version.
+AIR311_names.py:27:1: AIR311 `airflow.datasets.DatasetAlias` is removed in Airflow 3.0; It still works in Airflow 3.0 but is expected to be removed in a future version.
    |
-19 | # airflow.datasets
+25 | # airflow.datasets
-20 | Dataset()
+26 | Dataset()
-21 | DatasetAlias()
+27 | DatasetAlias()
    | ^^^^^^^^^^^^ AIR311
-22 | DatasetAll()
+28 | DatasetAll()
-23 | DatasetAny()
+29 | DatasetAny()
    |
    = help: Use `airflow.sdk.AssetAlias` instead
 
-AIR311_names.py:22:1: AIR311 `airflow.datasets.DatasetAll` is removed in Airflow 3.0; It still works in Airflow 3.0 but is expected to be removed in a future version.
+AIR311_names.py:28:1: AIR311 `airflow.datasets.DatasetAll` is removed in Airflow 3.0; It still works in Airflow 3.0 but is expected to be removed in a future version.
    |
-20 | Dataset()
+26 | Dataset()
-21 | DatasetAlias()
+27 | DatasetAlias()
-22 | DatasetAll()
+28 | DatasetAll()
    | ^^^^^^^^^^ AIR311
-23 | DatasetAny()
+29 | DatasetAny()
-24 | Metadata()
+30 | Metadata()
    |
    = help: Use `airflow.sdk.AssetAll` instead
 
-AIR311_names.py:23:1: AIR311 `airflow.datasets.DatasetAny` is removed in Airflow 3.0; It still works in Airflow 3.0 but is expected to be removed in a future version.
+AIR311_names.py:29:1: AIR311 `airflow.datasets.DatasetAny` is removed in Airflow 3.0; It still works in Airflow 3.0 but is expected to be removed in a future version.
    |
-21 | DatasetAlias()
+27 | DatasetAlias()
-22 | DatasetAll()
+28 | DatasetAll()
-23 | DatasetAny()
+29 | DatasetAny()
    | ^^^^^^^^^^ AIR311
-24 | Metadata()
+30 | Metadata()
-25 | expand_alias_to_datasets()
+31 | expand_alias_to_datasets()
    |
    = help: Use `airflow.sdk.AssetAny` instead
 
-AIR311_names.py:24:1: AIR311 `airflow.datasets.metadata.Metadata` is removed in Airflow 3.0; It still works in Airflow 3.0 but is expected to be removed in a future version.
+AIR311_names.py:30:1: AIR311 `airflow.datasets.metadata.Metadata` is removed in Airflow 3.0; It still works in Airflow 3.0 but is expected to be removed in a future version.
    |
-22 | DatasetAll()
+28 | DatasetAll()
-23 | DatasetAny()
+29 | DatasetAny()
-24 | Metadata()
+30 | Metadata()
    | ^^^^^^^^ AIR311
-25 | expand_alias_to_datasets()
+31 | expand_alias_to_datasets()
    |
    = help: Use `airflow.sdk.Metadata` instead
 
-AIR311_names.py:25:1: AIR311 `airflow.datasets.expand_alias_to_datasets` is removed in Airflow 3.0; It still works in Airflow 3.0 but is expected to be removed in a future version.
+AIR311_names.py:31:1: AIR311 `airflow.datasets.expand_alias_to_datasets` is removed in Airflow 3.0; It still works in Airflow 3.0 but is expected to be removed in a future version.
    |
-23 | DatasetAny()
+29 | DatasetAny()
-24 | Metadata()
+30 | Metadata()
-25 | expand_alias_to_datasets()
+31 | expand_alias_to_datasets()
    | ^^^^^^^^^^^^^^^^^^^^^^^^ AIR311
-26 |
+32 |
-27 | # airflow.models.baseoperator
+33 | # airflow.decorators
    |
    = help: Use `airflow.sdk.expand_alias_to_assets` instead
 
-AIR311_names.py:28:1: AIR311 `airflow.models.baseoperator.chain` is removed in Airflow 3.0; It still works in Airflow 3.0 but is expected to be removed in a future version.
+AIR311_names.py:34:1: AIR311 `airflow.decorators.dag` is removed in Airflow 3.0; It still works in Airflow 3.0 but is expected to be removed in a future version.
    |
-27 | # airflow.models.baseoperator
+33 | # airflow.decorators
-28 | chain()
+34 | dag()
+   | ^^^ AIR311
+35 | task()
+36 | task_group()
+   |
+   = help: Use `airflow.sdk.dag` instead
+
+AIR311_names.py:35:1: AIR311 `airflow.decorators.task` is removed in Airflow 3.0; It still works in Airflow 3.0 but is expected to be removed in a future version.
+   |
+33 | # airflow.decorators
+34 | dag()
+35 | task()
+   | ^^^^ AIR311
+36 | task_group()
+37 | setup()
+   |
+   = help: Use `airflow.sdk.task` instead
+
+AIR311_names.py:36:1: AIR311 `airflow.decorators.task_group` is removed in Airflow 3.0; It still works in Airflow 3.0 but is expected to be removed in a future version.
+   |
+34 | dag()
+35 | task()
+36 | task_group()
+   | ^^^^^^^^^^ AIR311
+37 | setup()
+38 | teardown()
+   |
+   = help: Use `airflow.sdk.task_group` instead
+
+AIR311_names.py:37:1: AIR311 `airflow.decorators.setup` is removed in Airflow 3.0; It still works in Airflow 3.0 but is expected to be removed in a future version.
+   |
+35 | task()
+36 | task_group()
+37 | setup()
    | ^^^^^ AIR311
-29 | chain_linear()
+38 | teardown()
-30 | cross_downstream()
+   |
+   = help: Use `airflow.sdk.setup` instead
+
+AIR311_names.py:38:1: AIR311 `airflow.decorators.teardown` is removed in Airflow 3.0; It still works in Airflow 3.0 but is expected to be removed in a future version.
+   |
+36 | task_group()
+37 | setup()
+38 | teardown()
+   | ^^^^^^^^ AIR311
+39 |
+40 | # airflow.io
+   |
+   = help: Use `airflow.sdk.teardown` instead
+
+AIR311_names.py:41:1: AIR311 `airflow.io.path.ObjectStoragePath` is removed in Airflow 3.0; It still works in Airflow 3.0 but is expected to be removed in a future version.
+   |
+40 | # airflow.io
+41 | ObjectStoragePath()
+   | ^^^^^^^^^^^^^^^^^ AIR311
+42 | attach()
+   |
+   = help: Use `airflow.sdk.ObjectStoragePath` instead
+
+AIR311_names.py:42:1: AIR311 `airflow.io.storage.attach` is removed in Airflow 3.0; It still works in Airflow 3.0 but is expected to be removed in a future version.
+   |
+40 | # airflow.io
+41 | ObjectStoragePath()
+42 | attach()
+   | ^^^^^^ AIR311
+43 |
+44 | # airflow.models
+   |
+   = help: Use `airflow.sdk.io.attach` instead
+
+AIR311_names.py:45:1: AIR311 `airflow.models.DAG` is removed in Airflow 3.0; It still works in Airflow 3.0 but is expected to be removed in a future version.
+   |
+44 | # airflow.models
+45 | DAGFromModel()
+   | ^^^^^^^^^^^^ AIR311
+46 |
+47 | # airflow.models.baseoperator
+   |
+   = help: Use `airflow.sdk.DAG` instead
+
+AIR311_names.py:48:1: AIR311 `airflow.models.baseoperator.chain` is removed in Airflow 3.0; It still works in Airflow 3.0 but is expected to be removed in a future version.
+   |
+47 | # airflow.models.baseoperator
+48 | chain()
+   | ^^^^^ AIR311
+49 | chain_linear()
+50 | cross_downstream()
    |
    = help: Use `airflow.sdk.chain` instead
 
-AIR311_names.py:29:1: AIR311 `airflow.models.baseoperator.chain_linear` is removed in Airflow 3.0; It still works in Airflow 3.0 but is expected to be removed in a future version.
+AIR311_names.py:49:1: AIR311 `airflow.models.baseoperator.chain_linear` is removed in Airflow 3.0; It still works in Airflow 3.0 but is expected to be removed in a future version.
    |
-27 | # airflow.models.baseoperator
+47 | # airflow.models.baseoperator
-28 | chain()
+48 | chain()
-29 | chain_linear()
+49 | chain_linear()
    | ^^^^^^^^^^^^ AIR311
-30 | cross_downstream()
+50 | cross_downstream()
    |
    = help: Use `airflow.sdk.chain_linear` instead
 
-AIR311_names.py:30:1: AIR311 `airflow.models.baseoperator.cross_downstream` is removed in Airflow 3.0; It still works in Airflow 3.0 but is expected to be removed in a future version.
+AIR311_names.py:50:1: AIR311 `airflow.models.baseoperator.cross_downstream` is removed in Airflow 3.0; It still works in Airflow 3.0 but is expected to be removed in a future version.
    |
-28 | chain()
+48 | chain()
-29 | chain_linear()
+49 | chain_linear()
-30 | cross_downstream()
+50 | cross_downstream()
    | ^^^^^^^^^^^^^^^^ AIR311
-31 |
+51 |
-32 | # airflow.models.baseoperatolinker
+52 | # airflow.models.baseoperatolinker
    |
    = help: Use `airflow.sdk.cross_downstream` instead
 
-AIR311_names.py:36:1: AIR311 `airflow.timetables.datasets.DatasetOrTimeSchedule` is removed in Airflow 3.0; It still works in Airflow 3.0 but is expected to be removed in a future version.
+AIR311_names.py:56:1: AIR311 `airflow.models.dag.DAG` is removed in Airflow 3.0; It still works in Airflow 3.0 but is expected to be removed in a future version.
    |
-35 | # airflow.timetables.datasets
+55 | # airflow.models.dag
-36 | DatasetOrTimeSchedule()
+56 | DAGFromDag()
+   | ^^^^^^^^^^ AIR311
+57 | # airflow.timetables.datasets
+58 | DatasetOrTimeSchedule()
+   |
+   = help: Use `airflow.sdk.DAG` instead
+
+AIR311_names.py:58:1: AIR311 `airflow.timetables.datasets.DatasetOrTimeSchedule` is removed in Airflow 3.0; It still works in Airflow 3.0 but is expected to be removed in a future version.
+   |
+56 | DAGFromDag()
+57 | # airflow.timetables.datasets
+58 | DatasetOrTimeSchedule()
    | ^^^^^^^^^^^^^^^^^^^^^ AIR311
-37 |
+59 |
-38 | # airflow.utils.dag_parsing_context
+60 | # airflow.utils.dag_parsing_context
    |
    = help: Use `airflow.timetables.assets.AssetOrTimeSchedule` instead
 
-AIR311_names.py:39:1: AIR311 `airflow.utils.dag_parsing_context.get_parsing_context` is removed in Airflow 3.0; It still works in Airflow 3.0 but is expected to be removed in a future version.
+AIR311_names.py:61:1: AIR311 `airflow.utils.dag_parsing_context.get_parsing_context` is removed in Airflow 3.0; It still works in Airflow 3.0 but is expected to be removed in a future version.
    |
-38 | # airflow.utils.dag_parsing_context
+60 | # airflow.utils.dag_parsing_context
-39 | get_parsing_context()
+61 | get_parsing_context()
    | ^^^^^^^^^^^^^^^^^^^ AIR311
    |
    = help: Use `airflow.sdk.get_parsing_context` instead