From 37fffa618f9ec2fdc09abde6faa6b40b28151e0e Mon Sep 17 00:00:00 2001 From: Polina Bungina <27892524+hughcapet@users.noreply.github.com> Date: Fri, 26 May 2023 15:13:04 +0200 Subject: [PATCH] Refactor daemon entrypoints (#2697) - abstract_main only creates Config object using the passed configfile and instantiates the passed daemon class - common args parser is extracted into a separate func that is called from daemons' main funcs (specific args can be added afterwards) Co-authored-by: Alexander Kukushkin --- patroni/__main__.py | 33 ++++++++++++++++++++----- patroni/daemon.py | 50 ++++++++++++++++---------------------- patroni/raft_controller.py | 7 ++++-- tests/test_patroni.py | 12 ++++----- 4 files changed, 59 insertions(+), 43 deletions(-) diff --git a/patroni/__main__.py b/patroni/__main__.py index 9c7f5cea..2a669d3f 100644 --- a/patroni/__main__.py +++ b/patroni/__main__.py @@ -1,11 +1,13 @@ import logging import os import signal +import sys import time +from argparse import Namespace from typing import Any, Dict, Optional, TYPE_CHECKING -from patroni.daemon import AbstractPatroniDaemon, abstract_main +from patroni.daemon import AbstractPatroniDaemon, abstract_main, get_base_arg_parser if TYPE_CHECKING: # pragma: no cover from .config import Config @@ -133,21 +135,40 @@ class Patroni(AbstractPatroniDaemon): logger.exception('Exception during Ha.shutdown') -def patroni_main() -> None: +def patroni_main(configfile: str) -> None: from multiprocessing import freeze_support - from patroni.validator import schema freeze_support() - abstract_main(Patroni, schema) + abstract_main(Patroni, configfile) + + +def process_arguments() -> Namespace: + parser = get_base_arg_parser() + parser.add_argument('--validate-config', action='store_true', help='Run config validator and exit') + args = parser.parse_args() + + if args.validate_config: + from patroni.validator import schema + from patroni.config import Config, ConfigParseError + + try: + Config(args.configfile, validator=schema) + sys.exit() + except ConfigParseError as e: + sys.exit(e.value) + + return args def main() -> None: from patroni import check_psycopg + args = process_arguments() + check_psycopg() if os.getpid() != 1: - return patroni_main() + return patroni_main(args.configfile) # Patroni started with PID=1, it looks like we are in the container from types import FrameType @@ -180,7 +201,7 @@ def main() -> None: signal.signal(signal.SIGTERM, passtochild) import multiprocessing - patroni = multiprocessing.Process(target=patroni_main) + patroni = multiprocessing.Process(target=patroni_main, args=(args.configfile,)) patroni.start() pid = patroni.pid patroni.join() diff --git a/patroni/daemon.py b/patroni/daemon.py index 88f6b133..e160ce46 100644 --- a/patroni/daemon.py +++ b/patroni/daemon.py @@ -6,6 +6,7 @@ Currently it is only used for the main "Thread" of ``patroni`` and ``patroni_raf from __future__ import print_function import abc +import argparse import os import signal import sys @@ -15,7 +16,22 @@ from typing import Any, Optional, Type, TYPE_CHECKING if TYPE_CHECKING: # pragma: no cover from .config import Config - from .validator import Schema + + +def get_base_arg_parser() -> argparse.ArgumentParser: + """Create a basic argument parser with the arguments used for both patroni and raft controller daemon. + + :returns: 'argparse.ArgumentParser' object + """ + from .config import Config + from .version import __version__ + + parser = argparse.ArgumentParser() + parser.add_argument('--version', action='version', version='%(prog)s {0}'.format(__version__)) + parser.add_argument('configfile', nargs='?', default='', + help='Patroni may also read the configuration from the {0} environment variable' + .format(Config.PATRONI_CONFIG_VARIABLE)) + return parser class AbstractPatroniDaemon(abc.ABC): @@ -141,41 +157,17 @@ class AbstractPatroniDaemon(abc.ABC): self.logger.shutdown() -def abstract_main(cls: Type[AbstractPatroniDaemon], validator: Optional['Schema'] = None) -> None: +def abstract_main(cls: Type[AbstractPatroniDaemon], configfile: str) -> None: """Create the main entry point of a given daemon process. - Expose a basic argument parser, parse the command-line arguments, and run the given daemon process. - :param cls: a class that should inherit from :class:`AbstractPatroniDaemon`. - :param validator: used to validate the daemon configuration schema, if requested by the user through - ``--validate-config`` CLI option. + :param configfile: """ - import argparse - from .config import Config, ConfigParseError - from .version import __version__ - - parser = argparse.ArgumentParser() - parser.add_argument('--version', action='version', version='%(prog)s {0}'.format(__version__)) - if validator: - parser.add_argument('--validate-config', action='store_true', help='Run config validator and exit') - parser.add_argument('configfile', nargs='?', default='', - help='Patroni may also read the configuration from the {0} environment variable' - .format(Config.PATRONI_CONFIG_VARIABLE)) - args = parser.parse_args() - validate_config = validator and args.validate_config try: - if validate_config: - Config(args.configfile, validator=validator) - sys.exit() - - config = Config(args.configfile) + config = Config(configfile) except ConfigParseError as e: - if e.value: - print(e.value, file=sys.stderr) - if not validate_config: - parser.print_help() - sys.exit(1) + sys.exit(e.value) controller = cls(config) try: diff --git a/patroni/raft_controller.py b/patroni/raft_controller.py index 2f9f7858..a9d7424c 100644 --- a/patroni/raft_controller.py +++ b/patroni/raft_controller.py @@ -1,7 +1,7 @@ import logging from .config import Config -from .daemon import AbstractPatroniDaemon, abstract_main +from .daemon import AbstractPatroniDaemon, abstract_main, get_base_arg_parser from .dcs.raft import KVStoreTTL logger = logging.getLogger(__name__) @@ -27,4 +27,7 @@ class RaftController(AbstractPatroniDaemon): def main() -> None: - abstract_main(RaftController) + parser = get_base_arg_parser() + args = parser.parse_args() + + abstract_main(RaftController, args.configfile) diff --git a/tests/test_patroni.py b/tests/test_patroni.py index ad8e775c..cb40b232 100644 --- a/tests/test_patroni.py +++ b/tests/test_patroni.py @@ -15,7 +15,7 @@ from patroni.exceptions import DCSError from patroni.postgresql import Postgresql from patroni.postgresql.config import ConfigHandler from patroni import check_psycopg -from patroni.__main__ import Patroni, main as _main, patroni_main +from patroni.__main__ import Patroni, main as _main from threading import Thread from . import psycopg_connect, SleepException @@ -52,14 +52,14 @@ class TestPatroni(unittest.TestCase): @patch('sys.argv', ['patroni.py']) def test_no_config(self): - self.assertRaises(SystemExit, patroni_main) + self.assertRaises(SystemExit, _main) @patch('sys.argv', ['patroni.py', '--validate-config', 'postgres0.yml']) @patch('socket.socket.connect_ex', Mock(return_value=1)) def test_validate_config(self): - self.assertRaises(SystemExit, patroni_main) + self.assertRaises(SystemExit, _main) with patch.object(config.Config, '__init__', Mock(return_value=None)): - self.assertRaises(SystemExit, patroni_main) + self.assertRaises(SystemExit, _main) @patch('pkgutil.iter_importers', Mock(return_value=[MockFrozenImporter()])) @patch('sys.frozen', Mock(return_value=True), create=True) @@ -94,11 +94,11 @@ class TestPatroni(unittest.TestCase): with patch('subprocess.call', Mock(return_value=1)): with patch.object(Patroni, 'run', Mock(side_effect=SleepException)): os.environ['PATRONI_POSTGRESQL_DATA_DIR'] = 'data/test0' - self.assertRaises(SleepException, patroni_main) + self.assertRaises(SleepException, _main) with patch.object(Patroni, 'run', Mock(side_effect=KeyboardInterrupt())): with patch('patroni.ha.Ha.is_paused', Mock(return_value=True)): os.environ['PATRONI_POSTGRESQL_DATA_DIR'] = 'data/test0' - patroni_main() + _main() @patch('os.getpid') @patch('multiprocessing.Process')