#!/usr/bin/perl

use strict;
use warnings;

=head1 DESCRIPTION

=encoding utf8

    Скрипт следит за запросами, выполняемыми mysql-сервером, и прибивает те, которые могут тормозить alter-ы. 

      
=head1 OPTIONS

    -H, --host <hostname>
        хост

    -P, --port <port>
        порт

    -u, --user <username>
        пользователь для коннекта к mysql

    -p, --pass <password>
        пароль для коннекта к mysql
        
    --check
        проверка на повторный запуск
        
    --db <database>
        название охраняемой базы, по умолчанию будет охранять все базы
        Использовать с осторожностью, только если точно знаешь, какую database затронет альтер

=head1 EXAMPLES

  Самое простое: 

  direct-query-guard -H ppctest-devtest-mysql.ppc.yandex.ru --port 3346 -u adiuser -p utro
  direct-query-guard -H ppctest-devtest-mysql.ppc.yandex.ru --port 3346 -u adiuser -p utro --check
  direct-query-guard -H ppctest-devtest-mysql.ppc.yandex.ru --port 3346 -u adiuser -p utro --check --db ppc

  ... TODO ...


=head1 TODO

    - !!! проверять, что есть привилегия PROCESS
    + охранять только запросы в waiting metadata lock
    - не убивать "молодые" запросы -- может, он быстренько закончится? -- но если их слишком много, то надо прибивать
    - не убивать запросы, которые сами waiting metadata lock -- непонятно, правильно ли
    - не убивать запросы моложе охраняемых (возможно, это повторение предыдущего пункта) -- не надо так делать; 
    select, запущенный позже альтера, может успеть взять metadata lock и заблокировать альтер 
    проверяется так: select sleep(100), alter, select sleep(200), kill <sleep 100> -- альтер продолжает висеть в waiting metadata lock 


=cut


use Data::Dumper;
use Getopt::Long qw(:config no_ignore_case);
use POSIX qw(strftime);
use YAML;
use List::Util qw(any);

use Yandex::DBTools;
use Yandex::DBShards;


run() unless caller();


