[airflow]: Import modules that has been moved to airflow providers (A…

…IR303) (#14764) ## Summary Many core Airflow features have been deprecated and moved to Airflow Providers since users might need to install an additional package (e.g., `apache-airflow-provider-fab==1.0.0`); a separate rule (AIR303) is created for this. As some of the changes only relate to the module/package moved, instead of listing out all the functions, variables, and classes in a module or a package, it warns the user to import from the new path instead of the specific name. The following is the ones that has been moved to `apache-airflow-provider-fab==1.0.0` * module moved * `airflow.api.auth.backend.basic_auth` → `airflow.providers.fab.auth_manager.api.auth.backend.basic_auth` * `airflow.api.auth.backend.kerberos_auth` → `airflow.providers.fab.auth_manager.api.auth.backend.kerberos_auth` * `airflow.auth.managers.fab.api.auth.backend.kerberos_auth` → `airflow.providers.fab.auth_manager.api.auth.backend.kerberos_auth` * `airflow.auth.managers.fab.security_manager.override` → `airflow.providers.fab.auth_manager.security_manager.override` * classes (e.g., functions, classes) moved * `airflow.www.security.FabAirflowSecurityManagerOverride` → `airflow.providers.fab.auth_manager.security_manager.override.FabAirflowSecurityManagerOverride` * `airflow.auth.managers.fab.fab_auth_manager.FabAuthManager` → `airflow.providers.fab.auth_manager.security_manager.FabAuthManager` ## Test Plan A test fixture has been included for the rule.
astral-sh · Dec 13, 2024 · dfd7f38 · dfd7f38
1 parent e96b13c
commit dfd7f38
Show file tree

Hide file tree

Showing 8 changed files with 283 additions and 0 deletions.
diff --git a/crates/ruff_linter/resources/test/fixtures/airflow/AIR303.py b/crates/ruff_linter/resources/test/fixtures/airflow/AIR303.py
@@ -0,0 +1,16 @@
+from airflow.api.auth.backend import basic_auth, kerberos_auth
+from airflow.api.auth.backend.basic_auth import auth_current_user
+from airflow.auth.managers.fab.api.auth.backend import (
+    kerberos_auth as backend_kerberos_auth,
+)
+from airflow.auth.managers.fab.fab_auth_manager import FabAuthManager
+from airflow.auth.managers.fab.security_manager import override as fab_override
+from airflow.www.security import FabAirflowSecurityManagerOverride
+
+basic_auth, kerberos_auth
+auth_current_user
+backend_kerberos_auth
+fab_override
+
+FabAuthManager
+FabAirflowSecurityManagerOverride
diff --git a/crates/ruff_linter/src/checkers/ast/analyze/expression.rs b/crates/ruff_linter/src/checkers/ast/analyze/expression.rs
@@ -223,6 +223,9 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
                     if checker.enabled(Rule::Airflow3Removal) {
                         airflow::rules::removed_in_3(checker, expr);
                     }
+                    if checker.enabled(Rule::Airflow3MovedToProvider) {
+                        airflow::rules::moved_to_provider_in_3(checker, expr);
+                    }
 
                     // Ex) List[...]
                     if checker.any_enabled(&[

diff --git a/crates/ruff_linter/src/codes.rs b/crates/ruff_linter/src/codes.rs
@@ -1046,6 +1046,7 @@ pub fn code_to_rule(linter: Linter, code: &str) -> Option<(RuleGroup, Rule)> {
         (Airflow, "001") => (RuleGroup::Stable, rules::airflow::rules::AirflowVariableNameTaskIdMismatch),
         (Airflow, "301") => (RuleGroup::Preview, rules::airflow::rules::AirflowDagNoScheduleArgument),
         (Airflow, "302") => (RuleGroup::Preview, rules::airflow::rules::Airflow3Removal),
+        (Airflow, "303") => (RuleGroup::Preview, rules::airflow::rules::Airflow3MovedToProvider),
 
         // perflint
         (Perflint, "101") => (RuleGroup::Stable, rules::perflint::rules::UnnecessaryListCast),

diff --git a/crates/ruff_linter/src/rules/airflow/mod.rs b/crates/ruff_linter/src/rules/airflow/mod.rs
@@ -16,6 +16,7 @@ mod tests {
     #[test_case(Rule::AirflowDagNoScheduleArgument, Path::new("AIR301.py"))]
     #[test_case(Rule::Airflow3Removal, Path::new("AIR302_args.py"))]
     #[test_case(Rule::Airflow3Removal, Path::new("AIR302_names.py"))]
+    #[test_case(Rule::Airflow3MovedToProvider, Path::new("AIR303.py"))]
     fn rules(rule_code: Rule, path: &Path) -> Result<()> {
         let snapshot = format!("{}_{}", rule_code.noqa_code(), path.to_string_lossy());
         let diagnostics = test_path(

diff --git a/crates/ruff_linter/src/rules/airflow/rules/mod.rs b/crates/ruff_linter/src/rules/airflow/rules/mod.rs
@@ -1,7 +1,9 @@
 pub(crate) use dag_schedule_argument::*;
+pub(crate) use moved_to_provider_in_3::*;
 pub(crate) use removal_in_3::*;
 pub(crate) use task_variable_name::*;
 
 mod dag_schedule_argument;
+mod moved_to_provider_in_3;
 mod removal_in_3;
 mod task_variable_name;
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
@@ -0,0 +1,185 @@
+use ruff_diagnostics::{Diagnostic, Violation};
+use ruff_macros::{derive_message_formats, ViolationMetadata};
+use ruff_python_ast::{Expr, ExprAttribute};
+use ruff_python_semantic::Modules;
+use ruff_text_size::Ranged;
+
+use crate::checkers::ast::Checker;
+
+#[derive(Debug, Eq, PartialEq)]
+enum Replacement {
+    ProviderName {
+        name: &'static str,
+        provider: &'static str,
+        version: &'static str,
+    },
+    ImportPathMoved {
+        original_path: &'static str,
+        new_path: &'static str,
+        provider: &'static str,
+        version: &'static str,
+    },
+}
+
+/// ## What it does
+/// Checks for uses of Airflow functions and values that have been moved to it providers.
+/// (e.g., apache-airflow-providers-fab)
+///
+/// ## Why is this bad?
+/// Airflow 3.0 moved various deprecated functions, members, and other
+/// values to its providers. The user needs to install the corresponding provider and replace
+/// the original usage with the one in the provider
+///
+/// ## Example
+/// ```python
+/// from airflow.auth.managers.fab.fab_auth_manage import FabAuthManager
+/// ```
+///
+/// Use instead:
+/// ```python
+/// from airflow.providers.fab.auth_manager.fab_auth_manage import FabAuthManager
+/// ```
+#[derive(ViolationMetadata)]
+pub(crate) struct Airflow3MovedToProvider {
+    deprecated: String,
+    replacement: Replacement,
+}
+
+impl Violation for Airflow3MovedToProvider {
+    #[derive_message_formats]
+    fn message(&self) -> String {
+        let Airflow3MovedToProvider {
+            deprecated,
+            replacement,
+        } = self;
+        match replacement {
+            Replacement::ProviderName {
+                name: _,
+                provider,
+                version: _,
+            } => {
+                format!("`{deprecated}` is moved into `{provider}` provider in Airflow 3.0;")
+            }
+            Replacement::ImportPathMoved {
+                original_path,
+                new_path: _,
+                provider,
+                version: _,
+            } => {
+                format!("Import path `{original_path}` is moved into `{provider}` provider in Airflow 3.0;")
+            }
+        }
+    }
+
+    fn fix_title(&self) -> Option<String> {
+        let Airflow3MovedToProvider { replacement, .. } = self;
+        if let Replacement::ProviderName {
+            name,
+            provider,
+            version,
+        } = replacement
+        {
+            Some(format!(
+                "Install `apache-airflow-provider-{provider}>={version}` and use `{name}` instead."
+            ))
+        } else if let Replacement::ImportPathMoved {
+            original_path: _,
+            new_path,
+            provider,
+            version,
+        } = replacement
+        {
+            Some(format!("Install `apache-airflow-provider-{provider}>={version}` and import from `{new_path}` instead."))
+        } else {
+            None
+        }
+    }
+}
+
+fn moved_to_provider(checker: &mut Checker, expr: &Expr, ranged: impl Ranged) {
+    let result =
+        checker
+            .semantic()
+            .resolve_qualified_name(expr)
+            .and_then(|qualname| match qualname.segments() {
+                // apache-airflow-providers-fab
+                ["airflow", "www", "security", "FabAirflowSecurityManagerOverride"] => Some((
+                    qualname.to_string(),
+                    Replacement::ProviderName {
+                        name: "airflow.providers.fab.auth_manager.security_manager.override.FabAirflowSecurityManagerOverride",
+                        provider: "fab",
+                        version: "1.0.0"
+                    },
+                )),
+                ["airflow","auth","managers","fab","fab_auth_manager", "FabAuthManager"] => Some((
+                    qualname.to_string(),
+                    Replacement::ProviderName{
+                        name: "airflow.providers.fab.auth_manager.security_manager.FabAuthManager",
+                        provider: "fab",
+                        version: "1.0.0"
+                },
+                )),
+                ["airflow", "api", "auth", "backend", "basic_auth", ..] => Some((
+                    qualname.to_string(),
+                    Replacement::ImportPathMoved{
+                        original_path: "airflow.api.auth.backend.basic_auth",
+                        new_path: "airflow.providers.fab.auth_manager.api.auth.backend.basic_auth",
+                        provider:"fab",
+                        version: "1.0.0"
+                },
+                )),
+                ["airflow", "api","auth","backend","kerberos_auth", ..] => Some((
+                    qualname.to_string(),
+                    Replacement::ImportPathMoved{
+                        original_path:"airflow.api.auth.backend.kerberos_auth",
+                        new_path: "airflow.providers.fab.auth_manager.api.auth.backend.kerberos_auth",
+                        provider: "fab",
+                        version:"1.0.0"
+                },
+                )),
+                ["airflow", "auth", "managers", "fab", "api", "auth", "backend", "kerberos_auth", ..] => Some((
+                    qualname.to_string(),
+                    Replacement::ImportPathMoved{
+                        original_path: "airflow.auth_manager.api.auth.backend.kerberos_auth",
+                        new_path: "airflow.providers.fab.auth_manager.api.auth.backend.kerberos_auth",
+                        provider: "fab",
+                        version: "1.0.0"
+                },
+                )),
+                ["airflow","auth","managers","fab","security_manager","override", ..] => Some((
+                    qualname.to_string(),
+                    Replacement::ImportPathMoved{
+                        original_path: "airflow.auth.managers.fab.security_manager.override",
+                        new_path: "airflow.providers.fab.auth_manager.security_manager.override",
+                        provider: "fab",
+                        version: "1.0.0"
+                },
+                )),
+
+                _ => None,
+            });
+    if let Some((deprecated, replacement)) = result {
+        checker.diagnostics.push(Diagnostic::new(
+            Airflow3MovedToProvider {
+                deprecated,
+                replacement,
+            },
+            ranged.range(),
+        ));
+    }
+}
+
+/// AIR303
+pub(crate) fn moved_to_provider_in_3(checker: &mut Checker, expr: &Expr) {
+    if !checker.semantic().seen_module(Modules::AIRFLOW) {
+        return;
+    }
+
+    match expr {
+        Expr::Attribute(ExprAttribute { attr: ranged, .. }) => {
+            moved_to_provider(checker, expr, ranged);
+        }
+        ranged @ Expr::Name(_) => moved_to_provider(checker, expr, ranged),
+        _ => {}
+    }
+}
diff --git a/...ter/src/rules/airflow/snapshots/ruff_linter__rules__airflow__tests__AIR303_AIR303.py.snap b/...ter/src/rules/airflow/snapshots/ruff_linter__rules__airflow__tests__AIR303_AIR303.py.snap
@@ -0,0 +1,74 @@
+---
+source: crates/ruff_linter/src/rules/airflow/mod.rs
+snapshot_kind: text
+---
+AIR303.py:10:1: AIR303 Import path `airflow.api.auth.backend.basic_auth` is moved into `fab` provider in Airflow 3.0;
+   |
+ 8 | from airflow.www.security import FabAirflowSecurityManagerOverride
+ 9 | 
+10 | basic_auth, kerberos_auth
+   | ^^^^^^^^^^ AIR303
+11 | auth_current_user
+12 | backend_kerberos_auth
+   |
+   = help: Install `apache-airflow-provider-fab>=1.0.0` and import from `airflow.providers.fab.auth_manager.api.auth.backend.basic_auth` instead.
+
+AIR303.py:10:13: AIR303 Import path `airflow.api.auth.backend.kerberos_auth` is moved into `fab` provider in Airflow 3.0;
+   |
+ 8 | from airflow.www.security import FabAirflowSecurityManagerOverride
+ 9 | 
+10 | basic_auth, kerberos_auth
+   |             ^^^^^^^^^^^^^ AIR303
+11 | auth_current_user
+12 | backend_kerberos_auth
+   |
+   = help: Install `apache-airflow-provider-fab>=1.0.0` and import from `airflow.providers.fab.auth_manager.api.auth.backend.kerberos_auth` instead.
+
+AIR303.py:11:1: AIR303 Import path `airflow.api.auth.backend.basic_auth` is moved into `fab` provider in Airflow 3.0;
+   |
+10 | basic_auth, kerberos_auth
+11 | auth_current_user
+   | ^^^^^^^^^^^^^^^^^ AIR303
+12 | backend_kerberos_auth
+13 | fab_override
+   |
+   = help: Install `apache-airflow-provider-fab>=1.0.0` and import from `airflow.providers.fab.auth_manager.api.auth.backend.basic_auth` instead.
+
+AIR303.py:12:1: AIR303 Import path `airflow.auth_manager.api.auth.backend.kerberos_auth` is moved into `fab` provider in Airflow 3.0;
+   |
+10 | basic_auth, kerberos_auth
+11 | auth_current_user
+12 | backend_kerberos_auth
+   | ^^^^^^^^^^^^^^^^^^^^^ AIR303
+13 | fab_override
+   |
+   = help: Install `apache-airflow-provider-fab>=1.0.0` and import from `airflow.providers.fab.auth_manager.api.auth.backend.kerberos_auth` instead.
+
+AIR303.py:13:1: AIR303 Import path `airflow.auth.managers.fab.security_manager.override` is moved into `fab` provider in Airflow 3.0;
+   |
+11 | auth_current_user
+12 | backend_kerberos_auth
+13 | fab_override
+   | ^^^^^^^^^^^^ AIR303
+14 | 
+15 | FabAuthManager
+   |
+   = help: Install `apache-airflow-provider-fab>=1.0.0` and import from `airflow.providers.fab.auth_manager.security_manager.override` instead.
+
+AIR303.py:15:1: AIR303 `airflow.auth.managers.fab.fab_auth_manager.FabAuthManager` is moved into `fab` provider in Airflow 3.0;
+   |
+13 | fab_override
+14 | 
+15 | FabAuthManager
+   | ^^^^^^^^^^^^^^ AIR303
+16 | FabAirflowSecurityManagerOverride
+   |
+   = help: Install `apache-airflow-provider-fab>=1.0.0` and use `airflow.providers.fab.auth_manager.security_manager.FabAuthManager` instead.
+
+AIR303.py:16:1: AIR303 `airflow.www.security.FabAirflowSecurityManagerOverride` is moved into `fab` provider in Airflow 3.0;
+   |
+15 | FabAuthManager
+16 | FabAirflowSecurityManagerOverride
+   | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ AIR303
+   |
+   = help: Install `apache-airflow-provider-fab>=1.0.0` and use `airflow.providers.fab.auth_manager.security_manager.override.FabAirflowSecurityManagerOverride` instead.
diff --git a/ruff.schema.json b/ruff.schema.json