sjyangkevin opened a new pull request, #74689:
URL: https://github.com/apache/airflow/pull/74689
<!-- SPDX-License-Identifier: Apache-2.0
https://www.apache.org/licenses/LICENSE-2.0 -->
<!--
Thank you for contributing!
Please provide above a brief description of the changes made in this pull
request.
Write a good git commit message following this guide:
https://chris.beams.io/posts/git-commit/
Please make sure that your code changes are covered with tests.
And in case of new features or big changes remember to adjust the
documentation.
For user-facing UI changes, please attach before/after screenshots (or a
short
screen recording) so reviewers can assess the visual impact.
Feel free to ping (in general) for the review if you do not see reaction for
a few days
(72 Hours is the minimum reaction time you can expect from volunteers) - we
sometimes miss notifications.
In case of an existing issue, reference it using one of the following:
* closes: #ISSUE
* related: #ISSUE
-->
Several Airflow providers offer `@task.<name>` decorators, such as
`@task.python`, `@task.kubernetes`, `@task.docker` and `@task.llm`. Each
decorator builds an operator behind the scenes, so the options it accepts are
defined in that operator's code. The type hints that editors and type checkers
show for these decorators come from somewhere else: a typing file in the Task
SDK, `task-sdk/src/airflow/sdk/definitions/decorators/__init__.pyi`. No hook
currently checks that the two agree, so as operators and decorators evolve, the
typing file can fall out of step without anyone noticing.
This PR adds a static check that compares the typing file with the source
code of every `@task.<name>` decorator registered in the repository, and brings
the typing file up to date.
closes: #48448
related: #51048
## Current state
Keeping the typing file in sync seems to be a manual step today. As the
issue notes, a step like this is easy to overlook, especially since it applies
to decorators spread across several providers. Comparing the typing file with
the current source shows where it has fallen behind:
- **12 registered decorators don't have a typing entry yet:** `agent`,
`analytics`, `llm`, `llm_batch`, `llm_branch`, `llm_file_analysis`,
`llm_schema_compare`, `llm_sql`, `sftp_sensor`, `snowpark`, `sql` and `stub`.
Editors can't suggest their options, and type checkers don't check them.
- **13 entries don't yet include every option their operators accept.** For
example, the two Kubernetes entries don't include about a dozen
`KubernetesPodOperator` options, such as `callbacks`, `logging_interval` and
`runtime_class_name`.
- **A few entries list options the decorators don't accept.** `@task.docker`
lists `use_dill`, which that decorator doesn't take. `@task.llm_schema_compare`
lists `output_type`, which the operator rejects with a `TypeError`, and
`serialize_output`, whose value the operator always sets itself. In these cases
a type checker accepts code that fails, or has no effect, at runtime.
## What this PR changes
1. Adds a prek hook, `check-task-decorator-arguments`, that fails whenever a
decorator's entry in the typing file doesn't match the options the decorator
actually accepts.
2. Brings the typing file up to date: adds entries for the 12 decorators
that don't have one, and updates 13 existing entries.
3. Updates the reminder comment in `KubernetesPodOperator` to point at the
typing file.
## How the check works
The check needs no list of decorators of its own. For each decorator it
follows the same trail a reader would follow by hand:
1. **Find the decorators.** It reads the `task-decorators` section of every
`provider.yaml`, which is the same list Airflow uses at runtime to find
`@task.<name>` decorators. A new provider or decorator is picked up
automatically.
2. **Find the operator each decorator builds.** It opens the registered
decorator function and finds the class passed as `decorated_operator_class`.
That is the operator created for every task the decorator produces.
3. **Follow the operator's parent classes.** It resolves each parent class
by reading the module's imports, including the compat modules that re-export
classes. This works across providers: the Snowflake decorator's operator, for
example, inherits from the standard provider's `PythonOperator`. It puts the
classes in the order Python itself would call their constructors, and stops at
`BaseOperator`. `BaseOperator` passes nothing further and rejects leftover
arguments, so no class after it can ever receive a user option.
4. **Collect every accepted option.** Each constructor along that path
contributes its parameters, as do the decorator function and Airflow's shared
decorator factory.
5. **Drop the options the decorator fills in itself.** Some options are
always set by the decorator, so users can't pass them:
- Kubernetes always sets `cmds`.
- Bash, SQL and the LLM decorators set `bash_command`, `sql` and `prompt`
from the value the decorated function returns.
- `LLMSchemaCompareOperator` always forces `serialize_output=True`.
The check recognises these because the value passed on is fixed, rather
than taken from the user's input. A value read from the user's arguments, such
as `name=kwargs.pop("name", default)`, stays a user option.
6. **Compare with the typing entry.** It reports options that are missing
from the entry and options the entry lists but the decorator doesn't accept.
Each message names the file and line where the option is defined or filled in.
Everything is read as source code with Python's `ast` module. Nothing is
imported, so the check runs on a clean checkout, without provider dependencies
or Breeze, in about 2 seconds. Example output:
```
__init__.pyi:542: @task.kubernetes is missing parameters: logging_interval.
Declared in providers/cncf/kubernetes/.../operators/pod.py:325.
__init__.pyi:382: @task.docker declares unexpected parameters: command.
command is internally supplied at providers/docker/.../decorators/docker.py:125.
```
## When the check fails
| Situation | What the check reports | What the contributor does |
|---|---|---|
| A new provider ships a decorator, or a provider adds one | The decorator
has no typing entry, or its entry is missing options, each with its source line
| Add the entry |
| A decorator is removed from `provider.yaml` | The typing file still has an
entry for it | Delete the entry |
| A decorator is renamed | The new name needs an entry, and the old entry is
left over | Rename the entry |
| `provider.yaml` points at a missing file or function | The missing file or
function | Fix the registration |
| An operator gains, loses or renames an option | Every decorator built on
it, including decorators in other providers (a `PythonOperator` change affects
9 decorators across 3 providers) | Update those entries |
| An operator gets a new parent class | The new parent's options are picked
up automatically | Update the entries |
| A decorator starts filling an option itself | The option is now
unexpected, with the line where it is filled | Remove it from the entry |
| Code is written in a way the check can't follow, e.g. picking the operator
class at runtime | An explanation of what it couldn't resolve. It never
silently passes | Use the usual pattern |
None of these cases require editing the check itself. Contributors only
touch the typing file, and the message says exactly what to change. The check
also requires plain `@task` to accept the same options as `@task.python`.
<details>
<summary><b>Example runs</b>: real output, one edit per scenario</summary>
Each scenario starts from this branch, makes one edit, and runs `prek run
check-task-decorator-arguments --all-files`. Long paths and option lists are
shortened with `…`.
On this branch as-is:
```
Check registered task decorator coverage and argument declarations....Passed
```
**1. A provider registers a new decorator without a typing entry.** This
scenario uses the `provider.yaml` entry and decorator module from open PR
#72033 (`@task.eks_pod`):
```
Check registered task decorator coverage and argument declarations....Failed
- hook id: check-task-decorator-arguments
- exit code: 1
@task.eks_pod: registered without a named typing declaration. Update
task-sdk/…/decorators/__init__.pyi.
@task.eks_pod: task-sdk/…/decorators/__init__.pyi: expected one
configuration overload TaskDecoratorCollection.eks_pod, found 0
```
After an empty `eks_pod` entry is added, the check lists every option to
declare and where each one comes from, here across two providers:
```
…/__init__.pyi:63: @task.eks_pod is missing parameters:
active_deadline_seconds, affinity, annotations, …, aws_conn_id, …,
cluster_name, …, region, …, xcom_sidecar_container_security_context. Declared
in providers/amazon/…/decorators/eks.py:72,
providers/amazon/…/operators/eks.py:1136,
providers/cncf/kubernetes/…/operators/pod.py:325. Update …/__init__.pyi.
```
**2. An operator gains an option.** Adding `cache_result: bool = False` to
`PythonOperator.__init__` flags every decorator built on it, across three
providers:
```
…/__init__.pyi:226: @task.branch is missing parameters: cache_result.
Declared in providers/standard/…/operators/python.py:183. Update …/__init__.pyi.
…/__init__.pyi:310: @task.branch_external_python is missing parameters:
cache_result. Declared in providers/standard/…/operators/python.py:183. Update
…/__init__.pyi.
…/__init__.pyi:246: @task.branch_virtualenv is missing parameters:
cache_result. Declared in providers/standard/…/operators/python.py:183. Update
…/__init__.pyi.
…/__init__.pyi:177: @task.external_python is missing parameters:
cache_result. Declared in providers/standard/…/operators/python.py:183. Update
…/__init__.pyi.
…/__init__.pyi:962: @task.pyspark is missing parameters: cache_result.
Declared in providers/standard/…/operators/python.py:183. Update …/__init__.pyi.
…/__init__.pyi:64: @task.python is missing parameters: cache_result.
Declared in providers/standard/…/operators/python.py:183. Update …/__init__.pyi.
…/__init__.pyi:360: @task.short_circuit is missing parameters:
cache_result. Declared in providers/standard/…/operators/python.py:183. Update
…/__init__.pyi.
…/__init__.pyi:1032: @task.snowpark is missing parameters: cache_result.
Declared in providers/standard/…/operators/python.py:183. Update …/__init__.pyi.
…/__init__.pyi:107: @task.virtualenv is missing parameters: cache_result.
Declared in providers/standard/…/operators/python.py:183. Update …/__init__.pyi.
```
**3. A typing entry lists an option the decorator always fills itself.**
Adding `cmds` to the `@task.kubernetes` entry:
```
…/__init__.pyi:542: @task.kubernetes declares unexpected parameters: cmds.
cmds is internally supplied at
providers/cncf/kubernetes/…/decorators/kubernetes.py:72. Update …/__init__.pyi
to match
providers/cncf/kubernetes/…/decorators/kubernetes.py:_KubernetesDecoratedOperator
and the constructors it inherits.
```
**4. A decorator is removed.** Deleting the `task-decorators` block from the
Docker provider's `provider.yaml`:
```
@task.docker: typing declaration has no provider registration. Update
task-sdk/…/__init__.pyi.
```
**5. Code the check can't follow.** Choosing the operator class at runtime
with `decorated_operator_class=_get_operator_class()`:
```
@task.docker: providers/docker/…/decorators/docker.py:213: cannot resolve
decorated_operator_class _get_operator_class()
```
**When the check deliberately passes:**
**6. A new option on `BaseOperator`.** Adding `new_shared_option` to
`BaseOperator.__init__` → `Passed`. Shared options go through each entry's
`**kwargs`, so the 25 entries don't need editing.
**7. A new operator option that the decorator fills itself.** Adding
`script_path` to `BashOperator`, with `_BashDecoratedOperator` passing
`script_path=SET_DURING_EXECUTION` → `Passed`. Users can't set the option, so
the entry must not list it. Without the decorator's line, the same edit fails
with `@task.bash is missing parameters: script_path`.
**8. Changing an option's type or default.** Changing `image: str | None =
None` to `image: str = "busybox"` in `KubernetesPodOperator` → `Passed`. Only
option names are compared; see below.
</details>
## What the check does not cover
- **Option types and default values.** Only option names are compared. A few
entries on `main` have types narrower than their operators. For example,
`@task.sensor` declares `poke_interval: float` although `timedelta` works.
Those can be fixed in a follow-up.
- **Options shared by every operator** (`task_id`, `retries`, `pool`, …).
The typing entries deliberately accept these through their trailing `**kwargs`
instead of listing about 55 options in every entry. The check keeps that
convention: it stops at `BaseOperator`, and it reports a shared option listed
in a single entry as unexpected.
- **`@task.run_if` and `@task.skip_if`.** They wrap another task rather than
creating an operator, and their signatures are defined in the Task SDK itself.
- **Decorators from providers outside this repository.** They still work at
runtime, but this repository's typing file can't describe them.
- **Options an operator reads from `**kwargs` without declaring them.** This
happens in 3 operator classes across the providers (two Databricks sensors and
one Cloud Composer sensor). None of them has a decorator today.
- **Other decorator families**, such as `@dag`, `@task_group` and `@asset`.
## Impact on CI and pre-commit
- **CI:** the static checks job already runs `prek --all-files`, so this is
one more hook taking about 2 seconds. It needs no Breeze image and imports no
provider code.
- **Local commits:** the hook uses `always_run: true`, so it runs on every
commit, adding about 2 seconds. This is needed because deleted files are not
passed to hooks, and removing a decorator's files would otherwise go unnoticed.
The check always scans all decorators, because a change to one shared file,
such as `PythonOperator`, the decorator factory or the typing file itself, can
affect decorators in several providers.
- **When it can block a PR:** when that PR leaves the typing file out of
sync, which is the intended case. It can also block a PR when two open PRs each
pass on their own but conflict once both are merged. For example, open PR
#72033 registers a new `@task.eks_pod` decorator without a typing entry, so
whichever of the two PRs merges second will need to add that entry. The error
message lists exactly what to add.
- **Coverage today:** all 25 `@task.<name>` decorators across 8 providers.
Every module that builds a decorator with the shared factory is registered in a
`provider.yaml`, so nothing is left unchecked.
## Question for reviewers: newsfragment?
The typing changes alter what type checkers accept for `@task.<name>` calls.
Twelve decorators now have typed signatures, and options the decorators don't
accept were removed. Runtime behaviour does not change. Since the Task SDK
ships with airflow-core, would you like an `airflow-core/newsfragments` entry
for this, and if so of which type? I haven't added one, following the guideline
to add newsfragments only when the change is clearly user-facing.
## Testing
- `uv run --project scripts pytest scripts/tests/`: 1,814 passed, 1 skipped.
This includes 133 tests for the new hook. Each rule in the check was also
disabled in turn to confirm that a test fails (13 of 13).
- `prek run --from-ref <base> --stage pre-commit`: passed.
- `prek run --from-ref <base> --stage manual --skip
upgrade-important-versions`: passed, except `mypy-providers`, which couldn't
run because my local Breeze CI image is out of date. The only provider file it
would check is `pod.py`, where the only change is a comment. CI runs it.
- `prek run mypy-task-sdk --all-files`: passed. This type-checks the typing
file.
- `breeze build-docs task-sdk`: passed. The Task SDK API docs are generated
from the typing file.
- Ran the derivation over all 1,262 operator and sensor classes in the
providers: every one resolved through to `BaseOperator`.
- Not run locally: the core and provider test suites `breeze verify` selects
(Kubernetes, Amazon, Google, …). They are selected because of the `.pyi` file
and a comment-only change in `pod.py`, neither of which changes runtime
behaviour. CI runs them.
---
##### Was generative AI tooling used to co-author this PR?
- [X] Yes — Claude Code (Opus 5.5)
Generated-by: Claude Code (Opus 5.5) following [the
guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions)
---
* Read the **[Pull Request
Guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#pull-request-guidelines)**
for more information. Note: commit author/co-author name and email in commits
become permanently public when merged.
* For fundamental code changes, an Airflow Improvement Proposal
([AIP](https://cwiki.apache.org/confluence/display/AIRFLOW/Airflow+Improvement+Proposals))
is needed.
* When adding dependency, check compliance with the [ASF 3rd Party License
Policy](https://www.apache.org/legal/resolved.html#category-x).
* For significant user-facing changes create newsfragment:
`{pr_number}.significant.rst`, in
[airflow-core/newsfragments](https://github.com/apache/airflow/tree/main/airflow-core/newsfragments).
You can add this file in a follow-up commit after the PR is created so you
know the PR number.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
