kaxil commented on code in PR #65587:
URL: https://github.com/apache/airflow/pull/65587#discussion_r3683480825
##########
airflow-core/src/airflow/api_fastapi/execution_api/app.py:
##########
@@ -379,6 +392,7 @@ class InProcessExecutionAPI:
needed so that we can use the sync httpx client
"""
+ request_scoped_server_context: bool = attrs.field(default=False,
kw_only=True)
Review Comment:
Separate from the `.app` vs `.transport` placement discussion above: why is
this opt-in at all? `InProcessExecutionAPI` is the server side of the Execution
API in every case, so forcing server context per request looks like it should
be the default.
As it stands `triggerer_job_runner.in_process_api_server()` builds
`InProcessExecutionAPI()` without the flag, and is safe only because
`TriggererJobRunner._execute` happens to set `_AIRFLOW_PROCESS_CONTEXT=server`
process-wide in a different file, with nothing tying the two together.
Defaulting to `True` would cover that instance and let the parameter go away.
##########
airflow-core/src/airflow/process_context.py:
##########
@@ -0,0 +1,59 @@
+# 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.
+from __future__ import annotations
+
+import os
+import sys
+from collections.abc import Generator
+from contextlib import contextmanager
+from contextvars import ContextVar
+from typing import Literal
+
+__all__ = [
+ "get_process_context",
+ "override_process_context",
+ "should_use_task_sdk_api_path",
+]
+
+_PROCESS_CONTEXT_OVERRIDE: ContextVar[str | None] = ContextVar(
+ "_AIRFLOW_PROCESS_CONTEXT_OVERRIDE",
+ default=None,
+)
+
+
+def get_process_context() -> str | None:
+ """Return the current process context, preferring request-scoped
overrides."""
+ return _PROCESS_CONTEXT_OVERRIDE.get() or
os.environ.get("_AIRFLOW_PROCESS_CONTEXT")
+
+
+@contextmanager
+def override_process_context(context: Literal["server", "client"]) ->
Generator[None, None, None]:
+ """Temporarily override the current process context for the active
execution flow."""
+ token = _PROCESS_CONTEXT_OVERRIDE.set(context)
+ try:
+ yield
+ finally:
+ _PROCESS_CONTEXT_OVERRIDE.reset(token)
+
+
+def should_use_task_sdk_api_path() -> bool:
+ """Return True when execution-context helpers should route through Task
SDK APIs."""
+ if get_process_context() == "server":
Review Comment:
This gives the process-wide env var precedence over `SUPERVISOR_COMMS`,
which is the reverse of what `ensure_secrets_backend_loaded()` in
`task-sdk/.../execution_time/supervisor.py` does (it checks `SUPERVISOR_COMMS`
first, specifically so a worker never falls back to `MetastoreBackend`).
That inversion is reachable: `action_cli(check_db=True)` in
`airflow/utils/cli.py` sets `_AIRFLOW_PROCESS_CONTEXT=server` around the whole
command body, and `dag_test` uses that decorator. So under `airflow dags test`,
user code calling core `airflow.models.Variable.get()` now reads the metastore
directly instead of going through the Task SDK: no deprecation warning, and no
`AirflowSecretsBackendAccessDenied` enforcement. `PythonVirtualenvOperator`
defaults to `inherit_env=True`, so the venv child gets `=server` too and would
need a working `sql_alchemy_conn` inside the venv to read a variable that used
to come back over comms.
Since the ASGI wrapper already sets the ContextVar for exactly the request
you're fixing, could the server short-circuit read `_PROCESS_CONTEXT_OVERRIDE`
only and leave `SUPERVISOR_COMMS` winning over the env var? The env arm only
changes behaviour in a process that has both, and that is the case the wrapper
covers.
##########
task-sdk/src/airflow/sdk/execution_time/supervisor.py:
##########
@@ -1922,10 +1922,10 @@ def _send_new_log_fd(self, req_id: int) -> None:
@functools.lru_cache(maxsize=1)
-def in_process_api_server():
+def in_process_api_server(*, request_scoped_server_context: bool = False):
Review Comment:
This function is `lru_cache(maxsize=1)`, so the new parameter becomes part
of the cache key. A call without the kwarg and a call with it now evict each
other, and each miss builds a fresh `InProcessExecutionAPI`: new FastAPI app,
new `create_dag_bag()`, and a new daemon event-loop thread once `.transport` is
touched.
`test_api_client_clears_dag_bag_override_when_dag_is_none` shows the trap,
it had to start passing the kwarg to keep looking at the same instance
`_api_client` uses. Setting the flag on the instance, or making it
unconditional, keeps the singleton a singleton.
##########
task-sdk/tests/task_sdk/execution_time/test_supervisor.py:
##########
@@ -3439,6 +3439,18 @@ def execute(self, context: Context):
assert isinstance(result.error, _Failure)
assert isinstance(collected[0], _Failure)
+ def test_api_client_uses_request_scoped_server_context(self):
+ api = mock.Mock()
+ api.transport = httpx.MockTransport(lambda request:
httpx.Response(status_code=200, json={}))
+
+ with patch(
+ "airflow.sdk.execution_time.supervisor.in_process_api_server",
return_value=api
+ ) as factory:
+ client = InProcessTestSupervisor._api_client()
+
+ factory.assert_called_once_with(request_scoped_server_context=True)
Review Comment:
The factory is mocked here, so nothing in this test runs
`_RequestScopedServerContextApp.__call__`. And the four new server-context
tests in `test_variables.py` / `test_connections.py` set
`_AIRFLOW_PROCESS_CONTEXT=server` in the environment and go through the
`TestClient(cached_app(...))` fixture, so they exercise the env branch of
`should_use_task_sdk_api_path()` rather than the wrapper.
That leaves the actual scenario from #65482 uncovered: `dag.test()` from a
script sets no env var, so the ContextVar is the only thing preventing the
re-entry. Picking up @jason810496's earlier ask for a test of the issue
scenario, could you add one that drives a request through
`InProcessExecutionAPI(request_scoped_server_context=True).transport` with
`SUPERVISOR_COMMS` present and asserts the SDK path isn't re-entered?
--
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]
