kaxil commented on code in PR #74150:
URL: https://github.com/apache/airflow/pull/74150#discussion_r4353453778
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -318,6 +341,21 @@ def get_conn(self) -> Model:
)
return self._model
+ async def aget_conn(self) -> Model:
+ """
+ Return the configured pydantic-ai ``Model``, for use inside an async
task.
+
+ The same model as :meth:`get_conn`, fallback chain included, resolved
in a worker
+ thread. Resolving it reads connections with ``get_connection`` and
``extra_dejson``,
+ and a subclass's :meth:`_get_provider_kwargs` may read more
(Snowflake's builds a
+ ``SnowflakeHook``); each of those sends to the supervisor
synchronously, which raises
+ ``DeadlockImminentError`` on the event loop while another async call
is in flight,
+ but only waits its turn from any other thread.
+ """
+ if self._model is not None:
+ return self._model
+ return await asyncio.to_thread(self.get_conn)
Review Comment:
Nothing serializes these thread hops, which the sync path never needed. With
`await asyncio.gather(hook.acreate_agent(instructions="a"),
hook.acreate_agent(instructions="b"))` on one hook, the second thread can find
`_conn` already set while the first is still in the `MaskSecret` round trip
inside `extra_dejson`, so it reads an empty extra and either raises "No model
specified" or drops the fallback chain. On 3.0 and 3.1 `send()` has no lock at
all (it arrived in 3.2.0), so gathered calls on two different hooks write to
the same socket at once; wrapping the call in `_call_locked`, as
`AirflowToolset.run_blocking` already does for that reason, would cover all
three methods.
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -206,7 +208,28 @@ def get_hook(cls, conn_id: str, hook_params: dict | None =
None):
f"one of {expected_types}. Change the connection type to
'pydanticai', put the API key "
"in the API Key field (stored as the password) and the model
(e.g. 'openai:gpt-5') in Model."
)
- return conn.get_hook(hook_params=hook_params)
+ hook = conn.get_hook(hook_params=hook_params)
+ if isinstance(hook, PydanticAIHook):
+ # Saves the second fetch of this same connection the hook's first
``get_conn``
+ # would make.
+ hook._seed_connection(conn)
Review Comment:
Small one: `_seed_connection`'s docstring still describes
`_resolve_fallback_models` as its caller, so `get_hook` should be named there
too.
--
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]
