[airflow] Add lint rule to show error for removed context variables in airflow

[sunank200]

Summary

Airflow 3.0 removes a set of deprecated context variables that were phased out in 2.x. This PR introduces lint rules to detect usage of these removed variables in various patterns, helping identify incompatibilities. The removed context variables include:

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

Detected Patterns and Examples

The linter now flags the use of removed context variables in the following scenarios:

1. Direct Subscript Access

   execution_date = context["execution_date"]  # Flagged

2. .get("key") Method Calls

   print(context.get("execution_date"))  # Flagged

3. Variables Assigned from get_current_context()
   If a variable is assigned from get_current_context() and then used to access a removed key:

   c = get_current_context()
   print(c.get("execution_date"))  # Flagged

4. Function Parameters in @task-Decorated Functions
   Parameters named after removed context variables in functions decorated with @task are flagged:

   from airflow.decorators import task
   
   @task
   def my_task(execution_date, **kwargs):  # Parameter 'execution_date' flagged
       pass

5. Removed Keys in Task Decorator kwargs and Other Scenarios
   Other similar patterns where removed context variables appear (e.g., as part of kwargs in a @task function) are also detected.

from airflow.decorators import task

@task
def process_with_execution_date(**context):
    execution_date = lambda: context["execution_date"]  # flagged
    print(execution_date)

@task(kwargs={"execution_date": "2021-01-01"})   # flagged
def task_with_kwargs(**context):  
    pass

Test Plan

Test fixtures covering various patterns of deprecated context usage are included in this PR. For example:

from airflow.decorators import task, dag, get_current_context
from airflow.models import DAG
from airflow.operators.dummy import DummyOperator
import pendulum
from datetime import datetime

@task
def access_invalid_key_task(**context):
    print(context.get("conf"))  # 'conf' flagged

@task
def print_config(**context):
    execution_date = context["execution_date"]  # Flagged
    prev_ds = context["prev_ds"]                # Flagged

@task
def from_current_context():
    context = get_current_context()
    print(context["execution_date"])            # Flagged

# Usage outside of a task decorated function
c = get_current_context()
print(c.get("execution_date"))                 # Flagged

@task
def some_task(execution_date, **kwargs):
    print("execution date", execution_date)     # Parameter flagged

@dag(
    start_date=pendulum.datetime(2021, 1, 1, tz="UTC")
)
def my_dag():
    task1 = DummyOperator(
        task_id="task1",
        params={
            "execution_date": "{{ execution_date }}",  # Flagged in template context
        },
    )

    access_invalid_key_task()
    print_config()
    from_current_context()
    
dag = my_dag()

Ruff will emit AIR302 diagnostics for each deprecated usage, with suggestions when applicable, aiding in code migration to Airflow 3.0.

related: apache/airflow#44409, apache/airflow#41641

[github-actions]

ruff-ecosystem results

Linter (stable)

✅ ecosystem check detected no linter changes.

Linter (preview)

✅ ecosystem check detected no linter changes.

[dhruvmanila]

Going to focus on reviewing this PR instead of #15240 for now as I think this one supersedes the other one but please correct me if I'm wrong.

[dhruvmanila]

Thank you for working on this. I've a couple of doubts which I've highlighted in the review comments and #15240 (comment).

[dhruvmanila]

Thank you for updating the PR, I plan on looking at it later today.

[dhruvmanila]

This is looking good, can you update the PR description to include all the checks that are being done? I'm mainly looking for all the structural matching that's being done here, not specific symbols or variables that's being checked. I'm having a hard time keeping track of them :)

[dhruvmanila]

Please re-request for review when it's ready :)

[Lee-W]

Please re-request for review when it's ready :)

I fix the CI failure. I think it's ready for review now. Thanks!

[dhruvmanila]

Thank you for updating the PR description with the patterns that are being checked here, it's helpful as a reference when reviewing.

Regarding 1, 2, 3 pattern, do they need to be in function / method which has the @task decorator or can they be in any function / method?

I think we're pretty close to finishing this up, thank you for your patience!

[dhruvmanila]

This seems to be highlighting the wrong range. I think you meant to highlight the parameter execution_date?

Can we add more test cases to check functions with multiple parameters which has both deprecated and non-deprecated parameters?

[dhruvmanila]

Are these related to the changes made in this PR? This is fine, I just want to confirm.

[dhruvmanila]

I don't think we can use 0 as the default value here. What would happen if it's a positional parameter which is expected to be at position 1 or 2? This will return incorrect argument.

I think we should also add test cases where there are multiple arguments that are deprecated in the same function intermixed with non-deprecated arguments.

[dhruvmanila]

What's the ** check is for? I only see ** being used in parameter in the tests and PR description but this is checking the call function and not the function parameter. Did you meant to look at the definition of value and check that instead? If so, I think that will require adding functionality to get the parameter definition for the binding which would be something like:

fn find_parameter(semantic: &SemanticModel, name: &ast::ExprName) -> Option<ast::AnyParameterRef> {
	let binding_id = semantic.only_binding(name_expr)?;
	let binding = semantic.binding(binding_id);
	let ast::StmtFunctionDef { parameters, .. } = binding.statement(semantic)?.as_function_def_stmt()?;
    parameters
        .iter()
        .find(|parameter| parameter.name().range() == binding.range())
}

[dhruvmanila]

This function doesn't have the @task decorator but the diagnostics is being raised here, is this correct?

[dhruvmanila]

Do we need to update the documentation for this function as it's only checking subscript expressions?