Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[airflow]: extend removed method calls (AIR302) #15054

Merged
merged 2 commits into from
Dec 20, 2024
+197 −0
Merged

[airflow]: extend removed method calls (AIR302) #15054

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
6799bdf
feat(AIR302): add the following rules for method name changes
Lee-W Dec 19, 2024
129dbba
fix: fix clippy warnig and rewrite method_name match
Lee-W Dec 20, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
31 changes: 31 additions & 0 deletions crates/ruff_linter/resources/test/fixtures/airflow/AIR302_names.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -253,3 +253,34 @@

# airflow.www.utils
get_sensitive_variables_fields, should_hide_value_for_key

from airflow.datasets.manager import DatasetManager

dm = DatasetManager()
dm.register_dataset_change()
dm.create_datasets()
dm.notify_dataset_created()
dm.notify_dataset_changed()
dm.notify_dataset_alias_created()


from airflow.lineage.hook import HookLineageCollector

hlc = HookLineageCollector()
hlc.create_dataset()
hlc.add_input_dataset()
hlc.add_output_dataset()
hlc.collected_datasets()


from airflow.providers.amazon.auth_manager.aws_auth_manager import AwsAuthManager

aam = AwsAuthManager()
aam.is_authorized_dataset()


from airflow.providers_manager import ProvidersManager

pm = ProvidersManager()
pm.initialize_providers_asset_uri_resources()
pm.dataset_factories
111 changes: 110 additions & 1 deletion crates/ruff_linter/src/rules/airflow/rules/removal_in_3.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
use ruff_diagnostics::{Diagnostic, Edit, Fix, FixAvailability, Violation};
use ruff_macros::{derive_message_formats, ViolationMetadata};
use ruff_python_ast::{name::QualifiedName, Arguments, Expr, ExprAttribute, ExprCall};
use ruff_python_ast::{name::QualifiedName, Arguments, Expr, ExprAttribute, ExprCall, Identifier};
use ruff_python_semantic::analyze::typing;
use ruff_python_semantic::Modules;
use ruff_text_size::Ranged;

Expand Down Expand Up @@ -154,6 +155,112 @@ fn removed_argument(checker: &mut Checker, qualname: &QualifiedName, arguments:
};
}

fn diagnostic_for_method(deprecated: &Identifier, replacement: Replacement) -> Option<Diagnostic> {
Lee-W marked this conversation as resolved.
Show resolved Hide resolved
let diagnostic = Diagnostic::new(
Airflow3Removal {
deprecated: deprecated.as_str().to_string(),
replacement,
},
deprecated.range(),
);

Some(diagnostic)
}

fn removed_method(checker: &mut Checker, func: &Box<Expr>) {
Lee-W marked this conversation as resolved.
Show resolved Hide resolved
let Expr::Attribute(ExprAttribute { attr, value, .. }) = func.as_ref() else {
return;
};

let Some(qualname) = typing::resolve_assignment(value, checker.semantic()) else {
return;
};

match *qualname.segments() {
["airflow", "datasets", "manager", "DatasetManager"] => {
match attr.as_str() {
"register_dataset_change" => {
checker.diagnostics.extend(diagnostic_for_method(
attr,
Replacement::Name("register_asset_change".to_string()),
));
}
"create_datasets" => {
checker.diagnostics.extend(diagnostic_for_method(
attr,
Replacement::Name("create_assets".to_string()),
));
}
"notify_dataset_created" => {
checker.diagnostics.extend(diagnostic_for_method(
attr,
Replacement::Name("notify_asset_created".to_string()),
));
}
"notify_dataset_changed" => {
checker.diagnostics.extend(diagnostic_for_method(
attr,
Replacement::Name("notify_asset_changed".to_string()),
));
}
"notify_dataset_alias_created" => {
checker.diagnostics.extend(diagnostic_for_method(
attr,
Replacement::Name("notify_asset_alias_created".to_string()),
));
}
&_ => {}
};
}
["airflow", "lineage", "hook", "HookLineageCollector"] => {
match attr.as_str() {
"create_dataset" => {
checker.diagnostics.extend(diagnostic_for_method(
attr,
Replacement::Name("create_asset".to_string()),
));
}
"add_input_dataset" => {
checker.diagnostics.extend(diagnostic_for_method(
attr,
Replacement::Name("add_input_asset".to_string()),
));
}
"add_output_dataset" => {
checker.diagnostics.extend(diagnostic_for_method(
attr,
Replacement::Name("add_output_asset".to_string()),
));
}
"collected_datasets" => {
checker.diagnostics.extend(diagnostic_for_method(
attr,
Replacement::Name("collected_assets".to_string()),
));
}
&_ => {}
};
}
["airflow", "providers", "amazon", "auth_manager", "aws_auth_manager", "AwsAuthManager"] => {
if attr.as_str() == "is_authorized_dataset" {
checker.diagnostics.extend(diagnostic_for_method(
attr,
Replacement::Name("is_authorized_asset".to_string()),
));
}
}
["airflow", "providers_manager", "ProvidersManager"] => {
if attr.as_str() == "initialize_providers_dataset_uri_resources" {
checker.diagnostics.extend(diagnostic_for_method(
attr,
Replacement::Name("initialize_providers_asset_uri_resources".to_string()),
));
}
}
_ => {}
}
}

