[airflow] Refine and add rules to capture deprecated arguments and a decorator (AIR301)

[sjyangkevin]

Summary

Related: apache/airflow#41641 (comment)

The PR add/refine rules for the following function signature change in Airflow 3.0.

Variable.get()

According to the following documents:
2.11.0: https://airflow.apache.org/docs/apache-airflow/2.11.0/core-concepts/variables.html
3.0.6: https://airflow.apache.org/docs/apache-airflow/3.0.6/core-concepts/variables.html

More context: apache/airflow#41641 (comment)

In Airflow 2.x, the class method get of Variable from airflow.models has a keyword argument named default_var.

from airflow.models import Variable

# Returns the value of default_var (None) if the variable is not set
baz = Variable.get("baz", default_var=None)

In AIR311, there is a rule to detect the migration from airflow.models.Variable to airflow.sdk.Variable. However, in Airflow 3, the change of the keyword argument from default_var to default is not yet captured. The syntax in Airflow 3 is as follow:

from airflow.sdk import Variable

# Returns the value of default (None) if the variable is not set
baz = Variable.get("baz", default=None)

Therefore, a new rule is added in AIR303 to detect the use of default_var when Variable is imported from airflow.sdk. The rule suggests a fix to use the default as argument name instead of default_var.

provide_context is deprecated from python operators

Second, the provide_context parameter is deprecated in Airflow 2 for python operators. It is still a valid syntax to set it in Airflow 2, but it is removed in Airflow 3. A rule is added to detect the presence of the keyword argument in PythonOperator and PythonVirtualenvOperator. The constructor call only accept keyword-only arguments. The rule will raise a warning message when this keyword argument is passed.

# This is ok in Airflow 2, but invalid in Airflow 3
PythonOperator(task_id="task", python_callable=lambda: None, provide_context=True)

Deprecated value for trigger_rule argument in TaskFlow API and Airflow operators

The current rule implemented in check_name only catches TriggerRule.NONE_FAILED_OR_SKIPPED via qualified name resolution. However, the argument is usually passed in the following way. (Here, PythonOperator is an example, all Airflow operators can be configured with a trigger_rule, since they inherit from BaseOperator).

Therefore, the existing rule is removed and re-implemented. There are two cases will be captured. A string value is passed into the @task decorator or the operator, the rule will fix none_failed_or_skipped to none_failed_min_one_success and TriggerRule.NONE_FAILED_OR_SKIPPED to TriggerRule.NONE_FAILED_MIN_ONE_SUCCESS

@task(trigger_rule="none_failed_or_skipped")
def invalid_trigger_rule_task():
    pass

@task(trigger_rule=TriggerRule.NONE_FAILED_OR_SKIPPED)
def invalid_trigger_rule_task():
    pass

@task(trigger_rule="all_success")
def valid_trigger_rule_task():
    pass

PythonOperator(task_id="invalid_trigger_rule_task", python_callable=lambda: None, trigger_rule="none_failed_or_skipped")

PythonOperator(task_id="invalid_trigger_rule_task", python_callable=lambda: None, trigger_rule=TriggerRule.NONE_FAILED_OR_SKIPPED)

PythonOperator(task_id="valid_trigger_rule_task", python_callable=lambda: None, trigger_rule="all_success")

Deprecated apply_defaults decorator

Similar to the above case, this decorator is also checked via qualified name resolution. The decorator is used in the following way. The existing rule is also removed, and re-implemented to capture the below use case, and suggest a fix to remove this decorator entirely.

class DecoratedOperator(BaseOperator):
    # this is deprecated in Airflow 3
    @apply_defaults
    def __init__(self, message, **kwargs):
        super().__init__(**kwargs)
        self.message = message

    def execute(self, context):
        print(self.message)

Test Plan

New test cases have been added to AIR301_args.py and AIR301_decorator.py. The test snapshots are updated.

[astral-sh-bot]

ruff-ecosystem results

Linter (stable)

✅ ecosystem check detected no linter changes.

Linter (preview)

✅ ecosystem check detected no linter changes.

[Lee-W]

Thanks! LGTM

[ntBre]

Thank you!