postgres.py

# coding=utf-8

"""
Collect metrics from postgresql

#### Dependencies

 * psycopg2

#### Example Configuration

#Section with defaults
enabled=True
password=default_password
port = 5432

#Instance specific configs
[instances]

[[postgres_a]]
host = db1.loc
password = instance_password_a

[[postgres_b]]
host = db2.loc
port = 5433
password = instance_password_b

"""

import diamond.collector
from diamond.collector import str_to_bool

try:
    import psycopg2
    import psycopg2.extras
except ImportError:
    psycopg2 = None


class PostgresqlCollector(diamond.collector.Collector):
    """
    PostgreSQL collector class
    """

    def get_default_config_help(self):
        """
        Return help text for collector
        """
        config_help = super(PostgresqlCollector,
                            self).get_default_config_help()

        config_help.update({
            'dbname': 'DB to connect to in order to get list of DBs in PgSQL',
            'user': 'Username',
            'password': 'Password',
            'port': 'Port number',
            'password_provider': "Whether to auth with supplied password or"
                                 " .pgpass file  <password|pgpass>",
            'sslmode': 'Whether to use SSL - <disable|allow|require|...>',
                       'underscore': 'Convert _ to .',
            'extended': 'Enable collection of extended database stats.',
            'metrics': 'List of enabled metrics to collect',
            'pg_version':   "The version of postgres that you'll be monitoring"
                            "eg. in format 9.2",
            'has_admin': 'Admin privileges are required to execute some'
                         ' queries.',
            'instances': 'A subcategory of postgres instances with a host '
                         'and port. Optionally all variables can be '
                         'overridden per instance (see example).',
        })
        return config_help

    def get_default_config(self):
        """
        Return default config.
        """
        config = super(PostgresqlCollector, self).get_default_config()

        config.update({
            'path': 'postgres',
            'host': 'localhost',
            'dbname': 'postgres',
            'user': 'postgres',
            'password': 'postgres',
            'port': 5432,
            'password_provider': 'password',
            'sslmode': 'disable',
            'underscore': False,
            'extended': False,
            'metrics': [],
            'pg_version': 9.2,
            'has_admin': True,
            'instances': {},
        })
        return config

    def collect(self):
        """
        Do pre-flight checks, get list of db names, collect metrics, publish
        """
        if psycopg2 is None:
            self.log.error('Unable to import module psycopg2')
            return {}

        instances = self.config.get('instances')

        # HACK: setting default with subcategory messes up merging of configs,
        # so we only set the default if one wasn't provided.
        if not instances:
            instances = {
                'default': {
                    'host': self.config['host'],
                }
            }

        for instance in instances:
            # Get list of databases
            dbs = self._get_db_names(instance)
            if len(dbs) == 0:
                self.log.error("I have 0 databases!")
                return {}

            if self._get_config(instance, 'metrics'):
                metrics = self._get_config(instance, 'metrics')
            elif str_to_bool(self._get_config(instance, 'extended')):
                metrics = registry['extended']
                if str_to_bool(self._get_config(instance, 'has_admin')) \
                        and 'WalSegmentStats' not in metrics:
                    metrics.append('WalSegmentStats')
            else:
                metrics = registry['basic']

            # Iterate every QueryStats class
            for metric_name in set(metrics):
                if metric_name not in metrics_registry:
                    self.log.error(
                        'metric_name %s not found in metric registry'
                        % metric_name)
                    continue

                for dbase in dbs:
                    conn = self._connect(instance, database=dbase)
                    try:
                        klass = metrics_registry[metric_name]
                        stat = klass(dbase, conn,
                                     underscore=self._get_config(instance,
                                                                 'underscore'))
                        stat.fetch(self._get_config(instance, 'pg_version'))
                        for metric, value in stat:
                            if value is not None:
                                self.publish("%s.%s" % (instance, metric),
                                             value)

                        # Setting multi_db to True will run this query on all
                        # known databases. This is bad for queries that hit
                        # views like pg_database, which are shared
                        # across databases.
                        #
                        # If multi_db is False, bail early after the first query
                        # iteration. Otherwise, continue to remaining databases.
                        if stat.multi_db is False:
                            break
                    finally:
                        conn.close()

    def _get_config(self, instance, name):
        """
        Return instance config value or value from default section
        if it is overriden or None
        """
        instance_config = self.config['instances'].get(instance)

        if instance_config:
            return instance_config.get(name, self.config.get(name, None)
                                       if name != 'host' else None)
        else:
            return None

    def _get_db_names(self, instance):
        """
        Try to get a list of db names
        """
        query = """
            SELECT datname FROM pg_database
            WHERE datallowconn AND NOT datistemplate
            AND NOT datname='postgres' AND NOT datname='rdsadmin' ORDER BY 1
        """
        conn = self._connect(instance,
                             self._get_config(instance, 'dbname'))
        cursor = conn.cursor(cursor_factory=psycopg2.extras.DictCursor)
        cursor.execute(query)
        datnames = [d['datname'] for d in cursor.fetchall()]
        conn.close()

        # Exclude `postgres` database list, unless it is the
        # only database available (required for querying pg_stat_database)
        if not datnames:
            datnames = ['postgres']

        return datnames

    def _connect(self, instance, database=None):
        """
        Connect to given database
        """
        conn_args = {
            'host': self._get_config(instance, 'host'),
            'user': self._get_config(instance, 'user'),
            'password': self._get_config(instance, 'password'),
            'port': self._get_config(instance, 'port'),
            'sslmode': self._get_config(instance, 'sslmode'),
        }

        if database:
            conn_args['database'] = database
        else:
            conn_args['database'] = 'postgres'

        # libpq will use ~/.pgpass only if no password supplied
        if self._get_config(instance, 'password_provider') == 'pgpass':
            del conn_args['password']

        try:
            conn = psycopg2.connect(**conn_args)
        except Exception, e:
            self.log.error(e)
            raise e

        # Avoid using transactions, set isolation level to autocommit
        conn.set_isolation_level(0)
        return conn


