[airflow] Add ruff rules to catch deprecated Airflow imports for Airflow 3.1 (AIR321)#22376
Merged
ntBre merged 18 commits intoastral-sh:mainfrom Feb 4, 2026
Merged
Conversation
|
| code | total | + violation | - violation | + fix | - fix |
|---|---|---|---|---|---|
| AIR321 | 462 | 462 | 0 | 0 | 0 |
sjyangkevin
commented
Jan 4, 2026
.../rules/airflow/snapshots/ruff_linter__rules__airflow__tests__AIR321_AIR321_names_fix.py.snap
Outdated
Show resolved
Hide resolved
38 tasks
Member
|
@Lee-W could you take a look at the PR if it catches the semantics you want. |
Lee-W
reviewed
Jan 5, 2026
Contributor
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! |
sjyangkevin
force-pushed
the
catch-deprecated-imports-airflow-3_1
branch
from
5f3ef77 to
1d8cbe3
Compare
amoghrajesh
reviewed
Jan 6, 2026
crates/ruff_linter/resources/test/fixtures/airflow/AIR321_names_fix.py
Outdated
Show resolved
Hide resolved
crates/ruff_linter/resources/test/fixtures/airflow/AIR321_names_fix.py
Outdated
Show resolved
Hide resolved
crates/ruff_linter/resources/test/fixtures/airflow/AIR321_names_fix.py
Outdated
Show resolved
Hide resolved
crates/ruff_linter/resources/test/fixtures/airflow/AIR321_names_fix.py
Outdated
Show resolved
Hide resolved
Contributor
|
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! |
Contributor
Author
|
@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
force-pushed
the
catch-deprecated-imports-airflow-3_1
branch
from
1d8cbe3 to
0a88188
Compare
sjyangkevin
force-pushed
the
catch-deprecated-imports-airflow-3_1
branch
from
6e876dc to
e786ba2
Compare
sjyangkevin
commented
Jan 14, 2026
Contributor
Author
There was a problem hiding this comment.
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.
- Exclude the followings
airflow.utils.task_group.get_task_group_children_getter
airflow.utils.task_group.task_group_to_dict
- Fix Import Path in AIR311
airflow.sensors.base.poke_mode_only → airflow.sdk.bases.sensor.poke_mode_only
- Move from AIR301 to AIR321
airflow.secrets.cache.SecretCache → from airflow.sdk import SecretCache
- Keep
_internalmodules and show a warning message to indicate these are internal API that may change without notice. - Update the other rules to adapt for the new Replacement
SourceModuleMovedToSDK
Thanks!!
.../src/rules/airflow/snapshots/ruff_linter__rules__airflow__tests__AIR301_AIR301_names.py.snap
Show resolved
Hide resolved
Contributor
A Diagnotic without fix is basically a warning |
Contributor
Author
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
force-pushed
the
catch-deprecated-imports-airflow-3_1
branch
from
e786ba2 to
75e00c6
Compare
sjyangkevin
commented
Jan 15, 2026
Contributor
Author
There was a problem hiding this comment.
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.
- Move
SecretCachefrom AIR301 to AIR321 (tested in Airflow 3.0.6 and the original import is still valid) - Move the test case for
SecretCachefrom AIR301 to AIR321 - For
_internalmodules, the rule is updated to useMessageto show a warning. Hence, thewarning_messagefield is removed from theSourceModuleMovedToSDK. This field is only used for the internal module cases, but will beNonefor majority. - Update the
fix_titleforSourceModuleMovedToSDKto "{name}has been moved to{module}since Airflow 3.x (with apache-airflow-task-sdk>={version}).", where version is1.0.6for rules in AIR301, and1.1.6for rules in AIR321. - For AIR321, update
preview_sinceto0.14.12 - Test snapshots are all updated.
- Comments in apache/airflow#54714 (comment) are resolved.
Thanks!
Lee-W
reviewed
Jan 16, 2026
sjyangkevin
force-pushed
the
catch-deprecated-imports-airflow-3_1
branch
from
8791001 to
6ef68e7
Compare
sjyangkevin
commented
Jan 16, 2026
Comment on lines
43
to
48
| // Symbols moved to internal module in Airflow 3. Used when we want to raise a warning. | ||
| // e.g., `airflow.utils.setup_teardown.BaseSetupTeardownContext` to `airflow.sdk.definitions._internal.setup_teardown.BaseSetupTeardownContext` | ||
| InternalModule { | ||
| module: &'static str, | ||
| name: String, | ||
| }, |
Contributor
Author
There was a problem hiding this comment.
add a new item in the enum which used to show warning message for internal module
.../src/rules/airflow/snapshots/ruff_linter__rules__airflow__tests__AIR321_AIR321_names.py.snap
Outdated
Show resolved
Hide resolved
sjyangkevin
force-pushed
the
catch-deprecated-imports-airflow-3_1
branch
from
6ef68e7 to
332b175
Compare
1 task
sjyangkevin
force-pushed
the
catch-deprecated-imports-airflow-3_1
branch
from
332b175 to
3e093db
Compare
sjyangkevin
commented
Jan 19, 2026
| name, | ||
| suggest_fix, | ||
| .. | ||
| } if *suggest_fix => (module, name.as_str()), |
Contributor
Author
There was a problem hiding this comment.
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
commented
Jan 19, 2026
Comment on lines
+43
to
+49
| // Symbols updated in Airflow 3 with only module changed. Used when we want to include custom message and optionally report diagnostics. | ||
| SourceModuleMovedWithMessage { | ||
| module: &'static str, | ||
| name: String, | ||
| message: &'static str, | ||
| suggest_fix: bool, | ||
| }, |
Contributor
Author
There was a problem hiding this comment.
Create this new enum item which can be used as follow:
- The
fix_titleformat is: "{name}has been moved to{module}since Airflow 3.1.{message}"; we can use themessageparameter to include a static string as additional message that we want to show to users. - The
suggest_fixparameter 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.,BaseHookis 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,
},
… and move SecretCache test case to AIR321
…report diagnostic
sjyangkevin
force-pushed
the
catch-deprecated-imports-airflow-3_1
branch
from
c948314 to
a513e55
Compare
sjyangkevin
commented
Jan 28, 2026
Contributor
Author
There was a problem hiding this comment.
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
commented
Jan 28, 2026
| 78 | ds_format("2026-01-01", "%Y-%m-%d", "%m-%d-%y") | ||
| 79 | datetime_diff_for_humans( | ||
| | | ||
| help: `ds_add` has been moved to `airflow.sdk.execution_time.macros` since Airflow 3.1. Requires `apache-airflow-task-sdk>=1.1.0,<=1.1.6`. For `apache-airflow-task-sdk>=1.1.7`, import from `airflow.sdk` instead. |
Contributor
Author
There was a problem hiding this comment.
the message is refined a bit, hope it is better than before. the version is updated to >=1.1.7
Lee-W
reviewed
Jan 29, 2026
| 78 | ds_format("2026-01-01", "%Y-%m-%d", "%m-%d-%y") | ||
| 79 | datetime_diff_for_humans( | ||
| | | ||
| help: `ds_add` has been moved to `airflow.sdk.execution_time.macros` since Airflow 3.1. Requires `apache-airflow-task-sdk>=1.1.0,<=1.1.6`. For `apache-airflow-task-sdk>=1.1.7`, import from `airflow.sdk` instead. |
Contributor
Author
|
Hi @MichaReiser , @ntBre , would appreciate if we could get your help to review again. Thanks! |
Contributor
|
Will do! I should be able to take a look this week :) |
ntBre
pushed a commit
that referenced
this pull request
Feb 9, 2026
…context key for Airflow 3.0 (`AIR301`) (#22850) ## Summary <!-- What's the purpose of the change? What does it do, and why? --> Context: apache/airflow#41641 1. apache/airflow#45961 * <strike>create_dagrun removed from airflow...DAG</strike> (This has already been implemented in AIR301) * context key dag_run.external_trigger removed 2. apache/airflow#45960 * context["inlet_events"]["url"] → context["inlet_events"][Asset("url")] 3. apache/airflow#41348 * context key triggering_dataset_events → triggering_asset_events The existing AIR301 rules can detect when users access a removed key such as `execution_date` through Airflow's `context`. For example, `context["execution_date"]`, or `context["dag_run"]`. However, if an attribute is deprecated from a context key, such as `context["dag_run"].external_trigger`, the current implement will not flag it. This PR adds the logic for such check, and add two rules to flag the deprecated attribute for the `"dag_run"` and `"inlet_events"` context key. In addition to this, `"triggering_dataset_events"` is a deprecated context key which can be handled by the existing rule. However, the existing rule doesn't raise a diagnostic. Hence, the rule logic is refactored a little bit, such that we can add this check and suggest a `Replacement::Rename`. ## Test Plan <!-- How was it tested? --> The test cases have been added to `AIR301_context.py`, and all the tests have been run locally and success. @Lee-W , could you please review it when you have time, thanks! ## Notes In #22376, we introduced some improvements to the AIR301 code. I will re-base this PR when we are all good on that, so it can pick up those code structure improvements, and updated rules.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
5 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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 ofAIR301. The rules check whether a removed Airflow name is used, and match onExpr::NameandExpr::Attribute.Test Plan
The following two test files are added:
AIR321_names.pyAll 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.AIR321_names_fix.pyAll the test cases in this file raise NO violation (i.e., all checks should pass). The snapshot file is empty.
Document Update