[airflow]: extend removed args (AIR302) (#14765)

## Summary Airflow 3.0 removes various deprecated functions, members, modules, and other values. They have been deprecated in 2.x, but the removal causes incompatibilities that we want to detect. This PR deprecates the following names. * in `DAG` * `sla_miss_callback` was removed * in `airflow.operators.trigger_dagrun.TriggerDagRunOperator` * `execution_date` was removed * in `airflow.operators.weekday.DayOfWeekSensor`, `airflow.operators.datetime.BranchDateTimeOperator` and `airflow.operators.weekday.BranchDayOfWeekOperator` * `use_task_execution_day` was removed in favor of `use_task_logical_date` The full list of rules we will extend apache/airflow#44556 ## Test Plan  A test fixture is included in the PR.
astral-sh · Dec 6, 2024 · 3ea14d7 · 3ea14d7
1 parent 4fdd4dd
commit 3ea14d7
Show file tree

Hide file tree

Showing 3 changed files with 166 additions and 27 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
@@ -1,13 +1,29 @@
 from airflow import DAG, dag
 from airflow.timetables.simple import NullTimetable
 
+from airflow.operators.trigger_dagrun import TriggerDagRunOperator
+from airflow.providers.standard.operators import trigger_dagrun
+
+from airflow.operators.datetime import BranchDateTimeOperator
+from airflow.providers.standard.operators import datetime
+
+from airflow.sensors.weekday import DayOfWeekSensor, BranchDayOfWeekOperator
+from airflow.providers.standard.sensors import weekday
+
 DAG(dag_id="class_schedule", schedule="@hourly")
 
 DAG(dag_id="class_schedule_interval", schedule_interval="@hourly")
 
 DAG(dag_id="class_timetable", timetable=NullTimetable())
 
 
+def sla_callback(*arg, **kwargs):
+    pass
+
+
+DAG(dag_id="class_sla_callback", sla_miss_callback=sla_callback)
+
+
 @dag(schedule="0 * * * *")
 def decorator_schedule():
     pass
@@ -21,3 +37,42 @@ def decorator_schedule_interval():
 @dag(timetable=NullTimetable())
 def decorator_timetable():
     pass
+
+
+@dag(sla_miss_callback=sla_callback)
+def decorator_sla_callback():
+    pass
+
+
+@dag()
+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"
+    )
+
+    branch_dt_op = datetime.BranchDateTimeOperator(
+        task_id="branch_dt_op", use_task_execution_day=True
+    )
+    branch_dt_op2 = BranchDateTimeOperator(
+        task_id="branch_dt_op2", 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_op2 = BranchDayOfWeekOperator(task_id="bdow_op2", use_task_execution_day=True)
+
+    trigger_dagrun_op >> trigger_dagrun_op2
+    branch_dt_op >> branch_dt_op2
+    dof_task_sensor >> dof_task_sensor2
+    bdow_op >> bdow_op2
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
@@ -96,6 +96,39 @@ fn removed_argument(checker: &mut Checker, qualname: &QualifiedName, arguments:
                 "timetable",
                 Some("schedule"),
             ));
+            checker.diagnostics.extend(diagnostic_for_argument(
+                arguments,
+                "sla_miss_callback",
+                None::<&str>,
+            ));
+        }
+        ["airflow", .., "operators", "trigger_dagrun", "TriggerDagRunOperator"] => {
+            checker.diagnostics.extend(diagnostic_for_argument(
+                arguments,
+                "execution_date",
+                Some("logical_date"),
+            ));
+        }
+        ["airflow", .., "operators", "datetime", "BranchDateTimeOperator"] => {
+            checker.diagnostics.extend(diagnostic_for_argument(
+                arguments,
+                "use_task_execution_day",
+                Some("use_task_logical_date"),
+            ));
+        }
+        ["airflow", .., "operators", "weekday", "DayOfWeekSensor"] => {
+            checker.diagnostics.extend(diagnostic_for_argument(
+                arguments,
+                "use_task_execution_day",
+                Some("use_task_logical_date"),
+            ));
+        }
+        ["airflow", .., "operators", "weekday", "BranchDayOfWeekOperator"] => {
+            checker.diagnostics.extend(diagnostic_for_argument(
+                arguments,
+                "use_task_execution_day",
+                Some("use_task_logical_date"),
+            ));
         }
         _ => {}
     };

diff --git a/...rc/rules/airflow/snapshots/ruff_linter__rules__airflow__tests__AIR302_AIR302_args.py.snap b/...rc/rules/airflow/snapshots/ruff_linter__rules__airflow__tests__AIR302_AIR302_args.py.snap
@@ -2,36 +2,87 @@
 source: crates/ruff_linter/src/rules/airflow/mod.rs
 snapshot_kind: text
 ---
