#!/usr/bin/perl

# This is a library of functions common to two or more programs within
# the Captrap network traffic usage monitor.

# Copyright 2010 Corey Hickey


# This file is part of Captrap.
#
# Captrap is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# Captrap is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with Captrap.  If not, see <http://www.gnu.org/licenses/>.

=head1 NAME

DBCache - a library for client-side caching of DBI queries

=head1 SYNOPSIS

  use DBCache;

  my $host     = 'localhost';
  my $db       = 'pmacct';
  my $user     = 'captrap';
  my $pass     = '6356134';
  my $dsn =  "DBI:mysql:database=$db;host=$host;mysql_auto_reconnect=1";

  my $dbh = DBCache->connect($dsn, $user, $pass,
    { RaiseError => 1, AutoCommit => 1 }) or die;

  my $sel = "select mac_dst, sum(bytes) from acct_eth1 group by mac_dst";
  my $sth = $dbh->prepare($sel);
  $sth->dbcache_execute();

  while (defined(my $row = $sth->fetchrow_arrayref())) {
    print "mac: $row->[0]  bytes: $row->[1]\n";
  }

  $dbh->disconnect();

=head1 DESCRIPTION

As of yet, this library is not intended to be used outside of Captrap; DBCache
is an incomplete wrapper around DBI, and methods which are not re-implemented
within DBCache may give incorrect results. To be safe, use only the methods
documented in this text.

Also, there is no error checking in DBCache! This is ok when used with
RaiseError -- DBI will cause the program to die if there is a database error.

No functions are exported by default.

=head2 Methods

The following methods are implemented in DBCache. Other methods from DBI are
available, but may not work or may give incorrect results when caching is
enabled. Use them with care.

=over

=cut


use strict;
use warnings;

use DBI;

# Note that "private" attributes in DBI must start with "private_"; see perdoc
# DBI.

# To debug:
# use Data::Dumper
# print Dumper($dbh->{private_dbcache); # all handles
# print Dumper($sth->{private_stcache); # this handle

package DBCache;
our @ISA = qw(DBI);

=item connect()

This is a wrapper for DBI::connect(), with no altered behavior aside from
making caching possible if dbcache_execute() is called.

=cut

sub connect {
  my $self = shift;
  my $dbh = $self->SUPER::connect(@_);
  $dbh->{private_dbcache} = {};
  return $dbh;
}


package DBCache::db;
our @ISA = qw(DBI::db);


package DBCache::st;
our @ISA = qw(DBI::st);

=item dbcache_execute(@params)

This method is a wrapper for DBI::db:execute() that will make the results be
cached. Only use this for read-only operations. Use the regular execute() for
UPDATE, DELETE, etc.

=cut

sub dbcache_execute {
  my $self = shift;
  my @params = @_;
  my $s = $self->{Statement};
  # use ';' for readability when printed out
  my $key = join(';', ($s, @params));
  my $dbh = $self->{Database};
  my $dbcache = $dbh->{private_dbcache};
  my $my_dbcache = $dbcache->{$key};
  $self->{private_stcache} = {
    pos     => 0,
    dbcache => $my_dbcache, # could be undef right now
  };
  if (defined($dbcache->{$key})) {
    ++$dbcache->{$key}->{use};
    return 1; # always success, since it's already cached
  }
  # all right, we have a new statement to execute and cache
  # print STDERR "executing new sth\n";
  my $ok = $self->execute(@params);
  if (! $ok) {
    return $ok;
  }
  $my_dbcache = {
    sth => $self,
    key => $key,
    use => 1,
    len => 0,
  };
  my $stcache = $self->{private_stcache};
  $dbcache->{$key} = $stcache->{dbcache} = $my_dbcache;
  return $ok;
}


=item dbcache_uncache()

Remove the cached item matching the statement and parameters with which this
sth was most recently executed. If there are any unfinished handles, this
method does nothing and returns undef; otherwise, this method returns 1.

=cut

sub dbcache_uncache {
  my $self = shift;
  my $stcache = $self->{private_stcache};
  if (! defined($stcache)) {
    die "dbcache: not a cached handle?";
  }
  my $dbh = $self->{Database};
  my $dbcache = $dbh->{private_dbcache};
  my $my_dbcache = $stcache->{dbcache};
  if ($my_dbcache->{use}) {
    return undef;
  }
  $dbcache->{$my_dbcache->{key}} = undef;
  $self->{private_stcache} = undef;
  return 1;
}


=item finish()

This method doesn't actually finish the statement handle, but just makes it
"unused" and detached from the dbcache.

On handles that weren't executed with dbcache_execute(), this method passes
through to the original DBI method.
=cut

sub finish {
  my $self = shift;
  my @params = @_; # shouldn't be any, but pass them anyway
  my $stcache = $self->{private_stcache};
  # if this isn't cached, just pass through
  if (! defined($stcache)) {
    return $self->SUPER::finish(@params);
  }
  my $my_dbcache = $stcache->{dbcache};
  --$my_dbcache->{use};
  $self->{private_stcache} = undef;
  return 1;
}


=item fetchrow_arrayref()

Return a reference to an array of values from the next row. This will be pulled
from the cache; if the cache doesn't yet have that row, it will be fetched
on-demand from the db. Unlike the original DBI fetchrow_arrayref() method, this
method will return a separate reference for each row, so it is safe to store
the references for later read-only use.

On handles that weren't executed with dbcache_execute(), this method passes
through to the original DBI method.

=cut

sub fetchrow_arrayref {
  my $self = shift;
  my @params = @_; # shouldn't be any, but pass them anyway
  my $stcache = $self->{private_stcache};
  # if this isn't cached, just pass through
  if (! defined($stcache)) {
    return $self->SUPER::fetchrow_arrayref(@params);
  }
  my $my_dbcache = $stcache->{dbcache};
  if (! defined($my_dbcache)) {
    # someone keeps calling fetchrow_arrayref() after getting undef already
    return undef;
  }
  if ($stcache->{pos} > $my_dbcache->{len}) {
    die "dbcache internal error: stcache overran dbcache";
  }
  if ($stcache->{pos} == $my_dbcache->{len}) {
    my $sth = $my_dbcache->{sth}; # $sth is not necessarily the same as $self
    if (! defined($sth)) {
      # got all data from db earlier
      $self->finish();
      return undef;
    }
    # We need to fetch a row from the db. According to DBI doc,
    # fetchrow_arrayref() reuses the same ref, so we make our own ref to
    # store in cache.
    # print STDERR "fetching; pos: $stcache->{pos}  len: $my_dbcache->{len}\n";
    my @line = $sth->SUPER::fetchrow_array();
    if (! @line) {
      # got all data from db, so clear dbcache's sth
      $my_dbcache->{sth} = undef;
      $self->finish();
      return undef;
    }
    push(@{$my_dbcache->{data}}, \@line);
    ++$my_dbcache->{len};
    # Ok, we've fetched the next line into the cache, so continue on.
  }
  # print STDERR "returning; pos: $stcache->{pos}\n";
  return $my_dbcache->{data}->[$stcache->{pos}++];
}


=item fetch()

This method passes through to fetchrow_arrayref().

=cut

sub fetch {
  my $self = shift;
  return $self->fetchrow_arrayref(@_);
}

1;

=back

=head1 AUTHOR

Corey Hickey <bugfood-c@fatooh.org>

This library is free software; you may redistribute and/or modify it under the
terms of the GNU General Public License, version 3. See the source file for the
usual GPL preamble and the COPYING file for a copy of the GPL.

=head1 SEE ALSO

DBI, Captrap

=cut