Add more checks for removed context variables

Author: sunank200

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

@@ -1,4 +1,9 @@
-from airflow.decorators import task
+from airflow.models import DAG
+from airflow.operators.dummy import DummyOperator
+from datetime import datetime
+from airflow.plugins_manager import AirflowPlugin
+from airflow.decorators import task, get_current_context
+from airflow.models.baseoperator import BaseOperator
 
 @task
 def print_config(**context):
@@ -17,3 +22,57 @@ def print_config(**context):
     tomorrow_ds = context["tomorrow_ds"]
     yesterday_ds = context["yesterday_ds"]
     yesterday_ds_nodash = context["yesterday_ds_nodash"]
+
+with DAG(
+    dag_id="example_dag",
+    schedule_interval="@daily",
+    start_date=datetime(2023, 1, 1),
+    template_searchpath=["/templates"],
+) as dag:
+    task1 = DummyOperator(
+        task_id="task1",
+        params={
+            # Removed variables in template
+            "execution_date": "{{ execution_date }}",
+            "next_ds": "{{ next_ds }}",
+            "prev_ds": "{{ prev_ds }}"
+        },
+    )
+
+class CustomMacrosPlugin(AirflowPlugin):
+    name = "custom_macros"
+    macros = {
+        "execution_date_macro": lambda context: context["execution_date"],
+        "next_ds_macro": lambda context: context["next_ds"]
+    }
+
+@task
+def print_config():
+    context = get_current_context()
+    execution_date = context["execution_date"]
+    next_ds = context["next_ds"]
+    next_ds_nodash = context["next_ds_nodash"]
+    next_execution_date = context["next_execution_date"]
+    prev_ds = context["prev_ds"]
+    prev_ds_nodash = context["prev_ds_nodash"]
+    prev_execution_date = context["prev_execution_date"]
+    prev_execution_date_success = context["prev_execution_date_success"]
+    tomorrow_ds = context["tomorrow_ds"]
+    yesterday_ds = context["yesterday_ds"]
+    yesterday_ds_nodash = context["yesterday_ds_nodash"]
+
+class CustomOperator(BaseOperator):
+    def execute(self, context):
+        execution_date = context["execution_date"]
+        next_ds = context["next_ds"]
+        next_ds_nodash = context["next_ds_nodash"]
+        next_execution_date = context["next_execution_date"]
+        prev_ds = context["prev_ds"]
+        prev_ds_nodash = context["prev_ds_nodash"]
+        prev_execution_date = context["prev_execution_date"]
+        prev_execution_date_success = context["prev_execution_date_success"]
+        tomorrow_ds = context["tomorrow_ds"]
+        yesterday_ds = context["yesterday_ds"]
+        yesterday_ds_nodash = context["yesterday_ds_nodash"]
+
+

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

@@ -15,6 +15,7 @@ enum Replacement {
     None,
     Name(&'static str),
     Message(&'static str),
+    ContextRemovalMessage(&'static str),
 }
 
 /// ## What it does
@@ -44,7 +45,6 @@ enum Replacement {
 pub(crate) struct Airflow3Removal {
     deprecated: String,
     replacement: Replacement,
-    is_context_variable: bool,
 }
 
 impl Violation for Airflow3Removal {
@@ -55,19 +55,17 @@ impl Violation for Airflow3Removal {
         let Airflow3Removal {
             deprecated,
             replacement,
-            is_context_variable,
         } = self;
         match replacement {
             Replacement::None => format!("`{deprecated}` is removed in Airflow 3.0"),
             Replacement::Name(_) => {
                 format!("`{deprecated}` is removed in Airflow 3.0")
             }
             Replacement::Message(message) => {
-                if *is_context_variable {
-                    format!("`{deprecated}` {message}")
-                } else {
-                    format!("`{deprecated}` is removed in Airflow 3.0; {message}")
-                }
+                format!("`{deprecated}` is removed in Airflow 3.0; {message}")
+            }
+            Replacement::ContextRemovalMessage(message) => {
+                format!("`{deprecated}` {message}")
             }
         }
     }
@@ -95,7 +93,6 @@ fn diagnostic_for_argument(
                 Some(name) => Replacement::Name(name),
                 None => Replacement::None,
             },
-            is_context_variable: false,
         },
         keyword
             .arg
@@ -213,7 +210,6 @@ fn removed_method(checker: &mut Checker, expr: &Expr) {
             Airflow3Removal {
                 deprecated: attr.to_string(),
                 replacement,
-                is_context_variable: false,
             },
             attr.range(),
         ));
@@ -651,7 +647,6 @@ fn removed_name(checker: &mut Checker, expr: &Expr, ranged: impl Ranged) {
             Airflow3Removal {
                 deprecated,
                 replacement,
-                is_context_variable: false,
             },
             ranged.range(),
         ));
@@ -667,31 +662,32 @@ fn extract_name_from_slice(slice: &Expr) -> Option<String> {
 }
 
 pub(crate) fn removed_context_variable(checker: &mut Checker, expr: &Expr) {
+    const REMOVED_CONTEXT_KEYS: [&str; 12] = [
+        "conf",
+        "execution_date",
+        "next_ds",
+        "next_ds_nodash",
+        "next_execution_date",
+        "prev_ds",
+        "prev_ds_nodash",
+        "prev_execution_date",
+        "prev_execution_date_success",
+        "tomorrow_ds",
+        "yesterday_ds",
+        "yesterday_ds_nodash",
+    ];
+
     if let Expr::Subscript(ExprSubscript { value, slice, .. }) = expr {
         if let Expr::Name(ExprName { id, .. }) = &**value {
             if id.as_str() == "context" {
                 if let Some(key) = extract_name_from_slice(slice) {
-                    const REMOVED_CONTEXT_KEYS: [&str; 11] = [
-                        "execution_date",
-                        "next_ds",
-                        "next_ds_nodash",
-                        "next_execution_date",
-                        "prev_ds",
-                        "prev_ds_nodash",
-                        "prev_execution_date",
-                        "prev_execution_date_success",
-                        "tomorrow_ds",
-                        "yesterday_ds",
-                        "yesterday_ds_nodash",
-                    ];
                     if REMOVED_CONTEXT_KEYS.contains(&key.as_str()) {
                         checker.diagnostics.push(Diagnostic::new(
                             Airflow3Removal {
                                 deprecated: key,
-                                replacement: Replacement::Message(
-                                    "is removed in the Airflow context in Airflow 3.0",
+                                replacement: Replacement::ContextRemovalMessage(
+                                    "is removed in Airflow 3.0.",
                                 ),
-                                is_context_variable: true,
                             },
                             slice.range(),
                         ));
@@ -700,6 +696,23 @@ pub(crate) fn removed_context_variable(checker: &mut Checker, expr: &Expr) {
             }
         }
     }
+
+    if let Expr::StringLiteral(ExprStringLiteral { value, .. }) = expr {
+        let value_str = value.to_string();
+        for key in REMOVED_CONTEXT_KEYS {
+            if value_str.contains(&format!("{{{{ {key} }}}}")) {
+                checker.diagnostics.push(Diagnostic::new(
+                    Airflow3Removal {
+                        deprecated: key.to_string(),
+                        replacement: Replacement::ContextRemovalMessage(
+                            "is removed in Airflow 3.0.",
+                        ),
+                    },
+                    expr.range(),
+                ));
+            }
+        }
+    }
 }
 
 /// AIR302