From 14a5a2629e2392aec9a27dc129d58cdb55c12840 Mon Sep 17 00:00:00 2001 From: Wei Lee Date: Fri, 20 Dec 2024 16:30:30 +0900 Subject: [PATCH] [airflow]: extend removed method calls (AIR302) (#15054) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary Airflow 3.0 removes various deprecated functions, members, modules, and other values. They have been deprecated in 2.x, but the removal causes incompatibilities that we want to detect. This PR deprecates the following names and add a function for removed methods * `airflow.datasets.manager.DatasetManager.register_dataset_change` → `airflow.assets.manager.AssetManager.register_asset_change` * `airflow.datasets.manager.DatasetManager.create_datasets` → `airflow.assets.manager.AssetManager.create_assets` * `airflow.datasets.manager.DatasetManager.notify_dataset_created` → `airflow.assets.manager.AssetManager.notify_asset_created` * `airflow.datasets.manager.DatasetManager.notify_dataset_changed` → `airflow.assets.manager.AssetManager.notify_asset_changed` * `airflow.datasets.manager.DatasetManager.notify_dataset_alias_created` → `airflow.assets.manager.AssetManager.notify_asset_alias_created` * `airflow.providers.amazon.auth_manager.aws_auth_manager.AwsAuthManager.is_authorized_dataset` → `airflow.providers.amazon.auth_manager.aws_auth_manager.AwsAuthManager.is_authorized_asset` * `airflow.lineage.hook.HookLineageCollector.create_dataset` → `airflow.lineage.hook.HookLineageCollector.create_asset` * `airflow.lineage.hook.HookLineageCollector.add_input_dataset` → `airflow.lineage.hook.HookLineageCollector.add_input_asset` * `airflow.lineage.hook.HookLineageCollector.add_output_dataset` → `airflow.lineage.hook.HookLineageCollector.dd_output_asset` * `airflow.lineage.hook.HookLineageCollector.collected_datasets` → `airflow.lineage.hook.HookLineageCollector.collected_assets` * `airflow.providers_manager.ProvidersManager.initialize_providers_dataset_uri_resources` → `airflow.providers_manager.ProvidersManager.initialize_providers_asset_uri_resources` ## Test Plan A test fixture is included in the PR. --- .../test/fixtures/airflow/AIR302_names.py | 31 ++++++ .../src/rules/airflow/rules/removal_in_3.rs | 63 +++++++++++ ...irflow__tests__AIR302_AIR302_names.py.snap | 103 ++++++++++++++++++ 3 files changed, 197 insertions(+) diff --git a/crates/ruff_linter/resources/test/fixtures/airflow/AIR302_names.py b/crates/ruff_linter/resources/test/fixtures/airflow/AIR302_names.py index 5eac04ccd7..e3b4939d9a 100644 --- a/crates/ruff_linter/resources/test/fixtures/airflow/AIR302_names.py +++ b/crates/ruff_linter/resources/test/fixtures/airflow/AIR302_names.py @@ -253,3 +253,34 @@ has_access_dataset # 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 diff --git a/crates/ruff_linter/src/rules/airflow/rules/removal_in_3.rs b/crates/ruff_linter/src/rules/airflow/rules/removal_in_3.rs index efb44d70c6..5ed1a122af 100644 --- a/crates/ruff_linter/src/rules/airflow/rules/removal_in_3.rs +++ b/crates/ruff_linter/src/rules/airflow/rules/removal_in_3.rs @@ -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_semantic::analyze::typing; use ruff_python_semantic::Modules; use ruff_text_size::Ranged; @@ -154,6 +155,66 @@ fn removed_argument(checker: &mut Checker, qualname: &QualifiedName, arguments: }; } +fn removed_method(checker: &mut Checker, expr: &Expr) { + let Expr::Call(ExprCall { func, .. }) = expr else { + return; + }; + + let Expr::Attribute(ExprAttribute { attr, value, .. }) = func.as_ref() else { + return; + }; + + let Some(qualname) = typing::resolve_assignment(value, checker.semantic()) else { + return; + }; + + let replacement = match *qualname.segments() { + ["airflow", "datasets", "manager", "DatasetManager"] => match attr.as_str() { + "register_dataset_change" => { + Some(Replacement::Name("register_asset_change".to_string())) + } + "create_datasets" => Some(Replacement::Name("create_assets".to_string())), + "notify_dataset_created" => Some(Replacement::Name("notify_asset_created".to_string())), + "notify_dataset_changed" => Some(Replacement::Name("notify_asset_changed".to_string())), + "notify_dataset_alias_created" => { + Some(Replacement::Name("notify_asset_alias_created".to_string())) + } + &_ => None, + }, + ["airflow", "lineage", "hook", "HookLineageCollector"] => match attr.as_str() { + "create_dataset" => Some(Replacement::Name("create_asset".to_string())), + "add_input_dataset" => Some(Replacement::Name("add_input_asset".to_string())), + "add_output_dataset" => Some(Replacement::Name("add_output_asset".to_string())), + "collected_datasets" => Some(Replacement::Name("collected_assets".to_string())), + &_ => None, + }, + ["airflow", "providers", "amazon", "auth_manager", "aws_auth_manager", "AwsAuthManager"] => { + match attr.as_str() { + "is_authorized_dataset" => { + Some(Replacement::Name("is_authorized_asset".to_string())) + } + &_ => None, + } + } + ["airflow", "providers_manager", "ProvidersManager"] => match attr.as_str() { + "initialize_providers_dataset_uri_resources" => Some(Replacement::Name( + "initialize_providers_asset_uri_resources".to_string(), + )), + &_ => None, + }, + _ => None, + }; + if let Some(replacement) = replacement { + checker.diagnostics.push(Diagnostic::new( + Airflow3Removal { + deprecated: attr.as_str().to_string(), + replacement, + }, + attr.range(), + )); + } +} + fn removed_name(checker: &mut Checker, expr: &Expr, ranged: impl Ranged) { let result = checker @@ -607,6 +668,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, expr); } Expr::Attribute(ExprAttribute { attr: ranged, .. }) => removed_name(checker, expr, ranged), ranged @ Expr::Name(_) => removed_name(checker, expr, ranged), diff --git a/crates/ruff_linter/src/rules/airflow/snapshots/ruff_linter__rules__airflow__tests__AIR302_AIR302_names.py.snap b/crates/ruff_linter/src/rules/airflow/snapshots/ruff_linter__rules__airflow__tests__AIR302_AIR302_names.py.snap index fdf7d98ee9..2400af6ed7 100644 --- a/crates/ruff_linter/src/rules/airflow/snapshots/ruff_linter__rules__airflow__tests__AIR302_AIR302_names.py.snap +++ b/crates/ruff_linter/src/rules/airflow/snapshots/ruff_linter__rules__airflow__tests__AIR302_AIR302_names.py.snap @@ -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 @@ -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