fn removed_name(checker: &mut Checker, expr: &Expr, ranged: impl Ranged) {
let result =
checker
Expand Down Expand Up @@ -607,6 +714,8 @@ pub(crate) fn removed_in_3(checker: &mut Checker, expr: &Expr) {
if let Some(qualname) = checker.semantic().resolve_qualified_name(func) {
removed_argument(checker, &qualname, arguments);
};

removed_method(checker, func);
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we skip calling removed_method if the resolved_qualified_name call was successful or does it need both?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think both are needed.

Copy link
Contributor Author

@Lee-W Lee-W Dec 19, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

e.g.,

from airflow.datasets.manager import DatasetManager

d = DatasetManager()
d.register_dataset_change()

needs 2 warnings

  1. import error
  2. method call error

}
Expr::Attribute(ExprAttribute { attr: ranged, .. }) => removed_name(checker, expr, ranged),
ranged @ Expr::Name(_) => removed_name(checker, expr, ranged),
Expand Down
103 changes: 103 additions & 0 deletions ...c/rules/airflow/snapshots/ruff_linter__rules__airflow__tests__AIR302_AIR302_names.py.snap
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -894,6 +894,8 @@ AIR302_names.py:255:1: AIR302 `airflow.www.utils.get_sensitive_variables_fields`
254 | # airflow.www.utils
255 | get_sensitive_variables_fields, should_hide_value_for_key
| ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ AIR302
256 |
257 | from airflow.datasets.manager import DatasetManager
|
= help: Use `airflow.utils.log.secrets_masker.get_sensitive_variables_fields` instead

Expand All @@ -902,5 +904,106 @@ AIR302_names.py:255:33: AIR302 `airflow.www.utils.should_hide_value_for_key` is
254 | # airflow.www.utils
255 | get_sensitive_variables_fields, should_hide_value_for_key
| ^^^^^^^^^^^^^^^^^^^^^^^^^ AIR302
256 |
257 | from airflow.datasets.manager import DatasetManager
|
= help: Use `airflow.utils.log.secrets_masker.should_hide_value_for_key` instead

AIR302_names.py:260:4: AIR302 `register_dataset_change` is removed in Airflow 3.0
|
259 | dm = DatasetManager()
260 | dm.register_dataset_change()
| ^^^^^^^^^^^^^^^^^^^^^^^ AIR302
261 | dm.create_datasets()
262 | dm.notify_dataset_created()
|
= help: Use `register_asset_change` instead

AIR302_names.py:261:4: AIR302 `create_datasets` is removed in Airflow 3.0
|
259 | dm = DatasetManager()
260 | dm.register_dataset_change()
261 | dm.create_datasets()
| ^^^^^^^^^^^^^^^ AIR302
262 | dm.notify_dataset_created()
263 | dm.notify_dataset_changed()
|
= help: Use `create_assets` instead

AIR302_names.py:262:4: AIR302 `notify_dataset_created` is removed in Airflow 3.0
|
260 | dm.register_dataset_change()
261 | dm.create_datasets()
262 | dm.notify_dataset_created()
| ^^^^^^^^^^^^^^^^^^^^^^ AIR302
263 | dm.notify_dataset_changed()
264 | dm.notify_dataset_alias_created()
|
= help: Use `notify_asset_created` instead

AIR302_names.py:263:4: AIR302 `notify_dataset_changed` is removed in Airflow 3.0
|
261 | dm.create_datasets()
262 | dm.notify_dataset_created()
263 | dm.notify_dataset_changed()
| ^^^^^^^^^^^^^^^^^^^^^^ AIR302
264 | dm.notify_dataset_alias_created()
|
= help: Use `notify_asset_changed` instead

AIR302_names.py:264:4: AIR302 `notify_dataset_alias_created` is removed in Airflow 3.0
|
262 | dm.notify_dataset_created()
263 | dm.notify_dataset_changed()
264 | dm.notify_dataset_alias_created()
| ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ AIR302
|
= help: Use `notify_asset_alias_created` instead

AIR302_names.py:270:5: AIR302 `create_dataset` is removed in Airflow 3.0
|
269 | hlc = HookLineageCollector()
270 | hlc.create_dataset()
| ^^^^^^^^^^^^^^ AIR302
271 | hlc.add_input_dataset()
272 | hlc.add_output_dataset()
|
= help: Use `create_asset` instead

AIR302_names.py:271:5: AIR302 `add_input_dataset` is removed in Airflow 3.0
|
269 | hlc = HookLineageCollector()
270 | hlc.create_dataset()
271 | hlc.add_input_dataset()
| ^^^^^^^^^^^^^^^^^ AIR302
272 | hlc.add_output_dataset()
273 | hlc.collected_datasets()
|
= help: Use `add_input_asset` instead

AIR302_names.py:272:5: AIR302 `add_output_dataset` is removed in Airflow 3.0
|
270 | hlc.create_dataset()
271 | hlc.add_input_dataset()
272 | hlc.add_output_dataset()
| ^^^^^^^^^^^^^^^^^^ AIR302
273 | hlc.collected_datasets()
|
= help: Use `add_output_asset` instead

AIR302_names.py:273:5: AIR302 `collected_datasets` is removed in Airflow 3.0
|
271 | hlc.add_input_dataset()
272 | hlc.add_output_dataset()
273 | hlc.collected_datasets()
| ^^^^^^^^^^^^^^^^^^ AIR302
|
= help: Use `collected_assets` instead

AIR302_names.py:279:5: AIR302 `is_authorized_dataset` is removed in Airflow 3.0
|
278 | aam = AwsAuthManager()
279 | aam.is_authorized_dataset()
| ^^^^^^^^^^^^^^^^^^^^^ AIR302
|
= help: Use `is_authorized_asset` instead
Loading