-AIR302_args.py:6:39: AIR302 `schedule_interval` is removed in Airflow 3.0; use `schedule` instead
-  |
-4 | DAG(dag_id="class_schedule", schedule="@hourly")
-5 | 
-6 | DAG(dag_id="class_schedule_interval", schedule_interval="@hourly")
-  |                                       ^^^^^^^^^^^^^^^^^ AIR302
-7 | 
-8 | DAG(dag_id="class_timetable", timetable=NullTimetable())
-  |
-
-AIR302_args.py:8:31: AIR302 `timetable` is removed in Airflow 3.0; use `schedule` instead
-  |
-6 | DAG(dag_id="class_schedule_interval", schedule_interval="@hourly")
-7 | 
-8 | DAG(dag_id="class_timetable", timetable=NullTimetable())
-  |                               ^^^^^^^^^ AIR302
-  |
-
-AIR302_args.py:16:6: AIR302 `schedule_interval` is removed in Airflow 3.0; use `schedule` instead
-   |
-16 | @dag(schedule_interval="0 * * * *")
+AIR302_args.py:15:39: AIR302 `schedule_interval` is removed in Airflow 3.0; use `schedule` instead
+   |
+13 | DAG(dag_id="class_schedule", schedule="@hourly")
+14 | 
+15 | DAG(dag_id="class_schedule_interval", schedule_interval="@hourly")
+   |                                       ^^^^^^^^^^^^^^^^^ AIR302
+16 | 
+17 | DAG(dag_id="class_timetable", timetable=NullTimetable())
+   |
+
+AIR302_args.py:17:31: AIR302 `timetable` is removed in Airflow 3.0; use `schedule` instead
+   |
+15 | DAG(dag_id="class_schedule_interval", schedule_interval="@hourly")
+16 | 
+17 | DAG(dag_id="class_timetable", timetable=NullTimetable())
+   |                               ^^^^^^^^^ AIR302
+   |
+
+AIR302_args.py:24:34: AIR302 `sla_miss_callback` is removed in Airflow 3.0
+   |
+24 | DAG(dag_id="class_sla_callback", sla_miss_callback=sla_callback)
+   |                                  ^^^^^^^^^^^^^^^^^ AIR302
+   |
+
+AIR302_args.py:32:6: AIR302 `schedule_interval` is removed in Airflow 3.0; use `schedule` instead
+   |
+32 | @dag(schedule_interval="0 * * * *")
    |      ^^^^^^^^^^^^^^^^^ AIR302
-17 | def decorator_schedule_interval():
-18 |     pass
+33 | def decorator_schedule_interval():
+34 |     pass
    |
 
-AIR302_args.py:21:6: AIR302 `timetable` is removed in Airflow 3.0; use `schedule` instead
+AIR302_args.py:37:6: AIR302 `timetable` is removed in Airflow 3.0; use `schedule` instead
    |
-21 | @dag(timetable=NullTimetable())
+37 | @dag(timetable=NullTimetable())
    |      ^^^^^^^^^ AIR302
-22 | def decorator_timetable():
-23 |     pass
+38 | def decorator_timetable():
+39 |     pass
+   |
+
+AIR302_args.py:42:6: AIR302 `sla_miss_callback` is removed in Airflow 3.0
+   |
+42 | @dag(sla_miss_callback=sla_callback)
+   |      ^^^^^^^^^^^^^^^^^ AIR302
+43 | def decorator_sla_callback():
+44 |     pass
+   |
+
+AIR302_args.py:50:39: AIR302 `execution_date` is removed in Airflow 3.0; use `logical_date` instead
+   |
+48 | def decorator_deprecated_operator_args():
+49 |     trigger_dagrun_op = trigger_dagrun.TriggerDagRunOperator(
+50 |         task_id="trigger_dagrun_op1", execution_date="2024-12-04"
+   |                                       ^^^^^^^^^^^^^^ AIR302
+51 |     )
+52 |     trigger_dagrun_op2 = TriggerDagRunOperator(
+   |
+
+AIR302_args.py:53:39: AIR302 `execution_date` is removed in Airflow 3.0; use `logical_date` instead
+   |
+51 |     )
+52 |     trigger_dagrun_op2 = TriggerDagRunOperator(
+53 |         task_id="trigger_dagrun_op2", execution_date="2024-12-04"
+   |                                       ^^^^^^^^^^^^^^ AIR302
+54 |     )
+   |
+
+AIR302_args.py:57:33: AIR302 `use_task_execution_day` is removed in Airflow 3.0; use `use_task_logical_date` instead
+   |
+56 |     branch_dt_op = datetime.BranchDateTimeOperator(
+57 |         task_id="branch_dt_op", use_task_execution_day=True
+   |                                 ^^^^^^^^^^^^^^^^^^^^^^ AIR302
+58 |     )
+59 |     branch_dt_op2 = BranchDateTimeOperator(
+   |
+
+AIR302_args.py:60:34: AIR302 `use_task_execution_day` is removed in Airflow 3.0; use `use_task_logical_date` instead
+   |
+58 |     )
+59 |     branch_dt_op2 = BranchDateTimeOperator(
+60 |         task_id="branch_dt_op2", use_task_execution_day=True
+   |                                  ^^^^^^^^^^^^^^^^^^^^^^ AIR302
+61 |     )
    |