feat(AIR301): apply AutoImport to cases that names has been changed in Airflow 3

Author: Lee-W

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

@@ -9,9 +9,7 @@
     PY311,
     PY312,
 )
-from airflow.api_connexion.security import requires_access, requires_access_dataset
-
-from airflow.auth.managers.models.resource_details import DatasetDetails
+from airflow.api_connexion.security import requires_access
 from airflow.configuration import (
     as_dict,
     get,
@@ -24,32 +22,13 @@
 )
 from airflow.contrib.aws_athena_hook import AWSAthenaHook
 from airflow.datasets import DatasetAliasEvent
-
-
-
-
-
 from airflow.hooks.base_hook import BaseHook
-
-
-
 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
-from airflow.providers.common.io.datasets import file as common_io_file
-from airflow.providers.fab.auth_manager import fab_auth_manager
-from airflow.providers.google.datasets import bigquery, gcs
 from airflow.providers.mysql.datasets import mysql
-from airflow.providers.openlineage.utils.utils import (
-    DatasetInfo,
-    translate_airflow_dataset,
-)
 from airflow.providers.postgres.datasets import postgres
 from airflow.providers.trino.datasets import trino
 from airflow.secrets.local_filesystem import LocalFilesystemBackend
-from airflow.security.permissions import RESOURCE_DATASET
 from airflow.sensors.base_sensor_operator import BaseSensorOperator
-
 from airflow.triggers.external_task import TaskStateTrigger
 from airflow.utils import dates
 from airflow.utils.dag_cycle_tester import test_cycle
@@ -75,13 +54,9 @@
 
 # airflow root
 PY36, PY37, PY38, PY39, PY310, PY311, PY312
-DatasetFromRoot()
 
 # airflow.api_connexion.security
-requires_access, requires_access_dataset
-
-
-
+requires_access
 
 
 # airflow.configuration
@@ -96,77 +71,16 @@
 DatasetAliasEvent()
 
 
-
-
-
-
 # airflow.hooks
 BaseHook()
 
 
-
-
-
-
-
-
-
-
-
-
-
-
-
-# airflow.operators.branch_operator
-BaseBranchOperator()
-
-# airflow.operators.dagrun_operator
-TriggerDagRunLink()
-TriggerDagRunOperator()
-
-# airflow.operators.email_operator
-EmailOperator()
-
-# airflow.operators.latest_only_operator
-LatestOnlyOperator()
-
-# airflow.operators.python_operator
-BranchPythonOperator()
-PythonOperator()
-PythonVirtualenvOperator()
-ShortCircuitOperator()
-
 # airflow.operators.subdag.*
 SubDagOperator()
 
-# airflow.providers.amazon
-AvpEntities.DATASET
-s3.create_dataset
-s3.convert_dataset_to_openlineage
-s3.sanitize_uri
-
-# airflow.providers.common.io
-common_io_file.convert_dataset_to_openlineage
-common_io_file.create_dataset
-common_io_file.sanitize_uri
-
-# airflow.providers.fab
-fab_auth_manager.is_authorized_dataset
-
-# airflow.providers.google
-bigquery.sanitize_uri
-
-gcs.create_dataset
-gcs.sanitize_uri
-gcs.convert_dataset_to_openlineage
-
 # airflow.providers.mysql
 mysql.sanitize_uri
 
-# airflow.providers.openlineage
-DatasetInfo()
-translate_airflow_dataset
-
 # airflow.providers.postgres
 postgres.sanitize_uri
 
@@ -178,16 +92,10 @@
 LocalFilesystemBackend()
 
 
-
-
-
 # airflow.sensors.base_sensor_operator
 BaseSensorOperator()
 
 
-
-
-
 # airflow.triggers.external_task
 TaskStateTrigger()
 
@@ -235,10 +143,10 @@
 TriggerRule.DUMMY
 TriggerRule.NONE_FAILED_OR_SKIPPED
 
+
 # airflow.www.auth
 has_access
 
-
 # airflow.www.utils
 get_sensitive_variables_fields
 should_hide_value_for_key

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

@@ -1,51 +1,45 @@
 from __future__ import annotations
 
-from airflow.auth.managers.base_auth_manager import is_authorized_dataset
+from airflow.api_connexion.security import requires_access_dataset
+from airflow.auth.managers.models.resource_details import (
+    DatasetDetails,
+    is_authorized_dataset,
+)
 from airflow.datasets.manager import (
     DatasetManager,
     dataset_manager,
     resolve_dataset_manager,
 )
 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.secrets.local_filesystem import load_connections
