[airflow] Add ruff rules to catch deprecated Airflow imports for Airflow 3.1 (AIR321)

[sjyangkevin]

Summary

This PR is related to the discussion: apache/airflow#54714

This change creates a new code (AIR321) and implement ruff rules to catch, and/or fix deprecated imports in Airflow for Airflow 3.1. The rules are implemented by following the structure of AIR301. The rules check whether a removed Airflow name is used, and match on Expr::Name and Expr::Attribute.

Test Plan

The following two test files are added:

1. crates/ruff_linter/resources/test/fixtures/airflow/AIR321_names.py
2. crates/ruff_linter/resources/test/fixtures/airflow/AIR321_names_fix.py

AIR321_names.py
All the test cases in this file should raise violations and fixes should be suggested when running the test with --unsafe-fixes. The test results shown in the snapshot are expected.

cargo run -p ruff -- check crates/ruff_linter/resources/test/fixtures/airflow/AIR321_names.py --no-cache --preview --select AIR321 --unsafe-fixes

AIR321_names_fix.py
All the test cases in this file raise NO violation (i.e., all checks should pass). The snapshot file is empty.

Document Update

[astral-sh-bot]

ruff-ecosystem results

Linter (stable)

✅ ecosystem check detected no linter changes.

Linter (preview)

ℹ️ ecosystem check detected linter changes. (+462 -0 violations, +0 -0 fixes in 1 projects; 54 projects unchanged)

apache/airflow (+462 -0 violations, +0 -0 fixes)

ruff check --no-cache --exit-zero --no-fix --output-format concise --preview --select ALL

