File Coverage

blib/lib/Thrift/API/HiveClient2.pm

Criterion	Covered	Total	%
statement	35	180	19.4
branch	0	66	0.0
condition	0	18	0.0
subroutine	12	31	38.7
pod	6	8	75.0
total	53	303	17.4

package Thrift::API::HiveClient2;
$Thrift::API::HiveClient2::VERSION = '0.024';
{
  $Thrift::API::HiveClient2::DIST = 'Thrift-API-HiveClient2';
}

# ABSTRACT: Perl to HiveServer2 Thrift API wrapper

use 5.010;
use strict;
use warnings;
use Moo;
use Carp;
use Scalar::Util qw( reftype blessed );
use List::MoreUtils 'zip';

use Thrift;
use Thrift::Socket;
use Thrift::BufferedTransport;

# Protocol loading is done dynamically later.

use Thrift::API::HiveClient2::TCLIService;

# See https://msdn.microsoft.com/en-us/library/ms711683(v=vs.85).aspx
my @odbc_coldesc_fields = qw(
    TABLE_CAT
    TABLE_SCHEM
    TABLE_NAME
    COLUMN_NAME
    DATA_TYPE
    TYPE_NAME
    COLUMN_SIZE
    BUFFER_LENGTH
    DECIMAL_DIGITS
    NUM_PREC_RADIX
    NULLABLE
    REMARKS
    COLUMN_DEF
    SQL_DATA_TYPE
    SQL_DATETIME_SUB
    CHAR_OCTET_LENGTH
    ORDINAL_POSITION
    IS_NULLABLE
);

my @tabledesc_fields = qw(
    TABLE_CAT
    TABLE_SCHEM
    TABLE_NAME
    TABLE_TYPE
    REMARKS
);

# Don't use XS for now, fails initializing properly with BufferedTransport. See
# Thrift::XS documentation.
has use_xs => (
    is      => 'rwp',
    default => sub {0},
    lazy    => 1,
);

has host => (
    is      => 'ro',
    default => sub {'localhost'},
);
has port => (
    is      => 'ro',
    default => sub {10_000},
);
has sasl => (
    is      => 'ro',
    default => 0,
);

# Kerberos principal
# Usually in the format 'hive/{hostname}@REALM.COM';
has principal => (
    is => 'rw',
);

# 1 hour default recv socket timeout. Increase for longer-running queries
# called "timeout" for simplicity's sake, as this is how a user will experience
# it: a time after which the Thrift stack will throw an exception if not
# getting an answer from the server

has timeout => (
    is      => 'rw',
    default => sub {3_600},
);

# These exist to make testing with various other Thrift Implementation classes
# easier, eventually.

has _socket    => ( is => 'rwp', lazy => 1 );
has _transport => ( is => 'rwp', lazy => 1 );
has _protocol  => ( is => 'rwp', lazy => 1 );
has _client    => ( is => 'rwp', lazy => 1 );
has _sasl      => ( is => 'rwp', lazy => 1 );

# setters implied by the 'rwp' mode on the attrs above.

sub _set_socket    { $_[0]->{_socket}    = $_[1] }
sub _set_transport { $_[0]->{_transport} = $_[1] }
sub _set_protocol  { $_[0]->{_protocol}  = $_[1] }
sub _set_client    { $_[0]->{_client}    = $_[1] }

sub _set_sasl {
    my ( $self, $sasl ) = @_;
    return if !$sasl;

    # This normally selects XS first (hopefully)
    require Authen::SASL;
    Authen::SASL->import;

    require Thrift::SASL::Transport;
    Thrift::SASL::Transport->import;

    if ( $sasl == 1 ) {
        return $self->{_sasl} = Authen::SASL->new( mechanism => 'GSSAPI' );
    }
    elsif ( reftype $sasl eq "HASH" ) {
        return $self->{_sasl} = Authen::SASL->new(%$sasl);    #, debug => 8 );
    }
    die "Incorrect parameter passed to _set_sasl";
}

