udsy19 opened a new pull request, #76171:
URL: https://github.com/apache/airflow/pull/76171
## Why
`context_to_airflow_vars`
(`task-sdk/src/airflow/sdk/execution_time/context.py`) prefixes the keys
returned by `get_airflow_context_vars` (cluster policy /
`airflow_local_settings.py`):
```python
if in_env_var_format and not key_raw.startswith(ENV_VAR_FORMAT_PREFIX):
key = ENV_VAR_FORMAT_PREFIX + key_raw.upper()
elif not key_raw.startswith(DEFAULT_FORMAT_PREFIX):
key = DEFAULT_FORMAT_PREFIX + key_raw
```
In env-var format, a key that already starts with `AIRFLOW_CTX_` skips the
first branch and falls into the *default-format* one, so it gets the dotted
prefix:
```python
# get_airflow_context_vars returns {"AIRFLOW_CTX_CLUSTER": "main"}
context_to_airflow_vars(ctx, in_env_var_format=True)
# -> {"airflow.ctx.AIRFLOW_CTX_CLUSTER": "main"} (expected
{"AIRFLOW_CTX_CLUSTER": "main"})
```
The task runner exports this result to `os.environ` (`task_runner.py`,
`in_env_var_format=True`), so the task gets a dotted variable name instead of
the one the policy defined. Airflow 2 (`airflow/utils/operator_helpers.py` on
`v2-10-stable`) checked each format against its own prefix and kept such keys
as-is; this regressed when the logic moved to the Task SDK.
## Fix
Check the prefix only against the requested format: keep a key that already
has that format's prefix, otherwise add it.
## Tests
Extends `TestAirflowContextHelpers.test_context_to_airflow_vars_from_policy`
in `task-sdk/tests/task_sdk/execution_time/test_context.py` with
already-prefixed keys in both formats. On `main` the env-var case returns
`{'airflow.ctx.AIRFLOW_CTX_CLUSTER': ...}` and fails; with the fix it passes,
and all of `test_context.py` passes.
---
##### Was generative AI tooling used to co-author this PR?
- [X] Yes. I reviewed the change and ran the tests above.
--
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]