+ airflow-core/tests/unit/ti_deps/deps/test_task_concurrency.py:35:16: AIR321 `airflow.models.baseoperator.BaseOperator` is moved in Airflow 3.1
+ dev/airflow_perf/scheduler_dag_execution_timing.py:181:24: AIR321 [*] `airflow.utils.timezone.utcnow` is moved in Airflow 3.1
+ kubernetes-tests/tests/kubernetes_tests/test_kubernetes_pod_operator.py:59:20: AIR321 [*] `airflow.utils.timezone.datetime` is moved in Airflow 3.1
+ performance/src/performance_dags/performance_dag/performance_dag_utils.py:83:27: AIR321 [*] `airflow.utils.timezone.utcnow` is moved in Airflow 3.1
+ providers/alibaba/tests/unit/alibaba/cloud/log/test_oss_task_handler.py:59:26: AIR321 `airflow.utils.timezone.datetime` is moved in Airflow 3.1
+ providers/alibaba/tests/unit/alibaba/cloud/sensors/test_analyticdb_spark.py:26:16: AIR321 [*] `airflow.utils.timezone.datetime` is moved in Airflow 3.1
+ providers/amazon/src/airflow/providers/amazon/aws/executors/utils/exponential_backoff_retry.py:72:20: AIR321 [*] `airflow.utils.timezone.utcnow` is moved in Airflow 3.1
+ providers/amazon/tests/system/amazon/aws/example_mongo_to_s3.py:45:16: AIR321 `airflow.utils.timezone.datetime` is moved in Airflow 3.1
+ providers/amazon/tests/system/amazon/aws/example_mwaa.py:121:35: AIR321 [*] `airflow.utils.timezone.utc` is moved in Airflow 3.1
+ providers/amazon/tests/system/amazon/aws/example_mwaa.py:151:35: AIR321 [*] `airflow.utils.timezone.utc` is moved in Airflow 3.1
+ providers/amazon/tests/unit/amazon/aws/executors/ecs/test_ecs_executor.py:881:77: AIR321 `airflow.utils.timezone.utcnow` is moved in Airflow 3.1
+ providers/amazon/tests/unit/amazon/aws/hooks/test_s3.py:325:25: AIR321 `airflow.utils.timezone.datetime` is moved in Airflow 3.1
+ providers/amazon/tests/unit/amazon/aws/hooks/test_s3.py:326:23: AIR321 `airflow.utils.timezone.datetime` is moved in Airflow 3.1
+ providers/amazon/tests/unit/amazon/aws/hooks/test_s3.py:585:74: AIR321 `airflow.utils.timezone.datetime` is moved in Airflow 3.1
+ providers/amazon/tests/unit/amazon/aws/hooks/test_s3.py:586:75: AIR321 `airflow.utils.timezone.datetime` is moved in Airflow 3.1
+ providers/amazon/tests/unit/amazon/aws/hooks/test_s3.py:601:66: AIR321 `airflow.utils.timezone.datetime` is moved in Airflow 3.1
+ providers/amazon/tests/unit/amazon/aws/hooks/test_s3.py:602:67: AIR321 `airflow.utils.timezone.datetime` is moved in Airflow 3.1
+ providers/amazon/tests/unit/amazon/aws/hooks/test_s3.py:694:74: AIR321 `airflow.utils.timezone.datetime` is moved in Airflow 3.1
+ providers/amazon/tests/unit/amazon/aws/hooks/test_s3.py:695:75: AIR321 `airflow.utils.timezone.datetime` is moved in Airflow 3.1
+ providers/amazon/tests/unit/amazon/aws/hooks/test_s3.py:725:74: AIR321 `airflow.utils.timezone.datetime` is moved in Airflow 3.1
+ providers/amazon/tests/unit/amazon/aws/hooks/test_s3.py:726:75: AIR321 `airflow.utils.timezone.datetime` is moved in Airflow 3.1
+ providers/amazon/tests/unit/amazon/aws/hooks/test_s3.py:762:41: AIR321 `airflow.utils.timezone.datetime` is moved in Airflow 3.1
+ providers/amazon/tests/unit/amazon/aws/hooks/test_s3.py:800:74: AIR321 `airflow.utils.timezone.datetime` is moved in Airflow 3.1
+ providers/amazon/tests/unit/amazon/aws/hooks/test_s3.py:801:75: AIR321 `airflow.utils.timezone.datetime` is moved in Airflow 3.1
+ providers/amazon/tests/unit/amazon/aws/hooks/test_s3.py:928:41: AIR321 `airflow.utils.timezone.datetime` is moved in Airflow 3.1
+ providers/amazon/tests/unit/amazon/aws/hooks/test_s3.py:934:41: AIR321 `airflow.utils.timezone.datetime` is moved in Airflow 3.1
+ providers/amazon/tests/unit/amazon/aws/hooks/test_s3.py:945:41: AIR321 `airflow.utils.timezone.datetime` is moved in Airflow 3.1
+ providers/amazon/tests/unit/amazon/aws/hooks/test_s3.py:951:41: AIR321 `airflow.utils.timezone.datetime` is moved in Airflow 3.1
+ providers/amazon/tests/unit/amazon/aws/hooks/test_s3.py:995:41: AIR321 `airflow.utils.timezone.datetime` is moved in Airflow 3.1
+ providers/amazon/tests/unit/amazon/aws/log/test_cloudwatch_task_handler.py:141:26: AIR321 `airflow.utils.timezone.datetime` is moved in Airflow 3.1
+ providers/amazon/tests/unit/amazon/aws/log/test_cloudwatch_task_handler.py:195:16: AIR321 `airflow.utils.timezone.datetime` is moved in Airflow 3.1
+ providers/amazon/tests/unit/amazon/aws/log/test_cloudwatch_task_handler.py:261:26: AIR321 `airflow.utils.timezone.datetime` is moved in Airflow 3.1
+ providers/amazon/tests/unit/amazon/aws/log/test_cloudwatch_task_handler.py:334:17: AIR321 `airflow.utils.timezone.datetime` is moved in Airflow 3.1
+ providers/amazon/tests/unit/amazon/aws/log/test_cloudwatch_task_handler.py:335:42: AIR321 `airflow.utils.timezone.datetime` is moved in Airflow 3.1
+ providers/amazon/tests/unit/amazon/aws/log/test_cloudwatch_task_handler.py:389:33: AIR321 `airflow.utils.timezone.datetime` is moved in Airflow 3.1
+ providers/amazon/tests/unit/amazon/aws/operators/test_s3.py:651:24: AIR321 `airflow.utils.timezone.utcnow` is moved in Airflow 3.1
+ providers/amazon/tests/unit/amazon/aws/operators/test_s3.py:652:42: AIR321 `airflow.utils.timezone.datetime` is moved in Airflow 3.1
+ providers/amazon/tests/unit/amazon/aws/operators/test_s3.py:676:27: AIR321 `airflow.utils.timezone.utcnow` is moved in Airflow 3.1
+ providers/amazon/tests/unit/amazon/aws/operators/test_s3.py:707:15: AIR321 `airflow.utils.timezone.utcnow` is moved in Airflow 3.1
+ providers/amazon/tests/unit/amazon/aws/operators/test_s3.py:796:17: AIR321 `airflow.utils.timezone.datetime` is moved in Airflow 3.1
+ providers/amazon/tests/unit/amazon/aws/operators/test_s3.py:803:17: AIR321 `airflow.utils.timezone.datetime` is moved in Airflow 3.1
+ providers/amazon/tests/unit/amazon/aws/operators/test_s3.py:804:17: AIR321 `airflow.utils.timezone.datetime` is moved in Airflow 3.1
+ providers/amazon/tests/unit/amazon/aws/operators/test_s3.py:832:17: AIR321 `airflow.utils.timezone.datetime` is moved in Airflow 3.1
+ providers/amazon/tests/unit/amazon/aws/operators/test_s3.py:839:17: AIR321 `airflow.utils.timezone.datetime` is moved in Airflow 3.1
+ providers/amazon/tests/unit/amazon/aws/operators/test_s3.py:840:17: AIR321 `airflow.utils.timezone.datetime` is moved in Airflow 3.1
+ providers/amazon/tests/unit/amazon/aws/transfers/test_s3_to_sftp.py:47:16: AIR321 `airflow.utils.timezone.datetime` is moved in Airflow 3.1
+ providers/amazon/tests/unit/amazon/aws/transfers/test_sftp_to_s3.py:45:16: AIR321 `airflow.utils.timezone.datetime` is moved in Airflow 3.1
+ providers/apache/flink/tests/unit/apache/flink/operators/test_flink_kubernetes.py:200:51: AIR321 [*] `airflow.utils.timezone.datetime` is moved in Airflow 3.1
+ providers/apache/flink/tests/unit/apache/flink/sensors/test_flink_kubernetes.py:886:51: AIR321 [*] `airflow.utils.timezone.datetime` is moved in Airflow 3.1
+ providers/apache/hdfs/tests/unit/apache/hdfs/log/test_hdfs_task_handler.py:37:16: AIR321 `airflow.utils.timezone.datetime` is moved in Airflow 3.1
... 412 additional changes omitted for project

