Add get_async_extra_dejson to common.compat connection helpers by dabla · Pull Request #74147 · apache/airflow

dabla · 2026-10-03T12:44:44Z

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:

Add async aget_conn and acreate_embeddings to OpenAIHook #74149: async aget_conn() / acreate_embeddings() for OpenAIHook
Add async aget_conn and acreate_agent to PydanticAIHook #74150: async aget_conn() / acreate_agent() for PydanticAIHook
Read the msgraph connection extra with get_async_extra_dejson #74161: KiotaRequestAdapterHook (msgraph) reads its extra with it instead of json.loads(connection.extra)

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>

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>

* 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>

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>

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.

…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.

* 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() (apache#71890, Airflow 3.3.2+), or the raw extra via json.loads on older versions (as in apache#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 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> * 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>

* Add async aget_conn and acreate_agent to PydanticAIHook get_conn() resolves the hook's connection and every fallback 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, although the agent the hook builds is async-native. aget_conn() fetches the connection and its fallbacks with get_async_connection() (extra from Connection.aextra_dejson(), #71890, Airflow 3.3.2+, or the raw extra via json.loads on older versions as in #72130), seeds the existing connection cache and builds the model with the unchanged get_conn(); acreate_agent() does the same before create_agent(), spec_file included. 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> * Resolve the async PydanticAIHook paths in a worker thread Prefetching the connections only covered the base hook's own lookups: a subclass's _get_provider_kwargs (Snowflake's builds a SnowflakeHook) still made synchronous Task SDK calls on the event loop, and on Airflow 3.1 to 3.3.1 get_async_connection itself reads extra_dejson synchronously. aget_conn() and acreate_agent() now run the unchanged get_conn() / create_agent() in asyncio.to_thread: from any thread other than the loop's, a synchronous send() waits for the in-flight asend() to release the comms lock instead of raising DeadlockImminentError. This drops the prefetch, the prefetched fallback cache and the extra parameter of _seed_connection. aget_hook() is the async counterpart of get_hook(), with the same connection type check, so an async operator can dispatch to the pydanticai* subclasses. get_hook() now caches the connection it fetched on the hook it returns, so the hook's first get_conn() does not fetch it again. The tests refuse the synchronous lookups only on the event loop, drive a subclass that reads another connection as primary and as fallback, and check that every argument of acreate_agent reaches the agent. The hook docs gain a section on using it from an async task, and uv.lock is back to main's. * Publish the PydanticAIHook connection and its extra together _get_conn_and_extra set _conn before reading extra_dejson, whose secret masking is a supervisor round trip. A second acreate_agent gathered on the same hook found _conn set and read the empty default extra, so it raised "No model specified" or dropped the fallback chain; a failed extra read left the hook half-seeded the same way. Both are now published once both are in hand. get_hook's type check already guarantees a PydanticAIHook, so the isinstance guard before seeding is gone, and _seed_connection's docstring describes its contract instead of one caller. * Run the async PydanticAIHook lookups on a worker thread of their own asyncio.to_thread shared the default executor with asend(), which holds the comms lock while a pool thread reads its reply. Enough gathered hook calls filled the pool: every thread blocked in send() on the lock and nothing was left to run the read, so the task hung. The three async methods now run on a single-worker executor of the hook module, with the caller's context variables copied because task logging binds them. One worker also serializes the hook's sends on Airflow 3.0 and 3.1, where the channel to the supervisor has no lock of its own. The docs section states Airflow 3.2.1 and later (3.2.0 reads the reply outside the lock) and that the loop blocks before 3.3.1, where DeadlockImminentError arrived. * Test the PydanticAIHook context copy and the two-thread cache publish The context copy in _run_off_loop had no test: a plain functools.partial kept every async test green, while it is what carries the task instance fields the task runner binds for logging into the worker thread. A test now sets a context variable before each of the three async methods and checks the worker sees it. With one worker, the two gathered acreate_agent calls ran one after the other, so that test no longer exercised the publish order of _get_conn_and_extra: the old order kept it green. Two threads calling get_conn() on one hook replace it, and the note in _get_conn_and_extra names that case instead of gathered calls. --------- Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

dabla self-assigned this Oct 3, 2026

boring-cyborg Bot added area:providers provider:common-compat labels Oct 3, 2026

dabla requested review from eladkal and shahar1 October 3, 2026 12:44

This was referenced Oct 3, 2026

Add async aget_conn and acreate_embeddings to OpenAIHook #74149

Merged

Add async aget_conn and acreate_agent to PydanticAIHook #74150

Merged

shahar1 approved these changes Oct 3, 2026

View reviewed changes

dabla mentioned this pull request Oct 3, 2026

Read the msgraph connection extra with get_async_extra_dejson #74161

Merged

1 task done

dabla merged commit 484a87b into apache:main Oct 3, 2026
113 checks passed

This was referenced Oct 3, 2026

Test get_async_extra_dejson against the real supervisor comms #74166

Merged

Add AzureAISearchHook to the Microsoft Azure provider #74290

Merged

moomindani mentioned this pull request Oct 7, 2026

Fix AsyncToSync error in Databricks deferrable operators on Airflow 3.0 #74406

Merged

1 task done

potiuk mentioned this pull request Oct 8, 2026

Status of testing Providers that were prepared on October 06, 2026 #74459

Closed

This was referenced Oct 14, 2026

Resolve async connections off the event loop before Airflow 3.3.2 dabla/airflow#1

Closed

Resolve async connections off the event loop before Airflow 3.3.2 #74766

Merged

potiuk mentioned this pull request Oct 17, 2026

Prepare ad-hoc providers release 2026-10-17 #74933

Merged

1 task done

hussein-awala mentioned this pull request Oct 21, 2026

Status of testing Providers that were prepared on October 20, 2026 #75332

Closed

28 of 93 tasks

Navigation Menu

Repository navigation

Add get_async_extra_dejson to common.compat connection helpers - #74147

Add get_async_extra_dejson to common.compat connection helpers#74147
dabla merged 1 commit into
apache:mainfrom
dabla:feature/common-compat-get-async-extra-dejson

dabla commented Oct 3, 2026 •

edited

Loading

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

dabla commented Oct 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Was generative AI tooling used to co-author this PR?

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

dabla commented Oct 3, 2026 •

edited

Loading