From 8ce093655fa0815fb97ed5a64f6a4056885f755d Mon Sep 17 00:00:00 2001
From: Daniel Nichter <daniel@percona.com>
Date: Sun, 11 Sep 2011 10:22:40 -0600
Subject: [PATCH] Add dms arg, remove before_row callback, pass some args to
 callbacks, get next boundary if no rows (so exec_nibble callback can skip
 nibbles).

---
 lib/NibbleIterator.pm  | 51 ++++++++++++++++++++++--------------------
 t/lib/NibbleIterator.t | 18 ++++-----------
 2 files changed, 31 insertions(+), 38 deletions(-)

diff --git a/lib/NibbleIterator.pm b/lib/NibbleIterator.pm
index 0ffd89ef..c4b7683c 100644
--- a/lib/NibbleIterator.pm
+++ b/lib/NibbleIterator.pm
@@ -103,7 +103,7 @@ sub new {
    MKDEBUG && _d('Next upper boundary statement:', $ub_sql);
 
    my $nibble_sql
-      = "SELECT "
+      = ($args{dms} ? "$args{dms} " : "SELECT ")
       . ($args{select} ? $args{select}
                        : join(', ', map { $q->quote($_) } @{$asc->{cols}}))
       . " FROM $from "
@@ -117,7 +117,7 @@ sub new {
    # If the chunk size is >= number of rows in table, then we don't
    # need to chunk; we can just select all rows, in order, at once.
    my $one_nibble_sql
-      = "SELECT "
+      = ($args{dms} ? "$args{dms} " : "SELECT ")
       . ($args{select} ? $args{select}
                        : join(', ', map { $q->quote($_) } @{$asc->{cols}}))
       . " FROM $from "
@@ -166,39 +166,53 @@ sub next {
       if ( $row ) {
          $self->{rowno}++;
          MKDEBUG && _d('Row', $self->{rowno}, 'in nibble', $self->{nibbleno});
-         if ( my $callback = $self->{callbacks}->{before_row} ) {
-            $callback->();
-         }
          # fetchrow_arraryref re-uses its internal arrayref, so we must copy.
          return [ @$row ];
       }
       MKDEBUG && _d('No more rows in nibble', $self->{nibbleno});
       if ( my $callback = $self->{callbacks}->{after_nibble} ) {
-         $callback->();
+         $callback->(
+            dbh => $self->{dbh},
+            tbl => $self->{tbl},
+         );
       }
       $self->{rowno}     = 0;
       $self->{have_rows} = 0;
    }
 
    # If there's another boundary, fetch the rows within it.
-   if ( $self->_next_boundaries() ) {
+   BOUNDARY:
+   while ( $self->_next_boundaries() ) {
       $self->{nibbleno}++;
       MKDEBUG && _d($self->{nibble_sth}->{Statement}, 'params:',
          join(', ', (@{$self->{lb}}, @{$self->{ub}})));
-      $self->{nibble_sth}->execute(@{$self->{lb}}, @{$self->{ub}});
-      $self->{have_rows} = $self->{nibble_sth}->rows();
+      if ( my $callback = $self->{callbacks}->{exec_nibble} ) {
+         $self->{have_rows} = $callback->(
+            dbh => $self->{dbh},
+            tbl => $self->{tbl},
+            sth => $self->{nibble_sth},
+            lb  => $self->{lb},
+            ub  => $self->{ub},
+         );
+      }
+      else {
+         $self->{nibble_sth}->execute(@{$self->{lb}}, @{$self->{ub}});
+         $self->{have_rows} = $self->{nibble_sth}->rows();
+      }
       if ( $self->{have_rows} ) {
          MKDEBUG && _d($self->{have_rows}, 'rows in nibble', $self->{nibbleno});
-         if ( my $callback = $self->{callbacks}->{before_nibble} ) {
-            $callback->();
-         }
          return $self->next();
       }
+      MKDEBUG && _d('No rows in nibble or nibble skipped');
+      next BOUNDARY;
    }
 
    MKDEBUG && _d('Done nibbling');
    if ( my $callback = $self->{callbacks}->{done} ) {
-      $callback->();
+      $callback->(
+         dbh => $self->{dbh},
+         tbl => $self->{tbl},
+      );
    }
    return;
 }
@@ -208,17 +222,6 @@ sub nibble_number {
    return $self->{nibbleno};
 }
 
-sub number_of_rows {
-   my ($self) = @_;
-   return $self->{have_rows};
-}
-
-sub row_number {
-   my ($self) = @_;
-   return $self->{nibbleno};
-}
-
-
 sub _can_nibble_once {
    my ($self) = @_;
    my ($dbh, $tbl, $q) = @{$self}{qw(dbh tbl Quoter)};
diff --git a/t/lib/NibbleIterator.t b/t/lib/NibbleIterator.t
index fefc68d8..6b2becb9 100644
--- a/t/lib/NibbleIterator.t
+++ b/t/lib/NibbleIterator.t
@@ -20,7 +20,7 @@ use OptionParser;
 use MySQLDump;
 use TableParser;
 use TableNibbler;
-use TableChecksum;
+use RowChecksum;
 use NibbleIterator;
 use PerconaTest;
 
@@ -47,7 +47,7 @@ my $tp = new TableParser(Quoter=>$q);
 my $du = new MySQLDump();
 my $nb = new TableNibbler(TableParser=>$tp, Quoter=>$q);
 my $o  = new OptionParser(description => 'NibbleIterator');
-my $tc = new TableChecksum(OptionParser => $o, Quoter=>$q);
+my $rc = new RowChecksum(OptionParser => $o, Quoter=>$q);
 
 $o->get_specs("$trunk/bin/pt-table-checksum");
 
@@ -251,8 +251,6 @@ $ni = make_nibble_iter(
    argv      => [qw(--databases test --chunk-size 2)],
    callbacks => {
       init          => sub { print "init\n" },
-      before_nibble => sub { print "before nibble ".$ni->nibble_number()."\n" },
-      before_row    => sub { print "before row\n" },
       after_nibble  => sub { print "after nibble ".$ni->nibble_number()."\n" },
       done          => sub { print "done\n" },
    }
@@ -269,17 +267,8 @@ my $output = output(
 is(
    $output,
 "init
-before nibble 1
-before row
-before row
 after nibble 1
-before nibble 2
-before row
-before row
 after nibble 2
-before nibble 3
-before row
-before row
 after nibble 3
 done
 done
@@ -287,6 +276,7 @@ done
    "callbacks"
 );
 
+# TODO: test exec_nibble callback
 
 # ############################################################################
 # Nibble a larger table by numeric pk id
@@ -315,7 +305,7 @@ SKIP: {
       tbl_struct => $tp->parse(
          $du->get_create_table($dbh, $q, 'sakila', 'country')),
    };
-   my $chunk_checksum = $tc->make_chunk_checksum(
+   my $chunk_checksum = $rc->make_chunk_checksum(
       dbh => $dbh,
       tbl => $tbl,
    );