+from airflow.metrics.validators import AllowListValidator, BlockListValidator
+from airflow.secrets.local_filesystm import load_connections
 from airflow.security.permissions import RESOURCE_DATASET
-from airflow.timetables.simple import DatasetTriggeredTimetable
 from airflow.www.auth import has_access_dataset
 
-# airflow.datasets.manager
-DatasetManager()
-dataset_manager
-resolve_dataset_manager
+requires_access_dataset()
 
-# airflow.auth.managers
-is_authorized_dataset
 DatasetDetails()
+is_authorized_dataset()
+
+DatasetManager()
+dataset_manager()
+resolve_dataset_manager()
 
-# airflow.lineage.hook
 DatasetLineageInfo()
 
-# airflow.listeners.spec.dataset
-on_dataset_changed
-on_dataset_created
-# airflow.metrics.validators
 AllowListValidator()
 BlockListValidator()
 
-# airflow.secrets.local_filesystem
 load_connections()
 
-# airflow.security.permissions
 RESOURCE_DATASET
-# airflow.timetables
-DatasetTriggeredTimetable()
 
-# airflow.www.auth
-has_access_dataset
+has_access_dataset()
+
+from airflow.listeners.spec.dataset import (
+    on_dataset_changed,
+    on_dataset_created,
+)
+
+on_dataset_created()
+on_dataset_changed()

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

@@ -0,0 +1,54 @@
+from __future__ import annotations
+
+from airflow.providers.amazon.aws.auth_manager.avp.entities.AvpEntities import DATASET
+from airflow.providers.amazon.aws.datasets.s3 import (
+    convert_dataset_to_openlineage as s3_convert_dataset_to_openlineage,
+)
+from airflow.providers.amazon.aws.datasets.s3 import create_dataset as s3_create_dataset
+from airflow.providers.common.io.dataset.file import (
+    convert_dataset_to_openlineage as io_convert_dataset_to_openlineage,
+)
+from airflow.providers.common.io.dataset.file import create_dataset as io_create_dataset
+from airflow.providers.fab.auth_manager.fab_auth_manager import is_authorized_dataset as fab_is_authorized_dataset
+from airflow.providers.google.datasets.bigquery import (
+    create_dataset as bigquery_create_dataset,
+)
+from airflow.providers.google.datasets.gcs import (
+    convert_dataset_to_openlineage as gcs_convert_dataset_to_openlineage,
+)
+from airflow.providers.google.datasets.gcs import create_dataset as gcs_create_dataset
+from airflow.providers.openlineage.utils.utils import (
+    DatasetInfo,
+    translate_airflow_dataset,
+)
+
+DATASET
+
+s3_create_dataset()
+s3_convert_dataset_to_openlineage()
+
+io_create_dataset()
+io_convert_dataset_to_openlineage()
+
+fab_is_authorized_dataset()
+
+# airflow.providers.google.datasets.bigquery
+bigquery_create_dataset()
+# airflow.providers.google.datasets.gcs
+gcs_create_dataset()
+gcs_convert_dataset_to_openlineage()
+# airflow.providers.openlineage.utils.utils
+DatasetInfo()
+translate_airflow_dataset()
+#
+# airflow.secrets.local_filesystem
+load_connections()
+#
+# airflow.security.permissions
+RESOURCE_DATASET
+
+# airflow.timetables
+DatasetTriggeredTimetable()
+#
+# airflow.www.auth
+has_access_dataset

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

@@ -18,6 +18,7 @@ mod tests {
     #[test_case(Rule::Airflow3Removal, Path::new("AIR301_args.py"))]
     #[test_case(Rule::Airflow3Removal, Path::new("AIR301_names.py"))]
     #[test_case(Rule::Airflow3Removal, Path::new("AIR301_names_fix.py"))]
+    #[test_case(Rule::Airflow3Removal, Path::new("AIR301_provider_names_fix.py"))]
     #[test_case(Rule::Airflow3Removal, Path::new("AIR301_names_try.py"))]
     #[test_case(Rule::Airflow3Removal, Path::new("AIR301_class_attribute.py"))]
     #[test_case(Rule::Airflow3Removal, Path::new("AIR301_airflow_plugin.py"))]