class QueryStats(object):
    query = None
    path = None

    def __init__(self, dbname, conn, parameters=None, underscore=False):
        self.conn = conn
        self.dbname = dbname
        self.underscore = underscore
        self.parameters = parameters
        self.data = list()

    def _translate_datname(self, datname):
        """
        Replace '_' with '.'
        """
        if self.underscore:
            datname = datname.replace("_", ".")
        return datname

    def fetch(self, pg_version):
        if float(pg_version) >= 9.6 and hasattr(self, 'post_96_query'):
            q = self.post_96_query
        elif float(pg_version) >= 9.2 and hasattr(self, 'post_92_query'):
            q = self.post_92_query
        else:
            q = self.query

        cursor = self.conn.cursor(cursor_factory=psycopg2.extras.DictCursor)
        try:
            cursor.execute(q, self.parameters)
            rows = cursor.fetchall()
            for row in rows:
                # If row is length 2, assume col1, col2 forms key: value
                if len(row) == 2:
                    self.data.append({
                        'datname': self._translate_datname(self.dbname),
                        'metric': row[0],
                        'value': row[1],
                    })

                # If row > length 2, assume each column name maps to
                # key => value
                else:
                    for key, value in row.iteritems():
                        if key in ('datname', 'schemaname', 'relname',
                                   'indexrelname', 'funcname',):
                            continue

                        self.data.append({
                            'datname': self._translate_datname(row.get(
                                'datname', self.dbname)),
                            'schemaname': row.get('schemaname', None),
                            'relname': row.get('relname', None),
                            'indexrelname': row.get('indexrelname', None),
                            'funcname': row.get('funcname', None),
                            'metric': key,
                            'value': value,
                        })

        # Clean up
        finally:
            cursor.close()

    def __iter__(self):
        for data_point in self.data:
            yield (self.path % data_point, data_point['value'])


class DatabaseStats(QueryStats):
    """
    Database-level summary stats
    """
    path = "database.%(datname)s.%(metric)s"
    multi_db = False
    post_92_query = """
        SELECT pg_stat_database.datname as datname,
               pg_stat_database.numbackends as numbackends,
               pg_stat_database.xact_commit as xact_commit,
               pg_stat_database.xact_rollback as xact_rollback,
               pg_stat_database.blks_read as blks_read,
               pg_stat_database.blks_hit as blks_hit,
               pg_stat_database.tup_returned as tup_returned,
               pg_stat_database.tup_fetched as tup_fetched,
               pg_stat_database.tup_inserted as tup_inserted,
               pg_stat_database.tup_updated as tup_updated,
               pg_stat_database.tup_deleted as tup_deleted,
               pg_database_size(pg_database.datname) AS size
        FROM pg_database
        JOIN pg_stat_database
        ON pg_database.datname = pg_stat_database.datname
        WHERE pg_stat_database.datname
        NOT IN ('template0','template1','postgres', 'rdsadmin')
    """
    query = post_92_query.replace(
        'pg_stat_database.temp_files as temp_files,',
        '').replace(
        'pg_stat_database.temp_bytes as temp_bytes,',
        '')


