pg_dbsizes

#!/usr/bin/env python
"""
Rudimentary implementation of a multigraph for postgresql DB sizes.

This is very much in development and I cannot guarantee everything works as
intended! Use at your own risk!
"""

import psycopg2
import atexit
import sys
from os import getenv
from textwrap import dedent


def cleanup(connection):
    def fun():
        connection.close()
    return fun


def construct_dsn(dbname, user, password='', host='', port=0):
    elements = [
        'dbname=%s' % dbname,
        'user=%s' % user
    ]
    if password:
        elements.append('password=%s' % password)
    if host:
        elements.append('host=%s' % host)
    if port:
        elements.append('port=%d' % port)

    return ' '.join(elements)


def get_stats(connection):
    cursor = connection.cursor()
    cursor.execute('SELECT datname, pg_database_size(datname) '
                   'FROM pg_database WHERE datistemplate=false')
    sizes = cursor.fetchall()
    cursor.execute('SELECT datname FROM pg_stat_activity')
    cursor.close()
    return {
        'sizes': {row[0]: row[1] for row in sizes}
    }


def render_config(stats):
    print(dedent(
        '''\
        multigraph pg_sizes
        graph_title PostgreSQL Database Sizes
        graph_args --base 1024
        graph_vlabel Size in Bytes
        graph_category PostgreSQL

        size.info Size in Bytes
        size.label Size in Bytes
        size.draw AREA
        size.min 0
        '''))

    template = dedent(
        '''\
        multigraph pg_sizes.{dbname}
        graph_title Size for {dbname}
        graph_args --base 1024
        graph_category PostgreSQL
        graph_order size
        size.min 0
        size.draw AREA
        size.label Size (in Bytes)
        ''')
    for dbname, value in stats['sizes'].items():
        print(template.format(dbname=dbname))


def render_stats(stats):

    print('multigraph pg_sizes')
    print('size.value {value}'.format(value=sum(stats['sizes'].values())))
    for dbname, value in stats['sizes'].items():
        print('multigraph pg_sizes.{dbname}'.format(dbname=dbname))
        print('size.value {value}'.format(dbname=dbname, value=value))


if __name__ == '__main__':
    dbname = getenv('PG_DBNAME', 'template1')
    user = getenv('PG_USER', 'postgres')
    password = getenv('PG_PASSWORD', '')
    host = getenv('PG_HOST', '')
    port = int(getenv('PG_PORT', 0))

    connection = psycopg2.connect(construct_dsn(
        dbname,
        user,
        password,
        host,
        port)
    )
    atexit.register(cleanup(connection))

    try:
        action = sys.argv[1]
    except IndexError:
        action = 'fetch'

    if action not in ('config', 'fetch'):
        raise ValueError("Action not recognized")

    stats = get_stats(connection)
    if action == 'config':
        render_config(stats)
    elif action == 'fetch':
        render_stats(stats)