sub run
{
    my %O = %{parse_options()};

    $Yandex::DBTools::DONT_SEND_LETTERS = 1;
    %Yandex::DBTools::DB_CONFIG = (
        CHILDS => {
            my_db => {
                'AutoCommit' => '1',
                'CHILDS' => {
                    '_' => {},
                },
                'connect_timeout' => '4',
                'host' => $O{mysql_host},
                'pass' => $O{mysql_password},
                'port' => $O{mysql_port},
                'user' => $O{mysql_user},
                'utf8' => '1',
                'db'   => 'information_schema',
            },
        },
    );
    
    #проверяем, работает ли уже guard
    my $check_copy_query;
    my $mysql_error_msg = "";
    open my $temp_fh, '>', \$mysql_error_msg;
    {
         local *STDERR = $temp_fh;
         $Yandex::DBTools::DB_CONFIG{CHILDS}{my_db}{db} = $O{db} if $O{db};
         $check_copy_query = eval{do_sql("my_db", "SELECT /* dbs-guard-kill-me */ SLEEP(5);")};
    }
    
    $Yandex::DBTools::DB_CONFIG{CHILDS}{my_db}{db} = 'information_schema';
    
    if ($@) {
        die "Неожиданная ошибка при выполнении запроса:\n$mysql_error_msg" if $mysql_error_msg!~/Lost connection to MySQL/;
        if ($O{check}) {
            exit 0;
        } else {
            die "error: direct-query-guard уже запущен\n";
        }
    } else {
        if ($O{check}) {
            die "error: direct-query-guard не запущен\n";
        }
    }
    
    #проверяем наши привилегии
    my @all_dbs = map { values $_ } @{get_all_sql("my_db", "SHOW DATABASES;")};
    my @all_grants = map { values $_ } @{get_all_sql("my_db", "SHOW GRANTS FOR CURRENT_USER;")};

    # если когда-то понадобится явный список игорируемых БД — можно сделать список и регвыр собирать через join
    my @need_dbs = $O{db} ? grep { $_ eq $O{db} } @all_dbs : grep { $_ !~ /^(information_schema|db1|heartbeat|.*lost.?found|mysql|sys|performance_schema|xtrabackup_backupfiles)$/ } @all_dbs;

    my %granted_dbs = ();

    # костыль после отрыва super-привилегий direct-sql: https://st.yandex-team.ru/DIRECT-153260
    if ($O{mysql_user} ne "ppc" 
            && $O{mysql_user} ne "direct-test"
            # пользователь, использующийся в unit-тестах
            && $O{mysql_user} ne "adiuser") {
        die "ошибка: нужно запускать с таким же пользователем, с каким работает приложение (ppc в продакшене, direct-test на ТС)\nЕсли запускается через dbs-guard, для этого нужно не указывать параметр -u\n";
    }

    # закомментированный кусок был актуален, когда у direct-sql были привилегии прибивать все запросы, TODO удалить
#    # нам нужно, чтобы для всех охраняемых баз у текущего пользователя были гранты ALL PRIVILEGES или SUPER
#    foreach my $grant (@all_grants) {
#        if ($grant =~ /^GRANT \s+ .* \b(ALL|SUPER)\b .*\s ON \s(?:(?<db>\*)|(?:\`(?<db>[\w\d-]+)\`))\.\*\s TO \s+ \S+ .*$/ix) {
#            my $granted_db = $+{db};
#            if ($granted_db eq '*') {
#                p("info: granted all privileges on all databases");
#                %granted_dbs = map { $_ => 1 } @all_dbs;
#            } else {
#                p("info: granted all privileges on $granted_db");
#                $granted_dbs{$granted_db} = 1;
#            }
#        } else {
#            # пропускаем гранты, которые мы не умеем парсить
#            p("warn: skip grant '$grant' (can't parse 'em, but probably it's OK)");
#        }
#    }
#
#    if (any { !exists($granted_dbs{$_}) } @need_dbs) {
#        my $missing_dbs = join(',', grep { !exists($granted_dbs{$_}) } @need_dbs);
#        die "error: нет нужных привилегий у текущего пользователя для баз: $missing_dbs\nСОВЕТ: возможно, параметр --db <ppc|ppcdict> сделает хорошо. Прочитай direct-query-guard --help\n.";
#    }
    p("info: starting guard loop");

    my $orig_procname = $0;
    my $i;
    my $killed_cnt = 0;
    while ( 1 ){
        $i++;
        $0 = "($killed_cnt / $i) $orig_procname";
        my $pl = get_all_sql("my_db", ['select * from information_schema.processlist', where => {command__ne => "Sleep", ($O{db} ? (db => $O{db}) : () ) }]); 

        my $to_kill = find_queries_to_kill($pl);
        
        my $killed_cnt_on_iter = kill_queries($to_kill);
        
        $killed_cnt += $killed_cnt_on_iter;
        
        #проверяем наличие убитых запросов на текущей итерации
        if ($killed_cnt_on_iter == 0) {
            sleep 1;
        };
    }

    exit 0;
}


sub p
{
    print localtime(time)." $_\n" for @_;
}


