fix(AIR303): fix ImportPathMoved / ProviderName misuse

astral-sh · Feb 7, 2025 · be0b104 · be0b104
1 parent 7db5a92
commit be0b104
Show file tree

Hide file tree

Showing 6 changed files with 1,640 additions and 1,472 deletions.
diff --git a/crates/ruff_linter/resources/test/fixtures/airflow/AIR302_names.py b/crates/ruff_linter/resources/test/fixtures/airflow/AIR302_names.py
@@ -45,8 +45,6 @@
 from airflow.listeners.spec.dataset import on_dataset_changed, on_dataset_created
 from airflow.metrics.validators import AllowListValidator, BlockListValidator
 from airflow.operators import dummy_operator
-from airflow.operators.bash import BashOperator
-from airflow.operators.bash_operator import BashOperator as LegacyBashOperator
 from airflow.operators.branch_operator import BaseBranchOperator
 from airflow.operators.dagrun_operator import TriggerDagRunLink, TriggerDagRunOperator
 from airflow.operators.dummy import DummyOperator, EmptyOperator
@@ -171,10 +169,6 @@
 dummy_operator.EmptyOperator()
 dummy_operator.DummyOperator()
 
-# airflow.operators.bash / airflow.operators.bash_operator
-BashOperator()
-LegacyBashOperator()
-
 # airflow.operators.branch_operator
 BaseBranchOperator()
 

