Rough pgbouncer plugin for netdata

pgbouncer.chart.py

#!/usr/bin/env python
"""
Collects stats about pgbouncer client and postgres server connections
If psycopg2 is available it will use it instead of psql and will
query pgbouncer both for client and postgres server connections.

test out using:

    /opt/netdata/usr/libexec/netdata/plugins.d/python.d.plugin 1 debug trace pgbouncer

"""
from collections import OrderedDict, deque
import six

# present in the Netdata environment
try:
    from bases.FrameworkServices.SimpleService import SimpleService
    from bases.charts import ChartError, CHART_TYPES
except ImportError:
    from fake_netdata import SimpleService, ChartError, CHART_TYPES

try:
    import psycopg2
    import psycopg2.extensions
    import psycopg2.extras

    PSYCOPG2 = True
except ImportError:
    PSYCOPG2 = False


def _chart_f(
    chart_id,
    name=None,
    title=None,
    units="count",
    family=None,
    context=None,
    chart_type="line",
    lines=None,
):
    """Emits chart options with defaults for netdata plugin use"""
    # https://github.com/firehol/netdata/wiki/External-Plugins#chart
    assert chart_type in CHART_TYPES
    chart_vals = {"options": [name, title, units, family, context, chart_type], "lines": lines}
    return chart_id, chart_vals


CHARTS = OrderedDict(
    [
        _chart_f(
            "pgbouncer_client_connections",
            title="Current client connections to pgbouncer",
            context="show clients",
            family="pgbouncer statistics",
            lines=[["pgb_client_conns", "connections", "absolute"]],
        ),
        _chart_f(
            "pgbouncer_server_connections",
            title="Current server connections from pgbouncer",
            context="show servers",
            family="pgbouncer statistics",
            lines=[["pgb_server_conns", "connections", "absolute"]],
        ),
        _chart_f(
            "pgbouncer_pool_clients",
            title="Current pool client connections from pgbouncer",
            context="show pool clients",
            family="pgbouncer statistics",
            lines=deque(),
        ),
        _chart_f(
            "pgbouncer_pool_servers",
            title="Current pool server connections from pgbouncer",
            context="show pool servers",
            family="pgbouncer statistics",
            lines=deque(),
        ),
        _chart_f(
            "pgbouncer_pool_conns_avail",
            title="Available conns in pgbouncer pool",
            context="show pool conns available",
            family="pgbouncer statistics",
            units="%",
            lines=deque(),
        ),
    ]
)
ORDER = list(CHARTS.keys())

# default module values (can be overridden per job in `config`)
update_every = 2
priority = 90000
retries = 10


class PgConn(object):
    """wraps pg connections with an easy iterator"""

    conn = None

    def __init__(self, error_function, info_function, **params):
        self.params = params
        # logging hacks
        self.error = error_function
        self.info = info_function

    def __repr__(self):
        return "<{self.__class__.__name__}>{self.params!r}".format(self=self)

    def connect(self):
        """grab a new connection if needed."""
        if self.conn:
            return True
        try:
            conn = psycopg2.connect(**self.params)
            conn.set_isolation_level(psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT)
            self.info("Connected %s" % conn.dsn)
            self.conn = conn
            return True
        except psycopg2.OperationalError as error:
            self.error("Failed to connect %r. Error: %s" % (self, error))
            self.conn = None
            return False

    def query_iter(self, query):
        """iterate queries, handle exceptions"""
        try:
            self.connect()
            with self.conn.cursor(cursor_factory=psycopg2.extras.NamedTupleCursor) as cursor:
                cursor.execute(query)
                for record in cursor:
                    yield record
        except psycopg2.OperationalError as error:
            self.error("Error: %s" % error)
            self.conn = None
        except AttributeError:
            self.error("No connection.")


class Service(SimpleService):
    """
    Roughly Netdata will call `check()`, `create()`, and then periodically `get_data()`.

    https://github.com/firehol/netdata/blob/master/python.d/python_modules/bases/FrameworkServices/SimpleService.py
    """

    _pool_chart_names = tuple(chart for chart in CHARTS if chart.startswith("pgbouncer_pool_"))
    pgb_conn = None

    def __init__(self, configuration=None, name=None):
        super(Service, self).__init__(configuration=configuration, name=name)
        self.order = ORDER
        self.definitions = CHARTS

        params = dict(user="postgres", password=None, port="6432")
        params.update(self.configuration.get("db_params", {}))
        self.db_params = params

        self.known_entities = set()
        self.data = dict()

    def check(self):
        """netdata calls this first to see if the module is available"""
        if not PSYCOPG2:
            self.error("'python-psycopg2' module is needed to use postgres.chart.py")
            return False

        self.pgb_conn = PgConn(
            database="pgbouncer",
            info_function=self.info,
            error_function=self.error,
            **self.db_params
        )
        if not (self.pgb_conn and self.pgb_conn.connect()):
            return False

        # attempt to pre-load dbname dimensions/variables
        for r in self.pgb_conn.query_iter("SHOW DATABASES"):
            self._handle_entity_dimension(r.database)

        return True

    def get_data(self):
        """netdata calls this method for data"""
        # avoid allocation cost
        self.data.clear()
        data = self.data

        # overall client/server stats
        for r in self.pgb_conn.query_iter("SHOW LISTS"):
            if r.list == "used_clients":
                data["pgb_client_conns"] = r.items
            elif r.list == "used_servers":
                data["pgb_server_conns"] = r.items
            else:
                continue

        # get client/server pool stats, update lines/dimensions
        for r in self.pgb_conn.query_iter("SHOW POOLS"):
            dbname = r.database
            c_conns = r.cl_active + r.cl_waiting
            s_conns = r.sv_active + r.sv_idle + r.sv_used + r.sv_tested + r.sv_login
            c_name = "pool_clients_" + dbname
            s_name = "pool_servers_" + dbname
            # add data point
            data[c_name] = c_conns
            data[s_name] = s_conns
            # update chart dimensions if needed
            if dbname not in self.known_entities:
                self._handle_entity_dimension(dbname)

        # get pgbouncer_pool_conns_avail
        for r in self.pgb_conn.query_iter("SHOW DATABASES"):
            if r.current_connections > 0:
                conns_avail = float(r.pool_size - r.current_connections)
                percent_avail = conns_avail / r.pool_size * 100.0
            else:
                percent_avail = 100.0
            dimension_name = "pool_conns_avail_" + r.database
            data[dimension_name] = percent_avail
        return data

    def _handle_entity_dimension(self, dbname):
        """handles dynamic entity dimensions in "pool" charts"""
        for chart_name in self._pool_chart_names:
            stat_name = chart_name[10:] + "_" + dbname  # "pgbouncer_" is 10 chars
            dimension = [stat_name, dbname, "absolute"]

            if hasattr(self, "definitions"):  # this attr is removed after `create()`
                self.definitions[chart_name]["lines"].append(dimension)
                continue
            try:
                self.charts[chart_name].add_dimension(dimension)
            except ChartError as error:
                self.error(
                    "[SKIPPED] (dimension='{dimension!r}': {error})".format(
                        dimension=dimension, error=error
                    )
                )
        self.known_entities.add(dbname)


# Run standalone test outside of netdata plugin context
if __name__ == "__main__":
    svc = Service(dict())
    # print svc._get_entity_info()
    # print svc._get_pools_info_from_pgbouncer()
    assert svc.check()
    print(svc.get_data())