kaxil commented on code in PR #72156:
URL: https://github.com/apache/airflow/pull/72156#discussion_r4042189467
##########
providers/common/ai/tests/unit/common/ai/hooks/test_pydantic_ai.py:
##########
@@ -240,6 +246,585 @@ def test_get_conn_caches_model(self, mock_infer_model):
mock_infer_model.assert_called_once()
+class _ConnRegistry:
+ """
+ In-memory stand-in for connection and hook lookup.
+
+ ``_resolve_fallback_models`` goes through ``BaseHook.get_hook``, which
needs both the
+ metadata DB and provider discovery; this resolves both from a dict instead.
+ """
+
+ def __init__(self) -> None:
+ self.conns: dict[str, Connection] = {}
+ self.hook_classes: dict[str, type[PydanticAIHook]] = {}
+
+ def add(
+ self,
+ conn_id: str,
+ *,
+ conn_type: str = "pydanticai",
+ hook_class: type[PydanticAIHook] = PydanticAIHook,
+ password: str | None = None,
+ extra: dict | None = None,
+ ) -> None:
+ self.conns[conn_id] = Connection(
+ conn_id=conn_id,
+ conn_type=conn_type,
+ password=password,
+ extra=json.dumps(extra) if extra else None,
+ )
+ self.hook_classes[conn_id] = hook_class
+
+ def get_connection(self, conn_id: str) -> Connection:
+ try:
+ return self.conns[conn_id]
+ except KeyError:
+ raise AirflowNotFoundException(f"The conn_id `{conn_id}` isn't
defined") from None
+
+ def get_hook(self, conn_id: str, hook_params: dict | None = None):
+ if conn_id not in self.conns:
+ raise AirflowNotFoundException(f"The conn_id `{conn_id}` isn't
defined")
+ hook_class = self.hook_classes[conn_id]
+ return hook_class(llm_conn_id=conn_id, **(hook_params or {}))
+
+
[email protected]
+def registry():
+ """Patch connection and hook lookup onto a registry the test populates."""
+ reg = _ConnRegistry()
+ with (
+ patch.object(PydanticAIHook, "get_connection",
side_effect=reg.get_connection),
Review Comment:
Following my own thread here, because the production side already pays this
rather than only risking it on a regression. `hooks/pydantic_ai.py:371` calls
`PydanticAIHook.get_hook(conn_id)`, and `BaseHook.get_hook` fetches the
connection then throws it away -- `Connection.get_hook` constructs the hook
with `**{connection_id_attribute_name: self.conn_id}` only
(`task-sdk/src/airflow/sdk/definitions/connection.py:288`) -- so each hook
refetches the same connection on its first `_get_conn_and_extra`, and a chain
of N costs 2N lookups. In a worker each one is a supervisor round trip, since
`secrets.use_cache` defaults to off during task execution
(`execution_time/cache.py:87`). `_preset_connections` already exists for
exactly this and `_get_connection` checks it before any backend
(`execution_time/context.py:216`, and `Connection.test_connection` sets it the
same way), so priming it around the loop would halve the lookups without
reaching into hook privates.
--
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]
