apache · dabla · Jul 16, 2026 · Jul 16, 2026 · Jul 16, 2026 · Jul 16, 2026
diff --git a/docs/spelling_wordlist.txt b/docs/spelling_wordlist.txt
@@ -1366,6 +1366,7 @@ Quantiles
 quantiles
 Qubole
 qubole
+queryable
 querybook
 Querybooks
 querybooks

@@ -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
+        correct behavior on earlier Airflow versions too: providers are independently
+        releasable, so upgrading the provider alone should be sufficient to unblock users
+        regardless of which Airflow minor version they run.
+
+        :param conn_id: Connection ID as supplied by the caller; may be a secrets-backend alias.
+        :param hook_params: Extra keyword arguments forwarded to the hook constructor.
+        :return: Configured :class:`PydanticAIHook` (or the matching subclass) with
+            ``llm_conn_id`` set to *conn_id*.
+        """
+        hook = super().get_hook(conn_id, hook_params=hook_params)
+        # On Airflow < 3.4 super() does not rekey; setattr is a no-op on >= 3.4.
+        setattr(hook, hook.conn_name_attr, conn_id)
+        return hook
+
     def _get_provider_kwargs(
         self,
         api_key: str | None,

@@ -833,3 +833,47 @@ def test_get_conn_uses_explicit_project(self, mock_infer_provider_class, mock_in
         factory = mock_infer_model.call_args[1]["provider_factory"]
         factory("google-vertex")
         mock_provider_cls.assert_called_with(project="my-project", location="europe-west4")
+
+
+class TestPydanticAIHookGetHook:
+    """Tests for alias-rewriting secrets-backend behaviour specific to PydanticAI hooks."""
+
+    def test_get_conn_uses_alias_not_canonical_after_rewrite(self):
+        """
+        After get_hook() restores the alias, get_conn() must call get_connection()
+        with the alias — not the canonical id that the secrets backend wrote onto
+        conn.conn_id.  This verifies the alias re-fetch succeeds in production.
+        """
+        alias = "openai.my_alias"
+        canonical = "openai.my_alias.gpt_5_3_chat.prod"
+
+        canonical_conn = Connection(
+            conn_id=canonical,
+            conn_type="pydanticai-azure",
+            password="secret-key",
+            host="https://myresource.openai.azure.com",
+            extra='{"model": "azure:gpt-4o", "api_version": "2024-07-01-preview"}',
+        )
+        hook_with_canonical = PydanticAIAzureHook(llm_conn_id=canonical)
+
+        with (
+            patch.object(PydanticAIHook, "get_connection", return_value=canonical_conn),
+            patch.object(canonical_conn, "get_hook", return_value=hook_with_canonical),
+        ):
+            hook = PydanticAIHook.get_hook(alias, hook_params={"model_id": None})
+
+        # get_conn() must look up by the alias, not by the canonical id.
+        with (
+            patch(
+                "airflow.providers.common.ai.hooks.pydantic_ai.infer_model",
+                return_value=MagicMock(),
+            ),
+            patch(
+                "airflow.providers.common.ai.hooks.pydantic_ai.infer_provider_class",
+                return_value=MagicMock(return_value=MagicMock()),
+            ),
+            patch.object(hook, "get_connection", return_value=canonical_conn) as mock_get_conn,
+        ):
+            hook.get_conn()
+
+        mock_get_conn.assert_called_once_with(alias)
@@ -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 _build_hook_from_connection(cls, connection: Connection, conn_id: str, hook_params: dict | None = None) -> BaseHook:
+        """
+        Build a hook from *connection* and re-key it to *conn_id*.
+
+        Secrets backends may rewrite the ``conn_id`` on the
+        :class:`~airflow.sdk.definitions.connection.Connection` they return (e.g. a
+        Vault alias resolved to a canonical, environment-specific id).  If that
+        rewritten id is left on the hook, any subsequent ``get_connection()`` call
+        inside the hook will try to fetch by the canonical id — a key the secrets
+        backend does not recognize.
+
+        Re-keying to the caller's original *conn_id* ensures every secrets-backend
+        lookup continues to use the registered alias.
+
+        :param connection: Connection object returned by the secrets backend.
+        :param conn_id: Original connection ID as supplied by the caller (the registered alias).
+        :param hook_params: Extra keyword arguments forwarded to the hook constructor.
+        :return: Hook instance keyed to *conn_id*.
+        """
+        hook = connection.get_hook(hook_params=hook_params)
+        if hasattr(hook, "conn_name_attr"):
+            setattr(hook, hook.conn_name_attr, conn_id)
+        return hook
+
     @classmethod
     def get_hook(cls, conn_id: str, hook_params: dict | None = None):
         """
