#!/usr/bin/perl -w
# -*- perl -*-
#
=head1 NAME

postgresql_space_ - Plugin to monitor PostgreSQL disk usage.

=head1 INSTALLATION

Link the plugin to /etc/munin/plugins/ :
  ln -s /usr/share/munin/plugins/postgres_space_ /etc/munin/plugins/postgres_space_dbname

The database name from the link is only used in the graph title. You have to
configure the name of the database to connect in the configuration file. 

=head1 CONFIGURATION

The plugin supports the following parameters:
	
 config 	(required)
 autoconf	(optional - used by munin-node-configure)


To configure the plugin set configuration variables in
/etc/munin/plugin-conf.d/postgres :

 PGHOST		- Database server to use. Defaults to using ident
		  authentication with the local server.
 PGPORT		- Port to connect to. Defaults to '5432'.
 PGDATABASE	- Database to connect to. Defaults to 'template1'.
		  Note that this can be different from the database to
		  monitor, which is defined via the symlink.
 PGUSER		- User to connect as, if necessary.
 PGPASSWORD	- Corresponding password to use, if necessary.

(See libpq documentation for more.)

Note that PGDATABASE will default to 'template1' in this plugin, and
without PGHOST it will try ident authentication with the local server,
as the user that the plugin is running as.

Configuration example:

# Use local server, ident authentication with the 'postgres' user.
 [postgres_*]
 user postgres
 env.PGDATABASE myappdb

# Use local server, TCP authentication with a username and password.
 [postgres_*]
 env.PGHOST localhost
 env.PGDATABASE myappdb
 env.PGUSER someuser
 env.PGPASSWORD somepassword

=head1 AUTHOR

 Written by Bjrn Ruberg (bjorn@linpro.no) 2006
 Rewritten by Moses Moore 2006-04-08  moc.iazom@sesom

=head1 LICENSE

Licenced under GPL


=head1 MAGIC MARKERS

 #%# family=auto
 #%# capabilities=suggest

=cut

use strict;
use DBI;
use vars qw ( $debug $suggest $configure $dbh );

# Default to template1 database.
$ENV{'PGDATABASE'} ||= 'template1';

if (exists $ARGV[0]) {
    if ($ARGV[0] eq 'autoconf') {
        # Check for DBD::Pg
        if (! eval "require DBD::Pg;") {
            print "no (DBD::Pg not found)";
            exit 0;
        }
        # Then we try to detect PostgreSQL presence.
        my $tempdbh = DBI->connect ('dbi:Pg:', '', '', { PrintError => 0 });
        if ($tempdbh) {
            print "yes\n";
            exit 0;
        } else {
            print "no (Can't connect to given host, please check environment settings)\n";
            exit 0;
        }
    } elsif ($ARGV[0] and $ARGV[0] eq 'debug') {
        # Set config flag
        $debug = 1;
    } elsif ($ARGV[0] and $ARGV[0] eq 'config') {
        # Set config flag
        $configure = 1;
    } elsif ($ARGV[0] eq 'suggest') {
        # doesn't always work
        my @datasources = DBI->data_sources ('Pg');
        foreach my $dsn (grep !/\=template\d$/, @datasources) {
            (my $db = $dsn) =~ s/^.*=//;
            print "$db\n";
        }
        exit 0;
    }
}

# Must do this here, after checking for autoconf/suggest/etc, because the
# plugin must be able to run before it is linked to the databases.

# Note that $dbname is not necessarily the same as $ENV{'PGDATABASE'}.
my (undef, undef, $dbname) = split (/_/, $0, 3);
die "No dbname configured (did you make the proper symlink?)" unless $dbname;

my @datasources = DBI->data_sources ('Pg')
    or die ("Can't read any possible data sources: $?");

my $dbh = DBI->connect ('dbi:Pg:', '', '', {RaiseError =>1});
unless($dbh) {
    die("Error connecting to database. (". $DBI::errstr .")\n");
}

if ($configure) {
    print <<_EOM;
graph_title Postgres database $dbname
graph_args -l 0 --base 1024
graph_vlabel bytes
graph_category PostgreSQL
graph_info Size
size.label Database size (bytes)
size.info Database size
size.type GAUGE
size.draw AREA
indexsize.label Index size (bytes)
indexsize.info Index size
indexsize.type GAUGE
indexsize.draw STACK
metasize.label Meta database size (bytes)
metasize.info Meta database size
metasize.type GAUGE
metasize.draw STACK
metaindexsize.label Meta index size (bytes)
metaindexsize.info Meta index size
metaindexsize.type GAUGE
metaindexsize.draw STACK
_EOM
} else {
    my $database_pages = 0;
    my $database_indexes = 0;
    my $metadatabase_pages = 0;
    my $metadatabase_indexes = 0;
    my @names = $dbh->tables;
    
    # Find relfilenode and relpages from the given table
    my $q_ind = "SELECT relkind, relfilenode, relpages FROM pg_class
                     WHERE relname = ?
                     UNION
                     SELECT relkind, relfilenode, relpages FROM pg_class
                     WHERE relfilenode IN (SELECT indexrelid FROM pg_index 
                     WHERE indrelid IN (SELECT relfilenode FROM pg_class
                     WHERE relname = ?))";
    my $sth = $dbh->prepare ($q_ind) or die $dbh->errstr;
    
    # Iterate over the tables in the database
    foreach my $table (@names) {
        my $meta = 1;
        print "#TABLE: $table\n" if $debug;
        my $table_pages = 0;
        my $table_indexes = 0;
        my $metatable_pages = 0;
        my $metatable_indexes = 0;
        # "public" tables are the user data
        $meta = 0 if $table =~ /^public\./;
        $table =~ s/^.*\.//;
        
        # Call the query with $table twice for each side of the UNION
        $sth->execute ($table, $table) or die $dbh->errstr;
        while (my ($relkind, $relfilenode, $relpages) = $sth->fetchrow_array) {
            if ($relkind eq 'r') {
                $table_pages     += $relpages if $meta == 0;
                $metatable_pages += $relpages if $meta == 1;
            } elsif ($relkind eq 'i') {
                $table_indexes     += $relpages if $meta == 0;
                $metatable_indexes += $relpages if $meta == 1;
            }
            # Define the query
            my $q2 = "SELECT SUM(relpages) 
                      FROM pg_class 
                      WHERE relname IN (?, ?)";
            my $sth2 = $dbh->prepare ($q2);
            $sth2->execute ("pg_toast_${relfilenode}",
                            "pg_toast_${relfilenode}_index");
            my $relpages = $sth2->fetchrow_array;
            if ($relkind eq 'r') {
                $table_pages     += $relpages if defined($relpages) and $meta == 0;
                $metatable_pages += $relpages if defined($relpages) and $meta == 1;
            } elsif ($relkind eq 'i') {
                $table_indexes     += $relpages if defined($relpages) and $meta == 0;
                $metatable_indexes += $relpages if defined($relpages) and $meta == 1;
            }
            print "#\tR:$relfilenode\tP:$table_pages\tI:$table_indexes\n" if $debug;
        }
        $database_pages       += $table_pages;
        $database_indexes     += $table_indexes;
        $metadatabase_pages   += $metatable_pages;
        $metadatabase_indexes += $metatable_indexes;
    }    
    $sth->finish;
    $dbh->disconnect;
    print "size\.value " . $database_pages * 8192 . "\n";
    print "indexsize\.value " . $database_indexes * 8192 . "\n";
    print "metasize\.value " . $metadatabase_pages * 8192 . "\n";
    print "metaindexsize\.value " . $metadatabase_indexes * 8192 . "\n";
}
