Repository navigation
Add async aget_conn and acreate_embeddings to OpenAIHook - #74149
Merged
Merged
Conversation
kaxil
reviewed
Oct 8, 2026
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>
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>
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.
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.
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.
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.
_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.
dabla
force-pushed
the
feature/openai-hook-async-conn
branch
from
October 12, 2026 09:48
47f89bb to
c0c1940
Compare
kaxil
reviewed
Oct 12, 2026
YAshhh29
reviewed
Oct 14, 2026
1 task done
….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.
1 task done
Patching aget_conn let the test supply distinct clients itself, so it kept passing if aget_conn cached one client on the hook again.
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.
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.
Only release managers raise provider dependency floors; the marker lets them bump it to the release that ships get_async_extra_dejson.
mypy checks the annotated _async_client helper and rejects the dict passed to CreateEmbeddingResponse.usage, which is typed Usage.
kaxil
reviewed
Oct 15, 2026
kaxil
approved these changes
Oct 15, 2026
Co-authored-by: Kaxil Naik <kaxilnaik@gmail.com>
Co-authored-by: Kaxil Naik <kaxilnaik@gmail.com>
Co-authored-by: Kaxil Naik <kaxilnaik@gmail.com>
Co-authored-by: Kaxil Naik <kaxilnaik@gmail.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
OpenAIHookonly has a synchronous path:get_conn()resolves the connection withget_connection()and readsextra_dejson, whose secret masking sendsMaskSecretto the supervisor synchronously. Inside an async task with another async SDK call in flight (e.g. items running concurrently), both raiseDeadlockImminentError, so the hook can't be used from async code.This adds an async path that mirrors the sync one:
aget_conn()returns anAsyncOpenAIclient configured exactly likeget_conn()'s client (openai_client_kwargs,base_url,api_key,auth_typeapi_key/workload_identity). The connection comes fromget_async_connection(hook=self), so a subclass override ofaget_connection()applies asget_connection()does forget_conn(), and its extra fromget_async_extra_dejson()(common.compat, Add get_async_extra_dejson to common.compat connection helpers #74147):Connection.aextra_dejson()on Airflow 3.3.2+ (Fix DeadlockImminentError when a connection is resolved inside an async task #71890),extra_dejsonfrom a worker thread on older versions. Resolving the connection without any synchronous Task SDK call is guaranteed from Airflow 3.3.2: before that, core'saget_connection()still masks the connection's secrets synchronously, which a separate common.compat change would have to cover._client_kwargs()helper shared byget_conn()andaget_conn(), so the two paths can't drift.acreate_embeddings()is the async counterpart ofcreate_embeddings(), same overloads and return shape (one vector for a single input, one per item for a batch), on a client built byaget_conn()for the call and closed when it returns. Eachaget_conn()call builds a new client, which the caller closes withasync withorawait client.close(), asAnthropicHook.get_async_conn()does.Further async counterparts (chat completions, responses, ...) can follow the same pattern.
Tests: the existing
get_conntests pass unchanged; new tests coveraget_conn(API key in password/extra, client kwargs, workload identity, invalidauth_type, a subclass overridingaget_connection(), and a connection withoutaextra_dejson()as on Airflow < 3.3.2, whoseextra_dejsonmay only run off the event loop) with a connection whoseextra_dejsonraises on the loop and whoseaextra_dejson()must have been awaited, a syncget_connectionthat raises, andacreate_embeddingsfor text, token and batch input, plus three concurrent calls on one hook each closing their own client.Was generative AI tooling used to co-author this PR?
Generated-by: Claude Opus 5.5 following the guidelines