sub find_queries_to_kill
{
    my ($pl) = @_;

    my @to_kill;

    my %to_protect;
    my %drop_trigger;

    # ищем, кого надо охранять
    for my $p ( @$pl ){
        next unless $p->{INFO};
        
        #проверяем наличие специального комментария
        if ( $p->{INFO} =~m!/\*\s*dbs-guard-kill-me\s*\*/! ){ 
            my $table = "/* dbs-guard-kill-me */";
            push @to_kill, {
                process => $p,
                because_of => {
                    table => $table,
                    processes => [$p],
                }

            };
        }
        
        # охраняе только запросы, ждущие лока на метаданные
        next unless $p->{STATE} && ($p->{STATE} eq 'Waiting for table metadata lock');
        if ( $p->{INFO} =~ /RENAME\s+TABLE\s+(.*)/i ){
            my $tables_str = $1;
            my @tables = map { s/\s+TO\s+.*$//ri } split /\s*,\s*/, $tables_str;
            $to_protect{$p->{ID}} = {
                tables => \@tables,
                process => $p,
            };
        } elsif( $p->{INFO} =~ /ALTER\s+TABLE\s+([^\s]+)/i ){
            my $table = $1;
            $to_protect{$p->{ID}} = {
                tables => [$table],
                process => $p,
            };
        } elsif( $p->{INFO} =~ /TRUNCATE\s+TABLE\s+([^\s]+)/i ){
            my $table = $1;
            $to_protect{$p->{ID}} = {
                tables => [$table],
                process => $p,
            };
        } elsif ( $p->{INFO} =~ /CREATE\s+TRIGGER.*ON\s+([^\s]+)/i ){
            my $table = $1;
            $to_protect{$p->{ID}} = {
                tables => [$table],
                process => $p,
            };
        } elsif ( $p->{INFO} =~ /DROP\s+TRIGGER\s+(?:IF\s+EXISTS)?\s*([^\s]+)/i ){
            my $trigger = $1;
            $trigger =~ s/^.*\.//;
            $trigger =~ s/`//g;
            $drop_trigger{$trigger} = $p;
        } elsif ($p->{INFO} =~ /(?:CREATE|DROP)\s+(?:\S+\s+)*INDEX.*\s+ON\s+([^\s\(]+)/i) {
            my $table = $1;
            $to_protect{$p->{ID}} = {
                tables => [$table],
                process => $p,
            };
        }
    }
    # по drop trigger найти, какие таблицы затронуты
    if ( scalar keys %drop_trigger ){
        my $triggers = get_all_sql("my_db", 'select * from information_schema.triggers;'); 
        for my $tr ( @$triggers ){
            next unless $drop_trigger{$tr->{TRIGGER_NAME}};
            my $p = $drop_trigger{TRIGGER_NAME}->{process};
            $to_protect{$p->{ID}} = {
                tables => [ $tr->{EVENT_OBJECT_TABLE} ],
                process => $p,
            };
        }
    }

    return \@to_kill unless keys %to_protect > 0;

    my %table2queries;
    for my $pr ( values %to_protect ){
        # таблицы могли оказаться в обратных кавычках или с указанием базы
        for my $t (map {s/(^`|`$)//gr} map {s/^.*\.//r} @{$pr->{tables}} ){
            push @{$table2queries{$t}}, $pr->{process};
        }
    }

    return \@to_kill unless keys %table2queries > 0;
    
    my $tables_to_protect_regexp = join "|", keys %table2queries;

    for my $p ( @$pl ){
        # не пытаться прибить сами охраняемые запросы
        next if $to_protect{$p->{ID}};
        if($p->{INFO} && $p->{INFO} =~ /\b($tables_to_protect_regexp)\b/ ){
            my $table = $1;
            push @to_kill, {
                process => $p,
                because_of => {
                    table => $table,
                    processes => $table2queries{$table},
                }

            };
        }
    }

    return \@to_kill;
}


sub kill_queries
{
    my ($to_kill) = @_;

    my $cnt = 0;
    for my $p (@$to_kill){
        my $query_snippet = $p->{process}->{INFO} =~ s/\n/ /r;
        $query_snippet =~ s/\s{2,/ /;
        $query_snippet = substr($query_snippet, 0, 120);
        my $id = $p->{process}->{ID};
        my $table_snippet = "$p->{because_of}->{table} (".substr($p->{because_of}->{processes}->[0]->{INFO}, 0, 30).")";
        p("to protect table $table_snippet going to kill session $id ($query_snippet)");
        my $res = eval {do_sql("my_db", "kill ?", $id)};
        my $error = $@;
        if ( $error ){
            p("ERROR: $@");
        } else {
            p "success (code ".($res+0).")" ;
            $cnt++;
        } 
    }

    return $cnt;
}


sub parse_options
{
    my %O = (
    );

    GetOptions(
        "help" => sub {
            system("podselect -section NAME -section DESCRIPTION -section OPTIONS -section EXAMPLES $0 | pod2text"); 
            exit 0;
        },
        "H|host=s" => \$O{mysql_host},
        "P|port=s" => \$O{mysql_port},
        "u|user=s" => \$O{mysql_user},
        "p|password=s" => \$O{mysql_password},
        "check" => \$O{check},
        "db=s" => \$O{db},
    ) || die "can't parse options, stop";

    return \%O;
}

