From 8e4824813f0ec0f12bb28352f751b66d0f53884c Mon Sep 17 00:00:00 2001
From: Ilaria Migliozzi <87898948+migliozziilaria@users.noreply.github.com>
Date: Mon, 6 Nov 2023 16:22:14 +0100
Subject: [PATCH] added new hook before_die (#509)

* added new hook before_die

before die, the script calls this hook in oder to execute extra user's
operations

* PR-509 -  added new hook before_die

- Added test cases

---------

Co-authored-by: Sveta Smirnova <sveta.smirnova@percona.com>
---
 bin/pt-online-schema-change                   |  20 ++-
 t/pt-online-schema-change/plugin_before_die.t | 114 ++++++++++++++++++
 .../samples/plugins/all_hooks.pm              |   5 +
 .../samples/plugins/before_die.pm             |  20 +++
 .../plugins/before_die_after_create.pm        |  33 +++++
 .../plugins/before_die_after_nibble.pm        |  32 +++++
 6 files changed, 222 insertions(+), 2 deletions(-)
 create mode 100644 t/pt-online-schema-change/plugin_before_die.t
 create mode 100644 t/pt-online-schema-change/samples/plugins/before_die.pm
 create mode 100644 t/pt-online-schema-change/samples/plugins/before_die_after_create.pm
 create mode 100644 t/pt-online-schema-change/samples/plugins/before_die_after_nibble.pm

diff --git a/bin/pt-online-schema-change b/bin/pt-online-schema-change
index 38311546..6ab96703 100755
--- a/bin/pt-online-schema-change
+++ b/bin/pt-online-schema-change
@@ -8481,7 +8481,7 @@ my $term_readkey = eval {
 
 use sigtrap 'handler', \&sig_int, 'normal-signals';
 
-
+my $plugin;
 my $exit_status = 0;
 my $oktorun     = 1;
 my $dont_interrupt_now = 0;
@@ -8535,6 +8535,9 @@ sub _die {
    $exit_status ||= 255;
    chomp ($msg);
    print "$msg\n";
+   if ( $plugin && $plugin->can('before_die') ) {
+     $plugin->before_die(exit_status => $exit_status);
+   }
    exit $exit_status;
 }
 
@@ -8853,7 +8856,7 @@ sub main {
    # ########################################################################
    # Create --plugin.
    # ########################################################################
-   my $plugin;
+
    if ( my $file = $o->get('plugin') ) {
       _die("--plugin file $file does not exist", INVALID_PLUGIN_FILE) unless -f $file;
       eval {
@@ -9590,6 +9593,9 @@ sub main {
          $cxn->dbh()->do($sql);
       };
       if ( $EVAL_ERROR ) {
+        if ( $plugin && $plugin->can('before_die') ) {
+          $plugin->before_die(exit_status => $EVAL_ERROR);
+        }
          # this is trapped by a signal handler. Don't replace it with _die
          die "Error altering new table $new_tbl->{name}: $EVAL_ERROR\n";
       }
@@ -10172,6 +10178,10 @@ sub main {
       1 while $nibble_iter->next();
    };
    if ( $EVAL_ERROR ) {
+      if ( $plugin && $plugin->can('before_die') ) {
+        $plugin->before_die(exit_status => $EVAL_ERROR);
+      }
+
       die ts("Error copying rows from $orig_tbl->{name} to "
         . "$new_tbl->{name}: $EVAL_ERROR");
    }
@@ -13197,6 +13207,7 @@ These hooks, in this order, are called if defined:
    before_drop_old_table
    after_drop_old_table
    before_drop_triggers
+   before_die
    before_exit
    get_slave_lag
 
@@ -13313,6 +13324,11 @@ Here's a plugin file template for all hooks:
       print "PLUGIN before_drop_triggers\n";
    }
 
+   sub before_die {
+      my ($self, %args) = @_;
+      print "PLUGIN before_die\n";
+   }
+
    sub before_exit {
       my ($self, %args) = @_;
       print "PLUGIN before_exit\n";
diff --git a/t/pt-online-schema-change/plugin_before_die.t b/t/pt-online-schema-change/plugin_before_die.t
new file mode 100644
index 00000000..fbf6e941
--- /dev/null
+++ b/t/pt-online-schema-change/plugin_before_die.t
@@ -0,0 +1,114 @@
+#!/usr/bin/env perl
+
+BEGIN {
+   die "The PERCONA_TOOLKIT_BRANCH environment variable is not set.\n"
+      unless $ENV{PERCONA_TOOLKIT_BRANCH} && -d $ENV{PERCONA_TOOLKIT_BRANCH};
+   unshift @INC, "$ENV{PERCONA_TOOLKIT_BRANCH}/lib";
+};
+
+use strict;
+use warnings FATAL => 'all';
+use English qw(-no_match_vars);
+use Test::More;
+
+use Data::Dumper;
+use PerconaTest;
+use Sandbox;
+
+require "$trunk/bin/pt-online-schema-change";
+
+my $dp = new DSNParser(opts=>$dsn_opts);
+my $sb = new Sandbox(basedir => '/tmp', DSNParser => $dp);
+my $master_dbh = $sb->get_dbh_for('master');
+
+if ( !$master_dbh ) {
+   plan skip_all => 'Cannot connect to sandbox master';
+}
+
+# The sandbox servers run with lock_wait_timeout=3 and it's not dynamic
+# so we need to specify --set-vars innodb_lock_wait_timeout=3 else the
+# tool will die.
+my $master_dsn = 'h=127.1,P=12345,u=msandbox,p=msandbox';
+my @args       = (qw(--set-vars innodb_lock_wait_timeout=3));
+my $sample     = "t/pt-online-schema-change/samples/";
+my $plugin     = "$trunk/$sample/plugins";
+my $output;
+my $exit_status;
+
+# ############################################################################
+# https://bugs.launchpad.net/percona-toolkit/+bug/1171653
+# 
+# ############################################################################
+$sb->load_file('master', "$sample/basic_no_fks.sql");
+
+# Should be greater than chunk-size and big enough, so plugin will trigger few times
+my $num_rows = 5000;
+diag("Loading $num_rows into the table. This might take some time.");
+diag(`util/mysql_random_data_load --host=127.0.0.1 --port=12345 --user=msandbox --password=msandbox pt_osc t $num_rows`);
+
+($output, $exit_status) = full_output(
+   sub { pt_online_schema_change::main(@args,
+      "$master_dsn,D=pt_osc,t=t",
+      "--alter", "CHARACTER SET utf9",
+      '--plugin', "$plugin/before_die.pm",
+      '--execute') },
+);
+
+like(
+   $output,
+   qr/PLUGIN before_die/s,
+   'Plugin before_die called for invalid ALTER command'
+);
+
+like(
+   $output,
+   qr/Exit status: .*Unknown character set: 'utf9'/s,
+   'Expected exit status for invalid ALTER command'
+);
+
+($output, $exit_status) = full_output(
+   sub { pt_online_schema_change::main(@args,
+      "$master_dsn,D=pt_osc,t=t",
+      "--alter", "ENGINE=InnoDB",
+      '--plugin', "$plugin/before_die_after_nibble.pm",
+      '--execute') },
+);
+
+like(
+   $output,
+   qr/PLUGIN before_die/s,
+   'Plugin before_die called for error while copying nibble'
+);
+
+like(
+   $output,
+   qr/Exit status: .*You have an error in your SQL syntax/s,
+   'Expected exit status for error while copying nibble'
+);
+
+($output, $exit_status) = full_output(
+   sub { pt_online_schema_change::main(@args,
+      "$master_dsn,D=pt_osc,t=t",
+      "--alter", "ENGINE=InnoDB",
+      '--plugin', "$plugin/before_die_after_create.pm",
+      '--execute') },
+);
+
+like(
+   $output,
+   qr/PLUGIN before_die/s,
+   'Plugin before_die called for _die call'
+);
+
+like(
+   $output,
+   qr/Exit status: 4/s,
+   'Expected exit status for table definition error handled in the _die call'
+);
+
+# #############################################################################
+# Done.
+# #############################################################################
+$sb->wipe_clean($master_dbh);
+ok($sb->ok(), "Sandbox servers") or BAIL_OUT(__FILE__ . " broke the sandbox");
+done_testing;
diff --git a/t/pt-online-schema-change/samples/plugins/all_hooks.pm b/t/pt-online-schema-change/samples/plugins/all_hooks.pm
index 3aa37bd6..d97ee48a 100644
--- a/t/pt-online-schema-change/samples/plugins/all_hooks.pm
+++ b/t/pt-online-schema-change/samples/plugins/all_hooks.pm
@@ -96,6 +96,11 @@ sub before_drop_triggers {
    print "PLUGIN before_drop_triggers\n";
 }
  
+sub before_die {
+   my ($self, %args) = @_;
+   print "PLUGIN before_die\n";
+}
+
 sub before_exit {
    my ($self, %args) = @_;
    print "PLUGIN before_exit\n";
diff --git a/t/pt-online-schema-change/samples/plugins/before_die.pm b/t/pt-online-schema-change/samples/plugins/before_die.pm
new file mode 100644
index 00000000..2d99e6fc
--- /dev/null
+++ b/t/pt-online-schema-change/samples/plugins/before_die.pm
@@ -0,0 +1,20 @@
+package pt_online_schema_change_plugin;
+
+use strict;
+use warnings FATAL => 'all';
+use English qw(-no_match_vars);
+use constant PTDEBUG => $ENV{PTDEBUG} || 0;
+
+sub new {
+   my ($class, %args) = @_;
+   my $self = { %args };
+   return bless $self, $class;
+}
+
+sub before_die {
+   my ($self, %args) = @_;
+   print "PLUGIN before_die\n";
+   print "Exit status: $args{exit_status}\n";
+}
+
+1;
diff --git a/t/pt-online-schema-change/samples/plugins/before_die_after_create.pm b/t/pt-online-schema-change/samples/plugins/before_die_after_create.pm
new file mode 100644
index 00000000..ce089bfc
--- /dev/null
+++ b/t/pt-online-schema-change/samples/plugins/before_die_after_create.pm
@@ -0,0 +1,33 @@
+package pt_online_schema_change_plugin;
+
+use strict;
+use warnings FATAL => 'all';
+use English qw(-no_match_vars);
+use constant PTDEBUG => $ENV{PTDEBUG} || 0;
+
+sub new {
+   my ($class, %args) = @_;
+   my $self = { %args };
+   return bless $self, $class;
+}
+
+sub before_die {
+   my ($self, %args) = @_;
+   print "PLUGIN before_die\n";
+   print "Exit status: $args{exit_status}\n";
+}
+
+sub after_create_new_table {
+   my ($self, %args) = @_;
+
+   print "PLUGIN after_create_new_table\n";
+
+   my $dbh     = $self->{aux_cxn}->dbh;
+   my $new_tbl = $args{new_tbl}->{name};
+
+   # Remove PRIMARY KEY, so pt-osc fails with an error and handles 
+   # it in the _die call
+   $dbh->do("ALTER TABLE ${new_tbl} MODIFY COLUMN id INT NOT NULL, DROP PRIMARY KEY");
+}
+
+1;
diff --git a/t/pt-online-schema-change/samples/plugins/before_die_after_nibble.pm b/t/pt-online-schema-change/samples/plugins/before_die_after_nibble.pm
new file mode 100644
index 00000000..2e285342
--- /dev/null
+++ b/t/pt-online-schema-change/samples/plugins/before_die_after_nibble.pm
@@ -0,0 +1,32 @@
+package pt_online_schema_change_plugin;
+
+use strict;
+use warnings FATAL => 'all';
+use English qw(-no_match_vars);
+use constant PTDEBUG => $ENV{PTDEBUG} || 0;
+
+sub new {
+   my ($class, %args) = @_;
+   my $self = { %args };
+   return bless $self, $class;
+}
+
+sub before_die {
+   my ($self, %args) = @_;
+   print "PLUGIN before_die\n";
+   print "Exit status: $args{exit_status}\n";
+}
+
+sub on_copy_rows_after_nibble {
+   my ($self, %args) = @_;
+   my $tbl = $args{tbl};
+   print "PLUGIN on_copy_rows_after_nibble\n";
+   if ($tbl->{row_cnt} > 1000) {
+      my $dbh      = $self->{aux_cxn}->dbh;
+
+      # Run invalid query to get error
+      $dbh->do("SELECT * FRO " . $tbl->{name});
+   }
+}
+
+1;