Added config reader to the project.

pgbouncemgr/config.py

@@ -0,0 +1,163 @@
+# -*- coding: utf-8 -*-
+"""This module encapsulates the configuration for pgbouncemgr.
+   The configuration is read from a YAML file."""
+
+import yaml
+
+class ConfigException(Exception):
+    """Used for all exceptions that are raised from pgbouncemgr.config."""
+
+class ConfigNotFound(ConfigException):
+    """Raised when the requested configuration file does not exist."""
+    def __init__(self, path):
+        super().__init__(
+            "Configuration file not found (path=%s)" % path)
+
+class InvalidConfigSection(ConfigException):
+    """Raised when the configuration file contains an invalid section."""
+    def __init__(self, section):
+        super().__init__(
+            "Configuration file contains an invalid section (section=%s)" % section)
+
+class InvalidConfigKey(ConfigException):
+    """Raised when the configuration file contains an invalid key."""
+    def __init__(self, section, key):
+        super().__init__(
+            "Configuration file contains an invalid key (key=%s.%s)" %
+            (section, key))
+
+class InvalidConfigValue(ConfigException):
+    """Raised when the configuration file contains an invalid value."""
+    def __init__(self, reason, section, key, value):
+        super().__init__(
+            "Configuration file contains an invalid value: %s (key=%s.%s, value=%s)" %
+            (reason, section, key, value))
+
+
+class Config():
+    def __init__(self, path):
+        self._set_defaults()
+        parser = self._read(path)
+        self._build(parser)
+        self._resolve_db_connection_defaults()
+
+    def _set_defaults(self):
+        self.run_user = "postgres"
+        self.run_group = "postgres"
+        self.state_file = "/var/lib/pgbouncemgr/state.json"
+        self.poll_interval_in_sec = 2
+        self.db_connection_defaults = {
+            "host": "0.0.0.0",
+            "port": 5432,
+            "connect_timeout": 1,
+            "user": "pgbouncemgr",
+            "password": None,
+            "name": "template1",
+        }
+        self.pgbouncer = {
+            "pgbouncer_config": "/etc/pgbouncer/pgbouncer.ini",
+            "port": 6432,
+        }
+        self.nodes = {}
+
+    def _read(self, path):
+        try:
+            with open(path, 'r') as stream:
+                return yaml.safe_load(stream)
+        except FileNotFoundError:
+            raise ConfigNotFound(path)
+
+    def _build(self, parser):
+        if not hasattr(parser, "items"):
+            return
+        for section, data in parser.items():
+            if section == "main":
+                self._build_main(data)
+            elif section == "pgbouncer":
+                self._build_pgbouncer(data)
+            elif section == "db_connection_defaults":
+                self._build_db_connection_defaults(data)
+            elif section == "nodes":
+                self._build_nodes(data)
+            else:
+                raise InvalidConfigSection(section)
+
+    def _build_main(self, data):
+        if not data:
+            return
+        for key in data:
+            if not hasattr(self, key):
+                raise InvalidConfigKey("main", key)
+            value = data[key]
+            if key == "poll_interval_in_sec":
+                value = self._to_int("main", key, value)
+            else:
+                self._check_for_empty_value("main", key, value)
+            setattr(self, key, value)
+
+    def _build_pgbouncer(self, data):
+        for key in data:
+            if key not in self.pgbouncer:
+                raise InvalidConfigKey("pgbouncer", key)
+            value = data[key]
+            if key in ["port", "connect_timeout"]:
+                value = self._to_int("pgbouncer", key, value)
+            else:
+                self._check_for_empty_value("pgbouncer", key, value)
+            self.pgbouncer[key] = value
+
+
+    def _build_db_connection_defaults(self, data):
+        if not data:
+            return
+        for key in data:
+            if key not in self.db_connection_defaults:
+                raise InvalidConfigKey("db_connection_defaults", key)
+            value = data[key]
+            if key in ["port", "connect_timeout"]:
+                value = self._to_int("db_connection_defaults", key, value)
+            else:
+                self._check_for_empty_value("db_connection_defaults", key, value)
+            self.db_connection_defaults[key] = value
+
+    def _build_nodes(self, data):
+        if not hasattr(data, "items"):
+            return
+        for node_name, node_data in data.items():
+            self._build_node(node_name, node_data)
+
+    def _build_node(self, node_name, data):
+        node = dict()
+        if not data:
+          return
+        for key in data:
+            value = data[key]
+            if key in ["port", "connect_timeout"]:
+                value = self._to_int("node", key, value)
+            else:
+                self._check_for_empty_value("node[%s]" % node_name, key, value)
+            node[key] = value
+        self.nodes[node_name] = node
+
+    def _resolve_db_connection_defaults(self):
+        # Resolve defaults for the node connections.
+        for node in self.nodes.values():
+            for key, default_value in self.db_connection_defaults.items():
+                if key not in node:
+                    node[key] = default_value
+
+        # Resolve defaults for the pgbouncer console connection.
+        for key, default_value in self.db_connection_defaults.items():
+            if key not in self.pgbouncer or self.pgbouncer[key] is None:
+                self.pgbouncer[key] = default_value
+
+    def _to_int(self, section, key, value):
+        try:
+            return int(value)
+        except (TypeError, ValueError):
+            raise InvalidConfigValue("it must be an integer", section, key, value)
+
+    def _check_for_empty_value(self, section, key, value):
+        if value is None or (str(value)).strip() == "":
+            value = value if value is None else repr(value)
+            raise InvalidConfigValue("it must not be an empty value", section, key, value)

