kaxil commented on code in PR #74626:
URL: https://github.com/apache/airflow/pull/74626#discussion_r4258122477
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -165,8 +181,26 @@ def get_hook(cls, conn_id: str, hook_params: dict | None =
None):
Airflow 3's ``BaseHook.get_hook`` already takes ``hook_params``;
Airflow 2's does
not, so this mirrors the Airflow 3 body.
+
+ :raises ValueError: if the connection's type is registered to a hook
that is not a
+ ``PydanticAIHook``.
"""
- return cls.get_connection(conn_id).get_hook(hook_params=hook_params)
+ conn = cls.get_connection(conn_id)
+ hooks = _get_registered_hooks()
+ # Resolved from the registry rather than from ``cls.conn_type`` or the
subclasses
+ # imported so far: operators call this on the base class, and a
subclass from another
+ # provider (e.g. Snowflake's ``pydanticai_snowflake``) is only
imported on lookup.
+ # A missing or unregistered type is left to ``conn.get_hook``, whose
error explains a
+ # URI that parsed with no type, or the '-'/'_' spelling a renamed type
needs.
+ hook_info = hooks.get(conn.conn_type) if conn.conn_type else None
+ if hook_info is not None and not
issubclass(import_string(hook_info.hook_class_name), PydanticAIHook):
Review Comment:
Both providers managers register `generic` and `email` as
`HookInfo(hook_class_name=None, ...)`, so `hooks.get("generic")` isn't None and
`import_string(None)` raises `AttributeError: 'NoneType' object has no
attribute 'rsplit'` here. Not a regression, since `Connection.get_hook` hits
the same thing, but Generic with the key in the password is a likely wrong
pick, so it'd be good for it to get the new message too. Something like
`hook_info is not None and (hook_info.hook_class_name is None or not
issubclass(...))`, plus a `generic` case in the rejection test?
##########
providers/common/ai/tests/unit/common/ai/hooks/test_pydantic_ai.py:
##########
@@ -333,10 +350,21 @@ def registry():
# argument -- without it, `conn.get_hook()` calls the mock with zero
arguments and
# `reg.get_hook` would have no way to know which connection dispatched
it.
patch.object(Connection, "get_hook", side_effect=reg.get_hook,
autospec=True),
+ patch(
+
"airflow.providers.common.ai.hooks.pydantic_ai._get_registered_hooks",
Review Comment:
Every test patches `_get_registered_hooks`, so nothing runs the real
function or its `ImportError` fallback. Could you add one unpatched test, e.g.
`assert _get_registered_hooks()["pydanticai"].hook_class_name ==
qualname(PydanticAIHook)`? It shouldn't need the DB, and the compat matrix
would then cover the `ProvidersManager` branch on 2.11/3.0/3.1.
##########
providers/common/ai/tests/unit/common/ai/hooks/test_pydantic_ai.py:
##########
@@ -351,6 +379,68 @@ def test_builds_the_connection_hook_with_hook_params(self,
registry):
assert hook.model_id == "openai:gpt-5"
assert hook.fallback_conn_ids == []
+ def test_rejects_a_connection_of_another_type_before_dispatching(self,
registry):
+ registry.add("llm", conn_type="openai")
+ registry.register_hook("openai", qualname(_OtherProviderHook))
+ expected_types = ["pydanticai", "pydanticai_azure",
"pydanticai_bedrock", "pydanticai_vertex"]
+
+ with pytest.raises(
+ ValueError,
+ match=re.escape(
+ f"Connection 'llm' has type 'openai', but a Pydantic AI hook
needs one of {expected_types}"
+ ),
+ ):
+ PydanticAIHook.get_hook("llm", hook_params={"model_id":
"openai:gpt-5"})
+
+ Connection.get_hook.assert_not_called()
+
+ def test_accepts_a_subclass_from_a_module_not_imported_yet(self, registry,
tmp_path, monkeypatch):
+ """Another provider's subclass (e.g. Snowflake's) is only imported
when its type is looked up."""
+ module_name = "pydanticai_subclass_not_imported_yet"
+ (tmp_path / f"{module_name}.py").write_text(
+ "from airflow.providers.common.ai.hooks.pydantic_ai import
PydanticAIHook\n\n\n"
+ "class PydanticAIElsewhereHook(PydanticAIHook):\n"
+ " conn_type = 'pydanticai_elsewhere'\n"
+ )
+ monkeypatch.syspath_prepend(tmp_path)
+ registry.add("llm", conn_type="pydanticai_elsewhere")
Review Comment:
Because the subclass's type starts with `pydanticai`, this test also passes
if the guard is replaced by a plain `conn.conn_type.startswith("pydanticai")`
check. Registering it under a type without the prefix (and asserting
`module_name in sys.modules` after the call) would pin that the class lookup is
what accepts it.
##########
providers/common/ai/tests/unit/common/ai/hooks/test_pydantic_ai.py:
##########
@@ -351,6 +379,68 @@ def test_builds_the_connection_hook_with_hook_params(self,
registry):
assert hook.model_id == "openai:gpt-5"
assert hook.fallback_conn_ids == []
+ def test_rejects_a_connection_of_another_type_before_dispatching(self,
registry):
+ registry.add("llm", conn_type="openai")
+ registry.register_hook("openai", qualname(_OtherProviderHook))
+ expected_types = ["pydanticai", "pydanticai_azure",
"pydanticai_bedrock", "pydanticai_vertex"]
+
+ with pytest.raises(
+ ValueError,
+ match=re.escape(
+ f"Connection 'llm' has type 'openai', but a Pydantic AI hook
needs one of {expected_types}"
+ ),
+ ):
+ PydanticAIHook.get_hook("llm", hook_params={"model_id":
"openai:gpt-5"})
+
+ Connection.get_hook.assert_not_called()
+
+ def test_accepts_a_subclass_from_a_module_not_imported_yet(self, registry,
tmp_path, monkeypatch):
+ """Another provider's subclass (e.g. Snowflake's) is only imported
when its type is looked up."""
+ module_name = "pydanticai_subclass_not_imported_yet"
+ (tmp_path / f"{module_name}.py").write_text(
+ "from airflow.providers.common.ai.hooks.pydantic_ai import
PydanticAIHook\n\n\n"
+ "class PydanticAIElsewhereHook(PydanticAIHook):\n"
+ " conn_type = 'pydanticai_elsewhere'\n"
+ )
+ monkeypatch.syspath_prepend(tmp_path)
+ registry.add("llm", conn_type="pydanticai_elsewhere")
+ registry.register_hook("pydanticai_elsewhere",
f"{module_name}.PydanticAIElsewhereHook")
+ assert module_name not in sys.modules
+
+ try:
+ PydanticAIHook.get_hook("llm", hook_params={"model_id":
"openai:gpt-5"})
+ finally:
+ sys.modules.pop(module_name, None)
+
+ Connection.get_hook.assert_called_once_with(
+ registry.conns["llm"], hook_params={"model_id": "openai:gpt-5"}
+ )
+
+ def
test_leaves_an_unregistered_connection_type_to_connection_get_hook(self,
registry):
+ """``Connection.get_hook`` raises for it with a hint about the '-'/'_'
spelling."""
+ registry.add("llm", conn_type="pydanticai-azure")
+
+ PydanticAIHook.get_hook("llm")
+
+ Connection.get_hook.assert_called_once_with(registry.conns["llm"],
hook_params=None)
+
+ @pytest.mark.parametrize(
+ ("conn_type", "hook_class"),
+ [
+ ("pydanticai", PydanticAIHook),
+ ("pydanticai_azure", PydanticAIAzureHook),
+ ("pydanticai_bedrock", PydanticAIBedrockHook),
+ ("pydanticai_vertex", PydanticAIVertexHook),
+ ],
+ )
+ def test_accepts_every_pydanticai_connection_type(self, registry,
conn_type, hook_class):
+ registry.add("llm", conn_type=conn_type, hook_class=hook_class)
+
+ hook = PydanticAIHook.get_hook("llm")
+
+ assert type(hook) is hook_class
Review Comment:
This assertion can't fail: the patched `Connection.get_hook` returns
`hook_classes[conn_id]`, which `registry.add` set from this same `hook_class`.
`Connection.get_hook.assert_called_once_with(registry.conns["llm"],
hook_params=None)` checks what the test is about, and parametrizing over
`get_provider_info()["connection-types"]` would keep the list in step with
provider.yaml.
--
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]
