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

<!--
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?
-->

## Summary

<!-- What's the purpose of the change? What does it do, and why? -->

`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

<!-- How was it tested? -->

The test fixture has been updated as some cases can now be fixed
This commit is contained in:
Wei Lee 2025-05-16 04:03:02 +08:00 committed by GitHub
parent d3a7cb3fe4
commit 2ceba6ae67
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 103 additions and 52 deletions

View file

@ -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,

View file

@ -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;

View file

@ -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;

View file

@ -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,7 +395,21 @@ 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()
@ -350,3 +417,17 @@ AIR302_common_sql.py:129:1: AIR302 `airflow.operators.sqlite_operator.SqliteOper
| ^^^^^^^^^^^^^^ 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()