feat(AIR302): argument appbuilder is now removed in BaseAuthManager a…

…nd its subclasses
astral-sh · Dec 26, 2024 · 6e665c4 · 6e665c4
1 parent d1c9770
commit 6e665c4
Show file tree

Hide file tree

Showing 3 changed files with 221 additions and 163 deletions.
diff --git a/crates/ruff_linter/resources/test/fixtures/airflow/AIR302_args.py b/crates/ruff_linter/resources/test/fixtures/airflow/AIR302_args.py
@@ -6,6 +6,7 @@
 from airflow.providers.amazon.aws.log.s3_task_handler import S3TaskHandler
 from airflow.providers.apache.hdfs.log.hdfs_task_handler import HdfsTaskHandler
 from airflow.providers.elasticsearch.log.es_task_handler import ElasticsearchTaskHandler
+from airflow.providers.fab.auth_manager.fab_auth_manager import FabAuthManager
 from airflow.providers.google.cloud.log.gcs_task_handler import GCSTaskHandler
 from airflow.providers.standard.operators import datetime, trigger_dagrun
 from airflow.providers.standard.sensors import weekday
@@ -51,7 +52,9 @@ def decorator_deprecated_operator_args():
     trigger_dagrun_op = trigger_dagrun.TriggerDagRunOperator(
         task_id="trigger_dagrun_op1", execution_date="2024-12-04"
     )
-    trigger_dagrun_op2 = TriggerDagRunOperator(task_id="trigger_dagrun_op2", execution_date="2024-12-04")
+    trigger_dagrun_op2 = TriggerDagRunOperator(
+        task_id="trigger_dagrun_op2", execution_date="2024-12-04"
+    )
 
     branch_dt_op = datetime.BranchDateTimeOperator(
         task_id="branch_dt_op", use_task_execution_day=True, task_concurrency=5
@@ -62,10 +65,16 @@ def decorator_deprecated_operator_args():
         sla=timedelta(seconds=10),
     )
 
-    dof_task_sensor = weekday.DayOfWeekSensor(task_id="dof_task_sensor", use_task_execution_day=True)
-    dof_task_sensor2 = DayOfWeekSensor(task_id="dof_task_sensor2", use_task_execution_day=True)
+    dof_task_sensor = weekday.DayOfWeekSensor(
+        task_id="dof_task_sensor", use_task_execution_day=True
+    )
+    dof_task_sensor2 = DayOfWeekSensor(
+        task_id="dof_task_sensor2", use_task_execution_day=True
+    )
 
-    bdow_op = weekday.BranchDayOfWeekOperator(task_id="bdow_op", use_task_execution_day=True)
+    bdow_op = weekday.BranchDayOfWeekOperator(
+        task_id="bdow_op", use_task_execution_day=True
+    )
     bdow_op2 = BranchDayOfWeekOperator(task_id="bdow_op2", use_task_execution_day=True)
 
     trigger_dagrun_op >> trigger_dagrun_op2
@@ -79,3 +88,5 @@ def decorator_deprecated_operator_args():
 HdfsTaskHandler(filename_template="/tmp/test")
 ElasticsearchTaskHandler(filename_template="/tmp/test")
 GCSTaskHandler(filename_template="/tmp/test")
+
+FabAuthManager(None)
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
@@ -35,6 +35,30 @@ fn is_airflow_secret_backend(segments: &[&str]) -> bool {
     }
 }
 
+fn is_airflow_auth_manager(segments: &[&str]) -> bool {
+    match segments {
+        ["airflow", "auth", "manager", rest @ ..] => {
+            if let Some(last_element) = rest.last() {
+                last_element.ends_with("AuthManager")
+            } else {
+                false
+            }
+        }
+
+        ["airflow", "providers", rest @ ..] => {
+            if let (Some(pos), Some(last_element)) =
+                (rest.iter().position(|&s| s == "auth_manager"), rest.last())
+            {
+                pos + 1 < rest.len() && last_element.ends_with("AuthManager")
+            } else {
+                false
+            }
+        }
+
+        _ => false,
+    }
+}
+
 fn is_airflow_hook(qualname: &QualifiedName) -> bool {
     match qualname.segments() {
         ["airflow", "hooks", rest @ ..] => {
@@ -217,6 +241,20 @@ fn removed_argument(checker: &mut Checker, qualname: &QualifiedName, arguments:
                 None::<&str>,
             ));
         }
+        _ if is_airflow_auth_manager(qualname.segments()) => {
+            if !arguments.is_empty() {
+                checker.diagnostics.push(Diagnostic::new(
+                    Airflow3Removal {
+                        // deprecated: (*arguments).to_string(),
+                        deprecated: "appbuilder".to_string(),
+                        replacement: Replacement::Message(
+                            "The constructor takes no parameter now.",
+                        ),
+                    },
+                    arguments.range(),
+                ));
+            }
+        }
         _ if is_airflow_task_handler(qualname.segments()) => {
             checker.diagnostics.extend(diagnostic_for_argument(
                 arguments,