class UserFunctionStats(QueryStats):
    # http://www.pateldenish.com/2010/11/postgresql-track-functions-to-tune.html
    path = "%(datname)s.functions.%(funcname)s.%(metric)s"
    multi_db = True
    query = """
        SELECT funcname,
               calls,
               total_time/calls as time_per_call
        FROM pg_stat_user_functions
        WHERE calls <> 0
    """


class UserTableStats(QueryStats):
    path = "%(datname)s.tables.%(schemaname)s.%(relname)s.%(metric)s"
    multi_db = True
    query = """
        SELECT relname,
               schemaname,
               seq_scan,
               seq_tup_read,
               idx_scan,
               idx_tup_fetch,
               n_tup_ins,
               n_tup_upd,
               n_tup_del,
               n_tup_hot_upd,
               n_live_tup,
               n_dead_tup
        FROM pg_stat_user_tables
    """


class UserIndexStats(QueryStats):
    path = "%(datname)s.indexes.%(schemaname)s.%(relname)s." \
           "%(indexrelname)s.%(metric)s"
    multi_db = True
    query = """
        SELECT relname,
               schemaname,
               indexrelname,
               idx_scan,
               idx_tup_read,
               idx_tup_fetch
        FROM pg_stat_user_indexes
    """


class UserTableIOStats(QueryStats):
    path = "%(datname)s.tables.%(schemaname)s.%(relname)s.%(metric)s"
    multi_db = True
    query = """
        SELECT relname,
               schemaname,
               heap_blks_read,
               heap_blks_hit,
               idx_blks_read,
               idx_blks_hit
        FROM pg_statio_user_tables
    """


class UserIndexIOStats(QueryStats):
    path = "%(datname)s.indexes.%(schemaname)s.%(relname)s." \
           "%(indexrelname)s.%(metric)s"
    multi_db = True
    query = """
        SELECT relname,
               schemaname,
               indexrelname,
               idx_blks_read,
               idx_blks_hit
        FROM pg_statio_user_indexes
    """


class ConnectionStateStats(QueryStats):
    path = "%(datname)s.connections.%(metric)s"
    multi_db = True
    query = """
        SELECT tmp.state AS key,COALESCE(count,0) FROM
               (VALUES ('active'),
                       ('waiting'),
                       ('idle'),
                       ('idletransaction'),
                       ('unknown')
                ) AS tmp(state)
        LEFT JOIN
             (SELECT CASE WHEN waiting THEN 'waiting'
                          WHEN current_query = '<IDLE>' THEN 'idle'
                          WHEN current_query = '<IDLE> in transaction'
                              THEN 'idletransaction'
                          WHEN current_query = '<insufficient privilege>'
                              THEN 'unknown'
                          ELSE 'active' END AS state,
                     count(*) AS count
               FROM pg_stat_activity
               WHERE procpid != pg_backend_pid()
               GROUP BY CASE WHEN waiting THEN 'waiting'
                             WHEN current_query = '<IDLE>' THEN 'idle'
                             WHEN current_query = '<IDLE> in transaction'
                                 THEN 'idletransaction'
                             WHEN current_query = '<insufficient privilege>'
                                 THEN 'unknown' ELSE 'active' END
             ) AS tmp2
        ON tmp.state=tmp2.state ORDER BY 1
    """
    post_92_query = """
        SELECT tmp.mstate AS state,COALESCE(count,0) FROM
               (VALUES ('active'),
                       ('waiting'),
                       ('idle'),
                       ('idletransaction'),
                       ('unknown')
               ) AS tmp(mstate)
        LEFT JOIN
             (SELECT CASE WHEN waiting THEN 'waiting'
                          WHEN state = 'idle' THEN 'idle'
                          WHEN state LIKE 'idle in transaction%'
                              THEN 'idletransaction'
                          WHEN state = 'disabled'
                              THEN 'unknown'
                          WHEN query = '<insufficient privilege>'
                              THEN 'unknown'
                          ELSE 'active' END AS mstate,
                     count(*) AS count
               FROM pg_stat_activity
               WHERE pid != pg_backend_pid()
               GROUP BY CASE WHEN waiting THEN 'waiting'
                             WHEN state = 'idle' THEN 'idle'
                             WHEN state LIKE 'idle in transaction%'
                                 THEN 'idletransaction'
                             WHEN state = 'disabled'
                                 THEN 'unknown'
                             WHEN query = '<insufficient privilege>'
                                 THEN 'unknown'
                             ELSE 'active'
                        END
             ) AS tmp2
        ON tmp.mstate=tmp2.mstate ORDER BY 1
    """
    post_96_query = """
        SELECT tmp.state AS key,COALESCE(count,0) FROM
               (VALUES ('active'),
                       ('waiting'),
                       ('idle'),
                       ('idletransaction'),
                       ('unknown')
                ) AS tmp(state)
        LEFT JOIN
             (SELECT CASE WHEN wait_event IS NOT NULL THEN 'waiting'
                          WHEN state= 'idle' THEN 'idle'
                          WHEN state= 'idle in transaction'
                          THEN 'idletransaction'
                          WHEN state = 'active' THEN 'active'
                          ELSE 'unknown' END AS state,
                     count(*) AS count
               FROM pg_stat_activity
               WHERE pid != pg_backend_pid()
               GROUP BY CASE WHEN wait_event IS NOT NULL THEN 'waiting'
                          WHEN state= 'idle' THEN 'idle'
                          WHEN state= 'idle in transaction'
                          THEN 'idletransaction'
                          WHEN state = 'active' THEN 'active'
                          ELSE 'unknown' END
             ) AS tmp2
        ON tmp.state=tmp2.state ORDER BY 1
    """


class LockStats(QueryStats):
    path = "%(datname)s.locks.%(metric)s"
    multi_db = False
    query = """
        SELECT lower(mode) AS key,
               count(*) AS value
        FROM pg_locks
        WHERE database IS NOT NULL
        GROUP BY mode ORDER BY 1
    """


class RelationSizeStats(QueryStats):
    path = "%(datname)s.sizes.%(schemaname)s.%(relname)s.%(metric)s"
    multi_db = True
    query = """
        SELECT pg_class.relname,
               pg_namespace.nspname as schemaname,
               pg_relation_size(pg_class.oid) as relsize
        FROM pg_class
        INNER JOIN
          pg_namespace
        ON pg_namespace.oid = pg_class.relnamespace
        WHERE reltype != 0
        AND relkind != 'S'
        AND nspname NOT IN ('pg_catalog', 'information_schema')
    """


class BackgroundWriterStats(QueryStats):
    path = "bgwriter.%(metric)s"
    multi_db = False
    query = """
        SELECT checkpoints_timed,
               checkpoints_req,
               buffers_checkpoint,
               buffers_clean,
               maxwritten_clean,
               buffers_backend,
               buffers_alloc
        FROM pg_stat_bgwriter
    """


class WalSegmentStats(QueryStats):
    path = "wals.%(metric)s"
    multi_db = False
    query = """
        SELECT count(*) AS segments
        FROM pg_ls_dir('pg_xlog') t(fn)
        WHERE fn ~ '^[0-9A-Z]{24}$'
    """


class TransactionCount(QueryStats):
    path = "transactions.%(metric)s"
    multi_db = False
    query = """
        SELECT 'commit' AS type,
               sum(pg_stat_get_db_xact_commit(oid))
        FROM pg_database
        UNION ALL
        SELECT 'rollback',
               sum(pg_stat_get_db_xact_rollback(oid))
        FROM pg_database
    """


class IdleInTransactions(QueryStats):
    path = "%(datname)s.idle_in_tranactions.%(metric)s"
    multi_db = True
    base_query = """
        SELECT 'idle_in_transactions',
               max(COALESCE(ROUND(EXTRACT(epoch FROM now()-query_start)),0))
                   AS idle_in_transaction
        FROM pg_stat_activity
        WHERE %s
        GROUP BY 1
    """
    query = base_query % ("current_query = '<IDLE> in transaction'", )
    post_92_query = base_query % ("state LIKE 'idle in transaction%'", )


class LongestRunningQueries(QueryStats):
    path = "%(datname)s.longest_running.%(metric)s"
    multi_db = True
    base_query = """
        SELECT 'query',
            COALESCE(max(extract(epoch FROM CURRENT_TIMESTAMP-query_start)),0)
        FROM pg_stat_activity
        WHERE %s
        UNION ALL
        SELECT 'transaction',
            COALESCE(max(extract(epoch FROM CURRENT_TIMESTAMP-xact_start)),0)
        FROM pg_stat_activity
        WHERE 1=1
    """
    query = base_query % ("current_query NOT LIKE '<IDLE%'", )
    post_92_query = base_query % ("state NOT LIKE 'idle%'", )


class UserConnectionCount(QueryStats):
    path = "%(datname)s.user_connections.%(metric)s"
    multi_db = True
    query = """
        SELECT usename,
               count(*) as count
        FROM pg_stat_activity
        WHERE procpid != pg_backend_pid()
        GROUP BY usename
        ORDER BY 1
    """
    post_92_query = query.replace('procpid', 'pid')


class DatabaseConnectionCount(QueryStats):
    path = "database.%(metric)s.connections"
    multi_db = False
    query = """
        SELECT datname,
               count(datname) as connections
        FROM pg_stat_activity
        GROUP BY pg_stat_activity.datname
    """


class TableScanStats(QueryStats):
    path = "%(datname)s.scans.%(metric)s"
    multi_db = True
    query = """
        SELECT 'relname' AS relname,
               COALESCE(sum(seq_scan),0) AS sequential,
               COALESCE(sum(idx_scan),0) AS index
        FROM pg_stat_user_tables
    """


class TupleAccessStats(QueryStats):
    path = "%(datname)s.tuples.%(metric)s"
    multi_db = True
    query = """
        SELECT COALESCE(sum(seq_tup_read),0) AS seqread,
               COALESCE(sum(idx_tup_fetch),0) AS idxfetch,
               COALESCE(sum(n_tup_ins),0) AS inserted,
               COALESCE(sum(n_tup_upd),0) AS updated,
               COALESCE(sum(n_tup_del),0) AS deleted,
               COALESCE(sum(n_tup_hot_upd),0) AS hotupdated
        FROM pg_stat_user_tables
    """


class DatabaseReplicationStats(QueryStats):
    path = "database.replication.%(metric)s"
    multi_db = False
    query = """
        SELECT EXTRACT(epoch FROM
            current_timestamp - pg_last_xact_replay_timestamp()) as replay_lag
    """


class DatabaseXidAge(QueryStats):
    path = "%(datname)s.datfrozenxid.%(metric)s"
    multi_db = False
    query = """
        SELECT datname, age(datfrozenxid) AS age
        FROM pg_database WHERE datallowconn = TRUE
    """


metrics_registry = {
    'DatabaseStats': DatabaseStats,
    'DatabaseConnectionCount': DatabaseConnectionCount,
    'UserFunctionStats': UserFunctionStats,
    'UserTableStats': UserTableStats,
    'UserIndexStats': UserIndexStats,
    'UserTableIOStats': UserTableIOStats,
    'UserIndexIOStats': UserIndexIOStats,
    'ConnectionStateStats': ConnectionStateStats,
    'LockStats': LockStats,
    'RelationSizeStats': RelationSizeStats,
    'BackgroundWriterStats': BackgroundWriterStats,
    'WalSegmentStats': WalSegmentStats,
    'TransactionCount': TransactionCount,
    'IdleInTransactions': IdleInTransactions,
    'LongestRunningQueries': LongestRunningQueries,
    'UserConnectionCount': UserConnectionCount,
    'TableScanStats': TableScanStats,
    'TupleAccessStats': TupleAccessStats,
    'DatabaseReplicationStats': DatabaseReplicationStats,
    'DatabaseXidAge': DatabaseXidAge,
}

registry = {
    'basic': [
        'DatabaseStats',
        'DatabaseConnectionCount',
    ],
    'extended': [
        'DatabaseStats',
        'DatabaseConnectionCount',
        'DatabaseReplicationStats',
        'DatabaseXidAge',
        'UserFunctionStats',
        'UserTableStats',
        'UserIndexStats',
        'UserTableIOStats',
        'UserIndexIOStats',
        'ConnectionStateStats',
        'LockStats',
        'RelationSizeStats',
        'BackgroundWriterStats',
        'TransactionCount',
        'IdleInTransactions',
        'LongestRunningQueries',
        'UserConnectionCount',
        'TableScanStats',
        'TupleAccessStats',
    ],
}