# after constructon is complete, initialize any attributes that
# weren't set in the constructor.
sub BUILD {
    my $self = shift;

    $self->_set_socket( Thrift::Socket->new( $self->host, $self->port ) )
        unless $self->_socket;
    $self->_socket->setRecvTimeout( $self->timeout * 1000 );

    $self->_set_sasl( $self->sasl ) if ( $self->sasl && !$self->_sasl );

    if ( !$self->_transport ) {
        my $transport = Thrift::BufferedTransport->new( $self->_socket );
        if ( $self->_sasl ) {
            my $debug = 0;
            $self->_set_transport( Thrift::SASL::Transport->new( $transport, $self->_sasl, $debug, $self->principal ) );
        }
        else {
            $self->_set_transport($transport);
        }
    }

    $self->_set_protocol( $self->_init_protocol( $self->_transport ) )
        unless $self->_protocol;

    $self->_set_client( Thrift::API::HiveClient2::TCLIServiceClient->new( $self->_protocol ) )
        unless $self->_client;
}

sub _init_protocol {
    my $self = shift;
    my $err;
    my $protocol = eval {
        $self->use_xs
            && require Thrift::XS::BinaryProtocol;
        Thrift::XS::BinaryProtocol->new( $self->_transport );
    } or do { $err = $@; 0 };
    $protocol
        ||= do { require Thrift::BinaryProtocol; Thrift::BinaryProtocol->new( $self->_transport ) };
    $self->_set_use_xs(0) if ref($protocol) !~ /XS/;

    # TODO Add warning when XS was asked but failed to load
    return $protocol;
}

sub connect {
    my ($self) = @_;
    $self->_transport->open;
}

has _session => (
    is  => 'rwp',
    isa => sub {
        my($val) = @_;
        if (   !blessed( $val )
            || !$val->isa('Thrift::API::HiveClient2::TOpenSessionResp')
        ) {
            die sprintf "Session `%s` isn't a Thrift::API::HiveClient2::TOpenSessionResp",
                $val // '[undefined]'
            ;
        }
    },
    lazy    => 1,
    builder => '_build_session',
);

has username => (
    is      => 'rwp',
    lazy    => 1,
    default => sub { $ENV{USER} },
);

has password => (
    is      => 'rwp',
    lazy    => 1,
    default => sub {''},
);

sub _build_session {
    my $self = shift;
    $self->_transport->open if !$self->_transport->isOpen;
    return $self->_client->OpenSession(
        Thrift::API::HiveClient2::TOpenSessionReq->new(
            {   username => $self->username,
                password => $self->password,
            }
        )
    );
}

has _session_handle => (
    is  => 'rwp',
    isa => sub {
        my($val) = @_;
        if (   !blessed( $val )
            || !$val->isa('Thrift::API::HiveClient2::TSessionHandle')
        ) {
            die sprintf "Session handle `%s` isn't a Thrift::API::HiveClient2::TSessionHandle",
                            $val // '[undefined]'
            ;
        }
    },
    lazy    => 1,
    builder => '_build_session_handle',
    predicate => '_has_session_handle',
);

sub _build_session_handle {
    my $self = shift;
    return $self->_session->{sessionHandle};
}

has _operation => (
    is  => "rwp",
    isa => sub {
        my($val) = @_;
        if ( defined $val
            && (
                !blessed( $val )
                || (   !$val->isa('Thrift::API::HiveClient2::TExecuteStatementResp')
                    && !$val->isa('Thrift::API::HiveClient2::TGetColumnsResp')
                    && !$val->isa('Thrift::API::HiveClient2::TGetTablesResp')
                    )
            )
        ) {
            die "Operation `%s` isn't a Thrift::API::HiveClient2::T*Resp",
                    $val // '[undefined]'
            ;
        }
    },
    lazy => 1,
);

has _operation_handle => (
    is  => 'rwp',
    isa => sub {
        my($val) = @_;
        if (
            defined $val
            && (   !blessed( $val )
                || !$val->isa('Thrift::API::HiveClient2::TOperationHandle')
                )
        ) {
            die sprintf "Operation handle isn't a Thrift::API::HiveClient2::TOperationHandle",
                            $val // '[undefined]'
            ;
        }
    },
    lazy => 1,
);

