#!/usr/bin/env perl
# ex:ts=4

# This script computes disk usage of a table, counting both data disk
# usage and indexes disk usage, if any.  I think it could use the
# pg_relation_size() function and friends, but those don't exist in
# old PostgreSQL versions (7.4.x for instance).
#
use strict;
use warnings;

use DBI;
use File::Basename;
use Getopt::Long;
use Number::Bytes::Human qw(format_bytes);

sub get_db_tables {
	my ($dbh, $namespace) = @_;

	my $sth = $dbh->prepare("SELECT c.relname " .
		"FROM pg_class c, pg_namespace n " .
		"WHERE c.relkind = 'r' AND n.oid = c.relnamespace AND n.nspname = ?");
	$sth->execute($namespace) or die $dbh->errstr;
	my @tables = ();
	while (my $row = $sth->fetchrow_hashref) {
		push(@tables, $row->{'relname'});
	}
	return @tables;
}

sub get_page_size {
	my ($dbh) = @_;

	my $sth = $dbh->prepare("SHOW block_size");
	if (!$sth->execute) {
		print STDERR "WARN:  no \"block_size\" configuration parameter, " .
			"assuming 8192\n";
		return 8192;
	}
	my $row = $sth->fetchrow_arrayref;
	return $row->[0];
}

sub get_table_usage {
	my ($dbh, $pagesize, $namespace, $table) = @_;

	my $usage = {};
	my $sth = $dbh->prepare("SELECT c.oid, c.relpages, c.relhasindex, " .
			"c.reltuples, c.relfilenode, c.reltoastrelid " .
		"FROM pg_class c, pg_namespace n " .
		"WHERE c.relname = ? AND c.relnamespace = n.oid AND n.nspname = ?");
	$sth->execute($table, $namespace) or die $dbh->errstr;
	my $row = $sth->fetchrow_hashref;
	if (!defined $row) {
		die "FATAL:  table \"$namespace.$table\" doesn't exist\n";
	}

	$usage->{'data'} = $row->{'relpages'} * $pagesize;
	$usage->{'tuples'} = $row->{'reltuples'};
	$usage->{'filenode'} = $row->{'relfilenode'};

	my $total = $usage->{'data'};
	# Now grab the indexes usage, if any.
	if ($row->{'relhasindex'}) {
		$usage->{'indexes'} = {};
		$sth = $dbh->prepare("SELECT c.relname, c.relpages " .
			"FROM pg_class c, pg_index i " .
			"WHERE c.oid = i.indexrelid AND " .
				"c.relkind = 'i' AND i.indrelid = ?");
		$sth->execute($row->{'oid'}) or die $dbh->errstr;
		while (my $idx = $sth->fetchrow_hashref) {
			my $size = $idx->{'relpages'} * $pagesize;
			$usage->{'indexes'}->{$idx->{'relname'}} = $size;
			$total += $size;
		}
	}

	# Finally, get the TOAST tables usage.
	if ($row->{'reltoastrelid'} != 0) {
		$usage->{'toast'} = {};
		$sth = $dbh->prepare("SELECT c.relname, c.reltoastidxid, c.relpages " .
			"FROM pg_class c WHERE c.oid = ? AND c.relkind = 't'");
		$sth->execute($row->{'reltoastrelid'}) or die $dbh->errstr;
		my $toast = $sth->fetchrow_hashref;
		my $size = $toast->{'relpages'} * $pagesize;
		$usage->{'toast'}->{'name'} = $toast->{'relname'};
		$usage->{'toast'}->{'data'} = $size;
		$total += $size;
		$sth = $dbh->prepare("SELECT c.relname, c.relpages " .
			"FROM pg_class c WHERE c.oid = ? AND c.relkind = 'i'");
		$sth->execute($toast->{'reltoastidxid'}) or die $dbh->errstr;
		my $idx = $sth->fetchrow_hashref;
		$usage->{'toast'}->{'index'} = $idx->{'relname'};
		$size = $idx->{'relpages'} * $pagesize;
		$usage->{'toast'}->{'idxsize'} = $size;
		$total += $size;
	}

	$usage->{'total'} = $total;
	return $usage;
}

