diff --git a/Dockerfile-dev b/Dockerfile-dev index dba6d16..e1250e7 100644 --- a/Dockerfile-dev +++ b/Dockerfile-dev @@ -1,8 +1,9 @@ FROM python:3.6-alpine # Install dependencies -COPY requirements.txt /prometheus-pgbouncer-exporter/ +COPY . /prometheus-pgbouncer-exporter RUN apk add --update --no-cache postgresql-dev gcc musl-dev && \ pip install -r /prometheus-pgbouncer-exporter/requirements.txt WORKDIR /prometheus-pgbouncer-exporter +RUN python setup.py install diff --git a/prometheus_pgbouncer_exporter/cli.py b/prometheus_pgbouncer_exporter/cli.py index 7ff8eac..b81a5da 100644 --- a/prometheus_pgbouncer_exporter/cli.py +++ b/prometheus_pgbouncer_exporter/cli.py @@ -11,8 +11,29 @@ from .collector import PgbouncersMetricsCollector +def read_config_file(config_file): + # Read config file + config = Config() + try: + config.read(config_file) + logging.getLogger().info("Config file successfully read from {file}".format(file=config_file)) + except Exception as error: + logging.getLogger().fatal("Unable to read config file from {file}".format(file=config_file), extra={"exception": str(error)}) + raise Exception('Cannot read config file') + + # Validate config + try: + config.validate() + except Exception as error: + logging.getLogger().fatal("The config file {file} is invalid: {error}".format(file=config_file, error=str(error))) + raise Exception('Invalid config file') + + return config + + def main(): shutdown = False + sighup_received_count = 0 # Parse arguments parser = argparse.ArgumentParser() @@ -31,6 +52,13 @@ def main(): # Register signal handler def _on_sighup(signal, frame): + # we want to do minimum work in the signal handling for HUP, so that a HUP received + # while handling another HUP won't change the execution order of the two HUPs + # (second and then the first) and leave us in a broken state. + nonlocal sighup_received_count + sighup_received_count += 1 + logging.getLogger().info("Received SIGHUP - Incrementing HUP counter to %s", + sighup_received_count) if args.log_file is not "stdout": logging.getLogger().info("Received SIGHUP - log file is closing") logHandler.close() @@ -46,30 +74,24 @@ def _on_sigterm(signal, frame): signal.signal(signal.SIGTERM, _on_sigterm) # Read config file - config = Config() - try: - config.read(args.config) - logging.getLogger().info("Config file successfully read from {file}".format(file=args.config)) - except Exception as error: - logging.getLogger().fatal("Unable to read config file from {file}".format(file=args.config), extra={"exception": str(error)}) - sys.exit(1) - - # Validate config - try: - config.validate() - except Exception as error: - logging.getLogger().fatal("The config file {file} is invalid: {error}".format(file=args.config, error=str(error))) - sys.exit(1) + config = read_config_file(args.config) # Register our custom collector - REGISTRY.register(PgbouncersMetricsCollector(config.getPgbouncers())) + pgbcollector = PgbouncersMetricsCollector(config.getPgbouncers()) + REGISTRY.register(pgbcollector) # Start server start_http_server(config.getExporterPort(), config.getExporterHost()) logging.getLogger().info("Exporter listening on {host}:{port}".format(host=config.getExporterHost(), port=config.getExporterPort())) while not shutdown: - time.sleep(1) + # check if HUPs were received in the last cycle and subsequently reload the pgbouncer hosts + if sighup_received_count: + logging.getLogger().info("Processing SIGHUP - Reloading pgbouncer hosts") + sighup_received_count = 0 + config = read_config_file(args.config) + pgbcollector.update(config.getPgbouncers()) + time.sleep(4) logging.getLogger().info("Exporter has shutdown") diff --git a/prometheus_pgbouncer_exporter/collector.py b/prometheus_pgbouncer_exporter/collector.py index 885c0b8..1e04b45 100644 --- a/prometheus_pgbouncer_exporter/collector.py +++ b/prometheus_pgbouncer_exporter/collector.py @@ -8,6 +8,9 @@ class PgbouncersMetricsCollector(): def __init__(self, configs: List[PgbouncerConfig]): + self.update(configs) + + def update(self, configs: List[PgbouncerConfig]): self.collectors = list(map(lambda config: PgbouncerMetricsCollector(config), configs)) def collect(self): diff --git a/tests/test_cli.py b/tests/test_cli.py new file mode 100644 index 0000000..3a39971 --- /dev/null +++ b/tests/test_cli.py @@ -0,0 +1,98 @@ +import unittest +import os +import signal +import threading +import time + +from prometheus_pgbouncer_exporter import cli +from unittest import mock + + +CURR_DIR=os.path.dirname(os.path.realpath(__file__)) + + +class TestCLI(unittest.TestCase): + + @mock.patch('prometheus_pgbouncer_exporter.cli.logging') + def testReadConfigFile(self, mock_logging): + mock_logger = mock.MagicMock() + mock_logging.getLogger.return_value = mock_logger + + # should read and validate this config file without any issues + config = cli.read_config_file(CURR_DIR + '/fixtures/config-with-one-pgbouncer.yml') + + self.assertEqual(mock_logger.info.call_count, 1) + self.assertEqual(config.getExporterHost(), "0.0.0.0") + self.assertEqual(config.getExporterPort(), 1234) + self.assertEqual(len(config.getPgbouncers()), 1) + + self.assertEqual(config.getPgbouncers()[0].getDsn(), "postgresql://user:password@host:6431/pgbouncer") + self.assertEqual(config.getPgbouncers()[0].getDsnWithMaskedPassword(), "postgresql://user:***@host:6431/pgbouncer") + self.assertEqual(config.getPgbouncers()[0].getConnectTimeout(), 2) + self.assertEqual(config.getPgbouncers()[0].getIncludeDatabases(), ["one", "two"]) + self.assertEqual(config.getPgbouncers()[0].getExcludeDatabases(), ["three"]) + self.assertEqual(config.getPgbouncers()[0].getExtraLabels(), {"first": "1", "second": "2"}) + + # Should raise exception saying 'Cannot read config file' + with self.assertRaises(Exception) as err: + cli.read_config_file(CURR_DIR + '/path/to/none.yml') + + self.assertEqual('Cannot read config file', str(err.exception)) + self.assertEqual(mock_logger.fatal.call_count, 1) + + # Should raise exception saying 'Invalid config file' + with self.assertRaises(Exception) as err: + cli.read_config_file(CURR_DIR + '/fixtures/config-empty.yml') + + self.assertEqual('Invalid config file', str(err.exception)) + self.assertEqual(mock_logger.fatal.call_count, 2) + + @mock.patch('prometheus_pgbouncer_exporter.cli.os._exit') + @mock.patch('prometheus_pgbouncer_exporter.collector.PgbouncersMetricsCollector.update') + @mock.patch('prometheus_pgbouncer_exporter.collector.PgbouncersMetricsCollector.collect') + @mock.patch('prometheus_pgbouncer_exporter.cli.argparse.ArgumentParser') + def testSignalHandling(self, mock_argparser, mock_collect, mock_update, mock_exit): + pid = os.getpid() + hup_count = 0 + + def send_signal(): + nonlocal hup_count + # Keep sending HUP until PgbouncersMetricsCollector.update is called at least twice. + # The first call is made when PgbouncersMetricsCollector class is initialized + # Any subsequent call happens as a result of config reload from HUP + while mock_update.call_count < 2: + os.kill(pid, signal.SIGHUP) + time.sleep(0.5) + hup_count += 1 + os.kill(pid, signal.SIGTERM) + + # initialize siguhup handler to ignore sighup, this will be overridden by cli.main + def _on_sighup(signal, frame): + pass + + signal.signal(signal.SIGHUP, _on_sighup) + + args = mock.MagicMock() + args.config = CURR_DIR + '/fixtures/config-with-one-pgbouncer.yml' + args.log_file = 'stdout' + args.log_level = 'CRITICAL' + + parser = mock.MagicMock() + parser.parse_args.return_value = args + mock_argparser.return_value = parser + + thread = threading.Thread(target=send_signal) + thread.daemon = True + thread.start() + + self.assertEqual(mock_update.call_count, 0) + cli.main() + self.assertEqual(mock_update.call_count, 2) + mock_exit.assert_called_once_with(0) + + # test if hup handler aggregates hup counts and 'update' is not called on every HUP + self.assertGreater(hup_count, mock_update.call_count) + + +if __name__ == '__main__': + unittest.main()