dabla commented on code in PR #74150:
URL: https://github.com/apache/airflow/pull/74150#discussion_r4362074796
##########
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:
The half-seeded read was real: `_get_conn_and_extra` published `_conn`
before reading `extra_dejson`, so a second gathered call saw the empty default.
It now publishes both together (897e742b1e), and a test gathers two
`acreate_agent` calls on one hook with a fallback and a slow `extra_dejson`,
and checks both get the complete `FallbackModel`; a failed extra read no longer
leaves a half-seeded hook either.
I left the lock out on purpose: `BaseAsyncOperator` and the comms
`_thread_lock` both arrived in 3.2.0, so every version that can run an async
task already serializes `send()` across threads, and `get_conn` shares nothing
else between threads (fallback hooks are new instances, `_model` is assigned
once). On 3.0/3.1 only a sync task running its own `asyncio.gather` of hook
calls is exposed. If you want that corner covered too, routing the three hops
through `AirflowToolset.run_blocking` is a one-line change each; I just didn't
want to serialize hook resolution with every tool call for it.
##########
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:
Reworded to the contract: it primes the cache with a connection the caller
already fetched, for the fallback hooks and for `get_hook` (897e742b1e).
--
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]
