root/lang/perl/DBIx-ShardManager-Pacific/trunk/pacific_divide @ 35192

#! /usr/bin/perl

use strict;
use warnings;

use DBI;
use Getopt::Long;
use File::Temp qw(tempfile);
use IO::File::AtomicChange;
use JSON;
use Storable qw(dclone);

our $VERSION = 0.01;

my $READ_PRIVS = "SELECT";
my $WRITE_PRIVS = "INSERT,UPDATE,DELETE,CREATE TEMPORARY TABLES,EXECUTE,LOCK TABLES";

my %DEFAULT_PORT = (
    mysql => 3306,
);

my (
    $opt_rdbms, $opt_database, $opt_user, $opt_password, $opt_shard_def,
    $opt_new_host, $opt_new_port, $opt_from_id, $opt_incline_source,
    $opt_drop_old_user_after, $opt_help, $opt_version,
);

sub usage {
    print <<"EOT";
Usage: $0 [options]

Options:
  --rdbms=driver      driver (default: mysql)
  --database=db_name  database name (mandatory)
  --user=user         administrator user of the rdb nodes (default: root)
  --password=pass     (default: empty)
  --shard-def=file    shard definition file (mandatory)
  --new-host=host     hostname of the new primary node (mandatory)
  --new-port=port     (optional)
  --from=lowest_id    lowest id of the data to be moved to the new node
  --incline-source=file
                      if set, re-installs incline triggers (default: empty)
  --drop-old-user-after=seconds
                      drops read privileges granted to old user after given
                      seconds (default: 60)
  --help              print this help
  --version           print version
EOT
    exit 0;
}

sub db_connect {
    my ($host, $port) = @_;
    $port ||= $DEFAULT_PORT{$opt_rdbms};
    my $dbh = DBI->connect(
        "DBI:$opt_rdbms:dbname=$opt_database;host=$host;port=$port",
        $opt_user,
        $opt_password,
        { RaiseError => 1 },
    ) or die $DBI::errstr;
    $dbh;
}

sub get_master_status {
    my $dbh = shift;
    $dbh->selectall_arrayref('SHOW MASTER STATUS', { Slice => {} })->[0];
}

sub get_slave_status {
    my $dbh = shift;
    $dbh->selectall_arrayref('SHOW SLAVE STATUS', { Slice => {} })->[0];
}

sub master_is {
    my ($dbh, $host, $port) = @_;
    $port ||= $DEFAULT_PORT{$opt_rdbms};
    my $status = get_slave_status($dbh);
    $status->{Master_Host} eq $host && $status->{Master_Port} == $port;
}

sub build_grant_username {
    my $u = shift;
    "'$u'\@'localhost','$u'\@'\%'";
}

GetOptions(
    'rdbms=s'               => \$opt_rdbms,
    'database=s'            => \$opt_database,
    'user=s'                => \$opt_user,
    'password=s'            => \$opt_password,
    'shard-def=s'           => \$opt_shard_def,
    'new-host=s'            => \$opt_new_host,
    'new-port=i'            => \$opt_new_port,
    'from-id=i'             => \$opt_from_id,
    'incline-source=s'      => \$opt_incline_source,
    'drop-old-user-after=i' => \$opt_drop_old_user_after,
    help                    => \$opt_help,
    version                 => \$opt_version,
) or exit 1;
usage()
    if $opt_help;
if ($opt_version) {
    print "$VERSION\n";
    exit 0;
}

# validate options
$opt_rdbms ||= 'mysql';
die "--rdbms=driver set to unknown rdbms:$opt_rdbms\n"
    unless $DEFAULT_PORT{$opt_rdbms || ''};
die "--database=db_name not specified\n"
    unless $opt_database;
$opt_user ||= 'root';
$opt_password ||= '';
die "--shard-def=file not specified\n"
    unless $opt_shard_def;
die "--new-host=host not specified\n"
    unless $opt_new_host;
