Lee-W commented on code in PR #74626:
URL: https://github.com/apache/airflow/pull/74626#discussion_r4262227659
##########
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:
It now also rejects connections whose registered `HookInfo` has
`hook_class_name=None` (`generic`, `email`) with the same "needs one of [...]"
message, instead of raising `AttributeError` from `import_string`.
--
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]
