#!/usr/bin/perl -w

use strict;
use warnings;
use Getopt::Long;
use RRDs;

sub usage() {
  print <<EOF;
rrdquery.pl selects data from an RRD file with a SQL like approach where each
consolidation function (CF) represents its own table, and each datasource plus
the timestamp are the columns.
Most commandline options are shared with rrdfetch(1). Additional parameters
are explained below.

Usage:
  rrdquery.pl <rrdfile> --info|-i
  rrdquery.pl <rrdfile> <cf> <filter> [--resolution|-r <resolution>]
              [--start|-s <start>] [--end|-e <end>] [--align-start|-a]
              [--daemon|-d <address>] [--utc|-u]
  rrdquery.pl --help|-h

Description:
  --info, -i   Output the names of all round robin archives (RRAs), the names
               of all datasources, as well as the oldest and latest update
               timestamp found in the round robin database (RRD) <rrdfile>
  --utc, -u    Display each timestamp in UTC time instead of local time
  --help, -h   Display this help ;-)
  <filter>     Select a datasource and optionally filter rows

Filter syntax:
  FUNC(<dsname>) [<comparator> <value>|IS NULL]

  FUNC(...)     Optional aggregate function, one of COUNT(), SUM(), AVG(),
                MIN(), or MAX()
  <dsname>      The name of a datasource in <rrdfile>
  <comparator>  Any of =, !=, <> (same as !=), <, <=, >=, or >
  <value>       An integer or float value
  IS NULL       Output any row where <dsname> is null or NaN

  Whitespace between <dsname>, <comparator>, and/or <value> may be omitted.
  Check for occurences where <dsname> has a non-NULL value can be achieved by
  specifying <dsname> only.
EOF
}

sub format_time($$) {
  my ($epoch, $utc) = @_;
  my ($sec, $min, $hour, $day, $mon, $year) = ($utc ?
                                               gmtime($epoch) :
                                               localtime($epoch));

  return sprintf("%04d-%02d-%02d %02d:%02d:%02d",
                 $year + 1900, $mon + 1, $day, $hour, $min, $sec);
}

my ($resolution, $start, $end, $alignstart, $daemon, $utc, $showinfo, $help);

GetOptions(
  "resolution|r=s" => \$resolution,
  "start|s=s"      => \$start,
  "end|e=s"        => \$end,
  "align-start|a"  => \$alignstart,
  "daemon|d=s"     => \$daemon,
  "utc|u"          => \$utc,
  "info|i"         => \$showinfo,
  "help|h"         => \$help,
) || die;

my ($rrdfile, $cf, $filter) = @ARGV;

if ($help || !$rrdfile || !($showinfo || ($cf && $filter))) {
  # show help, or missing rrdfile and/or any command
  usage();
  exit(!$help);
}

if ($showinfo) {
  # show info
  my $info = RRDs::info($rrdfile);
  my $err = RRDs::error;
  die "Error while getting info from $rrdfile: $err" if $err;

  my %rras;
  my %datasources;
  my $lastupdate;
  my $firstupdate;

  # get the name of all round robin archives and datasources,
  # the oldest update timestamp of all rras, and the latest update
  # timestamp of the database
  foreach (keys %$info) {
    if (/^rra\[(\d+)\]\.cf$/) {
      my $rraidx = $1;
      my $first = RRDs::first($rrdfile, "--rraindex", $rraidx);
      $err = RRDs::error;
      if ($err) {
        die "Error while getting first update of rra $rraidx from $rrdfile: " .
            $err;
      }

      if (!defined $firstupdate || ($first < $firstupdate)) {
        $firstupdate = $first;
      }

      $rras{$info->{$_}} = 0;
    } elsif (/^ds\[([^\]]+)\]\./) {
      $datasources{$1} = 0;
    } elsif ($_ eq "last_update") {
      $lastupdate = $info->{$_};
    }
  }

  print "Round robin archives (RRAs)\n";
  foreach my $rra(sort keys %rras) {
    print "  $rra\n";
  }

  print "\nData sources (DSs):\n";
  foreach my $ds(sort keys %datasources) {
    print "  $ds\n";
  }

  print "\nFirst update: ", $firstupdate,
        " (", format_time($firstupdate, $utc), ")",
        "\nLast update:  ", $lastupdate,
        " (", format_time($lastupdate, $utc), ")\n";
} else {
  # build options for rrdfetch
  my @cmd = ($rrdfile, $cf);
  push @cmd, "-r $resolution" if defined $resolution;
  push @cmd, "-s $start" if defined $start;
  push @cmd, "-e $end" if defined $end;
  push @cmd, "-a" if $alignstart;
  push @cmd, "-d $daemon" if $daemon;

  # SQL-ish aggregates
  my %aggregates = (
    "identity" => sub {
      my ($ts, $utc, $curval, $result) = @_;
      print $ts, " (", format_time($ts, $utc), ") ",
            (defined $curval ? $curval : "NULL"), "\n";
      return $result;
    },
    "count" => sub {
      my ($ts, $utc, $curval, $result) = @_;
      return ($result ? $result + 1 : 1);
    },
    "sum" => sub {
      my ($ts, $utc, $curval, $result) = @_;
      if (!defined $curval) {
        return $result;
      }
      return ($result ? ($result + $curval) : $curval);
    },
    "avg" => sub {
      my ($ts, $utc, $curval, $result) = @_;
      if (!defined $curval) {
        return $result;
      }
      if (defined $result) {
        $result->[0]++;
        $result->[1] += $curval;
      } else {
        $result = [1, $curval];
      }
      return $result;
    },
    "min" => sub {
      my ($ts, $utc, $curval, $result) = @_;
      if (defined $curval && (!defined $result || ($result->[1] > $curval))) {
        $result = [$ts, $curval];
      }
      return $result;
    },
    "max" => sub {
      my ($ts, $utc, $curval, $result) = @_;
      if (defined $curval && (!defined $result || ($result->[1] < $curval))) {
        $result = [$ts, $curval];
      }
      return $result;
    },
  );

  my ($ts, $step, $names, $data) = RRDs::fetch(@cmd);
  my $err = RRDs::error;
  die "Error while fetching data from $rrdfile: $err" if $err;

  # e.g. "sum(outBytes)", "count(inBytes) > 10", "errBytes > 0.1",
  # "inBytes iS NulL", or just "discOutBits"
  if ($filter !~ /^\s*((?<aggregate>count|sum|avg|min|max)\s*\(\s*)?(?<dsname>[^\s=!<>()]+)(\s*\))?((\s*(?<cmp>=|!=|<>|<|<=|>=|>)\s*(?<value>\d+(\.\d+)?))|\s+(<?isnull>is\s+null))?\s*$/i) {
    die "Invalid filter: $filter";
  }

  # $value is fixed float given by user filter
  my ($dsname, $cmp, $value, $isnull, $aggregate) = ($+{dsname}, $+{cmp},
                                                     $+{value}, $+{isnull},
                                                     $+{aggregate});
  my $aggregate_func = $aggregates{$aggregate ? $aggregate : "identity"};

  # dsidx holds column index of selected datasource
  my $dsidx;
  for ($dsidx = 0; $dsidx <= $#$names; $dsidx++) {
    last if ($dsname eq $names->[$dsidx]);
  }
  die "No such data source: $dsname" if ($dsidx > $#$names);

  # reference to comparator
  my %cmps = (
    "="  => sub { return defined $_[0] && ($_[0] == $_[1]); },
    "!=" => sub { return defined $_[0] && ($_[0] != $_[1]); },
    "<>" => sub { return defined $_[0] && ($_[0] != $_[1]); },
    "<"  => sub { return defined $_[0] && ($_[0] <  $_[1]); },
    "<=" => sub { return defined $_[0] && ($_[0] <= $_[1]); },
    ">=" => sub { return defined $_[0] && ($_[0] >= $_[1]); },
    ">"  => sub { return defined $_[0] && ($_[0] >  $_[1]); },
  );

  my $cmpfunc;
  if ($isnull) {
    $cmpfunc = sub { return !defined $_[0]; };
  } elsif ($cmp) {
    $cmpfunc = $cmps{$cmp};
  } else {
    $cmpfunc = sub { return defined $_[0]; };
  }

  my $result;
  foreach my $line(@$data) {
    my $curval = $line->[$dsidx];
    if (&$cmpfunc($curval, $value)) {
      $result = &$aggregate_func($ts, $utc, $curval, $result);
    }
    $ts += $step;
  }

  if (defined $result) {
    if ($aggregate eq "avg") {
      my $avg = $result->[1] / $result->[0];
      print $avg, "\n";
    } elsif (($aggregate eq "min") || ($aggregate eq "max")) {
      print $result->[0], " (", format_time($result->[0], $utc), ") ",
            $result->[1], "\n";
    } else {
      # count(), sum()
      print $result, "\n";
    }
  }
}