sub _cleanup_previous_operation {
    my $self = shift;

    # We seeem to have some memory leaks in the Hive server, let's try freeing the
    # operation handle explicitely
    if ( $self->_operation_handle ) {
        $self->_client->CloseOperation(
            Thrift::API::HiveClient2::TCloseOperationReq->new(
                { operationHandle => $self->_operation_handle, }
            )
        );
        $self->_set__operation(undef);
        $self->_set__operation_handle(undef);
    }
}

sub execute {
    my $self = shift;
    my ($query) = @_;    # make this a bit more flexible

    $self->_cleanup_previous_operation;

    my $rh = $self->_client->ExecuteStatement(
        Thrift::API::HiveClient2::TExecuteStatementReq->new(
            { sessionHandle => $self->_session_handle, statement => $query, confOverlay => {} }
        )
    );
    if ( $rh->{status}{errorCode} ) {
        die __PACKAGE__ . "::execute: $rh->{status}{errorMessage}; HQL was: \"$query\"";
    }
    $self->_set__operation($rh);
    $self->_set__operation_handle( $rh->{operationHandle} );
    return $rh;
}

{
    # cache the column names we need to extract from the bloated data structure
    # (keyed on query)
    my ( $column_keys, $column_names );

    sub fetch_hashref {
        my $self = shift;
        my ( $rh, $rows_at_a_time ) = @_;
        return $self->fetch( $rh, $rows_at_a_time, 1 );
    }

    sub fetch {
        my $self = shift;
        my ( $rh, $rows_at_a_time, $use_hashref ) = @_;

        # if $rh looks like a number, use it instead of $rows_at_a_time
        # it means we're using the new form for this call, which takes only the
        # number of wanted rows, or even nothing (and relies on the defaults,
        # and a cached copy of the query $rh)
        $rows_at_a_time = $rh if ( $rh && !$rows_at_a_time && $rh =~ /^[1-9][0-9]*$/ );
        $rows_at_a_time ||= 10_000;

        my $result = [];
        my $has_more_rows;

        # NOTE we don't use the provided $rh any more, maybe we should leave
        # that possibility open for parallel queries, but that would need a lot
        # more testing. Patches welcome.
        my $cached_rh = $self->_operation_handle;
        $rh = $self->_client->FetchResults(
            Thrift::API::HiveClient2::TFetchResultsReq->new(
                {   operationHandle => $cached_rh,
                    maxRows         => $rows_at_a_time,
                }
            )
        );
        if ( ref $rh eq 'Thrift::API::HiveClient2::TFetchResultsResp' ) {

            # NOTE that currently (july 2013) the hasMoreRows method is broken,
            # see the explanation in the POD
            $has_more_rows = $rh->hasMoreRows();

            for my $row ( @{ $rh->{results}{rows} || [] } ) {

                # Find which fields to extract from each row, only on the first iteration
                if ( !@{ $column_keys->{$cached_rh} || [] } ) {

                    # metadata for the query
                    if ($use_hashref) {
                        my $rh_meta = $self->_client->GetResultSetMetadata(
                            Thrift::API::HiveClient2::TGetResultSetMetadataReq->new(
                                { operationHandle => $cached_rh }
                            )
                        );
                        $column_names = [ map { $_->{columnName} }
                                @{ $rh_meta->{schema}{columns} || [] } ];
                    }

                    # TODO redo all this using the TGetResultSetMetadataResp object we retrieved
                    # above
                    for my $column ( @{ $row->{colVals} || [] } ) {

                        my $first_col = {%$column};

                        # Only 1 element of each TColumnValue is populated
                        # (although 7 keys are present, 1 for each possible data
                        # type) with a T*Value, and the rest is undef. Find out
                        # which is defined, and put the key (i.e. the data type) in
                        # cache, to reuse it to fetch the next rows faster.
                        # NOTE this data structure smells of Java and friends from
                        # miles away. Dynamically typed languages don't really need
                        # the bloat.
                        push @{ $column_keys->{$cached_rh} },
                            grep { ref $first_col->{$_} } keys %$first_col;
                    }
                }

                # TODO find something faster? (see comment above)

                my $idx    = 0;
                my $retval = [
                    map  { $_->value }
                    grep { defined $_ }
                    map  { $row->{colVals}[ $idx++ ]{$_} } @{ $column_keys->{$cached_rh} }
                ];
                if ($use_hashref) {
                    push @$result, { zip @$column_names, @$retval };
                }
                else {
                    push @$result, $retval;
                }
            }
        }
        return wantarray ? ( $result, $has_more_rows ) : ( @$result ? $result : undef );
    }
}

