pgbouncemgr/pgbouncemgr/config.py

# -*- coding: utf-8 -*-
# pylint: disable=no-self-use,too-few-public-methods
"""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():
    """This class is used to load the pgbouncemgr configuration from
       a yaml file and to enrich it with default values where needed."""
    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,
            "database": "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)