die "--from-id=lowest id not specified\n"
    unless defined $opt_from_id;
$opt_drop_old_user_after ||= 60;

# load definition file
my $shard_def = from_json(
    do {
        open my $fh, '<', $opt_shard_def
            or die "failed to open file:$opt_shard_def:$!";
        my $lines = join '', <$fh>;
        close $fh;
        $lines;
    },
    { utf8 => 1 },
);
die "failed to parse file:$opt_shard_def:root object is not a hash\n"
    unless ref $shard_def eq 'HASH';
die "algorithm is not \"range-int\" in file:$opt_shard_def\n"
    unless $shard_def->{algorithm} eq 'range-int';
die "cannot parse \"map\" attribute in file:$opt_shard_def\n"
    unless ref $shard_def->{map} eq 'HASH';
# get information of the original node to be divided
my ($orig_from_id, $orig_node) = do {
    my $cand_lb;
    for my $lb (sort { $a <=> $b } keys %{$shard_def->{map}}) {
        if ($lb <= $opt_from_id) {
            $cand_lb = $lb;
        }
    }
    die "--from-id=$opt_from_id is out of range\n"
        unless defined $cand_lb;
    my $orig_node = $shard_def->{map}->{$cand_lb};
    $orig_node = [ $orig_node ]
        unless ref $orig_node eq 'ARRAY';
    die "new node should be different from the original\n"
        if $opt_new_host eq $orig_node->[0]->{host}
            && ($opt_new_port || $DEFAULT_PORT{$opt_rdbms})
                == ($orig_node->[0]->{port} || $DEFAULT_PORT{$opt_rdbms});
    ($cand_lb, $orig_node);
};
print "INFO: dividing node: $orig_node->[0]->{host}:$orig_node->[0]->{port}\n";
die "no username defined for original host\n"
    unless $orig_node->[0]->{username};
my $orig_username = $orig_node->[0]->{username};

# connect to orig and new node
print "INFO: connecting to database nodes\n";
my $orig_dbh = db_connect($orig_node->[0]->{host}, $orig_node->[0]->{port});
my $new_dbh = db_connect($opt_new_host, $opt_new_port);

print "INFO: checking replication status\n";
die "new node is not currently slave of the original node\n"
    unless master_is(
        $new_dbh, $orig_node->[0]->{host}, $orig_node->[0]->{port},
    );
die "new node is not receiving events from master\n"
    unless get_slave_status($new_dbh)->{Slave_IO_State}
    eq 'Waiting for master to send event';

# drop write privileges
print "INFO: dropping write privileges to original node\n";
$orig_dbh->do(
    "REVOKE $WRITE_PRIVS ON $opt_database.* FROM "
        . build_grant_username($orig_username),
);
print "INFO: dropped write privileges from original node, to recover:\n";
print "INFO:   GRANT $WRITE_PRIVS ON $opt_database.* TO ", build_grant_username($orig_username), ";\n";
print "INFO:\n";

# TODO wait until all write queries on the original node stops running

print "INFO: waiting for new node to catch up\n";
# wait for the nodes to sync
while (1) {
    my $m = get_master_status($orig_dbh);
    my $s = get_slave_status($new_dbh);
    print "INFO:  $m->{Position} vs. $s->{Read_Master_Log_Pos}\n";
    last
        if $m->{File} eq $s->{Master_Log_File}
            && $m->{Position} eq $s->{Read_Master_Log_Pos};
    sleep 1;
}

# upgrade new node to master
print "INFO: nodes are in sync, stopping replication\n";
$new_dbh->do(
    'STOP SLAVE',
);

# drop old user from new node
print "INFO: dropping old user '$orig_username' from new node\n";
$new_dbh->do(
    'DROP USER ' . build_grant_username($orig_username),
);

my $new_username = "pac" . time;

