Skip to content

Commit fff5b97

Browse files
committed
refactor(AIR302): combine similiar cases
1 parent 21a6bc6 commit fff5b97

File tree

4 files changed

+53
-94
lines changed

4 files changed

+53
-94
lines changed

crates/ruff_linter/resources/test/fixtures/airflow/AIR302.py

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -139,18 +139,18 @@
139139
from airflow.operators.pig_operator import PigOperator
140140
from airflow.operators.postgres_operator import Mapping, PostgresOperator
141141
from airflow.operators.presto_check_operator import (
142-
PrestoCheckOperator,
143-
PrestoIntervalCheckOperator,
144-
PrestoValueCheckOperator,
142+
CheckOperator as SQLCheckOperator2,
145143
)
146144
from airflow.operators.presto_check_operator import (
147-
SQLCheckOperator as SQLCheckOperator2,
145+
IntervalCheckOperator as SQLIntervalCheckOperator2,
148146
)
149147
from airflow.operators.presto_check_operator import (
150-
SQLIntervalCheckOperator as SQLIntervalCheckOperator2,
148+
PrestoCheckOperator,
149+
PrestoIntervalCheckOperator,
150+
PrestoValueCheckOperator,
151151
)
152152
from airflow.operators.presto_check_operator import (
153-
SQLValueCheckOperator as SQLValueCheckOperator2,
153+
ValueCheckOperator as SQLValueCheckOperator2,
154154
)
155155
from airflow.operators.presto_to_mysql import (
156156
PrestoToMySqlOperator,

crates/ruff_linter/src/rules/airflow/rules/moved_to_provider_in_3.rs

Lines changed: 41 additions & 82 deletions
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,7 @@ impl Violation for Airflow3MovedToProvider {
6262
fn fix_title(&self) -> Option<String> {
6363
let Airflow3MovedToProvider { replacement, .. } = self;
6464
match replacement {
65+
ProviderReplacement::None => {None}
6566
ProviderReplacement::ProviderName {
6667
name,
6768
provider,
@@ -79,7 +80,6 @@ impl Violation for Airflow3MovedToProvider {
7980
} => {
8081
Some(format!("Install `apache-airflow-provider-{provider}>={version}` and use `{module}.{name}` instead."))
8182
} ,
82-
_ => None,
8383
}
8484
}
8585
}
@@ -187,68 +187,48 @@ fn check_names_moved_to_provider(checker: &Checker, expr: &Expr, ranged: TextRan
187187
provider: "common-sql",
188188
version: "1.0.0"
189189
},
190-
["airflow", "operators", "check_operator", rest] => match *rest {
191-
"SQLCheckOperator" | "CheckOperator" => ProviderReplacement::ProviderName {
190+
["airflow", "operators", "check_operator" | "sql", "SQLCheckOperator"]
191+
| ["airflow", "operators", "check_operator" | "druid_check_operator" | "presto_check_operator", "CheckOperator"]
192+
| ["airflow", "operators", "druid_check_operator", "DruidCheckOperator"]
193+
| ["airflow", "operators", "presto_check_operator", "PrestoCheckOperator"] => {
194+
ProviderReplacement::ProviderName {
192195
name: "airflow.providers.common.sql.operators.sql.SQLCheckOperator",
193196
provider: "common-sql",
194197
version: "1.1.0"
195-
},
196-
"SQLIntervalCheckOperator" | "IntervalCheckOperator" => ProviderReplacement::ProviderName {
197-
name: "airflow.providers.common.sql.operators.sql.SQLIntervalCheckOperator",
198-
provider: "common-sql",
199-
version: "1.1.0"
200-
},
201-
"SQLThresholdCheckOperator" | "ThresholdCheckOperator" => ProviderReplacement::ProviderName {
202-
name: "airflow.providers.common.sql.operators.sql.SQLThresholdCheckOperator",
203-
provider: "common-sql",
204-
version: "1.1.0"
205-
},
206-
"SQLValueCheckOperator" | "ValueCheckOperator" => ProviderReplacement::ProviderName {
207-
name: "airflow.providers.common.sql.operators.sql.SQLValueCheckOperator",
208-
provider: "common-sql",
209-
version: "1.1.0"
210-
},
211-
_ => return
198+
}
212199
},
213-
["airflow", "operators", "druid_check_operator", "DruidCheckOperator"] => ProviderReplacement::ProviderName {
214-
name: "airflow.providers.common.sql.operators.sql.SQLCheckOperator",
200+
["airflow", "operators", "check_operator" | "presto_check_operator", "IntervalCheckOperator"]
201+
| ["airflow", "operators", "check_operator" | "sql", "SQLIntervalCheckOperator"]
202+
| ["airflow", "operators", "presto_check_operator", "PrestoIntervalCheckOperator"] => ProviderReplacement::ProviderName {
203+
name: "airflow.providers.common.sql.operators.sql.SQLIntervalCheckOperator",
215204
provider: "common-sql",
216-
version: "4.0.0"
205+
version: "1.1.0"
206+
},
207+
["airflow", "operators", "check_operator" | "sql" , "SQLThresholdCheckOperator"]
208+
| ["airflow", "operators", "check_operator", "ThresholdCheckOperator"] => ProviderReplacement::ProviderName {
209+
name: "airflow.providers.common.sql.operators.sql.SQLThresholdCheckOperator",
210+
provider: "common-sql",
211+
version: "1.1.0"
212+
},
213+
["airflow", "operators", "check_operator" | "presto_check_operator", "ValueCheckOperator"]
214+
| ["airflow", "operators", "presto_check_operator", "PrestoValueCheckOperator"]
215+
| ["airflow", "operators", "check_operator" | "sql", "SQLValueCheckOperator"] => ProviderReplacement::ProviderName {
216+
name: "airflow.providers.common.sql.operators.sql.SQLValueCheckOperator",
217+
provider: "common-sql",
218+
version: "1.1.0"
217219
},
218-
["airflow", "operators", "presto_check_operator", rest] => match *rest {
219-
"SQLCheckOperator" | "PrestoCheckOperator" => ProviderReplacement::ProviderName {
220-
name: "airflow.providers.common.sql.operators.sql.SQLCheckOperator",
221-
provider: "common-sql",
222-
version: "1.1.0"
223-
},
224-
"SQLIntervalCheckOperator" | "PrestoIntervalCheckOperator" => ProviderReplacement::ProviderName {
225-
name: "airflow.providers.common.sql.operators.sql.SQLIntervalCheckOperator",
226-
provider: "common-sql",
227-
version: "1.1.0"
228-
},
229-
"SQLValueCheckOperator" | "PrestoValueCheckOperator" => ProviderReplacement::ProviderName {
230-
name: "airflow.providers.common.sql.operators.sql.SQLValueCheckOperator",
231-
provider: "common-sql",
232-
version: "1.1.0"
233-
},
234-
_ => return
235-
}
236220
["airflow", "operators", "sql", rest] => match *rest {
237-
"BaseSQLOperator" |
238-
"BranchSQLOperator" |
239-
"SQLCheckOperator" |
240-
"SQLIntervalCheckOperator" |
241-
"SQLTablecheckOperator" |
242-
"SQLThresholdCheckOperator" => ProviderReplacement::SourceModuleMovedToProvider {
221+
"BaseSQLOperator"
222+
| "BranchSQLOperator"
223+
| "SQLTablecheckOperator" => ProviderReplacement::SourceModuleMovedToProvider {
243224
name: (*rest).to_string(),
244225
module: "airflow.providers.common.sql.operators.sql",
245226
provider: "common-sql",
246227
version: "1.1.0"
247228
},
248-
"SQLColumnCheckOperator" |
249-
"SQLValueCheckOperator" |
250-
"_convert_to_float_if_possible" |
251-
"parse_boolean" => ProviderReplacement::SourceModuleMovedToProvider {
229+
"SQLColumnCheckOperator"
230+
| "_convert_to_float_if_possible"
231+
| "parse_boolean" => ProviderReplacement::SourceModuleMovedToProvider {
252232
name: (*rest).to_string(),
253233
module: "airflow.providers.common.sql.operators.sql",
254234
provider: "common-sql",
@@ -261,6 +241,16 @@ fn check_names_moved_to_provider(checker: &Checker, expr: &Expr, ranged: TextRan
261241
provider: "common-sql",
262242
version: "1.0.0"
263243
},
244+
["airflow", "operators", "jdbc_operator", "JdbcOperator"]
245+
| ["airflow", "operators", "mssql_operator", "MsSqlOperator"]
246+
| ["airflow", "operators", "mysql_operator", "MySqlOperator"]
247+
| ["airflow", "operators", "oracle_operator", "OracleOperator"]
248+
| ["airflow", "operators", "postgres_operator", "PostgresOperator"]
249+
| ["airflow", "operators", "sqlite_operator", "SqliteOperator"] => ProviderReplacement::ProviderName {
250+
name: "airflow.providers.common.sql.operators.sql.SQLExecuteQueryOperator",
251+
provider: "common-sql",
252+
version: "1.3.0"
253+
},
264254

265255
// apache-airflow-providers-daskexecutor
266256
["airflow", "executors", "dask_executor", "DaskExecutor"] => ProviderReplacement::ProviderName {
@@ -464,11 +454,6 @@ fn check_names_moved_to_provider(checker: &Checker, expr: &Expr, ranged: TextRan
464454
provider: "jdbc",
465455
version: "1.0.0"
466456
},
467-
["airflow", "operators", "jdbc_operator", "JdbcOperator"] => ProviderReplacement::ProviderName {
468-
name: "airflow.providers.common.sql.operators.sql.SQLExecuteQueryOperator",
469-
provider: "common-sql",
470-
version: "1.3.0"
471-
},
472457

473458
// apache-airflow-providers-cncf-kubernetes
474459
["airflow", "executors", "kubernetes_executor_types", rest @ (
@@ -632,23 +617,13 @@ fn check_names_moved_to_provider(checker: &Checker, expr: &Expr, ranged: TextRan
632617
provider: "microsoft-mssql",
633618
version: "1.0.0"
634619
},
635-
["airflow", "operators", "mssql_operator", "MsSqlOperator"] => ProviderReplacement::ProviderName {
636-
name: "airflow.providers.common.sql.operators.sql.SQLExecuteQueryOperator",
637-
provider: "common-sql",
638-
version: "1.3.0"
639-
},
640620

641621
// apache-airflow-providers-mysql
642622
["airflow", "hooks", "mysql_hook", "MySqlHook"] => ProviderReplacement::ProviderName {
643623
name: "airflow.providers.mysql.hooks.mysql.MySqlHook",
644624
provider: "mysql",
645625
version: "1.0.0"
646626
},
647-
["airflow", "operators", "mysql_operator", "MySqlOperator"] => ProviderReplacement::ProviderName {
648-
name: "airflow.providers.common.sql.operators.sql.SQLExecuteQueryOperator",
649-
provider: "common-sql",
650-
version: "1.3.0"
651-
},
652627
["airflow", "operators", "presto_to_mysql", "PrestoToMySqlOperator" | "PrestoToMySqlTransfer"] => ProviderReplacement::ProviderName {
653628
name: "airflow.providers.mysql.transfers.presto_to_mysql.PrestoToMySqlOperator",
654629
provider: "mysql",
@@ -661,11 +636,6 @@ fn check_names_moved_to_provider(checker: &Checker, expr: &Expr, ranged: TextRan
661636
provider: "oracle",
662637
version: "1.0.0"
663638
},
664-
["airflow", "operators", "oracle_operator", "OracleOperator"] => ProviderReplacement::ProviderName {
665-
name: "airflow.providers.common.sql.operators.sql.SQLExecuteQueryOperator",
666-
provider: "common-sql",
667-
version: "1.3.0"
668-
},
669639

670640
// apache-airflow-providers-papermill
671641
["airflow", "operators", "papermill_operator", "PapermillOperator"] => ProviderReplacement::ProviderName {
@@ -693,11 +663,6 @@ fn check_names_moved_to_provider(checker: &Checker, expr: &Expr, ranged: TextRan
693663
version: "1.0.0"
694664
},
695665
["airflow", "operators", "postgres_operator", "Mapping"] => ProviderReplacement::None,
696-
["airflow", "operators", "postgres_operator", "PostgresOperator"] => ProviderReplacement::ProviderName {
697-
name: "airflow.providers.common.sql.operators.sql.SQLExecuteQueryOperator",
698-
provider: "common-sql",
699-
version: "1.3.0"
700-
},
701666

702667
// apache-airflow-providers-presto
703668
["airflow", "hooks", "presto_hook", "PrestoHook"] => ProviderReplacement::ProviderName {
@@ -742,12 +707,6 @@ fn check_names_moved_to_provider(checker: &Checker, expr: &Expr, ranged: TextRan
742707
provider: "sqlite",
743708
version: "1.0.0"
744709
},
745-
["airflow", "operators", "sqlite_operator", "SqliteOperator"] => ProviderReplacement::ProviderName {
746-
name: "airflow.providers.common.sql.operators.sql.SQLExecuteQueryOperator",
747-
provider: "common-sql",
748-
version: "1.3.0"
749-
},
750-
751710

752711
// apache-airflow-providers-zendesk
753712
["airflow", "hooks", "zendesk_hook", "ZendeskHook"] =>

crates/ruff_linter/src/rules/airflow/rules/suggested_to_move_to_provider_in_3.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -65,6 +65,7 @@ impl Violation for Airflow3SuggestedToMoveToProvider {
6565
fn fix_title(&self) -> Option<String> {
6666
let Airflow3SuggestedToMoveToProvider { replacement, .. } = self;
6767
match replacement {
68+
ProviderReplacement::None => {None}
6869
ProviderReplacement::ProviderName {
6970
name,
7071
provider,
@@ -82,7 +83,6 @@ impl Violation for Airflow3SuggestedToMoveToProvider {
8283
} => {
8384
Some(format!("Install `apache-airflow-provider-{provider}>={version}` and use `{module}.{name}` instead."))
8485
},
85-
_ => None,
8686
}
8787
}
8888
}

crates/ruff_linter/src/rules/airflow/snapshots/ruff_linter__rules__airflow__tests__AIR302_AIR302.py.snap

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -304,7 +304,7 @@ AIR302.py:259:1: AIR302 `airflow.operators.check_operator.SQLCheckOperator` is m
304304
|
305305
= help: Install `apache-airflow-provider-common-sql>=1.1.0` and use `airflow.providers.common.sql.operators.sql.SQLCheckOperator` instead.
306306

307-
AIR302.py:260:1: AIR302 `airflow.operators.presto_check_operator.SQLCheckOperator` is moved into `common-sql` provider in Airflow 3.0;
307+
AIR302.py:260:1: AIR302 `airflow.operators.presto_check_operator.CheckOperator` is moved into `common-sql` provider in Airflow 3.0;
308308
|
309309
258 | PrestoValueCheckOperator()
310310
259 | SQLCheckOperator()
@@ -348,7 +348,7 @@ AIR302.py:263:1: AIR302 `airflow.operators.check_operator.SQLIntervalCheckOperat
348348
|
349349
= help: Install `apache-airflow-provider-common-sql>=1.1.0` and use `airflow.providers.common.sql.operators.sql.SQLIntervalCheckOperator` instead.
350350

351-
AIR302.py:264:1: AIR302 `airflow.operators.presto_check_operator.SQLIntervalCheckOperator` is moved into `common-sql` provider in Airflow 3.0;
351+
AIR302.py:264:1: AIR302 `airflow.operators.presto_check_operator.IntervalCheckOperator` is moved into `common-sql` provider in Airflow 3.0;
352352
|
353353
262 | SQLColumnCheckOperator2()
354354
263 | SQLIntervalCheckOperator()
@@ -403,7 +403,7 @@ AIR302.py:269:1: AIR302 `airflow.operators.check_operator.SQLValueCheckOperator`
403403
|
404404
= help: Install `apache-airflow-provider-common-sql>=1.1.0` and use `airflow.providers.common.sql.operators.sql.SQLValueCheckOperator` instead.
405405

406-
AIR302.py:270:1: AIR302 `airflow.operators.presto_check_operator.SQLValueCheckOperator` is moved into `common-sql` provider in Airflow 3.0;
406+
AIR302.py:270:1: AIR302 `airflow.operators.presto_check_operator.ValueCheckOperator` is moved into `common-sql` provider in Airflow 3.0;
407407
|
408408
268 | SQLThresholdCheckOperator2()
409409
269 | SQLValueCheckOperator()
@@ -423,7 +423,7 @@ AIR302.py:271:1: AIR302 `airflow.operators.sql.SQLValueCheckOperator` is moved i
423423
272 | SqlSensor()
424424
273 | SqlSensor2()
425425
|
426-
= help: Install `apache-airflow-provider-common-sql>=1.0.0` and use `airflow.providers.common.sql.operators.sql.SQLValueCheckOperator` instead.
426+
= help: Install `apache-airflow-provider-common-sql>=1.1.0` and use `airflow.providers.common.sql.operators.sql.SQLValueCheckOperator` instead.
427427

428428
AIR302.py:272:1: AIR302 `airflow.sensors.sql.SqlSensor` is moved into `common-sql` provider in Airflow 3.0;
429429
|
@@ -516,7 +516,7 @@ AIR302.py:287:1: AIR302 `airflow.operators.druid_check_operator.DruidCheckOperat
516516
288 |
517517
289 | # apache-airflow-providers-apache-hdfs
518518
|
519-
= help: Install `apache-airflow-provider-common-sql>=4.0.0` and use `airflow.providers.common.sql.operators.sql.SQLCheckOperator` instead.
519+
= help: Install `apache-airflow-provider-common-sql>=1.1.0` and use `airflow.providers.common.sql.operators.sql.SQLCheckOperator` instead.
520520

521521
AIR302.py:290:1: AIR302 `airflow.hooks.webhdfs_hook.WebHDFSHook` is moved into `apache-hdfs` provider in Airflow 3.0;
522522
|

0 commit comments

Comments
 (0)