Skip to content

Add async Variable methods (aget, aset, adelete, akeys) - #72329

Open
dabla wants to merge 10 commits into
apache:mainfrom
dabla:feature/add-async-accessors-variable
Open

dabla wants to merge 10 commits into
apache:mainfrom
dabla:feature/add-async-accessors-variable

Conversation

@dabla

@dabla dabla commented Aug 31, 2026 •

Copy link
Copy Markdown
Contributor

Add async Variable methods (aget, aset, adelete, akeys) with tests

This PR adds async counterparts to all synchronous Variable methods.

What changed

Variable class (task-sdk/src/airflow/sdk/definitions/variable.py)

Four async class-methods are now fully functional:

Async method Sync equivalent Notes
Variable.aget Variable.get Awaits _async_get_variable; falls back to default on VARIABLE_NOT_FOUND
Variable.aset Variable.set Awaits _async_set_variable; checks secrets backends for write conflicts
Variable.adelete Variable.delete Awaits _async_delete_variable; invalidates the secret cache
Variable.akeys Variable.keys Awaits _async_get_variable_keys eagerly (see bug fix below)

Context helpers (task-sdk/src/airflow/sdk/execution_time/context.py)

The four underlying _async_* functions were already implemented:

  • _async_get_variable — checks SecretCache, iterates backends via
    asyncio.to_thread, falls through to the Execution API via asend.
  • _async_set_variable — warns on backend write conflicts via
    asyncio.to_thread, sends PutVariable via asend, invalidates cache.
  • _async_delete_variable — sends DeleteVariable via asend, invalidates cache.
  • _async_get_variable_keys — paginates GetVariableKeys requests via asend.

Tests

test_variables.py

Added TestAsyncVariables and TestAsyncVariableKeys mirroring all existing
sync test cases:

  • test_avar_get — simple and JSON-deserialized values
  • test_avar_set — plain and JSON-serialized values; verifies PutVariable payload
  • test_avar_delete — verifies DeleteVariable is sent via asend
  • test_akeys — prefix filtering (all / with prefix / empty)
  • test_akeys_paginates_when_results_exceed_page_size
  • test_akeys_raises_on_error_response
  • test_akeys_raises_on_unexpected_response_type

test_context.py

Added TestAsyncVariableContext with direct unit tests for the four context helpers
(previously untested):

  • _async_get_variable: from API, from secrets backend, not-found raises
  • _async_set_variable: simple value, JSON-serialized value
  • _async_delete_variable: sends DeleteVariable
  • _async_get_variable_keys: prefix filtering, pagination, error, unexpected response

Was generative AI tooling used to co-author this PR?
  • [ x ] Yes (please specify the tool below)
    GitHub Copilot (Claude Sonnet 4.6)

  • Read the 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) is needed.
  • When adding dependency, check compliance with the ASF 3rd Party License Policy.
  • For significant user-facing changes create newsfragment: {pr_number}.significant.rst, in airflow-core/newsfragments. You can add this file in a follow-up commit after the PR is created so you know the PR number.

@dabla dabla self-assigned this Aug 31, 2026
Fix Variable.akeys: wrapping an async call inside lazy_object_proxy.Proxy
is invalid because await cannot be used inside a plain lambda. The method
now awaits _async_get_variable_keys directly.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
@dabla
dabla force-pushed the feature/add-async-accessors-variable branch from 93418a6 to 6c0e252 Compare August 31, 2026 14:43
@dabla dabla added this to the Airflow 3.4.0 milestone Sep 14, 2026

@amoghrajesh amoghrajesh left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Requesting changes to avoid accidental merge, directionally good, but there's work needed.

Comment thread task-sdk/src/airflow/sdk/execution_time/context.py
Comment thread task-sdk/tests/task_sdk/definitions/test_variables.py Outdated
Comment thread task-sdk/tests/task_sdk/definitions/test_variables.py Outdated
dabla and others added 3 commits September 15, 2026 12:41
…ctly

Avoid importing and reassigning AsyncMock in every test; the fixture's
asend is already an AsyncMock, so tests can set return_value/side_effect
on it directly, per top-level import convention.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…cMock imports

Same cleanup as test_variables.py: set return_value/side_effect directly
on the fixture's mock_supervisor_comms.asend instead of importing and
reassigning AsyncMock per test; move json import to the top of the file.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
@dabla
dabla force-pushed the feature/add-async-accessors-variable branch from af48dc4 to f8be671 Compare September 15, 2026 12:09
@dabla
dabla requested a review from amoghrajesh September 16, 2026 06:43
dabla added a commit to dabla/airflow that referenced this pull request Sep 30, 2026
The error an iterated async task gets for a synchronous SDK call already
points at Variable.aget/aset, which apache#72329 adds; the class docstring
still said Variable had no async equivalent.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

@amoghrajesh amoghrajesh left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looking better, some comments.

Comment thread task-sdk/src/airflow/sdk/execution_time/context.py Outdated
Comment thread task-sdk/src/airflow/sdk/execution_time/context.py
Comment thread task-sdk/src/airflow/sdk/definitions/variable.py
Comment thread task-sdk/src/airflow/sdk/execution_time/context.py
Comment thread task-sdk/src/airflow/sdk/execution_time/context.py
dabla and others added 4 commits September 30, 2026 13:12
json is already imported at module level in context.py, and
_VARIABLE_KEYS_PAGE_SIZE is now imported once at the top of each test module.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Mirror the sync VariableAccessor masking tests for _async_get_variable so
_async_mask_and_deserialize_variable and the SecretCache hit path are covered.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@dabla
dabla requested a review from amoghrajesh September 30, 2026 14:46

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants