dabla commented on code in PR #73414: URL: https://github.com/apache/airflow/pull/73414#discussion_r4307710100
########## airflow-core/docs/security/secrets/secrets-backend/index.rst: ########## @@ -161,6 +161,62 @@ section of ``airflow.cfg``. Additional arguments to your SecretsBackend can be configured in ``airflow.cfg`` by supplying a JSON string to ``backend_kwargs``, which will be passed to the ``__init__`` of your SecretsBackend. See :ref:`Configuration <secrets_backend_configuration>` for more details, and :ref:`SSM Parameter Store <ssm_parameter_store_secrets>` for an example. +Async server-side reads +^^^^^^^^^^^^^^^^^^^^^^^ + +When the Execution API serves a Variable or Connection, it resolves the value with +``airflow.secrets.resolver.resolve_variable`` or ``airflow.secrets.resolver.resolve_connection``. +Both functions search the configured backends in the same order as synchronous reads. There is no +separate async configuration. + +A backend can implement either or both of these optional methods: + +.. code-block:: python + + async def aget_variable(self, key: str, team_name: str | None = None) -> str | None: ... + + + async def aget_connection(self, conn_id: str, team_name: str | None = None) -> Connection | None: ... + +They take the same arguments and return the same values as ``get_variable`` and ``get_connection``. +``aget_connection`` returns the core :py:class:`airflow.models.connection.Connection` or ``None``. + +The resolver calls these methods as follows: + +* **Precedence:** the resolver uses the async method when it is defined in the same class as its Review Comment: **[warning]** This page states the `aget_*` contract for any backend, but the precedence rule (and `team_name` forwarding) only holds on the API server. This section sits on the general secrets-backend page, so a backend author will read it as the contract for `aget_variable` / `aget_connection` wherever the backend is configured. In a task process the same methods are already picked up by the Task SDK with a different rule. [`_async_get_connection`](https://github.com/apache/airflow/blob/9ad8d78c27779b3cc6660d1a496676a4a195c203/task-sdk/src/airflow/sdk/execution_time/context.py#L294) and [`_async_get_variable`](https://github.com/apache/airflow/blob/9ad8d78c27779b3cc6660d1a496676a4a195c203/task-sdk/src/airflow/sdk/execution_time/context.py#L421) take `getattr(backend, "aget_*")` whenever it exists and call it without `team_name`. So on a worker: - a subclass that overrides only `get_variable` still gets the inherited `aget_variable` (the reverse of the **Precedence** bullet); - a team-scoped lookup never reaches `aget_*` with its team. The last paragraph has the same gap: "Task processes read secrets through the Execution API" is true for the `ExecutionAPISecretsBackend`, but env vars and `[workers] secrets_backend` are read in the task process first, by the SDK path above. Either scope the section to the Execution API server and say what the worker side does today, or align the SDK lookups with `_get_native_async_method` (here or in a follow-up linked from the page). One rule for one method name is easier to document than two. ########## airflow-core/src/airflow/secrets/resolver.py: ########## @@ -0,0 +1,241 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Asynchronous server-side resolution through configured secrets backends.""" + +from __future__ import annotations + +import contextlib +import inspect +import logging +from typing import TYPE_CHECKING, Any + +from starlette.concurrency import run_in_threadpool + +from airflow._shared.secrets_backend.base import accepts_kwarg, call_secrets_backend_method +from airflow._shared.secrets_masker import mask_secret +from airflow.configuration import conf, ensure_secrets_loaded +from airflow.exceptions import AirflowNotFoundException +from airflow.models.connection import AirflowSecretsBackendAccessDenied, Connection +from airflow.models.crypto import get_fernet +from airflow.secrets.metastore import MetastoreBackend + +if TYPE_CHECKING: + from collections.abc import Callable + + from sqlalchemy.ext.asyncio import AsyncSession + + from airflow.secrets import BaseSecretsBackend + +log = logging.getLogger(__name__) + + +def _get_secret_cache() -> Any: + from airflow.sdk import SecretCache # noqa: SDK001 + + return SecretCache + + +def _find_method_owner(backend: Any, method_name: str) -> int | None: + if method_name in vars(backend): + return -1 + for index, cls in enumerate(type(backend).__mro__): + if method_name in vars(cls): + return index + return None + + +def _get_native_async_method(backend: Any, *, async_name: str, sync_name: str) -> Callable[..., Any] | None: + if inspect.getattr_static(backend, async_name, None) is None: + return None + async_method = getattr(backend, async_name, None) + if async_method is None: + return None + + async_owner = _find_method_owner(backend, async_name) + sync_owner = _find_method_owner(backend, sync_name) + if async_owner is not None and sync_owner is not None and sync_owner < async_owner: + return None + return async_method + + +async def _call_backend_method( + method: Callable[..., Any], + *, + team_name: str | None, + session: AsyncSession | None = None, + **kwargs: Any, +) -> Any: + if session is not None and accepts_kwarg(method, "session"): + kwargs["session"] = session + + if inspect.iscoroutinefunction(method): + result = call_secrets_backend_method(method, team_name=team_name, **kwargs) + else: + result = await run_in_threadpool( + call_secrets_backend_method, + method, + team_name=team_name, + **kwargs, + ) + if inspect.isawaitable(result): + return await result + return result + + +async def _call_secret_cache(method: Callable[..., Any], *args: Any, **kwargs: Any) -> Any: + # ``SecretCache.init()`` stores entries in a ``multiprocessing.Manager`` dict, so every access to an + # initialized cache is IPC to the manager process. An uninitialized cache only misses or no-ops. + if _get_secret_cache()._cache is None: + return method(*args, **kwargs) + return await run_in_threadpool(method, *args, **kwargs) + + +async def _load_fernet() -> None: + # The first ``get_fernet()`` reads ``core.fernet_key``, which may run ``fernet_key_cmd`` or query a + # secrets backend. Once loaded, encrypting and decrypting stored values is local work. A failed + # load is not cached, so its error still surfaces where a value needs the key, as before. + if not get_fernet.cache_info().currsize: + with contextlib.suppress(Exception): Review Comment: **[warning]** When the Fernet load fails, the error is suppressed here and the load runs again on the event loop. `get_fernet` is `@functools.cache`d, and `cache` stores no result for a call that raised. So after this `suppress`, the next `get_fernet()` re-reads `core.fernet_key` (and runs `fernet_key_cmd` / `fernet_key_secret`) on the loop thread. In this PR that call comes from: - `Connection(conn_id=..., uri=...)` on a cache hit (via `set_password`); - `on_db_load()` / `Variable.val` in the metastore `aget_*` methods; - `connection.get_uri()` after a sync backend returns. Probe on this head: `SecretCache` initialised with one cached connection URI with a password, `get_fernet` cache cleared, and `conf.get("core", "FERNET_KEY")` raising `AirflowConfigException`. Then `await resolve_connection("c")` reads the key twice: once in `AnyIO worker thread`, then on `MainThread` (the loop). Because failures are not cached, every request repeats both reads. When `fernet_key_cmd` is slow or hangs (a Vault CLI timing out), every route on that server stalls on each Connection or encrypted Variable read, task heartbeats included. The comment above and the docs bullet ("Successful initialization is cached; failed loads can be retried") describe the retry but not where it runs. Two ways out, either keeps the loop free: - re-raise the load error for the paths that need the key, so it surfaces once, from the worker; - or run the decrypting step (`Connection(uri=...)`, `_build_connection`, `_get_variable_value`, `get_uri()`) in the worker while `get_fernet.cache_info().currsize == 0`. A variant of `test_failed_fernet_key_load_does_not_stop_resolution` with an encrypted value, asserting the threads of the `FERNET_KEY` reads as `test_cold_fernet_key_is_loaded_off_the_event_loop_once` does, would cover it. ########## airflow-core/tests/unit/api_fastapi/execution_api/test_app.py: ########## @@ -96,6 +97,96 @@ def test_server_and_in_process_apps_isolate_auth(in_process_first, monkeypatch): assert authenticated_response.json() == response.json() [email protected]("failure", [None, "startup", "shutdown", "dispose", "configure"]) [email protected]("airflow.api_fastapi.execution_api.app.InProcessExecutionAPI", autospec=True) +def test_in_process_fixture_owns_engine_lifetime(mock_api, request, failure): + previous_engine, previous_factory = settings.async_engine, settings.AsyncSession + api = InProcessExecutionAPI() + mock_api.return_value = api + original_lifespan = api.app.router.lifespan_context + disposed = [] + original_dispose = AsyncEngine.dispose + original_configure = settings._configure_async_session + + @asynccontextmanager + async def lifespan(app): + async with original_lifespan(app): + if failure == "startup": + raise RuntimeError("startup failed") + yield + if failure == "shutdown": + raise RuntimeError("shutdown failed") + + api.app.router.lifespan_context = lifespan + + async def dispose(engine, *args, **kwargs): + disposed.append((engine, asyncio.get_running_loop())) + await original_dispose(engine, *args, **kwargs) + if engine is not previous_engine and failure == "dispose": + raise RuntimeError("dispose failed") + + def configure(): + original_configure() + if failure == "configure": + raise RuntimeError("configure failed") + + async def query(factory): + async with factory() as session: + return await session.scalar(text("SELECT 1")) + + @api.app.get("/fixture-db-read") + async def read(): + return await query(settings.AsyncSession) + + fixture = request._fixturemanager.getfixturedefs("in_process_execution_api", request.node)[-1].func Review Comment: **[nit]** The tests of the test fixtures reach into pytest internals. `request._fixturemanager.getfixturedefs(...)[-1].func` here, and `fresh_async_engine.__wrapped__()` in `test_secrets_resolver.py`, call fixture functions directly through private or undocumented attributes. pytest disallows calling fixtures directly ([deprecation note](https://docs.pytest.org/en/stable/deprecations.html#calling-fixtures-directly)), and `_fixturemanager` / `getfixturedefs` are not public API (the `getfixturedefs` signature already changed in pytest 8.1). A pytest or pytest-asyncio upgrade can therefore break these two tests while the product code is fine. Together they add about 130 lines that test scaffolding rather than the change. The fixture's guarantees (a fresh engine per server loop, the previous pool left intact, globals restored) are already exercised by its consumers, `test_processor.py` and `test_triggerer_job.py`. If you want to keep explicit checks of the fixture, they can live under `tests_common` with the fixture extracted to a plain context manager that both the fixture and the tests call. ########## airflow-core/src/airflow/secrets/metastore.py: ########## @@ -23,15 +23,40 @@ from sqlalchemy import or_, select +from airflow._shared.secrets_masker import mask_secret from airflow.secrets import BaseSecretsBackend -from airflow.utils.session import NEW_SESSION, provide_session +from airflow.utils.session import NEW_SESSION, create_session_async, provide_session if TYPE_CHECKING: + from sqlalchemy.engine import RowMapping + from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.orm import Session from airflow.models import Connection +def _build_connection(values: RowMapping) -> Connection: Review Comment: **[nit]** `_build_connection` and `_get_variable_value` re-implement what the ORM load already does in the sync twins. `get_connection` / `get_variable` above load the entity with `session.scalar(select(Connection)...)`, `session.expunge(...)` it, and let the `@reconstructor` (`on_db_load`) do the masking. The async versions select table columns and rebuild the object by hand: `Connection()` followed by a `setattr` per mapped column, then `on_db_load()`, and a `Variable` whose `_val` / `is_encrypted` are poked directly. That is a second loader per model that has to follow every future model change (a `column_property`, a deferred column, a new `__init__` side effect), while `AsyncSession` supports the same ORM select + expunge ([SQLAlchemy asyncio docs](https://docs.sqlalchemy.org/en/20/orm/extensions/asyncio.html)): ```python conn = await session.scalar(select(Connection).where(...).limit(1)) if conn: session.expunge(conn) return conn ``` Is there a reason the borrowed request session must not see the entity at all, beyond what `expunge` gives? If so, a one-line comment on the helpers saying why would help the next reader. -- 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]