sub updated_shard_def {
    my $r = dclone($shard_def);
    $orig_node->[0]->{username} = $new_username;
    $r->{map}->{$opt_from_id} = [
        {
            %{$orig_node->[0]},
            host => $opt_new_host,
            ($opt_new_port ? (port => $opt_new_port) : ()),
        },
    ];
    return $r;
}

# reinstall incline triggers
if ($opt_incline_source) {
    print "INFO: reinstalling incline triggers\n";
    my ($tmp_fh, $tmp_def) = tempfile();
    print $tmp_fh to_json(updated_shard_def(), { utf8 => 1, pretty => 1 });
    $tmp_fh->flush;
    my @incline_cmd = (
        q(incline),
        '--mode=shard',
        "--source=$opt_incline_source",
        "--shard-source=$tmp_def",
        "--rdbms=$opt_rdbms",
        "--database=$opt_database",
        ($opt_user ? "--user=$opt_user" : ()),
        ($opt_password ? "--password=$opt_password" : ()),
    );
    if ($orig_from_id != $opt_from_id) {
        my @opts = (
            "--host=$orig_node->[0]->{host}",
            "--port=" . ($orig_node->[0]->{port} || $DEFAULT_PORT{$opt_rdbms}),
        );
        print "INFO:   ", join(' ', @incline_cmd, @opts), " drop-trigger\n";
        system(@incline_cmd, @opts, 'drop-trigger') == 0
            or die "failed to drop incline triggers from the original node:$?";
        print "INFO:   ", join(' ', @incline_cmd, @opts), " create-trigger\n";
        system(@incline_cmd, @opts, 'create-trigger') == 0
            or die "failed to install incline triggers to the original node:$?";
    }
    my @opts = (
        "--host=$opt_new_host",
        "--port=" . ($opt_new_port || $DEFAULT_PORT{$opt_rdbms}),
    );
    print "INFO:   ", join(' ', @incline_cmd, @opts), " drop-trigger\n";
    system(@incline_cmd, @opts, 'drop-trigger') == 0
        or die "failed to drop incline triggers from the new node:$?";
    print "INFO:   ", join(' ', @incline_cmd, @opts), " create-trigger\n";
    system(@incline_cmd, @opts, 'create-trigger') == 0
        or die "failed to install incline triggers to the new node:$?";
}

# create new user on both nodes
print "INFO: creating user '$new_username' on the original node\n";
$orig_dbh->do(
    "GRANT $READ_PRIVS,$WRITE_PRIVS ON $opt_database.* TO "
        . build_grant_username($new_username),
);
print "INFO: creating user '$new_username' on the new node\n";
$new_dbh->do(
    "GRANT $READ_PRIVS,$WRITE_PRIVS ON $opt_database.* TO "
        . build_grant_username($new_username),
);

# update shard definition
$shard_def = updated_shard_def();

{ # link current shard def as a backup
    my $i = 1;
    while (! link($opt_shard_def, "$opt_shard_def.$i")) {
        $i++;
    }
    print "INFO: link(2)ed original shard definition file to:",
        "$opt_shard_def.$i\n";
}

{ # update shard def
    my $fh = IO::File::AtomicChange->new($opt_shard_def, 'w')
        or die "failed to open file:$opt_shard_def:$!";
    $fh->print(to_json($shard_def, { utf8 => 1, pretty => 1 }));
    $fh->close;
}
print <<"EOT";
INFO: updated the shard definition file
INFO: 
INFO: CONGRATULATIONS!!! The node has been successfully divided!
INFO: Now you have to copy the shard definition file to your servers.
INFO: 
INFO: After $opt_drop_old_user_after seconds, read privileges to old user: $orig_username
INFO: on the original node will be revoked.  Or you can type Ctrl-C and
INFO  manually drop the privileges, by:
INFO:   DROP USER '$orig_username'\@'localhost','$orig_username'\@'\%'
INFO: 
EOT

sleep $opt_drop_old_user_after;

$orig_dbh->do(
    'DROP USER ' . build_grant_username($orig_username),
);
print "INFO: dropped old user from original node\n";