From 2ceba6ae67854aa2f766f48bff56d9a314ce2e7f Mon Sep 17 00:00:00 2001 From: Wei Lee Date: Fri, 16 May 2025 04:03:02 +0800 Subject: [PATCH] [`airflow`] Add autofixes for `AIR302` and `AIR312` (#17942) ## Summary `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 The test fixture has been updated as some cases can now be fixed --- .../ruff_linter/src/rules/airflow/helpers.rs | 5 - .../airflow/rules/moved_to_provider_in_3.rs | 39 +++----- .../suggested_to_move_to_provider_in_3.rs | 18 +--- ...w__tests__AIR302_AIR302_common_sql.py.snap | 93 +++++++++++++++++-- 4 files changed, 103 insertions(+), 52 deletions(-) diff --git a/crates/ruff_linter/src/rules/airflow/helpers.rs b/crates/ruff_linter/src/rules/airflow/helpers.rs index 54a4d2bda5..90c99234d3 100644 --- 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 index b0b7e36f57..8ce561c62f 100644 --- 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 index a95bdeef84..d8807d9623 100644 --- 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 index 42f27ffd6b..693e7f44bd 100644 --- 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,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()