[airflow] branch main updated: Print configuration on scheduler startup. (#22588)

[po...@apache.org]

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 78586b45a0 Print configuration on scheduler startup. (#22588)
78586b45a0 is described below

commit 78586b45a0f6007ab6b94c35b33790a944856e5e
Author: Karthikeyan Singaravelan <ti...@gmail.com>
AuthorDate: Mon Apr 4 17:35:49 2022 +0530

    Print configuration on scheduler startup. (#22588)
---
 airflow/cli/commands/scheduler_command.py    | 12 ++++++++-
 airflow/utils/cli.py                         | 39 ++++++++++++++++++++++++++++
 docs/apache-airflow/howto/set-config.rst     |  4 +++
 tests/cli/commands/test_scheduler_command.py | 18 +++++++++++++
 4 files changed, 72 insertions(+), 1 deletion(-)

diff --git a/airflow/cli/commands/scheduler_command.py b/airflow/cli/commands/scheduler_command.py
index bc6e983ee5..c4ecdf411a 100644
--- a/airflow/cli/commands/scheduler_command.py
+++ b/airflow/cli/commands/scheduler_command.py
@@ -26,7 +26,15 @@ from daemon.pidfile import TimeoutPIDLockFile
 from airflow import settings
 from airflow.jobs.scheduler_job import SchedulerJob
 from airflow.utils import cli as cli_utils
-from airflow.utils.cli import process_subdir, setup_locations, setup_logging, sigint_handler, sigquit_handler
+from airflow.utils.cli import (
+    get_config_with_source,
+    process_subdir,
+    setup_locations,
+    setup_logging,
+    sigconf_handler,
+    sigint_handler,
+    sigquit_handler,
+)
 
 
 def _create_scheduler_job(args):
@@ -53,6 +61,7 @@ def _run_scheduler_job(args):
 def scheduler(args):
     """Starts Airflow Scheduler"""
     print(settings.HEADER)
+    print(get_config_with_source())
 
     if args.daemon:
         pid, stdout, stderr, log_file = setup_locations(
@@ -72,6 +81,7 @@ def scheduler(args):
         signal.signal(signal.SIGINT, sigint_handler)
         signal.signal(signal.SIGTERM, sigint_handler)
         signal.signal(signal.SIGQUIT, sigquit_handler)
+        signal.signal(signal.SIGUSR1, sigconf_handler)
         _run_scheduler_job(args=args)
 
 
diff --git a/airflow/utils/cli.py b/airflow/utils/cli.py
index a15e1b7f02..9629951a67 100644
--- a/airflow/utils/cli.py
+++ b/airflow/utils/cli.py
@@ -18,6 +18,7 @@
 #
 """Utilities module for cli"""
 import functools
+import io
 import json
 import logging
 import os
@@ -28,12 +29,18 @@ import threading
 import traceback
 import warnings
 from argparse import Namespace
+from contextlib import redirect_stdout
 from datetime import datetime
 from typing import TYPE_CHECKING, Callable, Optional, TypeVar, cast
 
+import pygments
+from pygments.lexers.configs import IniLexer
+
 from airflow import settings
+from airflow.configuration import conf
 from airflow.exceptions import AirflowException
 from airflow.utils import cli_action_loggers
+from airflow.utils.code_utils import get_terminal_formatter
 from airflow.utils.log.non_caching_file_handler import NonCachingFileHandler
 from airflow.utils.platform import getuser, is_terminal_support_colors
 from airflow.utils.session import provide_session
@@ -271,6 +278,13 @@ def sigint_handler(sig, frame):
     sys.exit(0)
 
 
+def sigconf_handler(sig, frame):
+    """Print configuration and source including default values."""
+    config = get_config_with_source(include_default=True)
+    log = logging.getLogger(__name__)
+    log.info(config)
+
+
 def sigquit_handler(sig, frame):
     """
     Helps debug deadlocks by printing stacktraces when this gets a SIGQUIT
@@ -328,3 +342,28 @@ def suppress_logs_and_warning(f: T) -> T:
                     logging.disable(logging.NOTSET)
 
     return cast(T, _wrapper)
+
+
+def get_config_with_source(include_default: bool = False) -> str:
+    """Return configuration along with source for each option."""
+    config_dict = conf.as_dict(display_source=True)
+
+    with io.StringIO() as buf, redirect_stdout(buf):
+        for section, options in config_dict.items():
+            if not include_default:
+                options = {
+                    key: (value, source) for key, (value, source) in options.items() if source != "default"
+                }
+
+            # Print the section only when there are options after filtering
+            if options:
+                print(f"[{section}]")
+                for key, (value, source) in options.items():
+                    print(f"{key} = {value} [{source}]")
+                print()
+        code = buf.getvalue()
+
+        if is_terminal_support_colors():
+            code = pygments.highlight(code=code, formatter=get_terminal_formatter(), lexer=IniLexer())
+
+        return code
diff --git a/docs/apache-airflow/howto/set-config.rst b/docs/apache-airflow/howto/set-config.rst
index a31885b2da..75762c1a74 100644
--- a/docs/apache-airflow/howto/set-config.rst
+++ b/docs/apache-airflow/howto/set-config.rst
@@ -126,6 +126,10 @@ the example below.
     work as expected. A good example for that is :ref:`secret_key<config:webserver__secret_key>` which
     should be same on the Webserver and Worker to allow Webserver to fetch logs from Worker.
 
+    During startup the scheduler prints configuration options whose values differ from the default
+    values along with the source from which the value was loaded (i.e. airflow.cfg, environment variable, etc).
+    You can print all configuration options by sending ``SIGUSR1`` signal to the scheduler.
+
     The webserver key is also used to authorize requests to Celery workers when logs are retrieved. The token
     generated using the secret key has a short expiry time though - make sure that time on ALL the machines
     that you run airflow components on is synchronized (for example using ntpd) otherwise you might get
diff --git a/tests/cli/commands/test_scheduler_command.py b/tests/cli/commands/test_scheduler_command.py
index 00b851f55a..6ef1013dcc 100644
--- a/tests/cli/commands/test_scheduler_command.py
+++ b/tests/cli/commands/test_scheduler_command.py
@@ -15,6 +15,9 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
+import contextlib
+import io
+import os
 import unittest
 from unittest import mock
 
@@ -87,3 +90,18 @@ class TestSchedulerCommand(unittest.TestCase):
                 scheduler_command.scheduler(args)
             finally:
                 mock_process().terminate.assert_called()
+
+    @mock.patch.dict(os.environ, {'AIRFLOW__CORE__PARALLELISM': '35'}, clear=True)
+    @mock.patch("airflow.cli.commands.scheduler_command.SchedulerJob")
+    @mock.patch("airflow.cli.commands.scheduler_command.Process")
+    def test_scheduler_print_config(self, mock_process, mock_scheduler_job):
+        args = self.parser.parse_args(['scheduler'])
+        with conf_vars({("core", "sql_alchemy_conn_cmd"): 'echo hello'}):
+            with contextlib.redirect_stdout(io.StringIO()) as temp_stdout:
+                scheduler_command.scheduler(args)
+                output = temp_stdout.getvalue()
+
+                assert "sql_alchemy_conn = < hidden > [cmd]" in output
+                assert "max_active_tasks_per_dag = 16 [airflow.cfg]" in output
+                assert "parallelism = < hidden > [env var]" in output
+                assert "max_active_runs_per_dag = 16 [default]" not in output