sub get_columns {
    my $self = shift;
    my ( $table, $schema ) = @_;

    # note that not specifying a table name would return all columns for all
    # tables we probably don't want that, but feel free to change this
    # behaviour. Same goes for the schema name: we probably want a default
    # value for the schema, which is what we use here.
    die "Unspecified table name" if !$table;
    $schema //= "default";

    $self->_cleanup_previous_operation;

    my $rh = $self->_client->GetColumns(
        Thrift::API::HiveClient2::TGetColumnsReq->new(
            {   sessionHandle => $self->_session_handle,
                catalogName   => undef,
                schemaName    => $schema,
                tableName     => $table,
                columnName    => undef,
                confOverlay   => {}
            }
        )
    );
    if ( $rh->{status}{errorCode} ) {
        die __PACKAGE__ . "::execute: $rh->{status}{errorMessage}";
    }
    $self->_set__operation($rh);
    $self->_set__operation_handle( $rh->{operationHandle} );
    my $columns;
    while ( my $res = $self->fetch($rh) ) {
        for my $line (@$res) {
            my $idx = 0;
            push @$columns, { map { $_ => $line->[ $idx++ ] } @odbc_coldesc_fields };
        }
    }
    return $columns;
}

sub get_tables {
    my $self = shift;
    my ( $schema, $table_pattern ) = @_;

    # note that not specifying a table name would return all columns for all
    # tables we probably don't want that, but feel free to change this
    # behaviour. Same goes for the schema name: we probably want a default
    # value for the schema, which is what we use here.
    $schema //= "default";

    $self->_cleanup_previous_operation;

    my $rh = $self->_client->GetTables(
        Thrift::API::HiveClient2::TGetTablesReq->new(
            {   sessionHandle => $self->_session_handle,
                catalogName   => undef,
                schemaName    => $schema,
                tableName     => $table_pattern,
                confOverlay   => {},
            }
        )
    );
    if ( $rh->{status}{errorCode} ) {
        die __PACKAGE__ . "::execute: $rh->{status}{errorMessage}";
    }
    $self->_set__operation($rh);
    $self->_set__operation_handle( $rh->{operationHandle} );
    my $tables;
    while ( my $res = $self->fetch($rh) ) {
        for my $line (@$res) {
            my $idx = 0;
            push @$tables, { map { $_ => $line->[ $idx++ ] } @tabledesc_fields };
        }
    }
    return $tables;
}

sub DEMOLISH {
    my $self = shift;

    $self->_cleanup_previous_operation;

    if ( $self->_has_session_handle ) {
        $self->_client->CloseSession(
            Thrift::API::HiveClient2::TCloseSessionReq->new(
                { sessionHandle => $self->_session_handle, }
            )
        );
    }

    if ( $self->_transport ) {
        $self->_transport->close;
    }
}

# when the user calls a method on an object of this class, see if that method
# exists on the TCLIService object. If so, create a sub that calls that method
# on the client object. If not, die horribly.
sub AUTOLOAD {
    my ($self) = @_;
    ( my $meth = our $AUTOLOAD ) =~ s/.*:://;
    return if $meth eq 'DESTROY';
    print STDERR "$meth\n";
    no strict 'refs';
    if ( $self->_client->can($meth) ) {
        *$AUTOLOAD = sub { shift->_client->$meth(@_) };
        goto &$AUTOLOAD;
    }
    croak "No such method exists: $AUTOLOAD";
}

1;

__END__