Merge pull request #21 from zumper/master

Reload pgbouncer hosts info from config yaml file on HUP

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

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")
 

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):

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()