pgbouncemgr/state.py

@@ -44,7 +44,7 @@ class NodeCannotBePromoted(StateException):
 class InvalidSystemId(StateException):
     """Raised when an invalid system_id is used."""
     def __init__(self, invalid):
-        display = "None" if invalid is None else "'%s'" % invalid
+        display = invalid if invalid is None else repr(invalid)
         super().__init__(
             "Invalid system_id provided; it must be a non-empty string " +
             "(invalid system_id=%s)" % display)

tests/test_config.py

@@ -0,0 +1,90 @@
+# -*- coding: utf-8 -*-
+
+import os
+import unittest
+import json
+from pgbouncemgr.config import *
+
+def load_test_config(name):
+    path = os.path.join(
+        os.path.dirname(os.path.realpath(__file__)),
+        "testfiles", name)
+    return Config(path)
+
+class ConfigTests(unittest.TestCase):
+    def test_ConfigThatDoesNotExist_RaisesException(self):
+        with self.assertRaises(ConfigNotFound) as context:
+            load_test_config("/path/doesnotexist.yaml") 
+        self.assertIn("path=/path/doesnotexist.yaml", str(context.exception))
+
+    def test_ConfigContainingInvalidSection_RaisesException(self):
+        with self.assertRaises(InvalidConfigSection) as context:
+            load_test_config("invalid_section.yaml") 
+        self.assertIn("section=unknown_section", str(context.exception))
+
+    def test_ConfigContainingInvalidKey_RaisesException(self):
+        with self.assertRaises(InvalidConfigKey) as context:
+            load_test_config("invalid_key.yaml") 
+        self.assertIn("key=main.improbability_drive_level", str(context.exception))
+
+    def test_ConfigContainingInvalidIntegerValue_RaisesException(self):
+        with self.assertRaises(InvalidConfigValue) as context:
+            load_test_config("invalid_main_interval_in_sec.yaml") 
+        self.assertIn("it must be an integer", str(context.exception))
+        self.assertIn("key=main.poll_interval_in_sec, value=ten seconds", str(context.exception))
+
+    def test_ConfigContainingEmptyValue_RaisesException(self):
+        with self.assertRaises(InvalidConfigValue) as context:
+            load_test_config("invalid_node_empty_host.yaml") 
+        self.assertIn("it must not be an empty value", str(context.exception))
+        self.assertIn(r"key=node[nodeA].host, value='  \t\r\n  '", str(context.exception))
+
+    def test_ConfigFromEmptyConfigurationFile_UsesDefaultValues(self):
+        config = load_test_config("empty.yaml") 
+        # main
+        self.assertEqual("postgres", config.run_user)
+        self.assertEqual("postgres", config.run_group)
+        self.assertEqual("/var/lib/pgbouncemgr/state.json", config.state_file)
+        self.assertEqual(2, config.poll_interval_in_sec)
+        # db connection defaults
+        self.assertEqual("0.0.0.0", config.db_connection_defaults["host"])
+        self.assertEqual(5432, config.db_connection_defaults["port"])
+        self.assertEqual(1, config.db_connection_defaults["connect_timeout"])
+        self.assertEqual("pgbouncemgr", config.db_connection_defaults["user"])
+        self.assertEqual(None, config.db_connection_defaults["password"])
+        self.assertEqual("template1", config.db_connection_defaults["name"])
+        # pgbouncer
+        self.assertEqual("/etc/pgbouncer/pgbouncer.ini", config.pgbouncer["pgbouncer_config"])
+        self.assertEqual(6432, config.pgbouncer["port"])
+        # nodes
+        self.assertEqual({}, config.nodes)
+
+    def test_ConfigFromBasicConfigurationFile_DbConnectionDefaultsAreApplied(self):
+        config = load_test_config("basic.yaml") 
+        self.assertEqual({
+            "connect_timeout": 1,
+            "host": "0.0.0.0",
+            "name": "template1",
+            "password": "Wilmaaaaa!!!",
+            "pgbouncer_config": "/etc/pgbouncer/pgbouncer.ini",
+            "port": 6432,
+            "user": "pgbouncemgr"
+        }, config.pgbouncer)
+        self.assertEqual({
+            "nodeA": {
+                "connect_timeout": 1,
+                "host": "1.2.3.4",
+                "name": "template1",
+                "password": "Wilmaaaaa!!!",
+                "port": 8888,
+                "user": "pgbouncemgr"
+            },
+            "nodeB": {
+                "connect_timeout": 1,
+                "host": "2.3.4.5",
+                "name": "template1",
+                "password": "Wilmaaaaa!!!",
+                "port": 7777,
+                "user": "pgbouncemgr"
+            }
+        }, config.nodes)

