#!/usr/bin/perl -w

# pgindexsize
# Report on PostgreSQL index sizes
# by Jon Jensen <jon@swelter.net>
# 2003-03-07

use strict;
use DBI;
use Getopt::Long;
#use Data::Dumper;
#$Data::Dumper::Terse = 1;

select STDERR;
$| = 1;
select STDOUT;
$| = 1;


my %opt = qw(
	report	1
);
GetOptions(\%opt, qw(
	dsn=s
	dbname|database|db=s
	user|username=s
	both=s
	password|pass=s
	pgdata=s
	rebuild
	report!
	quiet
));

my $starttime = time;

my $base = $opt{pgdata} || $ENV{PGDATA} || "/var/lib/pgsql/data";
$base =~ s:/+$::;
$base .= "/base" unless $base =~ m:/base$:;

$opt{dbname} = $opt{user} = $opt{both} if $opt{both};

if ($opt{dsn}) {
	$opt{dbname} = $1 if ! $opt{dbname} and $opt{dsn} =~ /\bdbname\s*=\s*(\S+)/i;
}
else {
	die "Please specify DSN or database name.\n" unless $opt{dbname};
	$opt{dsn} = "dbi:Pg:dbname=$opt{dbname}";
};


my $sizelen = 12;

unless ($opt{quiet}) {
	print STDERR "Connecting to $opt{dsn}";
	print STDERR " as user $opt{user}" if $opt{user};
	print STDERR "\n";
}
my $dbh = DBI->connect($opt{dsn}, $opt{user}, $opt{password},
	{ AutoCommit => 1, RaiseError => 1 });
my $sth;


$sth = $dbh->prepare("SELECT datid FROM pg_stat_database WHERE datname = ?");
$sth->execute($opt{dbname});
my ($dboid) = ($sth->fetchrow_array);
$sth->finish;
die "Error getting OID of database '$opt{dbname}'\n" unless $dboid;


my $datadir = "$base/$dboid";
my $sizes1 = getoidsizes($datadir);
die "No OID files found in $datadir!\n" unless %$sizes1;


$sth = $dbh->prepare(<<'EOQ');
SELECT i.oid, c.relname AS tablename, i.relname AS indexname,
	pg_get_indexdef(x.indexrelid) AS indexdef
FROM pg_index x, pg_class c, pg_class i
WHERE c.relkind = 'r'::"char" AND i.relkind = 'i'::"char" AND c.oid = x.indrelid
	AND i.oid = x.indexrelid AND c.relname NOT LIKE 'pg_%'::"text"
EOQ
$sth->execute;
my $result1 = $sth->fetchall_hashref('indexname');
die "No results found for indexes query!\n" unless %$result1;

my $namelen = 0;
for (keys %$result1) {
	my $len = length($result1->{$_}{indexname});
	$namelen = $len if $len > $namelen;
}


#print Dumper($result1, $sizes1);


unless ($opt{rebuild}) {
	exit unless $opt{report};

	my $headerformat = "%-${namelen}s %${sizelen}s\n";
	my $rowformat = "%-${namelen}s %${sizelen}lu\n";
	my $separator = "-" x ($namelen + $sizelen + 1) . "\n";
	printf $headerformat, "Index Name", "Bytes";
	print $separator;
	my $totalsize = 0;
	for (sort keys %$result1) {
		my $size = $sizes1->{$result1->{$_}{oid}};
		$totalsize += $size;
		printf $rowformat, $_, $size;
	}
	print $separator;
	printf $rowformat, "(total)", $totalsize;
	exit;
}


for (sort keys %$result1) {
	my $sql = $result1->{$_}{indexdef};
	die "Index definition missing for $_!\n"
		unless $sql =~ /^CREATE(?:\s+UNIQUE)?\s+INDEX\s+$_\s+/;
	print STDERR "Rebuilding $_\n" unless $opt{quiet};
	$dbh->begin_work;
	$dbh->do("DROP INDEX $_");
	$dbh->do($sql);
	$dbh->commit;
}

exit unless $opt{report};

sleep 1;


$sth->execute;
my $result2 = $sth->fetchall_hashref('indexname');


my $sizes2 = getoidsizes($datadir);

my $headerformat = "%-${namelen}s" . (" %${sizelen}s" x 3) . " %4s\n";
my $rowformat = "%-${namelen}s" . (" %${sizelen}lu" x 2) . " %${sizelen}li %3li%%\n";
my $separator = "-" x ($namelen + ($sizelen + 1) * 3 + 5) . "\n";
printf $headerformat, "Index Name", "Before", "After", "Bytes Less", "Pct.";
print $separator;
my $beforetotal = my $aftertotal = 0;
for (sort keys %$result1) {
	my $before = $sizes1->{$result1->{$_}{oid}}; 
	$beforetotal += $before;
	my $after = $sizes2->{$result2->{$_}{oid}};
	$aftertotal += $after;
	printf $rowformat, $_, $before, $after, $before - $after, $after / $before * 100;
}
print $separator;
printf $rowformat, "(total)", $beforetotal, $aftertotal,
	$beforetotal - $aftertotal, $aftertotal / $beforetotal * 100;

printf STDERR "Elapsed time: %u seconds\n", time - $starttime
	unless $opt{quiet};


END {
	$dbh and $dbh->disconnect;
}


sub getoidsizes {
	my ($datadir) = @_;
	opendir DIR, $datadir or die "Can't opendir $datadir: $!\n";
	my %sizes = map { $_ => -s "$datadir/$_" } grep ! /\D/, readdir(DIR);
	closedir DIR;
	return \%sizes;
}