(airflow) branch main updated: Fix airflow-scheduler exiting with code 0 on exceptions (#36800)

[potiuk]

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 1d5d5022b8 Fix airflow-scheduler exiting with code 0 on exceptions 
(#36800)
1d5d5022b8 is described below

commit 1d5d5022b8fc92f23f9fdc3b61269e5c7acfaf39
Author: Joao Amaral <7281460+joaopama...@users.noreply.github.com>
AuthorDate: Mon Jan 15 17:59:38 2024 -0300

    Fix airflow-scheduler exiting with code 0 on exceptions (#36800)
    
    * Fix airflow-scheduler exiting with code 0 on exceptions
    
    * Fix static check
---
 airflow/cli/commands/scheduler_command.py    | 11 +++++++++--
 tests/cli/commands/test_scheduler_command.py |  3 ++-
 2 files changed, 11 insertions(+), 3 deletions(-)

diff --git a/airflow/cli/commands/scheduler_command.py 
b/airflow/cli/commands/scheduler_command.py
index fef0b97b2d..ed51381a21 100644
--- a/airflow/cli/commands/scheduler_command.py
+++ b/airflow/cli/commands/scheduler_command.py
@@ -19,7 +19,7 @@ from __future__ import annotations
 
 import logging
 from argparse import Namespace
-from contextlib import contextmanager
+from contextlib import ExitStack, contextmanager
 from multiprocessing import Process
 
 from airflow import settings
@@ -44,11 +44,18 @@ def _run_scheduler_job(args) -> None:
     
ExecutorLoader.validate_database_executor_compatibility(job_runner.job.executor)
     InternalApiConfig.force_database_direct_access()
     enable_health_check = conf.getboolean("scheduler", "ENABLE_HEALTH_CHECK")
-    with _serve_logs(args.skip_serve_logs), 
_serve_health_check(enable_health_check):
+    with ExitStack() as stack:
+        stack.enter_context(_serve_logs(args.skip_serve_logs))
+        stack.enter_context(_serve_health_check(enable_health_check))
+
         try:
             run_job(job=job_runner.job, execute_callable=job_runner._execute)
         except Exception:
             log.exception("Exception when running scheduler job")
+            raise
+        finally:
+            # Ensure that the contexts are closed
+            stack.close()
 
 
 @cli_utils.action_cli
diff --git a/tests/cli/commands/test_scheduler_command.py 
b/tests/cli/commands/test_scheduler_command.py
index 570a84fb66..608a537c6e 100644
--- a/tests/cli/commands/test_scheduler_command.py
+++ b/tests/cli/commands/test_scheduler_command.py
@@ -169,7 +169,8 @@ class TestSchedulerCommand:
         mock_scheduler_job,
     ):
         args = self.parser.parse_args(["scheduler"])
-        scheduler_command.scheduler(args)
+        with pytest.raises(Exception, match="run_job failed"):
+            scheduler_command.scheduler(args)
 
         # Make sure that run_job is called, that the exception has been 
logged, and that the serve_logs
         # sub-process has been terminated