kaxil commented on code in PR #74150:
URL: https://github.com/apache/airflow/pull/74150#discussion_r4393931460
##########
providers/common/ai/tests/unit/common/ai/hooks/test_pydantic_ai.py:
##########
@@ -2225,3 +2236,304 @@ def test_hook_resolves_from_environment_variable(self,
hook_class, serializer, m
assert resolved.conn_type == hook_class.conn_type
assert type(resolved.get_hook()) is hook_class
+
+
+def _off_the_event_loop(lookup):
+ """
+ Wrap a synchronous Task SDK lookup so it fails the test when it runs on
the event loop.
+
+ ``get_connection`` and the secret masking of ``extra_dejson`` send to the
supervisor
+ synchronously. On the event loop, with another async call in flight, that
raises
+ ``DeadlockImminentError``; from a worker thread it only waits its turn.
+ """
+
+ def guarded(*args, **kwargs):
+ try:
+ asyncio.get_running_loop()
+ except RuntimeError:
+ return lookup(*args, **kwargs)
+ raise AssertionError("synchronous Task SDK call on the event loop")
+
+ return guarded
+
+
[email protected]
+def loop_guarded_registry():
+ """The ``registry`` fixture, with every synchronous lookup refused on the
event loop."""
+ reg = _ConnRegistry()
+ with (
+ patch.object(
+ PydanticAIHook, "get_connection",
side_effect=_off_the_event_loop(reg.get_connection)
+ ) as mock_get_connection,
+ patch.object(
+ Connection,
+ "extra_dejson",
+ new=property(_off_the_event_loop(lambda conn:
json.loads(conn.extra) if conn.extra else {})),
+ ),
+ patch.object(Connection, "get_hook", side_effect=reg.get_hook,
autospec=True),
+ patch(
+
"airflow.providers.common.ai.hooks.pydantic_ai._get_registered_hooks",
+ autospec=True,
+ return_value=reg.hooks,
+ ),
+ ):
+ reg.get_connection_mock = mock_get_connection
+ yield reg
+
+
+class _ReadsAnotherConnectionHook(PydanticAIHook):
+ """A subclass whose provider kwargs come from another connection, as
Snowflake's do."""
+
+ def _get_provider_kwargs(self, api_key, base_url, extra):
+ return {"api_key":
self.get_connection("credentials").extra_dejson["token"]}
+
+
+class TestPydanticAIHookAsync:
+ @pytest.mark.asyncio
+ async def test_aget_conn_resolves_the_model_off_the_event_loop(
+ self, loop_guarded_registry, infer_model_stub
+ ):
+ loop_guarded_registry.add("primary", extra={"model":
"openai:gpt-5.6-sol"})
+
+ model = await PydanticAIHook(llm_conn_id="primary").aget_conn()
+
+ assert model is infer_model_stub.models["openai:gpt-5.6-sol"]
+
loop_guarded_registry.get_connection_mock.assert_called_once_with("primary")
+
+ @pytest.mark.asyncio
+ async def test_aget_conn_resolves_the_fallback_chain_off_the_event_loop(
+ self, loop_guarded_registry, infer_model_stub
+ ):
+ loop_guarded_registry.add(
+ "primary", extra={"model": "openai:gpt-5.6-sol",
"fallback_conn_ids": ["second", "third"]}
+ )
+ loop_guarded_registry.add("second", extra={"model":
"anthropic:claude-opus-4-6"})
+ loop_guarded_registry.add("third", extra={"model": "groq:llama-4"})
+
+ model = await PydanticAIHook(llm_conn_id="primary").aget_conn()
+
+ assert isinstance(model, FallbackModel)
+ assert model.models == [
+ infer_model_stub.models["openai:gpt-5.6-sol"],
+ infer_model_stub.models["anthropic:claude-opus-4-6"],
+ infer_model_stub.models["groq:llama-4"],
+ ]
+ assert loop_guarded_registry.get_connection_mock.call_args_list == [
+ call("primary"),
+ call("second"),
+ call("third"),
+ ]
+
+ @pytest.mark.asyncio
+ @pytest.mark.parametrize("subclass_conn_id", ["primary", "second"],
ids=["primary", "fallback"])
+ async def
test_aget_conn_runs_a_subclass_provider_kwargs_off_the_event_loop(
+ self, loop_guarded_registry, infer_model_stub, subclass_conn_id
+ ):
+ loop_guarded_registry.add("credentials", extra={"token":
"from-another-connection"})
+ loop_guarded_registry.add(
+ "primary", extra={"model": "openai:gpt-5.6-sol",
"fallback_conn_ids": ["second"]}
+ )
+ loop_guarded_registry.add(
+ "second", hook_class=_ReadsAnotherConnectionHook, extra={"model":
"anthropic:claude-opus-4-6"}
+ )
+ hook_class = _ReadsAnotherConnectionHook if subclass_conn_id ==
"primary" else PydanticAIHook
+
+ await hook_class(llm_conn_id="primary").aget_conn()
+
+
loop_guarded_registry.get_connection_mock.assert_any_call("credentials")
+
+ @pytest.mark.asyncio
+ async def test_aget_conn_is_cached(self, loop_guarded_registry,
infer_model_stub):
+ loop_guarded_registry.add("primary", extra={"model":
"openai:gpt-5.6-sol"})
+ hook = PydanticAIHook(llm_conn_id="primary")
+
+ assert await hook.aget_conn() is await hook.aget_conn()
+ assert hook.get_conn() is infer_model_stub.models["openai:gpt-5.6-sol"]
+ loop_guarded_registry.get_connection_mock.assert_called_once()
+
+ @pytest.mark.asyncio
+ @patch("airflow.providers.common.ai.hooks.pydantic_ai.Agent",
autospec=True)
+ async def test_acreate_agent_forwards_every_argument(
+ self, mock_agent_cls, loop_guarded_registry, infer_model_stub
+ ):
+ loop_guarded_registry.add("primary", extra={"model":
"openai:gpt-5.6-sol"})
+
+ agent = await PydanticAIHook(llm_conn_id="primary").acreate_agent(
+ int, instructions="Count.", retries=3
+ )
+
+ assert agent is mock_agent_cls.return_value
+ mock_agent_cls.assert_called_once_with(
+ infer_model_stub.models["openai:gpt-5.6-sol"],
+ output_type=int,
+ instructions="Count.",
+ retries=3,
+ capabilities=[MaskingCapability()],
+ )
+
+ @pytest.mark.asyncio
+ @patch("airflow.providers.common.ai.hooks.pydantic_ai.Agent",
autospec=True)
+ async def test_acreate_agent_with_spec_file_forwards_every_argument(
+ self, mock_agent_cls, loop_guarded_registry, infer_model_stub
+ ):
+ loop_guarded_registry.add("primary", extra={"model":
"openai:gpt-5.6-sol"})
+
+ await PydanticAIHook(llm_conn_id="primary").acreate_agent(
+ int, spec_file="/path/to/agent.yaml", instructions="Be brief.",
retries=3
+ )
+
+ mock_agent_cls.from_file.assert_called_once_with(
+ "/path/to/agent.yaml",
+ output_type=int,
+ model=infer_model_stub.models["openai:gpt-5.6-sol"],
+ instructions="Be brief.",
+ retries=3,
+ capabilities=[MaskingCapability()],
+ )
+
+ @pytest.mark.asyncio
+ @patch("airflow.providers.common.ai.hooks.pydantic_ai.Agent",
autospec=True)
+ async def test_acreate_agent_twice_fetches_each_connection_once(
+ self, mock_agent_cls, loop_guarded_registry, infer_model_stub
+ ):
+ loop_guarded_registry.add(
+ "primary", extra={"model": "openai:gpt-5.6-sol",
"fallback_conn_ids": ["second"]}
+ )
+ loop_guarded_registry.add("second", extra={"model":
"anthropic:claude-opus-4-6"})
+ hook = PydanticAIHook(llm_conn_id="primary")
+
+ await hook.acreate_agent(instructions="Be helpful.")
+ await hook.acreate_agent(instructions="Be helpful.")
+
+ assert loop_guarded_registry.get_connection_mock.call_args_list ==
[call("primary"), call("second")]
+
+ @pytest.mark.asyncio
+ async def
test_acreate_agent_without_instructions_or_spec_file_raises_before_fetching(
+ self, loop_guarded_registry
+ ):
+ loop_guarded_registry.add("primary", extra={"model":
"openai:gpt-5.6-sol"})
+
+ with pytest.raises(ValueError, match="instructions is required when
spec_file is not provided"):
+ await PydanticAIHook(llm_conn_id="primary").acreate_agent()
+
+ loop_guarded_registry.get_connection_mock.assert_not_called()
+
+ @pytest.mark.asyncio
+ async def test_aget_hook_builds_the_connection_hook_off_the_event_loop(
+ self, loop_guarded_registry, infer_model_stub
+ ):
+ loop_guarded_registry.add("llm", extra={"model": "openai:gpt-5.6-sol"})
+
+ hook = await PydanticAIHook.aget_hook("llm", hook_params={"model_id":
"anthropic:claude-opus-4-6"})
+
+ assert hook.model_id == "anthropic:claude-opus-4-6"
+ assert await hook.aget_conn() is
infer_model_stub.models["anthropic:claude-opus-4-6"]
+
loop_guarded_registry.get_connection_mock.assert_called_once_with("llm")
+
+ @pytest.mark.asyncio
+ async def test_aget_hook_rejects_a_connection_of_another_type(self,
loop_guarded_registry):
+ loop_guarded_registry.add("llm", conn_type="openai")
+ loop_guarded_registry.register_hook("openai",
qualname(_OtherProviderHook))
+
+ with pytest.raises(ValueError, match="Connection 'llm' has type
'openai'"):
+ await PydanticAIHook.aget_hook("llm")
+
+ Connection.get_hook.assert_not_called()
+
+ @pytest.mark.asyncio
+ @patch("airflow.providers.common.ai.hooks.pydantic_ai.Agent",
autospec=True)
+ async def
test_concurrent_acreate_agent_calls_each_see_a_complete_connection(
Review Comment:
With one worker these two gathered calls now run one after the other, so
this no longer races the publish in `_get_conn_and_extra`: putting `self._conn
= ...` back before the `extra_dejson` read keeps this test green, and it only
fails again with `max_workers=2`. Calling `hook.get_conn()` from two
`threading.Thread`s with the slow `extra_dejson` would still pin it. Otherwise
this test and the "two gathered async calls on one hook" note in
`_get_conn_and_extra` could go.
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -95,6 +102,20 @@ def _looks_like_unrecognized_provider_prefix(prefix: str)
-> bool:
return bool(_PROVIDER_SLUG_RE.match(prefix))
+# The async methods run their blocking lookups here, never on asyncio's
default executor:
+# ``asend()`` reads its reply on that executor while it holds the comms lock,
so blocking
+# lookups that fill it leave nothing to run the read and the task hangs. One
worker also
+# serializes this process's hook lookups on Airflow 3.0 and 3.1, where the
channel to the
+# supervisor has no lock of its own.
+_executor = concurrent.futures.ThreadPoolExecutor(max_workers=1,
thread_name_prefix="pydantic-ai-hook")
+
+
+async def _run_off_loop(fn: Callable[P, R], /, *args: P.args, **kwargs:
P.kwargs) -> R:
+ """Run ``fn`` on the hook's own worker thread, with the caller's context
variables (task logging binds them)."""
+ call = functools.partial(contextvars.copy_context().run, fn, *args,
**kwargs)
Review Comment:
`run_in_executor` doesn't carry the caller's context the way `to_thread`
did, so this `copy_context()` is what keeps the TI fields the task runner binds
with `bind_contextvars` on the lookup's log lines. Nothing tests it, though:
replacing this line with `functools.partial(fn, *args, **kwargs)` keeps the
async tests green. Could one test set a `ContextVar` before each of the three
hops and check that the `get_connection` side effect sees it?
--
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]