sub format_size {
	my ($size) = @_;

	return format_bytes($size) . 'B';
}

sub usage {
	my (%options) = @_;
	$options{'exit'} = 1 unless exists $options{'exit'};
	$options{'help'} = 0 unless exists $options{'help'};

	my $progname = basename($0);
	print STDERR "Usage: $progname [OPTIONS] table...\n";
	print STDERR "       $progname [OPTIONS] -a\n";
	if ($options{'help'}) {
		print "\nOptions:\n";
		print "  -h, --host=HOSTNAME\t\tDatabase server host\n";
		print "  -U, --user=NAME\t\tDatabase user name\n";
		print "  -p, --password=PASSWORD\t" .
			"Specify password to use for connection\n";
		print "  -d, --database=DBNAME\t\t" .
			"Specify database name to connect to\n";
		print "  -s, --schema=SCHEMA\t\t" .
			"Consider the named schema only (default \"public\")\n";
		print "  -a, --all\t\t\tHandle all the tables in the database\n";
		print "  -?, --help\t\t\tDisplay this help screen\n";
	}
	exit $options{'exit'};
}

# Default values
my $user = "pgsql";
my $passwd = undef;
my $host = undef;
my $dbname = undef;
my $namespace = "public";
my $all = 0;
my $help = 0;

my %options = (
	"help|?" => \$help,
	"database|d=s" => \$dbname,
	"all|a" => \$all,
	"schema|s=s" => \$namespace,
	"host|h=s" => \$host,
	"user|U=s" => \$user,
	"password|p=s" => \$passwd
);
my $result = GetOptions(%options);
if (!$result) {
	usage();
}

if ($help) {
	usage(help => 1, exit => 0);
}

if (!$all && @ARGV < 1) {
	usage();
}

$dbname = $user unless defined $dbname;
my $dsn = "dbi:Pg:dbname=$dbname";
if (defined $host) {
	$dsn .= ";host=$host";
}
my $dbh = DBI->connect($dsn, $user, $passwd,
        { RaiseError => 0, PrintError => 0 });
die $DBI::errstr if not defined $dbh;

my @tables = ();
if ($all) {
	@tables = get_db_tables($dbh, $namespace);
} else {
	@tables = @ARGV;
}

my $pagesize = get_page_size($dbh);

my $first = 1;
foreach my $table (@tables) {
	my $usage = eval { get_table_usage($dbh, $pagesize, $namespace, $table) };
	if ($@) {
		print $@;
		next;
	}

	print "\n" if !$first;
	print "Disk usage for table \"$namespace.$table\":\n";
	print "Configured page size:\t\t$pagesize\n";
	print "Estimated row count:\t\t$usage->{'tuples'}\n";
	print "Name of the on-disk file:\t$usage->{'filenode'}\n";
	print "Data usage:\t\t\t" . format_size($usage->{'data'}) . "\n";

	if (exists $usage->{'indexes'}) {
		print "Indexes:\n";
		while (my ($name, $size) = each %{$usage->{'indexes'}}) {
			print "  " . format_size($size) . "\t($name)\n";
		}
	}

	if (exists $usage->{'toast'}) {
		print "TOAST table:\n";
		print "  " . format_size($usage->{'toast'}->{'data'}) .
			"\t($usage->{'toast'}->{'name'})\n";
		print "TOAST table index:\n";
		print "  " . format_size($usage->{'toast'}->{'idxsize'}) .
			"\t($usage->{'toast'}->{'index'})\n";
	}

	print "Total disk usage:\t\t" .  format_size($usage->{'total'}) . "\n";
	$first = 0;
}

$dbh->disconnect;