tests/test_state.py

@@ -309,7 +309,7 @@ class NodeTests(unittest.TestCase):
 
         with self.assertRaises(InvalidSystemId) as context:
             node.system_id = "  \t\r\n\t   "
-        self.assertIn("invalid system_id='  \t\r\n\t   '", str(context.exception))
+        self.assertIn(r"invalid system_id='  \t\r\n\t   '", str(context.exception))
         self.assertFalse(state.modified)
 
     def test_SetSystemId_SetsSystemId_AndNotifiesChangeToState(self):

tests/testfiles/basic.yaml

@@ -0,0 +1,11 @@
+---
+db_connection_defaults:
+  port: 8888
+  password: Wilmaaaaa!!!
+
+nodes:
+  nodeA:
+    host: 1.2.3.4
+  nodeB:
+    host: 2.3.4.5
+    port: 7777

tests/testfiles/empty.yaml

@@ -0,0 +1 @@
+# Empty configuration file.

tests/testfiles/invalid_key.yaml

@@ -0,0 +1,3 @@
+---
+main:
+  improbability_drive_level: 42 

tests/testfiles/invalid_main_interval_in_sec.yaml

@@ -0,0 +1,3 @@
+---
+main:
+  poll_interval_in_sec: ten seconds

tests/testfiles/invalid_node_empty_host.yaml

@@ -0,0 +1,4 @@
+---
+nodes:
+  nodeA:
+    host: "  \t\r\n  "

tests/testfiles/invalid_section.yaml

@@ -0,0 +1,5 @@
+---
+main:
+db_connection_defaults:
+nodes:
+unknown_section: