override storage to access sth details

[dbsrgits/DBIx-Class-RowCountStatistics.git] / lib / CtrlO / DBIC / Cursor / RowCountStatistics.pm

diff --git a/lib/CtrlO/DBIC/Cursor/RowCountStatistics.pm b/lib/CtrlO/DBIC/Cursor/RowCountStatistics.pm
index ba0a828..882cc54 100644 (file)
--- a/lib/CtrlO/DBIC/Cursor/RowCountStatistics.pm
+++ b/lib/CtrlO/DBIC/Cursor/RowCountStatistics.pm
@@ -4,26 +4,51 @@ use warnings;
 package CtrlO::DBIC::Cursor::RowCountStatistics;
 use Class::Method::Modifiers;
 
+use CtrlO::DBIC::Cursor::RowCountStatistics::Storage;
+
 use parent 'DBIx::Class::Storage::DBI::Cursor';
 
 our $VERSION = '0.000001'; # 0.0.1
 $VERSION = eval $VERSION;
 
+around new => sub {
+    my $orig = shift;
+    my ($class, $storage, @rest) = @_;
+    return $orig->(
+        $class,
+        CtrlO::DBIC::Cursor::RowCountStatistics::Storage->new($storage),
+        @rest,
+    );
+};
+
 after next => sub {
     my ($self) = @_;
     $self->{_ctrlo_rcs_count}++
         unless $self->{_done};
 };
 
+around all => sub {
+    my $orig = shift;
+    my ($self) = @_;
+    my @rows = $orig->(@_);
+    $self->_emit_query_complete(scalar(@rows));
+    return @rows;
+};
+
 before __finish_sth => sub {
     my ($self) = @_;
     my $sql = $self->sth->{Statement};
+    $self->_emit_query_complete($self->{_ctrlo_rcs_count} || 0);
+};
+
+sub _emit_query_complete {
+    my ($self, $count) = @_;
     $self->storage->debugobj->query_complete(
-        $self->{_ctrlo_rcs_count} || 0,
-        $sql,
+        $count,
+        $self->storage->_cached_sql,
         # TODO pass bind params
     ) if $self->storage->debug;
-};
+}
 
 1;