This is an automated email from the ASF dual-hosted git repository.
potiuk pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 46186d28912 Fix connection resolution in CLI by setting server process
context in decorators (#61036)
46186d28912 is described below
commit 46186d2891228d786cb05f1c6dc2251987cc1b2d
Author: SeonghwanLee <[email protected]>
AuthorDate: Thu Jul 2 01:21:08 2026 +0900
Fix connection resolution in CLI by setting server process context in
decorators (#61036)
* Fix connection resolution in CLI by setting server process context in
decorators
* Apply review of @uranusjr - Restore _AIRFLOW_PROCESS_CONTEXT in finally
block for CLI decorators
* Apply review of @potiuk - Reduce scope of server context
* Add test case - test_bundle_validation_runs_with_server_context
* Apply review of @potiuk - Update action_cli, test_case
---
airflow-core/src/airflow/utils/cli.py | 13 +++++++++++-
.../cli/commands/test_dag_processor_command.py | 23 ++++++++++++++++++++++
2 files changed, 35 insertions(+), 1 deletion(-)
diff --git a/airflow-core/src/airflow/utils/cli.py
b/airflow-core/src/airflow/utils/cli.py
index 6af75dd5b5e..d0105e86efa 100644
--- a/airflow-core/src/airflow/utils/cli.py
+++ b/airflow-core/src/airflow/utils/cli.py
@@ -110,7 +110,18 @@ def action_cli(func=None, check_db=True):
if conf.getboolean("database", "check_migrations"):
check_and_run_migrations()
synchronize_log_template()
- return f(*args, **kwargs)
+ # Set server context so validation logic (e.g.
DagBundlesManager) runs correctly
+ original_ctx = os.environ.pop("_AIRFLOW_PROCESS_CONTEXT",
None)
+ try:
+ os.environ["_AIRFLOW_PROCESS_CONTEXT"] = "server"
+ return f(*args, **kwargs)
+ finally:
+ if original_ctx is None:
+ os.environ.pop("_AIRFLOW_PROCESS_CONTEXT", None)
+ else:
+ os.environ["_AIRFLOW_PROCESS_CONTEXT"] =
original_ctx
+ else:
+ return f(*args, **kwargs)
except Exception as e:
metrics["error"] = e
raise
diff --git a/airflow-core/tests/unit/cli/commands/test_dag_processor_command.py
b/airflow-core/tests/unit/cli/commands/test_dag_processor_command.py
index 6224be2c5d5..abdf52ad125 100644
--- a/airflow-core/tests/unit/cli/commands/test_dag_processor_command.py
+++ b/airflow-core/tests/unit/cli/commands/test_dag_processor_command.py
@@ -57,6 +57,29 @@ class TestDagProcessorCommand:
dag_processor_command.dag_processor(args)
assert mock_runner.call_args.kwargs["processor"].bundle_names_to_parse
== ["testing"]
+ @conf_vars({("core", "load_examples"): "False"})
+
@mock.patch("airflow.cli.commands.dag_processor_command.DagProcessorJobRunner")
+ @mock.patch("airflow.utils.cli.DagBundlesManager", autospec=True)
+ def test_bundle_validation_runs_with_server_context(self,
mock_manager_cls, mock_runner, monkeypatch):
+ mock_runner.return_value.job_type = "DagProcessorJob"
+ captured_ctx = {}
+
+ mock_bundle = mock.MagicMock()
+ mock_bundle.name = "bundle1"
+
+ def capture_env_and_return_bundles():
+ captured_ctx["during_validation"] =
os.environ.get("_AIRFLOW_PROCESS_CONTEXT")
+ return [mock_bundle]
+
+ mock_manager_cls.return_value.get_all_dag_bundles.side_effect =
capture_env_and_return_bundles
+
+ monkeypatch.delenv("_AIRFLOW_PROCESS_CONTEXT", raising=False)
+ args = self.parser.parse_args(["dag-processor", "--bundle-name",
"bundle1"])
+ dag_processor_command.dag_processor(args)
+
+ assert captured_ctx["during_validation"] == "server"
+ assert "_AIRFLOW_PROCESS_CONTEXT" not in os.environ
+
@mock.patch("airflow.cli.hot_reload.run_with_reloader")
def test_dag_processor_with_dev_flag(self, mock_reloader):
"""Ensure that dag-processor with --dev flag uses hot-reload"""
