[airflow] extend and fix AIR302 rules (#15525)

astral-sh · Jan 16, 2025 · fc9dd63 · fc9dd63
1 parent 79e52c7
commit fc9dd63
Show file tree

Hide file tree

Showing 3 changed files with 588 additions and 565 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
@@ -66,7 +66,7 @@
 )
 from airflow.providers.postgres.datasets import postgres
 from airflow.providers.trino.datasets import trino
-from airflow.secrets.local_filesystem import get_connection, load_connections
+from airflow.secrets.local_filesystem import LocalFilesystemBackend, load_connections
 from airflow.security.permissions import RESOURCE_DATASET
 from airflow.sensors.base_sensor_operator import BaseSensorOperator
 from airflow.sensors.date_time_sensor import DateTimeSensor
@@ -84,6 +84,7 @@
 from airflow.triggers.external_task import TaskStateTrigger
 from airflow.utils import dates
 from airflow.utils.dag_cycle_tester import test_cycle
+from airflow.utils.dag_parsing_context import get_parsing_context
 from airflow.utils.dates import (
     date_range,
     datetime_to_nano,
@@ -225,7 +226,10 @@
 trino.sanitize_uri
 
 # airflow.secrets
-get_connection, load_connections
+# get_connection
+lfb = LocalFilesystemBackend()
+lfb.get_connections()
+load_connections
 
 # airflow.security.permissions
 RESOURCE_DATASET
@@ -272,6 +276,9 @@
 # airflow.utils.dag_cycle_tester
 test_cycle
 
+# airflow.utils.dag_parsing_context
+get_parsing_context
+
 # airflow.utils.decorators
 apply_defaults
 

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
@@ -294,6 +294,11 @@ fn check_method(checker: &mut Checker, call_expr: &ExprCall) {
             )),
             _ => None,
         },
+        ["airflow", "secrets", "local_filesystem", "LocalFilesystemBackend"] => match attr.as_str()
+        {
+            "get_connections" => Some(Replacement::Name("get_connection")),
+            _ => None,
+        },
         ["airflow", "datasets", ..] | ["airflow", "Dataset"] => match attr.as_str() {
             "iter_datasets" => Some(Replacement::Name("iter_assets")),
             "iter_dataset_aliases" => Some(Replacement::Name("iter_asset_aliases")),
@@ -543,12 +548,13 @@ fn check_name(checker: &mut Checker, expr: &Expr, range: TextRange) {
         ["airflow", "secrets", "local_filesystem", "load_connections"] => {
             Replacement::Name("airflow.secrets.local_filesystem.load_connections_dict")
         }
-        ["airflow", "secrets", "local_filesystem", "get_connection"] => {
-            Replacement::Name("airflow.secrets.local_filesystem.load_connections_dict")
-        }
 
+        // airflow.utils.dag_parsing_context
+        ["airflow", "utils", "dag_parsing_context", "get_parsing_context"] => {
+            Replacement::Name("airflow.sdk.get_parsing_context")
+        }
         // airflow.utils.dates
-        ["airflow", "utils", "dates", "date_range"] => Replacement::Name("airflow.timetables."),
+        ["airflow", "utils", "dates", "date_range"] => Replacement::None,
         ["airflow", "utils", "dates", "days_ago"] => {
             Replacement::Name("pendulum.today('UTC').add(days=-N, ...)")
         }