[airflow] Apply auto fixes to cases where the names have changed in Airflow 3 (AIR301)

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.base_auth_manager import is_authorized_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.datasets.manager import (
-    DatasetManager,
-    dataset_manager,
-    resolve_dataset_manager,
-)
 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.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, load_connections
-from airflow.security.permissions import RESOURCE_DATASET
+from airflow.secrets.local_filesystem import LocalFilesystemBackend
 from airflow.sensors.base_sensor_operator import BaseSensorOperator
-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
@@ -70,19 +49,15 @@
 from airflow.utils.log import secrets_masker
 from airflow.utils.state import SHUTDOWN, terminating_states
 from airflow.utils.trigger_rule import TriggerRule
-from airflow.www.auth import has_access, has_access_dataset
+from airflow.www.auth import has_access
 from airflow.www.utils import get_sensitive_variables_fields, should_hide_value_for_key
 
 # airflow root
 PY36, PY37, PY38, PY39, PY310, PY311, PY312
-DatasetFromRoot()
 
 # airflow.api_connexion.security
-requires_access, requires_access_dataset
+requires_access
 
-# airflow.auth.managers
-is_authorized_dataset
-DatasetDetails()
 
 # airflow.configuration
 get, getboolean, getfloat, getint, has_option, remove_option, as_dict, set
@@ -95,76 +70,17 @@
 # airflow.datasets
 DatasetAliasEvent()
 
-# airflow.datasets.manager
-DatasetManager()
-dataset_manager
-resolve_dataset_manager
 
 # airflow.hooks
 BaseHook()
 
-# airflow.lineage.hook
-DatasetLineageInfo()
-
-# airflow.listeners.spec.dataset
-on_dataset_changed
-on_dataset_created
-
-# airflow.metrics.validators
-AllowListValidator()
-BlockListValidator()
-
-
-# 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
 
@@ -174,18 +90,12 @@
 # airflow.secrets
 # get_connection
 LocalFilesystemBackend()
-load_connections
 
-# airflow.security.permissions
-RESOURCE_DATASET
 
 # airflow.sensors.base_sensor_operator
 BaseSensorOperator()
 
 
-# airflow.timetables
-DatasetTriggeredTimetable()
-
 # airflow.triggers.external_task
 TaskStateTrigger()
 
@@ -233,9 +143,9 @@
 TriggerRule.DUMMY
 TriggerRule.NONE_FAILED_OR_SKIPPED
 
+
 # airflow.www.auth
 has_access
-has_access_dataset
 
 # airflow.www.utils
 get_sensitive_variables_fields

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

@@ -0,0 +1,45 @@
+from __future__ import annotations
+
+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.metrics.validators import AllowListValidator, BlockListValidator
+from airflow.secrets.local_filesystm import load_connections
+from airflow.security.permissions import RESOURCE_DATASET
+from airflow.www.auth import has_access_dataset
+
+requires_access_dataset()
+
+DatasetDetails()
+is_authorized_dataset()
+
+DatasetManager()
+dataset_manager()
+resolve_dataset_manager()
+
+DatasetLineageInfo()
+
+AllowListValidator()
+BlockListValidator()
+
+load_connections()
+
+RESOURCE_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

@@ -17,6 +17,8 @@ mod tests {
     #[test_case(Rule::AirflowDagNoScheduleArgument, Path::new("AIR002.py"))]
     #[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"))]