Skip to content

Commit

Permalink
refactor(AIR302): refactor regex usage
Browse files Browse the repository at this point in the history
  • Loading branch information
@Lee-W
Lee-W committed Dec 25, 2024
1 parent e179b0a commit 7ea9b06
Show file tree
Hide file tree
Showing 3 changed files with 85 additions and 40 deletions.
13 changes: 12 additions & 1 deletion crates/ruff_linter/resources/test/fixtures/airflow/AIR302_names.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -317,12 +317,23 @@
csm_backend.get_conn_uri()
csm_backend.get_connections()

from airflow.providers.hashicorp.secrets.vault import VaultBackend
from airflow.providers.hashicorp.secrets.vault import VaultBackend, NotAir302SecretError

vault_backend = VaultBackend()
vault_backend.get_conn_uri()
vault_backend.get_connections()

not_an_error = NotAir302SecretError()
not_an_error.get_conn_uri()

from airflow.providers.apache.beam.hooks import BeamHook, NotAir302HookError

beam_hook = BeamHook()
beam_hook.get_conn_uri()

not_an_error = NotAir302HookError()
not_an_error.get_conn_uri()

from airflow.providers_manager import ProvidersManager

provider_manager = ProvidersManager()
Expand Down
66 changes: 50 additions & 16 deletions crates/ruff_linter/src/rules/airflow/rules/removal_in_3.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,16 +4,58 @@ use ruff_python_ast::{name::QualifiedName, Arguments, Expr, ExprAttribute, ExprC
use ruff_python_semantic::analyze::typing;
use ruff_python_semantic::Modules;
use ruff_text_size::Ranged;
use std::sync::LazyLock;

use crate::checkers::ast::Checker;
use regex::Regex;

static SECRET_BACKEND_REGEX: LazyLock<Regex> =
LazyLock::new(|| Regex::new(r"airflow\..*secrets\.\w+\.\w+Backend").unwrap());
fn is_airflow_secret_backend(qualname: &QualifiedName) -> bool {
match qualname.segments() {
["airflow", "secrets", rest @ ..] => {
if let Some(last_element) = rest.last() {
last_element.ends_with("Backend")
} else {
false
}
}

["airflow", "providers", rest @ ..] => {
// Ensure 'operators' exists somewhere in the middle
if let (Some(pos), Some(last_element)) =
(rest.iter().position(|&s| s == "secrets"), rest.last())
{
pos + 1 < rest.len() && last_element.ends_with("Backend") // Check that 'operators' is not the last element
} else {
false
}
}

static AIRFLOW_HOOK_REGEX: LazyLock<Regex> =
LazyLock::new(|| Regex::new(r"airflow\..*hooks\.\w+\.\w+Hook").unwrap());
_ => false,
}
}

fn is_airflow_hook(qualname: &QualifiedName) -> bool {
match qualname.segments() {
["airflow", "hooks", rest @ ..] => {
if let Some(last_element) = rest.last() {
last_element.ends_with("Hook")
} else {
false
}
}

["airflow", "providers", rest @ ..] => {
// Ensure 'operators' exists somewhere in the middle
if let (Some(pos), Some(last_element)) =
(rest.iter().position(|&s| s == "hooks"), rest.last())
{
pos + 1 < rest.len() && last_element.ends_with("Hook")
} else {
false
}
}

_ => false,
}
}

#[derive(Debug, Eq, PartialEq)]
enum Replacement {
Expand Down Expand Up @@ -172,14 +214,6 @@ fn removed_class_attribute(checker: &mut Checker, expr: &Expr) {
return;
};

// checker.diagnostics.push(Diagnostic::new(
// Airflow3Removal {
// deprecated: qualname.to_string(),
// replacement: Replacement::Name("te"),
// },
// attr.range(),
// ));

let replacement = match *qualname.segments() {
["airflow", "providers_manager", "ProvidersManager"] => match attr.as_str() {
"dataset_factories" => Some(Replacement::Name("asset_factories")),
Expand Down Expand Up @@ -252,12 +286,12 @@ fn removed_method(checker: &mut Checker, expr: &Expr) {
"iter_dataset_aliases" => Some(Replacement::Name("iter_asset_aliases")),
&_ => None,
},
_ if SECRET_BACKEND_REGEX.is_match(&qualname.segments().join(".")) => match attr.as_str() {
_ if is_airflow_secret_backend(&qualname) => match attr.as_str() {
"get_conn_uri" => Some(Replacement::Name("get_conn_value")),
"get_connections" => Some(Replacement::Name("get_connection")),
&_ => None,
},
_ if AIRFLOW_HOOK_REGEX.is_match(&qualname.segments().join(".")) => match attr.as_str() {
_ if is_airflow_hook(&qualname) => match attr.as_str() {
"get_connections" => Some(Replacement::Name("get_connection")),
&_ => None,
},
Expand Down
46 changes: 23 additions & 23 deletions ...c/rules/airflow/snapshots/ruff_linter__rules__airflow__tests__AIR302_AIR302_names.py.snap
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1177,7 +1177,7 @@ AIR302_names.py:318:13: AIR302 `get_connections` is removed in Airflow 3.0
318 | csm_backend.get_connections()
| ^^^^^^^^^^^^^^^ AIR302
319 |
320 | from airflow.providers.hashicorp.secrets.vault import VaultBackend
320 | from airflow.providers.hashicorp.secrets.vault import VaultBackend, NotAir302SecretError
|
= help: Use `get_connection` instead

Expand All @@ -1197,47 +1197,47 @@ AIR302_names.py:324:15: AIR302 `get_connections` is removed in Airflow 3.0
324 | vault_backend.get_connections()
| ^^^^^^^^^^^^^^^ AIR302
325 |
326 | from airflow.providers_manager import ProvidersManager
326 | not_an_error = NotAir302SecretError()
|
= help: Use `get_connection` instead

AIR302_names.py:329:18: AIR302 `dataset_factories` is removed in Airflow 3.0
AIR302_names.py:340:18: AIR302 `dataset_factories` is removed in Airflow 3.0
|
328 | provider_manager = ProvidersManager()
329 | provider_manager.dataset_factories
339 | provider_manager = ProvidersManager()
340 | provider_manager.dataset_factories
| ^^^^^^^^^^^^^^^^^ AIR302
330 | provider_manager.dataset_uri_handlers
331 | provider_manager.dataset_to_openlineage_converters
341 | provider_manager.dataset_uri_handlers
342 | provider_manager.dataset_to_openlineage_converters
|
= help: Use `asset_factories` instead

AIR302_names.py:330:18: AIR302 `dataset_uri_handlers` is removed in Airflow 3.0
AIR302_names.py:341:18: AIR302 `dataset_uri_handlers` is removed in Airflow 3.0
|
328 | provider_manager = ProvidersManager()
329 | provider_manager.dataset_factories
330 | provider_manager.dataset_uri_handlers
339 | provider_manager = ProvidersManager()
340 | provider_manager.dataset_factories
341 | provider_manager.dataset_uri_handlers
| ^^^^^^^^^^^^^^^^^^^^ AIR302
331 | provider_manager.dataset_to_openlineage_converters
342 | provider_manager.dataset_to_openlineage_converters
|
= help: Use `asset_uri_handlers` instead

AIR302_names.py:331:18: AIR302 `dataset_to_openlineage_converters` is removed in Airflow 3.0
AIR302_names.py:342:18: AIR302 `dataset_to_openlineage_converters` is removed in Airflow 3.0
|
329 | provider_manager.dataset_factories
330 | provider_manager.dataset_uri_handlers
331 | provider_manager.dataset_to_openlineage_converters
340 | provider_manager.dataset_factories
341 | provider_manager.dataset_uri_handlers
342 | provider_manager.dataset_to_openlineage_converters
| ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ AIR302
332 |
333 | from airflow.lineage.hook import DatasetLineageInfo
343 |
344 | from airflow.lineage.hook import DatasetLineageInfo
|
= help: Use `asset_to_openlineage_converters` instead

AIR302_names.py:335:11: AIR302 `airflow.lineage.hook.DatasetLineageInfo` is removed in Airflow 3.0
AIR302_names.py:346:11: AIR302 `airflow.lineage.hook.DatasetLineageInfo` is removed in Airflow 3.0
|
333 | from airflow.lineage.hook import DatasetLineageInfo
334 |
335 | dl_info = DatasetLineageInfo()
344 | from airflow.lineage.hook import DatasetLineageInfo
345 |
346 | dl_info = DatasetLineageInfo()
| ^^^^^^^^^^^^^^^^^^ AIR302
336 | dl_info.dataset
347 | dl_info.dataset
|
= help: Use `airflow.lineage.hook.AssetLineageInfo` instead

0 comments on commit 7ea9b06

Please sign in to comment.