Skip to content

Add get_async_extra_dejson to common.compat connection helpers - #74147

Merged
dabla merged 1 commit into
apache:mainfrom
dabla:feature/common-compat-get-async-extra-dejson
Oct 3, 2026
Merged

dabla merged 1 commit into
apache:mainfrom
dabla:feature/common-compat-get-async-extra-dejson

Conversation

@dabla

@dabla dabla commented Oct 3, 2026 •

Copy link
Copy Markdown
Contributor

Connection.extra_dejson masks the extra's secrets with a synchronous call to the supervisor. When a hook reads it on an event loop with another async call in flight (an async task running items concurrently, a trigger), that raises DeadlockImminentError.

Airflow 3.3.2+ has Connection.aextra_dejson() (#71890, backported in #72895), but providers that still support older Airflow versions can't call it directly. Async hooks therefore deserialize the raw extra themselves (json.loads(conn.extra), e.g. #72130 and the msgraph hook), which skips masking the individual values of the extra.

This adds get_async_extra_dejson(conn) next to get_async_connection() in airflow.providers.common.compat.connection:

from airflow.providers.common.compat.connection import get_async_connection, get_async_extra_dejson

conn = await get_async_connection(self.conn_id)
extras = await get_async_extra_dejson(conn)
Airflow Path Secrets masked
3.3.2+ await conn.aextra_dejson() yes, asynchronously
older (3.x, 2.x) conn.extra_dejson in a worker thread (asgiref.sync_to_async), where blocking on the supervisor is safe — the same fallback get_async_connection() uses for get_connection() yes

Used by these PRs, which are drafts until this one is merged:

Tests cover the aextra_dejson() path (while extra_dejson raises), the fallback running off the event-loop thread, and a real Connection with and without an extra.


Was generative AI tooling used to co-author this PR?
  • Yes — Claude Opus 5.5

Generated-by: Claude Opus 5.5 following the guidelines

Connection.extra_dejson masks the extra's secrets with a synchronous
call to the supervisor, which raises DeadlockImminentError when a hook
reads it on an event loop with another async call in flight. Airflow
3.3.2+ has Connection.aextra_dejson() (apache#71890), but providers that still
support older versions cannot call it directly, so async hooks fall back
to json.loads(conn.extra) and skip the masking (apache#72130).

get_async_extra_dejson(conn) awaits Connection.aextra_dejson() when it
exists, and otherwise runs extra_dejson in a worker thread, where
blocking on the supervisor is safe, as get_async_connection() does for
get_connection().

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@dabla
dabla merged commit 484a87b into apache:main Oct 3, 2026
113 checks passed
dabla added a commit to dabla/airflow that referenced this pull request Oct 3, 2026
Replace the local aextra_dejson fallback with the common.compat helper
added in apache#74147, which masks the extra's secrets asynchronously on
Airflow 3.3.2+ and reads extra_dejson from a worker thread on older
versions.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
dabla added a commit to dabla/airflow that referenced this pull request Oct 3, 2026
Replace the local aextra_dejson fallback with the common.compat helper
added in apache#74147, which masks the extra's secrets asynchronously on
Airflow 3.3.2+ and reads extra_dejson from a worker thread on older
versions.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
dabla added a commit to dabla/airflow that referenced this pull request Oct 12, 2026
Replace the local aextra_dejson fallback with the common.compat helper
added in apache#74147, which masks the extra's secrets asynchronously on
Airflow 3.3.2+ and reads extra_dejson from a worker thread on older
versions.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
dabla added a commit to dabla/airflow that referenced this pull request Oct 14, 2026
Replace the local aextra_dejson fallback with the common.compat helper
added in apache#74147, which masks the extra's secrets asynchronously on
Airflow 3.3.2+ and reads extra_dejson from a worker thread on older
versions.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
dabla added a commit that referenced this pull request Oct 16, 2026
* Add async aget_conn and acreate_embeddings to OpenAIHook

get_conn() resolves the connection with get_connection() and reads
extra_dejson, whose secret masking sends to the supervisor
synchronously; inside an async task with another async SDK call in
flight both raise DeadlockImminentError. aget_conn() returns an
AsyncOpenAI client configured exactly like get_conn()'s, from
get_async_connection() with the extra from Connection.aextra_dejson()
(#71890, Airflow 3.3.2+), or the raw extra via json.loads on older
versions (as in #72130). The client settings move into
_client_kwargs(), shared by both paths. acreate_embeddings() is the
async counterpart of create_embeddings().

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

* Use get_async_extra_dejson from common.compat

Replace the local aextra_dejson fallback with the common.compat helper
added in #74147, which masks the extra's secrets asynchronously on
Airflow 3.3.2+ and reads extra_dejson from a worker thread on older
versions.

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

* Spec the mocked AsyncOpenAI client in test_acreate_embeddings

MagicMock() without a spec accepts any attribute, so a typo on the async
path such as client.embedding would still pass. spec=AsyncOpenAI guards
the client's namespace the way the sync tests' spec=OpenAI does.

* Assert aget_conn reads the extra through aextra_dejson

The connection mock carried an extra string and no test checked that
aextra_dejson() was awaited, so a hook reading json.loads(conn.extra),
which skips the secret masking, kept every test green. Drop extra from
the mock's spec and assert the await in the tests that have it.

* Share the embeddings result shaping between the sync and async paths

acreate_embeddings copied the three lines that pick one vector for a
single input and sort one vector per item for a batch. Move them into
_embeddings_result, used by create_embeddings and acreate_embeddings,
so the two paths cannot drift, like _client_kwargs for the client.

* Resolve the async connection through the hook

get_conn resolves through self.get_connection, so a subclass override
applies there, but aget_conn called get_async_connection without hook=
and the lookup went through BaseHook. Pass hook=self, as
AnthropicHook.get_async_conn does, and test it with a subclass that
overrides aget_connection, through the real common.compat helper.

* Build the async client per call instead of caching it on the hook

_aget_cached_conn checked the cache before two awaits, so concurrent
calls on a fresh hook each built a client, dropped all but one
unclosed, and the kept one was never closed and stayed bound to the
first event loop. aget_conn now builds a client per call and says to
close it, as AnthropicHook.get_async_conn does; acreate_embeddings
closes its client with async with. Tested with three concurrent calls
on one hook.

* State that aget_conn avoids synchronous Task SDK calls from Airflow 3.3.2

Before 3.3.2, core's aget_connection still masks the connection's secrets
synchronously, so the docstring promised more than older cores give.

* Pin the per-call client in the concurrent embeddings test

Patching aget_conn let the test supply distinct clients itself, so it kept
passing if aget_conn cached one client on the hook again.

* Note the token exchange per client for workload identity in aget_conn

The SDK caches the exchanged token on each client, so callers making many
calls should know they can keep one client open and reuse it.

* Require common.compat 1.21.0 for get_async_extra_dejson

The hook imports get_async_extra_dejson at module level, so with an older
common.compat the whole hook, sync get_conn included, fails to import.

* Restore the use-next-version marker on the common.compat dependency

Only release managers raise provider dependency floors; the marker lets
them bump it to the release that ships get_async_extra_dejson.

* Build the embeddings usage as a Usage model in the async test client

mypy checks the annotated _async_client helper and rejects the dict
passed to CreateEmbeddingResponse.usage, which is typed Usage.

* Update providers/openai/src/airflow/providers/openai/hooks/openai.py

Co-authored-by: Kaxil Naik <kaxilnaik@gmail.com>

* Update providers/openai/src/airflow/providers/openai/hooks/openai.py

Co-authored-by: Kaxil Naik <kaxilnaik@gmail.com>

* Update providers/openai/src/airflow/providers/openai/hooks/openai.py

Co-authored-by: Kaxil Naik <kaxilnaik@gmail.com>

* Update providers/openai/tests/unit/openai/hooks/test_openai.py

Co-authored-by: Kaxil Naik <kaxilnaik@gmail.com>

* Drop test docstrings that repeat the test names

---------

Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Co-authored-by: Kaxil Naik <kaxilnaik@gmail.com>
dabla added a commit to dabla/airflow that referenced this pull request Oct 16, 2026
Replace the local aextra_dejson fallback with the common.compat helper
added in apache#74147, which masks the extra's secrets asynchronously on
Airflow 3.3.2+ and reads extra_dejson from a worker thread on older
versions.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
dabla added a commit to dabla/airflow that referenced this pull request Oct 18, 2026
get_async_extra_dejson (apache#74147) keeps the extra's secret masking off the
event loop on every version, but fetching the connection did not: before
3.3.2, BaseHook.aget_connection masks synchronously in core, so async
hooks could still raise DeadlockImminentError on 3.1 to 3.3.1.
shahar1 pushed a commit that referenced this pull request Oct 18, 2026
…4766)

get_async_extra_dejson (#74147) keeps the extra's secret masking off the
event loop on every version, but fetching the connection did not: before
3.3.2, BaseHook.aget_connection masks synchronously in core, so async
hooks could still raise DeadlockImminentError on 3.1 to 3.3.1.
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