Changes by rule (1 rules affected)

code	total	+ violation	- violation	+ fix	- fix
AIR321	462	462	0	0	0

[MichaReiser]

@Lee-W could you take a look at the PR if it catches the semantics you want.

[Lee-W]

@Lee-W could you take a look at the PR if it catches the semantics you want.

At a high level, yes, but some details will need some polish. Will let you know when it's at least ok from my side. Thanks!

[Lee-W]

The logic looks solid, but we’ll need to revisit the list. @sjyangkevin, could you please make this a draft? We can discuss the list further in relation to the Airflow issue. Thanks!

[sjyangkevin]

@Lee-W and @amoghrajesh , thanks for the feedback! I have converted it into a draft. I will also happy to work on the fix after the further discussion.

[sjyangkevin]

Hi @Lee-W and @amoghrajesh ,

I've refined the rules based on the feedback. I not sure if we can make a rule a "warning" in ruff; currently, I introduce a new Replacement SourceModuleMovedToSDK which we can use to embed a warning message. Let me know if it is the right approach.

Below is the summary of changes.

1. Exclude the followings

airflow.utils.task_group.get_task_group_children_getter
airflow.utils.task_group.task_group_to_dict

2. Fix Import Path in AIR311

airflow.sensors.base.poke_mode_only → airflow.sdk.bases.sensor.poke_mode_only

3. Move from AIR301 to AIR321

airflow.secrets.cache.SecretCache → from airflow.sdk import SecretCache

4. Keep _internal modules and show a warning message to indicate these are internal API that may change without notice.
5. Update the other rules to adapt for the new Replacement SourceModuleMovedToSDK

Thanks!!

[Lee-W]

I not sure if we can make a rule a "warning" in ruff

A Diagnotic without fix is basically a warning

[sjyangkevin]

I not sure if we can make a rule a "warning" in ruff

A Diagnotic without fix is basically a warning

Thanks! The current implementation proposes a fix and include a warning message. I can refactor this to use a simple message to notify the module move, and include a message to tell these APIs are subject to change (don't encourage using these). So, we don't suggest fix for it. Which one do you think could be a better idea?

Current Implementation suggest fixes and show warning message.

[sjyangkevin]

Hi @Lee-W and @amoghrajesh , I made a few adjustments to the PR and I think it is almost ready for open it again for review. Let me know if you have any feedback before opening it again. Below is the summary of changes made.

1. Move SecretCache from AIR301 to AIR321 (tested in Airflow 3.0.6 and the original import is still valid)
2. Move the test case for SecretCache from AIR301 to AIR321
3. For _internal modules, the rule is updated to use Message to show a warning. Hence, the warning_message field is removed from the SourceModuleMovedToSDK. This field is only used for the internal module cases, but will be None for majority.
4. Update the fix_title for SourceModuleMovedToSDK to "{name} has been moved to {module} since Airflow 3.x (with apache-airflow-task-sdk>={version}).", where version is 1.0.6 for rules in AIR301, and 1.1.6 for rules in AIR321.
5. For AIR321, update preview_since to 0.14.12
6. Test snapshots are all updated.
7. Comments in apache/airflow#54714 (comment) are resolved.

Thanks!

[sjyangkevin]

add a new item in the enum which used to show warning message for internal module

[sjyangkevin]

Here, we leveraged the suggest_fix parameter from SourceModuleMovedWithMessage to handle when we would like to show a warning and when we would like to suggest a fix with custom message.

[sjyangkevin]

Create this new enum item which can be used as follow:

1. The fix_title format is: "{name} has been moved to {module} since Airflow 3.1. {message}"; we can use the message parameter to include a static string as additional message that we want to show to users.
2. The suggest_fix parameter is used to optionally report diagnostics. In some cases, we might want to just show a warning (e.g., internal module usage), but in some case, we might want to report diagnostic and suggest fix (e.g., BaseHook is moved from one module to another)

Why message is a static string. Use the following example to explain. Making message a String allow us to use format! to templating the string and render it at runtime. However, only rest or variables defined in the accessible scope can be used to render the string. As module and name will be handled in fix_title. It might not be necessary to use those values to render the message. So, here use static string.

[
            "airflow",
            "utils",
            "setup_teardown",
            rest @ ("BaseSetupTeardownContext" | "SetupTeardownContext"),
        ] => Replacement::SourceModuleMovedWithMessage {
            module: "airflow.sdk.definitions._internal.setup_teardown",
            name: rest.to_string(),
            message: "This is an internal module which is not suggested to be used and is subject to change without notice.",
            suggest_fix: false,
        },

[sjyangkevin]

I also found modules moved to sdk here. Update the Rename/SourceModuleMoved to SourceModuleMovedToSDK, and the task-sdk version is 1.0.0.

Test Command:

breeze shell --use-airflow-version 3.0.0 --python 3.12 --backend postgres --db-reset

[sjyangkevin]

the message is refined a bit, hope it is better than before. the version is updated to >=1.1.7

[Lee-W]

it's better :)

[Lee-W]

LGTM :)

[Lee-W]

it's better :)

[sjyangkevin]

Hi @MichaReiser , @ntBre , would appreciate if we could get your help to review again. Thanks!

[ntBre]

Will do! I should be able to take a look this week :)

[ntBre]

Looks good, thank you!