diff --git a/crates/ruff_linter/resources/test/fixtures/airflow/AIR303.py b/crates/ruff_linter/resources/test/fixtures/airflow/AIR303.py
@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from airflow.api.auth.backend import basic_auth, kerberos_auth
 from airflow.api.auth.backend.basic_auth import auth_current_user
 from airflow.auth.managers.fab.api.auth.backend import (
@@ -41,7 +43,11 @@
 from airflow.hooks.webhdfs_hook import WebHDFSHook
 from airflow.hooks.zendesk_hook import ZendeskHook
 from airflow.kubernetes.k8s_model import K8SModel, append_to_pod
-from airflow.kubernetes.kube_client import _disable_verify_ssl, _enable_tcp_keepalive, get_kube_client
+from airflow.kubernetes.kube_client import (
+    _disable_verify_ssl,
+    _enable_tcp_keepalive,
+    get_kube_client,
+)
 from airflow.kubernetes.kubernetes_helper_functions import (
     add_pod_suffix,
     annotations_for_logging_task_metadata,
@@ -55,31 +61,47 @@
     PodDefaults,
     PodGenerator,
     PodGeneratorDeprecated,
-    add_pod_suffix as add_pod_suffix2,
     datetime_to_label_safe_datestring,
     extend_object_field,
     label_safe_datestring_to_datetime,
     make_safe_label_value,
     merge_objects,
+)
+from airflow.kubernetes.pod_generator import (
+    add_pod_suffix as add_pod_suffix2,
+)
+from airflow.kubernetes.pod_generator import (
     rand_str as rand_str2,
 )
 from airflow.kubernetes.pod_generator_deprecated import (
     PodDefaults as PodDefaults3,
+)
+from airflow.kubernetes.pod_generator_deprecated import (
     PodGenerator as PodGenerator2,
+)
+from airflow.kubernetes.pod_generator_deprecated import (
     make_safe_label_value as make_safe_label_value2,
 )
 from airflow.kubernetes.pod_launcher import PodLauncher, PodStatus
 from airflow.kubernetes.pod_launcher_deprecated import (
     PodDefaults as PodDefaults2,
+)
+from airflow.kubernetes.pod_launcher_deprecated import (
     PodLauncher as PodLauncher2,
+)
+from airflow.kubernetes.pod_launcher_deprecated import (
     PodStatus as PodStatus2,
+)
+from airflow.kubernetes.pod_launcher_deprecated import (
     get_kube_client as get_kube_client2,
 )
 from airflow.kubernetes.pod_runtime_info_env import PodRuntimeInfoEnv
 from airflow.kubernetes.secret import K8SModel2, Secret
 from airflow.kubernetes.volume import Volume
 from airflow.kubernetes.volume_mount import VolumeMount
 from airflow.macros.hive import closest_ds_partition, max_partition
+from airflow.operators.bash import BashOperator
+from airflow.operators.bash_operator import BashOperator as LegacyBashOperator
 from airflow.operators.check_operator import (
     CheckOperator,
     IntervalCheckOperator,
@@ -117,8 +139,14 @@
     PrestoCheckOperator,
     PrestoIntervalCheckOperator,
     PrestoValueCheckOperator,
+)
+from airflow.operators.presto_check_operator import (
     SQLCheckOperator as SQLCheckOperator2,
+)
+from airflow.operators.presto_check_operator import (
     SQLIntervalCheckOperator as SQLIntervalCheckOperator2,
+)
+from airflow.operators.presto_check_operator import (
     SQLValueCheckOperator as SQLValueCheckOperator2,
 )
 from airflow.operators.presto_to_mysql import (
@@ -139,14 +167,24 @@
 from airflow.operators.sql import (
     BaseSQLOperator,
     BranchSQLOperator,
+    SQLTableCheckOperator,
+    _convert_to_float_if_possible,
+    parse_boolean,
+)
+from airflow.operators.sql import (
     SQLCheckOperator as SQLCheckOperator3,
+)
+from airflow.operators.sql import (
     SQLColumnCheckOperator as SQLColumnCheckOperator2,
+)
+from airflow.operators.sql import (
     SQLIntervalCheckOperator as SQLIntervalCheckOperator3,
-    SQLTableCheckOperator,
+)
+from airflow.operators.sql import (
     SQLThresholdCheckOperator as SQLThresholdCheckOperator2,
+)
+from airflow.operators.sql import (
     SQLValueCheckOperator as SQLValueCheckOperator3,
-    _convert_to_float_if_possible,
-    parse_boolean,
 )
 from airflow.operators.sqlite_operator import SqliteOperator
 from airflow.operators.trigger_dagrun import TriggerDagRunOperator
@@ -193,6 +231,8 @@
 _convert_to_float_if_possible()
 parse_boolean()
 BaseSQLOperator()
+BashOperator()
+LegacyBashOperator()
 BranchSQLOperator()
 CheckOperator()
 ConnectorProtocol()

diff --git a/crates/ruff_linter/src/rules/airflow/rules/moved_to_provider_in_3.rs b/crates/ruff_linter/src/rules/airflow/rules/moved_to_provider_in_3.rs
@@ -176,15 +176,13 @@ fn check_names_moved_to_provider(checker: &mut Checker, expr: &Expr, ranged: Tex
         },
 
         // apache-airflow-providers-celery
-        ["airflow", "config_templates", "default_celery", "DEFAULT_CELERY_CONFIG"] => Replacement::ImportPathMoved{
-            original_path: "airflow.config_templates.default_celery.DEFAULT_CELERY_CONFIG",
-            new_path: "airflow.providers.celery.executors.default_celery.DEFAULT_CELERY_CONFIG",
+        ["airflow", "config_templates", "default_celery", "DEFAULT_CELERY_CONFIG"] => Replacement::ProviderName{
+            name: "airflow.providers.celery.executors.default_celery.DEFAULT_CELERY_CONFIG",
             provider: "celery",
             version: "3.3.0"
         },
-        ["airflow", "executors", "celery_executor", "app"] => Replacement::ImportPathMoved{
-            original_path: "airflow.executors.celery_executor.app",
-            new_path: "airflow.providers.celery.executors.celery_executor_utils.app",
+        ["airflow", "executors", "celery_executor", "app"] => Replacement::ProviderName{
+            name: "airflow.providers.celery.executors.celery_executor_utils.app",
             provider: "celery",
             version: "3.3.0"
         },
@@ -200,15 +198,13 @@ fn check_names_moved_to_provider(checker: &mut Checker, expr: &Expr, ranged: Tex
         },
 
         // apache-airflow-providers-common-sql
-        ["airflow", "hooks", "dbapi", "ConnectorProtocol"] => Replacement::ImportPathMoved{
-            original_path: "airflow.hooks.dbapi.ConnectorProtocol",
-            new_path: "airflow.providers.common.sql.hooks.sql.ConnectorProtocol",
+        ["airflow", "hooks", "dbapi", "ConnectorProtocol"] => Replacement::ProviderName{
+            name: "airflow.providers.common.sql.hooks.sql.ConnectorProtocol",
             provider: "common-sql",
             version: "1.0.0"
         },
-        ["airflow", "hooks", "dbapi", "DbApiHook"] => Replacement::ImportPathMoved{
-            original_path: "airflow.hooks.dbapi.DbApiHook",
-            new_path: "airflow.providers.common.sql.hooks.sql.DbApiHook",
+        ["airflow", "hooks", "dbapi", "DbApiHook"] => Replacement::ProviderName{
+            name: "airflow.providers.common.sql.hooks.sql.DbApiHook",
             provider: "common-sql",
             version: "1.0.0"
         },
@@ -590,15 +586,13 @@ fn check_names_moved_to_provider(checker: &mut Checker, expr: &Expr, ranged: Tex
         },
 
         // apache-airflow-providers-cncf-kubernetes
-        ["airflow", "executors", "kubernetes_executor_types", "ALL_NAMESPACES"] => Replacement::ImportPathMoved{
-                original_path: "airflow.executors.kubernetes_executor_types.ALL_NAMESPACES",
-                new_path: "airflow.providers.cncf.kubernetes.executors.kubernetes_executor_types.ALL_NAMESPACES",
+        ["airflow", "executors", "kubernetes_executor_types", "ALL_NAMESPACES"] => Replacement::ProviderName{
+                name: "airflow.providers.cncf.kubernetes.executors.kubernetes_executor_types.ALL_NAMESPACES",
             provider: "cncf-kubernetes",
             version: "7.4.0"
         },
-        ["airflow", "executors", "kubernetes_executor_types", "POD_EXECUTOR_DONE_KEY"] => Replacement::ImportPathMoved{
-                original_path: "airflow.executors.kubernetes_executor_types.POD_EXECUTOR_DONE_KEY",
-                new_path: "airflow.providers.cncf.kubernetes.executors.kubernetes_executor_types.POD_EXECUTOR_DONE_KEY",
+        ["airflow", "executors", "kubernetes_executor_types", "POD_EXECUTOR_DONE_KEY"] => Replacement::ProviderName{
+                name: "airflow.providers.cncf.kubernetes.executors.kubernetes_executor_types.POD_EXECUTOR_DONE_KEY",
             provider: "cncf-kubernetes",
             version: "7.4.0"
         },
@@ -903,74 +897,86 @@ fn check_names_moved_to_provider(checker: &mut Checker, expr: &Expr, ranged: Tex
         },
 
         // apache-airflow-providers-standard
-        ["airflow", "operators", "datetime"] => Replacement::ImportPathMoved{
+        ["airflow", "operators", "bash", ..] => Replacement::ImportPathMoved{
+            original_path: "airflow.operators.bash",
+            new_path: "airflow.providers.standard.operators.bash",
+            provider: "standard",
+            version: "0.0.1"
+        },
+        ["airflow", "operators", "bash_operator", ..] => Replacement::ImportPathMoved{
+            original_path: "airflow.operators.bash_operator",
+            new_path: "airflow.providers.standard.operators.bash",
+            provider: "standard",
+            version: "0.0.1"
+        },
+        ["airflow", "operators", "datetime", ..] => Replacement::ImportPathMoved{
             original_path: "airflow.operators.datetime",
             new_path: "airflow.providers.standard.time.operators.datetime",
             provider: "standard",
             version: "0.0.1"
         },
-        ["airflow", "operators", "weekday"] => Replacement::ImportPathMoved{
+        ["airflow", "operators", "weekday", ..] => Replacement::ImportPathMoved{
             original_path: "airflow.operators.weekday",
             new_path: "airflow.providers.standard.time.operators.weekday",
             provider: "standard",
             version: "0.0.1"
         },
-        ["airflow", "sensors", "date_time"] => Replacement::ImportPathMoved{
+        ["airflow", "sensors", "date_time", ..] => Replacement::ImportPathMoved{
             original_path: "airflow.sensors.date_time",
             new_path: "airflow.providers.standard.time.sensors.date_time",
             provider: "standard",
             version: "0.0.1"
         },
-        ["airflow", "sensors", "time_sensor"] => Replacement::ImportPathMoved{
+        ["airflow", "sensors", "time_sensor", ..] => Replacement::ImportPathMoved{
             original_path: "airflow.sensors.time_sensor",
             new_path: "airflow.providers.standard.time.sensors.time",
             provider: "standard",
             version: "0.0.1"
         },
-        ["airflow", "sensors", "time_delta"] => Replacement::ImportPathMoved{
+        ["airflow", "sensors", "time_delta", ..] => Replacement::ImportPathMoved{
             original_path: "airflow.sensors.time_delta",
             new_path: "airflow.providers.standard.time.sensors.time_delta",
             provider: "standard",
             version: "0.0.1"
         },
-        ["airflow", "sensors", "weekday"] => Replacement::ImportPathMoved{
+        ["airflow", "sensors", "weekday", ..] => Replacement::ImportPathMoved{
             original_path: "airflow.sensors.weekday",
             new_path: "airflow.providers.standard.time.sensors.weekday",
             provider: "standard",
             version: "0.0.1"
         },
-        ["airflow", "hooks", "filesystem"] => Replacement::ImportPathMoved{
+        ["airflow", "hooks", "filesystem", ..] => Replacement::ImportPathMoved{
             original_path: "airflow.hooks.filesystem",
             new_path: "airflow.providers.standard.hooks.filesystem",
             provider: "standard",
             version: "0.0.1"
         },
-        ["airflow", "hooks", "package_index"] => Replacement::ImportPathMoved{
+        ["airflow", "hooks", "package_index", ..] => Replacement::ImportPathMoved{
             original_path: "airflow.hooks.package_index",
             new_path: "airflow.providers.standard.hooks.package_index",
             provider: "standard",
             version: "0.0.1"
         },
-        ["airflow", "hooks", "subprocess"] => Replacement::ImportPathMoved{
+        ["airflow", "hooks", "subprocess", ..] => Replacement::ImportPathMoved{
             original_path: "airflow.hooks.subprocess",
             new_path: "airflow.providers.standard.hooks.subprocess",
             provider: "standard",
             version: "0.0.1"
         },
 
-        ["airflow", "triggers", "external_task"] => Replacement::ImportPathMoved{
+        ["airflow", "triggers", "external_task", ..] => Replacement::ImportPathMoved{
             original_path: "airflow.triggers.external_task",
             new_path: "airflow.providers.standard.triggers.external_task",
             provider: "standard",
             version: "0.0.3"
         },
-        ["airflow", "triggers", "file"] => Replacement::ImportPathMoved{
+        ["airflow", "triggers", "file", ..] => Replacement::ImportPathMoved{
             original_path: "airflow.triggers.file",
             new_path: "airflow.providers.standard.triggers.file",
             provider: "standard",
             version: "0.0.3"
         },
-        ["airflow", "triggers", "temporal"] => Replacement::ImportPathMoved{
+        ["airflow", "triggers", "temporal", ..] => Replacement::ImportPathMoved{
             original_path: "airflow.triggers.temporal",
             new_path: "airflow.providers.standard.triggers.temporal",
             provider: "standard",

diff --git a/crates/ruff_linter/src/rules/airflow/rules/removal_in_3.rs b/crates/ruff_linter/src/rules/airflow/rules/removal_in_3.rs
@@ -684,9 +684,6 @@ fn check_name(checker: &mut Checker, expr: &Expr, range: TextRange) {
         ["airflow", "operators", "subdag", ..] => {
             Replacement::Message("The whole `airflow.subdag` module has been removed.")
         }
-        ["airflow", "operators", "bash" | "bash_operator", "BashOperator"] => {
-            Replacement::Name("airflow.providers.standard.operators.bash.BashOperator")
-        }
         ["airflow", "operators", "branch_operator", "BaseBranchOperator"] => {
             Replacement::Name("airflow.operators.branch.BaseBranchOperator")
         }