Lee-W commented on code in PR #69969:
URL: https://github.com/apache/airflow/pull/69969#discussion_r3600290079
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -98,6 +98,38 @@ def get_ui_field_behaviour() -> dict[str, Any]:
# Core connection / agent API
# ------------------------------------------------------------------
+ @classmethod
+ def get_hook(cls, conn_id: str, hook_params: dict | None = None) ->
PydanticAIHook:
+ """
+ Return the :class:`PydanticAIHook` subclass matching the connection's
``conn_type``.
+
+ Overrides :meth:`~airflow.sdk.bases.hook.BaseHook.get_hook` to
guarantee that the
+ hook's ``llm_conn_id`` always equals the *caller's* ``conn_id``, even
when the
+ underlying secrets backend rewrites it.
+
+ Some secrets backends (e.g. a custom Vault provider that supports
+ environment-independent connection aliases) return a
+ :class:`~airflow.sdk.definitions.connection.Connection` whose
``conn_id`` has been
+ resolved to a canonical, environment-specific key (e.g.
``"openai.prod"``). Without
+ this override, the hook would be constructed with that canonical id,
and any subsequent
+ :meth:`get_conn` call would try to fetch it from the secrets backend —
a key it does
+ not recognize — raising
:class:`~airflow.exceptions.AirflowNotFoundException`.
+
+ ``BaseHook`` applies the same fix as of Airflow 3.4. This override
ensures
Review Comment:
```suggestion
``BaseHook`` applies the same fix as of Airflow 3.4. This override
ensures
```
##########
task-sdk/src/airflow/sdk/bases/hook.py:
##########
@@ -76,6 +76,31 @@ async def aget_connection(cls, conn_id: str) -> Connection:
log.debug("Connection Retrieved '%s' (via task-sdk)", conn.conn_id)
return conn
+ @classmethod
+ def _hook_from_connection(cls, connection: Connection, conn_id: str,
hook_params: dict | None = None):
Review Comment:
```suggestion
def _build_hook_from_connection(cls, connection: Connection, conn_id:
str, hook_params: dict | None = None) -> BaseHook:
```
--
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]
