feat(AIR301): warn if airflow.lineage.hook.HookLineageCollector.create_asset has positional arguments

Author: Lee-W

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

@@ -96,3 +96,9 @@
 # airflow.secrets.local_filesystem
 lfb = LocalFilesystemBackend()
 lfb.get_connections()
+
+
+hlc.create_asset("there")
+hlc.create_asset("should", "be", "no", "posarg")
+hlc.create_asset(name="but", uri="kwargs are ok")
+hlc.create_asset()

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

@@ -21,6 +21,11 @@ pub(crate) enum Replacement {
     Message(&'static str),
     // The attribute name of a class has been changed.
     AttrName(&'static str),
+    // The attribute name of a class has been changed with extra Message.
+    AttrNameWithMessage {
+        attr_name: &'static str,
+        message: &'static str,
+    },
     // Symbols updated in Airflow 3 with replacement
     // e.g., `airflow.datasets.Dataset` to `airflow.sdk.Asset`
     Rename {

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

@@ -58,8 +58,12 @@ impl Violation for Airflow3Removal {
         } = self;
         match replacement {
             Replacement::None
-            | Replacement::AttrName(_)
             | Replacement::Message(_)
+            | Replacement::AttrName(_)
+            | Replacement::AttrNameWithMessage {
+                attr_name: _,
+                message: _,
+            }
             | Replacement::Rename { module: _, name: _ }
             | Replacement::SourceModuleMoved { module: _, name: _ } => {
                 format!("`{deprecated}` is removed in Airflow 3.0")
@@ -71,8 +75,11 @@ impl Violation for Airflow3Removal {
         let Airflow3Removal { replacement, .. } = self;
         match replacement {
             Replacement::None => None,
-            Replacement::AttrName(name) => Some(format!("Use `{name}` instead")),
             Replacement::Message(message) => Some((*message).to_string()),
+            Replacement::AttrName(name) => Some(format!("Use `{name}` instead")),
+            Replacement::AttrNameWithMessage { attr_name, message } => {
+                Some(format!("Use `{attr_name}` instead; {message}"))
+            }
             Replacement::Rename { module, name } => {
                 Some(format!("Use `{name}` from `{module}` instead."))
             }
@@ -98,7 +105,7 @@ pub(crate) fn airflow_3_removal_expr(checker: &Checker, expr: &Expr) {
             if let Some(qualified_name) = checker.semantic().resolve_qualified_name(func) {
                 check_call_arguments(checker, &qualified_name, arguments);
             }
-            check_method(checker, call_expr);
+            check_method(checker, call_expr, arguments);
             check_context_key_usage_in_call(checker, call_expr);
         }
         Expr::Attribute(attribute_expr @ ExprAttribute { range, .. }) => {
@@ -467,7 +474,7 @@ fn is_kwarg_parameter(semantic: &SemanticModel, name: &ExprName) -> bool {
 /// manager = DatasetManager()
 /// manager.register_datsaet_change()
 /// ```
-fn check_method(checker: &Checker, call_expr: &ExprCall) {
+fn check_method(checker: &Checker, call_expr: &ExprCall, arguments: &Arguments) {
     let Expr::Attribute(ExprAttribute { attr, value, .. }) = &*call_expr.func else {
         return;
     };
@@ -486,10 +493,28 @@ fn check_method(checker: &Checker, call_expr: &ExprCall) {
             _ => return,
         },
         ["airflow", "lineage", "hook", "HookLineageCollector"] => match attr.as_str() {
-            "create_dataset" => Replacement::AttrName("create_asset"),
             "add_input_dataset" => Replacement::AttrName("add_input_asset"),
             "add_output_dataset" => Replacement::AttrName("add_output_asset"),
             "collected_datasets" => Replacement::AttrName("collected_assets"),
+            "create_dataset" => {
+                if arguments.len() > 0 {
+                    Replacement::AttrNameWithMessage {
+                        attr_name: "create_asset",
+                        message: "Calling ``HookLineageCollector.create_asset`` with positional argument should raise an error",
+                    }
+                } else {
+                    Replacement::AttrName("create_asset")
+                }
+            }
+            "create_asset" => {
+                if arguments.len() > 0 {
+                    Replacement::Message(
+                        "Calling ``HookLineageCollector.create_asset`` with positional argument should raise an error",
+                    )
+                } else {
+                    return;
+                }
+            }
             _ => return,
         },
         ["airflow", "providers_manager", "ProvidersManager"] => match attr.as_str() {

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

@@ -54,8 +54,12 @@ impl Violation for Airflow3SuggestedUpdate {
         } = self;
         match replacement {
             Replacement::None
-            | Replacement::AttrName(_)
             | Replacement::Message(_)
+            | Replacement::AttrName(_)
+            | Replacement::AttrNameWithMessage {
+                attr_name: _,
+                message: _,
+            }
             | Replacement::Rename { module: _, name: _ }
             | Replacement::SourceModuleMoved { module: _, name: _ } => {
                 format!(
@@ -70,8 +74,11 @@ impl Violation for Airflow3SuggestedUpdate {
         let Airflow3SuggestedUpdate { replacement, .. } = self;
         match replacement {
             Replacement::None => None,
-            Replacement::AttrName(name) => Some(format!("Use `{name}` instead")),
             Replacement::Message(message) => Some((*message).to_string()),
+            Replacement::AttrName(name) => Some(format!("Use `{name}` instead")),
+            Replacement::AttrNameWithMessage { attr_name, message } => {
+                Some(format!("Use `{attr_name}` instead; {message}"))
+            }
             Replacement::Rename { module, name } => {
                 Some(format!("Use `{name}` from `{module}` instead."))
             }

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

@@ -631,8 +631,43 @@ AIR301 [*] `get_connections` is removed in Airflow 3.0
    |     ^^^^^^^^^^^^^^^
    |
 help: Use `get_connection` instead
-95 | 
-96 | # airflow.secrets.local_filesystem
-97 | lfb = LocalFilesystemBackend()
-   - lfb.get_connections()
-98 + lfb.get_connection()
+95  | 
+96  | # airflow.secrets.local_filesystem
+97  | lfb = LocalFilesystemBackend()
+    - lfb.get_connections()
+98  + lfb.get_connection()
+99  | 
+100 | 
+101 | hlc.create_asset("there")
+
+AIR301 `create_asset` is removed in Airflow 3.0
+   --> AIR301_class_attribute.py:101:5
+    |
+101 | hlc.create_asset("there")
+    |     ^^^^^^^^^^^^
+102 | hlc.create_asset("should", "be", "no", "posarg")
+103 | hlc.create_asset(name="but", uri="kwargs are ok")
+    |
+help: Calling ``HookLineageCollector.create_asset`` with positional argument should raise an error
+
+AIR301 `create_asset` is removed in Airflow 3.0
+   --> AIR301_class_attribute.py:102:5
+    |
+101 | hlc.create_asset("there")
+102 | hlc.create_asset("should", "be", "no", "posarg")
+    |     ^^^^^^^^^^^^
+103 | hlc.create_asset(name="but", uri="kwargs are ok")
+104 | hlc.create_asset()
+    |
+help: Calling ``HookLineageCollector.create_asset`` with positional argument should raise an error
+
+AIR301 `create_asset` is removed in Airflow 3.0
+   --> AIR301_class_attribute.py:103:5
+    |
+101 | hlc.create_asset("there")
+102 | hlc.create_asset("should", "be", "no", "posarg")
+103 | hlc.create_asset(name="but", uri="kwargs are ok")
+    |     ^^^^^^^^^^^^
+104 | hlc.create_asset()
+    |
+help: Calling ``HookLineageCollector.create_asset`` with positional argument should raise an error