[airflow] Add autofixes for AIR302 and AIR312 (#17942)

Lee-W · web-flow · commit 2ceba6ae6785 · 2025-05-15T16:03:02.000-04:00
&lt;!--
Thank you for contributing to Ruff! To help us out with reviewing,
please consider the following:

- Does this pull request include a summary of the change? (See below.)
- Does this pull request include a descriptive title?
- Does this pull request include references to any relevant issues?
--&gt;

## Summary

&lt;!-- What's the purpose of the change? What does it do, and why? --&gt;

`ProviderReplacement::Name` was designed back when we only wanted to do
linting. Now we also want to fix the user code. It would be easier for
us to replace them with better AutoImport struct.

## Test Plan

&lt;!-- How was it tested? --&gt;

The test fixture has been updated as some cases can now be fixed
diff --git a/crates/ruff_linter/src/rules/airflow/helpers.rs b/crates/ruff_linter/src/rules/airflow/helpers.rs
@@ -31,11 +31,6 @@ pub(crate) enum Replacement {
 #[derive(Clone, Debug, Eq, PartialEq)]
 pub(crate) enum ProviderReplacement {
     None,
-    ProviderName {
-        name: &'static str,
-        provider: &'static str,
-        version: &'static str,
-    },
     AutoImport {
         module: &'static str,
         name: &'static str,
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
@@ -46,12 +46,7 @@ impl Violation for Airflow3MovedToProvider {
             ProviderReplacement::None => {
                 format!("`{deprecated}` is removed in Airflow 3.0")
             }
-            ProviderReplacement::ProviderName {
-                name: _,
-                provider,
-                version: _,
-            }
-            | ProviderReplacement::AutoImport {
+            ProviderReplacement::AutoImport {
                 name: _,
                 module: _,
                 provider,
@@ -72,15 +67,6 @@ impl Violation for Airflow3MovedToProvider {
         let Airflow3MovedToProvider { replacement, .. } = self;
         match replacement {
             ProviderReplacement::None => {None}
-            ProviderReplacement::ProviderName {
-                name,
-                provider,
-                version,
-            } => {
-                Some(format!(
-                    "Install `apache-airflow-providers-{provider}>={version}` and use `{name}` instead."
-                ))
-            },
             ProviderReplacement::AutoImport {
                 name,
                 module,
@@ -122,7 +108,6 @@ fn check_names_moved_to_provider(checker: &Checker, expr: &Expr, ranged: TextRan
     };
 
     let replacement = match qualified_name.segments() {
-        // ProviderName: for cases that only one name has been moved
         // apache-airflow-providers-amazon
         ["airflow", "hooks", "S3_hook", rest @ ("S3Hook" | "provide_bucket_name")] => {
             ProviderReplacement::SourceModuleMovedToProvider {
@@ -238,8 +223,9 @@ fn check_names_moved_to_provider(checker: &Checker, expr: &Expr, ranged: TextRan
         }
         ["airflow", "operators", "druid_check_operator", "DruidCheckOperator"]
         | ["airflow", "operators", "presto_check_operator", "PrestoCheckOperator"] => {
-            ProviderReplacement::ProviderName {
-                name: "airflow.providers.common.sql.operators.sql.SQLCheckOperator",
+            ProviderReplacement::AutoImport {
+                module: "airflow.providers.common.sql.operators.sql",
+                name: "SQLCheckOperator",
                 provider: "common-sql",
                 version: "1.1.0",
             }
@@ -255,8 +241,9 @@ fn check_names_moved_to_provider(checker: &Checker, expr: &Expr, ranged: TextRan
             }
         }
         ["airflow", "operators", "presto_check_operator", "PrestoIntervalCheckOperator"] => {
-            ProviderReplacement::ProviderName {
-                name: "airflow.providers.common.sql.operators.sql.SQLIntervalCheckOperator",
+            ProviderReplacement::AutoImport {
+                module: "airflow.providers.common.sql.operators.sql",
+                name: "SQLIntervalCheckOperator",
                 provider: "common-sql",
                 version: "1.1.0",
             }
@@ -281,8 +268,9 @@ fn check_names_moved_to_provider(checker: &Checker, expr: &Expr, ranged: TextRan
             }
         }
         ["airflow", "operators", "presto_check_operator", "PrestoValueCheckOperator"] => {
-            ProviderReplacement::ProviderName {
-                name: "airflow.providers.common.sql.operators.sql.SQLValueCheckOperator",
+            ProviderReplacement::AutoImport {
+                module: "airflow.providers.common.sql.operators.sql",
+                name: "SQLValueCheckOperator",
                 provider: "common-sql",
                 version: "1.1.0",
             }
@@ -320,8 +308,9 @@ fn check_names_moved_to_provider(checker: &Checker, expr: &Expr, ranged: TextRan
         | ["airflow", "operators", "oracle_operator", "OracleOperator"]
         | ["airflow", "operators", "postgres_operator", "PostgresOperator"]
         | ["airflow", "operators", "sqlite_operator", "SqliteOperator"] => {
-            ProviderReplacement::ProviderName {
-                name: "airflow.providers.common.sql.operators.sql.SQLExecuteQueryOperator",
+            ProviderReplacement::AutoImport {
+                module: "airflow.providers.common.sql.operators.sql",
+                name: "SQLExecuteQueryOperator",
                 provider: "common-sql",
                 version: "1.3.0",
             }
@@ -943,7 +932,7 @@ fn check_names_moved_to_provider(checker: &Checker, expr: &Expr, ranged: TextRan
         ProviderReplacement::SourceModuleMovedToProvider { module, name, .. } => {
             Some((module, name.as_str()))
         }
-        _ => None,
+        ProviderReplacement::None => None,
     } {
         if is_guarded_by_try_except(expr, module, name, semantic) {
             return;
diff --git a/crates/ruff_linter/src/rules/airflow/rules/suggested_to_move_to_provider_in_3.rs b/crates/ruff_linter/src/rules/airflow/rules/suggested_to_move_to_provider_in_3.rs
@@ -47,12 +47,7 @@ impl Violation for Airflow3SuggestedToMoveToProvider {
             ProviderReplacement::None => {
                 format!("`{deprecated}` is removed in Airflow 3.0")
             }
-            ProviderReplacement::ProviderName {
-                name: _,
-                provider,
-                version: _,
-            }
-            | ProviderReplacement::AutoImport {
+            ProviderReplacement::AutoImport {
                 name: _,
                 module: _,
                 provider,
@@ -76,15 +71,6 @@ impl Violation for Airflow3SuggestedToMoveToProvider {
         let Airflow3SuggestedToMoveToProvider { replacement, .. } = self;
         match replacement {
             ProviderReplacement::None => None,
-            ProviderReplacement::ProviderName {
-                name,
-                provider,
-                version,
-            } => {
-                Some(format!(
-                    "Install `apache-airflow-providers-{provider}>={version}` and use `{name}` instead."
-                ))
-            },
             ProviderReplacement::AutoImport {
                 module,
                 name,
@@ -285,7 +271,7 @@ fn check_names_moved_to_provider(checker: &Checker, expr: &Expr, ranged: TextRan
         ProviderReplacement::SourceModuleMovedToProvider { module, name, .. } => {
             Some((module, name.as_str()))
         }
-        _ => None,
+        ProviderReplacement::None => None,
     } {
         if is_guarded_by_try_except(expr, module, name, semantic) {
             return;
diff --git a/crates/ruff_linter/src/rules/airflow/snapshots/ruff_linter__rules__airflow__tests__AIR302_AIR302_common_sql.py.snap b/crates/ruff_linter/src/rules/airflow/snapshots/ruff_linter__rules__airflow__tests__AIR302_AIR302_common_sql.py.snap
@@ -289,7 +289,7 @@ AIR302_common_sql.py:114:1: AIR302 `airflow.sensors.sql_sensor.SqlSensor` is mov
     |
     = help: Install `apache-airflow-providers-common-sql>=1.0.0` and use `airflow.providers.common.sql.sensors.sql.SqlSensor` instead.
 
-AIR302_common_sql.py:124:1: AIR302 `airflow.operators.jdbc_operator.JdbcOperator` is moved into `common-sql` provider in Airflow 3.0;
+AIR302_common_sql.py:124:1: AIR302 [*] `airflow.operators.jdbc_operator.JdbcOperator` is moved into `common-sql` provider in Airflow 3.0;
     |
 122 | from airflow.operators.sqlite_operator import SqliteOperator
 123 |
@@ -300,7 +300,19 @@ AIR302_common_sql.py:124:1: AIR302 `airflow.operators.jdbc_operator.JdbcOperator
     |
     = help: Install `apache-airflow-providers-common-sql>=1.3.0` and use `airflow.providers.common.sql.operators.sql.SQLExecuteQueryOperator` instead.
 
-AIR302_common_sql.py:125:1: AIR302 `airflow.operators.mssql_operator.MsSqlOperator` is moved into `common-sql` provider in Airflow 3.0;
+ℹ Safe fix
+120 120 | from airflow.operators.oracle_operator import OracleOperator
+121 121 | from airflow.operators.postgres_operator import PostgresOperator
+122 122 | from airflow.operators.sqlite_operator import SqliteOperator
+    123 |+from airflow.providers.common.sql.operators.sql import SQLExecuteQueryOperator
+123 124 | 
+124     |-JdbcOperator()
+    125 |+SQLExecuteQueryOperator()
+125 126 | MsSqlOperator()
+126 127 | MySqlOperator()
+127 128 | OracleOperator()
+
+AIR302_common_sql.py:125:1: AIR302 [*] `airflow.operators.mssql_operator.MsSqlOperator` is moved into `common-sql` provider in Airflow 3.0;
     |
 124 | JdbcOperator()
 125 | MsSqlOperator()
@@ -310,7 +322,20 @@ AIR302_common_sql.py:125:1: AIR302 `airflow.operators.mssql_operator.MsSqlOperat
     |
     = help: Install `apache-airflow-providers-common-sql>=1.3.0` and use `airflow.providers.common.sql.operators.sql.SQLExecuteQueryOperator` instead.
 
-AIR302_common_sql.py:126:1: AIR302 `airflow.operators.mysql_operator.MySqlOperator` is moved into `common-sql` provider in Airflow 3.0;
+ℹ Safe fix
+120 120 | from airflow.operators.oracle_operator import OracleOperator
+121 121 | from airflow.operators.postgres_operator import PostgresOperator
+122 122 | from airflow.operators.sqlite_operator import SqliteOperator
+    123 |+from airflow.providers.common.sql.operators.sql import SQLExecuteQueryOperator
+123 124 | 
+124 125 | JdbcOperator()
+125     |-MsSqlOperator()
+    126 |+SQLExecuteQueryOperator()
+126 127 | MySqlOperator()
+127 128 | OracleOperator()
+128 129 | PostgresOperator()
+
+AIR302_common_sql.py:126:1: AIR302 [*] `airflow.operators.mysql_operator.MySqlOperator` is moved into `common-sql` provider in Airflow 3.0;
     |
 124 | JdbcOperator()
 125 | MsSqlOperator()
@@ -321,7 +346,21 @@ AIR302_common_sql.py:126:1: AIR302 `airflow.operators.mysql_operator.MySqlOperat
     |
     = help: Install `apache-airflow-providers-common-sql>=1.3.0` and use `airflow.providers.common.sql.operators.sql.SQLExecuteQueryOperator` instead.
 
-AIR302_common_sql.py:127:1: AIR302 `airflow.operators.oracle_operator.OracleOperator` is moved into `common-sql` provider in Airflow 3.0;
+ℹ Safe fix
+120 120 | from airflow.operators.oracle_operator import OracleOperator
+121 121 | from airflow.operators.postgres_operator import PostgresOperator
+122 122 | from airflow.operators.sqlite_operator import SqliteOperator
+    123 |+from airflow.providers.common.sql.operators.sql import SQLExecuteQueryOperator
+123 124 | 
+124 125 | JdbcOperator()
+125 126 | MsSqlOperator()
+126     |-MySqlOperator()
+    127 |+SQLExecuteQueryOperator()
+127 128 | OracleOperator()
+128 129 | PostgresOperator()
+129 130 | SqliteOperator()
+
+AIR302_common_sql.py:127:1: AIR302 [*] `airflow.operators.oracle_operator.OracleOperator` is moved into `common-sql` provider in Airflow 3.0;
     |
 125 | MsSqlOperator()
 126 | MySqlOperator()
@@ -332,7 +371,21 @@ AIR302_common_sql.py:127:1: AIR302 `airflow.operators.oracle_operator.OracleOper
     |
     = help: Install `apache-airflow-providers-common-sql>=1.3.0` and use `airflow.providers.common.sql.operators.sql.SQLExecuteQueryOperator` instead.
 
-AIR302_common_sql.py:128:1: AIR302 `airflow.operators.postgres_operator.PostgresOperator` is moved into `common-sql` provider in Airflow 3.0;
+ℹ Safe fix
+120 120 | from airflow.operators.oracle_operator import OracleOperator
+121 121 | from airflow.operators.postgres_operator import PostgresOperator
+122 122 | from airflow.operators.sqlite_operator import SqliteOperator
+    123 |+from airflow.providers.common.sql.operators.sql import SQLExecuteQueryOperator
+123 124 | 
+124 125 | JdbcOperator()
+125 126 | MsSqlOperator()
+126 127 | MySqlOperator()
+127     |-OracleOperator()
+    128 |+SQLExecuteQueryOperator()
+128 129 | PostgresOperator()
+129 130 | SqliteOperator()
+
+AIR302_common_sql.py:128:1: AIR302 [*] `airflow.operators.postgres_operator.PostgresOperator` is moved into `common-sql` provider in Airflow 3.0;
     |
 126 | MySqlOperator()
 127 | OracleOperator()
@@ -342,11 +395,39 @@ AIR302_common_sql.py:128:1: AIR302 `airflow.operators.postgres_operator.Postgres
     |
     = help: Install `apache-airflow-providers-common-sql>=1.3.0` and use `airflow.providers.common.sql.operators.sql.SQLExecuteQueryOperator` instead.
 
-AIR302_common_sql.py:129:1: AIR302 `airflow.operators.sqlite_operator.SqliteOperator` is moved into `common-sql` provider in Airflow 3.0;
+ℹ Safe fix
+120 120 | from airflow.operators.oracle_operator import OracleOperator
+121 121 | from airflow.operators.postgres_operator import PostgresOperator
+122 122 | from airflow.operators.sqlite_operator import SqliteOperator
+    123 |+from airflow.providers.common.sql.operators.sql import SQLExecuteQueryOperator
+123 124 | 
+124 125 | JdbcOperator()
+125 126 | MsSqlOperator()
+126 127 | MySqlOperator()
+127 128 | OracleOperator()
+128     |-PostgresOperator()
+    129 |+SQLExecuteQueryOperator()
+129 130 | SqliteOperator()
+
+AIR302_common_sql.py:129:1: AIR302 [*] `airflow.operators.sqlite_operator.SqliteOperator` is moved into `common-sql` provider in Airflow 3.0;
     |
 127 | OracleOperator()
 128 | PostgresOperator()
 129 | SqliteOperator()
     | ^^^^^^^^^^^^^^ AIR302
     |
     = help: Install `apache-airflow-providers-common-sql>=1.3.0` and use `airflow.providers.common.sql.operators.sql.SQLExecuteQueryOperator` instead.
+
+ℹ Safe fix
+120 120 | from airflow.operators.oracle_operator import OracleOperator
+121 121 | from airflow.operators.postgres_operator import PostgresOperator
+122 122 | from airflow.operators.sqlite_operator import SqliteOperator
+    123 |+from airflow.providers.common.sql.operators.sql import SQLExecuteQueryOperator
+123 124 | 
+124 125 | JdbcOperator()
+125 126 | MsSqlOperator()
+126 127 | MySqlOperator()
+127 128 | OracleOperator()
+128 129 | PostgresOperator()
+129     |-SqliteOperator()
+    130 |+SQLExecuteQueryOperator()