@@ -86,7 +111,7 @@ def get_hook(cls, conn_id: str, hook_params: dict | None = None):
         :return: default hook for this connection
         """
         connection = cls.get_connection(conn_id)
-        return connection.get_hook(hook_params=hook_params)
+        return cls._hook_from_connection(connection, conn_id, hook_params=hook_params)
 
     @classmethod
     async def aget_hook(cls, conn_id: str, hook_params: dict | None = None):
@@ -103,7 +128,7 @@ async def aget_hook(cls, conn_id: str, hook_params: dict | None = None):
         :return: default hook for this connection
         """
         connection = await cls.aget_connection(conn_id)
-        return connection.get_hook(hook_params=hook_params)
+        return cls._hook_from_connection(connection, conn_id, hook_params=hook_params)
 
     def get_conn(self) -> Any:
         """Return connection for the hook."""

@@ -128,6 +128,60 @@ async def test_aget_hook(self, mock_supervisor_comms):
             )
         assert result is not None
 
+    def test_get_hook_preserves_original_conn_id_when_backend_rewrites_alias(self):
+        """
+        When a secrets backend rewrites alias → canonical on the returned Connection,
+        the hook's conn_name_attr must remain the original alias so that any future
+        get_connection() call stays on the registered key.
+        """
+        alias = "my-llm-alias"
+        canonical = "my-llm-alias.gpt-4o.prod"
+
+        class _FakeHook:
+            conn_name_attr = "my_conn_id"
+            my_conn_id = canonical
+
+        fake_hook = _FakeHook()
+        mock_conn = MagicMock()
+        mock_conn.get_hook.return_value = fake_hook
+
+        with patch.object(BaseHook, "get_connection", return_value=mock_conn):
+            result = BaseHook.get_hook(alias)
+
+        assert result.my_conn_id == alias
+
+    def test_get_hook_no_rewriting_leaves_conn_id_unchanged(self):
+        """When conn.conn_id matches the lookup key (no alias rewriting), the hook is unchanged."""
+        conn_id = "my_plain_conn"
+
+        class _FakeHook:
+            conn_name_attr = "my_conn_id"
+            my_conn_id = conn_id
+
+        fake_hook = _FakeHook()
+        mock_conn = MagicMock()
+        mock_conn.get_hook.return_value = fake_hook
+
+        with patch.object(BaseHook, "get_connection", return_value=mock_conn):
+            result = BaseHook.get_hook(conn_id)
+
+        assert result.my_conn_id == conn_id
+
+    def test_get_hook_without_conn_name_attr_is_returned_unchanged(self):
+        """Hooks without conn_name_attr are returned as-is — no AttributeError."""
+
+        class _MinimalHook:
+            pass
+
+        fake_hook = _MinimalHook()
+        mock_conn = MagicMock()
+        mock_conn.get_hook.return_value = fake_hook
+
+        with patch.object(BaseHook, "get_connection", return_value=mock_conn):
+            result = BaseHook.get_hook("some_conn")
+
+        assert result is fake_hook
+
     def test_get_connection_secrets_backend_configured(self, mock_supervisor_comms, tmp_path):
         path = tmp_path / "conn.env"
         path.write_text("CONN_A=mysql://host_a")