File Coverage

blib/lib/Riak/Client.pm

Criterion	Covered	Total	%
statement	47	260	18.0
branch	1	98	1.0
condition	0	37	0.0
subroutine	16	46	34.7
pod	16	17	94.1
total	80	458	17.4

#
# This file is part of Riak-Client
#
# This software is copyright (c) 2014 by Damien Krotkine.
#
# This is free software; you can redistribute it and/or modify it under
# the same terms as the Perl 5 programming language system itself.
#
package Riak::Client;
{
  $Riak::Client::VERSION = '1.95';
}

use 5.010;
use Riak::Client::PBC;
use Type::Params qw(compile);
use Types::Standard -types;
use Errno qw(EINTR);
use Scalar::Util qw(blessed);
use JSON::XS;
use Carp;
$Carp::Internal{ (__PACKAGE__) }++;
use Module::Runtime qw(use_module);
require bytes;
use Moo;

use IO::Socket::INET;
use IO::Socket::Timeout;

use Scalar::Util qw(weaken);

use constant {
    # error
    ERROR_RESPONSE_CODE            => 0,
    # ping
    PING_REQUEST_CODE              => 1,
    PING_RESPONSE_CODE             => 2,
    # get, get_raw
    GET_REQUEST_CODE               => 9,
    GET_RESPONSE_CODE              => 10,
    # put, put_raw
    PUT_REQUEST_CODE               => 11,
    PUT_RESPONSE_CODE              => 12,
    # del
    DEL_REQUEST_CODE               => 13,
    DEL_RESPONSE_CODE              => 14,
    # get_buckets
    GET_BUCKETS_REQUEST_CODE       => 15,
    GET_BUCKETS_RESPONSE_CODE      => 16,
    # get_keys
    GET_KEYS_REQUEST_CODE          => 17,
    GET_KEYS_RESPONSE_CODE         => 18,
    # get_bucket_props
    GET_BUCKET_PROPS_REQUEST_CODE  => 19,
    GET_BUCKET_PROPS_RESPONSE_CODE => 20,
    # set_bucket_props
    SET_BUCKET_PROPS_REQUEST_CODE  => 21,
    SET_BUCKET_PROPS_RESPONSE_CODE => 22,
    # map_reducd
    MAP_REDUCE_REQUEST_CODE        => 23,
    MAP_REDUCE_RESPONSE_CODE       => 24,
    # query_index
    QUERY_INDEX_REQUEST_CODE       => 25,
    QUERY_INDEX_RESPONSE_CODE      => 26,
};


# ABSTRACT: Fast and lightweight Perl client for Riak


has host    => ( is => 'ro', isa => Str,  required => 1 );
has port    => ( is => 'ro', isa => Int,  required => 1 );
has r       => ( is => 'ro', isa => Int,  default  => sub {2} );
has w       => ( is => 'ro', isa => Int,  default  => sub {2} );
has dw      => ( is => 'ro', isa => Int,  default  => sub {1} );
has connection_timeout => ( is => 'ro',                 isa => Num,  default  => sub {5} );
has read_timeout       => ( is => 'ro', predicate => 1, isa => Num,  default  => sub {5} );
has write_timeout      => ( is => 'ro', predicate => 1, isa => Num,  default  => sub {5} );
has no_delay           => ( is => 'ro',                 isa => Bool, default  => sub {0} );


has no_auto_connect => ( is => 'ro', isa => Bool,  default  => sub {0} );

has _on_connect_cb => ( is => 'rw' );

has _requests_lock => ( is => 'rw', default => sub { undef });

has _socket => ( is => 'ro', lazy => 1, builder => 1 );
sub _build__socket {
    my ($self) = @_;

    my $host = $self->host;
    my $port = $self->port;

    my $socket = IO::Socket::INET->new(
        PeerHost => $host,
        PeerPort => $port,
        Timeout  => $self->connection_timeout,
    );

    croak "Error ($!), can't connect to $host:$port"
      unless defined $socket;

    $self->has_read_timeout || $self->has_write_timeout
      or return $socket;

    # enable read and write timeouts on the socket
    IO::Socket::Timeout->enable_timeouts_on($socket);
    # setup the timeouts
    $self->has_read_timeout
      and $socket->read_timeout($self->read_timeout);
    $self->has_write_timeout
      and $socket->write_timeout($self->write_timeout);

    use Socket qw(IPPROTO_TCP TCP_NODELAY);
    $self->no_delay
      and $socket->setsockopt(IPPROTO_TCP, TCP_NODELAY, 1);

    return $socket;
}

sub BUILD {
    my ($self) = @_;
    $self->no_auto_connect
      and return;

    $self->connect();
}


sub connect {
    state $check = compile(Any, Optional[CodeRef]);
    my ( $self, $cb ) = $check->(@_);

    # that will perform connection
    $self->_socket();
    if ($cb) {
        $cb->();
        return;
    } else {
        return 1;
    }

}

has _getkeys_accumulator => (is => 'rw', init_arg => undef);
has _mapreduce_accumulator => (is => 'rw', init_arg => undef);


sub ping {
    state $check = compile(Any, Optional[CodeRef]);
    my ( $self, $cb ) = $check->(@_);
    $_[0]->_parse_response( {
        request_code   => PING_REQUEST_CODE,
        expected_code  => PING_RESPONSE_CODE,
        operation_name => 'ping',
        body_ref       => \'',
        cb             => $cb,
    } );
}


sub is_alive {
    state $check = compile(Any, Optional[CodeRef]);
    my ( $self, $cb ) = $check->(@_);
    my $res = eval { $self->ping; 1 };
    $cb and return $cb->($res);
    return $res;
}


sub get {
    state $check = compile(Any, Str, Str, Optional[CodeRef]);
    my ( $self, $bucket, $key, $cb ) = $check->(@_);
    $self->_fetch( $bucket, $key, 1, 0, $cb );
}


sub get_raw {
    state $check = compile(Any, Str, Str, Optional[CodeRef]);
    my ( $self, $bucket, $key, $cb ) = $check->(@_);
    $self->_fetch( $bucket, $key, 0, 0, $cb );
}


#my $LinksStructure = declare as ArrayRef[Dict[bucket => Str, key => Str, tag => Str]];
#coerce $LinksStructure, from HashRef[] Num, q{ int($_) };

sub put {
    my $cb = ref $_[-1] eq 'CODE' ? pop : undef;
    state $check = compile(Any, Str, Str, Any, Optional[Str],
                           Optional[HashRef[Str]], # indexes
                           Optional[ArrayRef[Dict[bucket => Str, key => Str, tag => Str]]], # links
                          );
    my ( $self, $bucket, $key, $value, $content_type, $indexes, $links ) = $check->(@_);

    ($content_type //= 'application/json')
      eq 'application/json'
        and $value = encode_json($value);

    $self->_store( $bucket, $key, $value, $content_type, $indexes, $links, $cb);
}



sub put_raw {
    my $cb = ref $_[-1] eq 'CODE' ? pop : undef;
    state $check = compile(Any, Str, Str, Any, Optional[Str],
                           Optional[HashRef[Str]], # indexes
                           Optional[ArrayRef[Dict[bucket => Str, key => Str, tag => Str]]], # links
                          );
    my ( $self, $bucket, $key, $value, $content_type, $indexes, $links ) = $check->(@_);

    $content_type ||= 'plain/text';
    $self->_store( $bucket, $key, $value, $content_type, $indexes, $links, $cb);
}


sub del {
    state $check = compile(Any, Str, Str, Optional[CodeRef]);
    my ( $self, $bucket, $key, $cb ) = $check->(@_);

    my $body = RpbDelReq->encode(
        {   key    => $key,
            bucket => $bucket,
            rw     => $self->dw
        }
    );

    $self->_parse_response( {
        request_code   => DEL_REQUEST_CODE,
        expected_code  => DEL_RESPONSE_CODE,
        operation_name => 'del',
        key            => $key,
        bucket         => $bucket,
        body_ref       => \$body,
        cb             => $cb,
    } );
}


sub get_keys {
    state $check = compile(Any, Str, Optional[CodeRef]);
    my ( $self, $bucket, $cb ) = $check->(@_);

    # reset accumulator
    $self->_getkeys_accumulator([]);
    my $body = RpbListKeysReq->encode( { bucket => $bucket } );
    $self->_parse_response( {
        request_code   => GET_KEYS_REQUEST_CODE,
        expected_code  => GET_KEYS_RESPONSE_CODE,
        operation_name => 'get_keys',
        key            => "*",
        bucket         => $bucket,
        body_ref       => \$body,
        cb             => $cb,
        handle_response => \&_handle_get_keys_response,
        lock_requests => 1,
    } );
}

sub _handle_get_keys_response {
    my ( $self, $encoded_message, $args ) = @_;

    # TODO: support for 1.4 (which provides 'stream', 'return_terms', and 'stream')
    my $obj = RpbListKeysResp->decode( $encoded_message );
    my @keys = @{$obj->keys // []};

    # case 1 : no user callback
    my $cb = $args->{cb};
    if (! $cb ) {
        # accumulate results
        push @{$self->_getkeys_accumulator}, @keys;

        # if more to come, return by saying so
        $obj->done
          or return (undef, 1);

        # all results are there, return the whole
        my $keys = $self->_getkeys_accumulator;
        $self->_getkeys_accumulator([]);
        return \$keys;
    }

    # case 2 : we have a user callback
    my $last_key;
    my $obj_done = $obj->done
      and $last_key = pop @keys;

    # no second arg = more to come
    $cb->($_) foreach @keys;

    # if more to come, return by saying so
    $obj->done
      or return (undef, 1);

    # process last keys if any
    defined $last_key and $cb->($last_key, 1);

    # means: nothing left to do, all results processed through callback
    return;
}


sub exists {
    state $check = compile(Any, Str, Str, Optional[CodeRef]);
    my ( $self, $bucket, $key, $cb ) = $check->(@_);
    $self->_fetch( $bucket, $key, 0, 1, $cb );
}

sub _fetch {
    my ( $self, $bucket, $key, $decode, $test_exist, $cb ) = @_;

    my $body = RpbGetReq->encode(
        {   r      => $self->r,
            key    => $key,
            bucket => $bucket,
            head   => $test_exist
        }
    );

    $self->_parse_response( {
        request_code  => GET_REQUEST_CODE,
        expected_code => GET_RESPONSE_CODE,
        operation_name => 'get',
        key       => $key,
        bucket    => $bucket,
        body_ref  => \$body,
        decode    => $decode,
        handle_response => \&_handle_get_response,
        test_exist => $test_exist,
        cb => $cb,
        cb_args => 1,
    } );
}

sub _handle_get_response {
    my ( $self, $encoded_message, $args ) = @_;

    defined $encoded_message
      or return _die_generic_error( "Undefined Message", 'get', $args );

    my $decoded_message = RpbGetResp->decode($encoded_message);
    my $content = $decoded_message->content;

    # empty content
    ref $content eq 'ARRAY'
      or return \undef;

    # if we just need to test existence
    $args->{test_exist}
      and return \1;

    # TODO: handle metadata
    my $value        = $content->[0]->value;
    my $content_type = $content->[0]->content_type;

    # if we need to decode
    $args->{decode} && ($content_type // '') eq 'application/json'
      and return \decode_json($value);

    # simply return the value
    return \$value;
}

sub _store {
    my ( $self, $bucket, $key, $encoded_value, $content_type, $indexes, $links, $cb ) = @_;

    my $body = RpbPutReq->encode(
        {   key     => $key,
            bucket  => $bucket,
            content => {
                value        => $encoded_value,
                content_type => $content_type,
                ( $indexes ?
                  ( indexes => [
                                map {
                                    { key => $_ , value => $indexes->{$_} }
                                } keys %$indexes
                               ])
                  : ()
                ),
                ( $links ? ( links => $links) : () ),
            },
        }
    );

    $self->_parse_response( {
        request_code   => PUT_REQUEST_CODE,
        expected_code  => PUT_RESPONSE_CODE,
        operation_name => 'put',
        key            => $key,
        bucket         => $bucket,
        body_ref       => \$body,
        cb             => $cb,
    } );
}


sub query_index {
    state $check = compile(Any, Str, Str, Str|ArrayRef, Optional[CodeRef]);
    my ( $self, $bucket, $index, $value_to_match, $cb ) = $check->(@_);

    my $query_type_is_eq = 0; # eq
    ref $value_to_match
      and $query_type_is_eq = 1; # range
    my $body = RpbIndexReq->encode(
        {   index    => $index,
            bucket   => $bucket,
            qtype    => $query_type_is_eq,
            $query_type_is_eq ?
            ( range_min => $value_to_match->[0],
              range_max => $value_to_match->[1] )
            : (key => $value_to_match ),
        }
    );

    $self->_parse_response( {
        request_code   => QUERY_INDEX_REQUEST_CODE,
        expected_code  => QUERY_INDEX_RESPONSE_CODE,
        operation_name => 'query_index',
        $query_type_is_eq ?
          (key => '2i query on ' . join('...', @$value_to_match) )
        : (key => $value_to_match ),
        bucket    => $bucket,
        body_ref  => \$body,
        handle_response => \&_handle_query_index_response,
        cb => $cb,
        lock_requests => 1,
    } );
}

sub _handle_query_index_response {
    my ( $self, $encoded_message, $args ) = @_;

    my $obj = RpbIndexResp->decode( $encoded_message );
    my @keys = @{$obj->keys // []};

    # case 1 : no user callback
    my $cb = $args->{cb}
      or return \\@keys;

    # case 2 : we have a user callback
    $cb->($_) foreach @keys;

    # means: nothing left to do, all results processed through callback
    return;

}


sub get_buckets {
    state $check = compile(Any, Optional[CodeRef]);
    my ( $self, $cb ) = $check->(@_);

    $self->_parse_response( {
        request_code    => GET_BUCKETS_REQUEST_CODE,
        expected_code   => GET_BUCKETS_RESPONSE_CODE,
        operation_name  => 'get_buckets',
        handle_response => \&_handle_get_buckets_response,
        cb              => $cb,
    } );
}

sub _handle_get_buckets_response {
    my ( $self, $encoded_message, $args ) = @_;
    my $obj = RpbListBucketsResp->decode( $encoded_message );
    return \($obj->buckets // []);
}


sub get_bucket_props {
    state $check = compile(Any, Str, Optional[CodeRef]);
    my ( $self, $bucket, $cb ) = $check->(@_);

    my $body = RpbGetBucketReq->encode( { bucket => $bucket } );
    $self->_parse_response( {
        request_code    => GET_BUCKET_PROPS_REQUEST_CODE,
        expected_code   => GET_BUCKET_PROPS_RESPONSE_CODE,
        bucket          => $bucket,
        body_ref        => \$body,
        handle_response => \&_handle_get_bucket_props_response,
        cb              => $cb,
    } );
}

sub _handle_get_bucket_props_response {
    my ( $self, $encoded_message, $args ) = @_;

    my $obj = RpbListBucketsResp->decode( $encoded_message );
    my $props = RpbBucketProps->decode($obj->buckets->[0]);
    return \{ %$props }; # unblessing variable
}


sub set_bucket_props {
    state $check = compile( Any, Str,
                            Dict[ n_val => Optional[Int],
                                  allow_mult => Optional[Bool] ],
                            Optional[CodeRef] );
    my ( $self, $bucket, $props, $cb ) = $check->(@_);
    $props->{n_val} && $props->{n_val} < 0 and croak 'n_val should be possitive integer';

    my $body = RpbSetBucketReq->encode({ bucket => $bucket, props => $props });
    $self->_parse_response( {
        request_code   => SET_BUCKET_PROPS_REQUEST_CODE,
        expected_code  => SET_BUCKET_PROPS_RESPONSE_CODE,
        bucket         => $bucket,
        body_ref       => \$body,
    } );
}


sub map_reduce {
  state $check = compile(Any, Any, Optional[CodeRef]);
  my ( $self, $request, $cb) = $check->(@_);

  my @args;

  push @args, ref($request) ? encode_json($request): $request;
  push @args, 'application/json';
  push @args, $cb if $cb;

  map_reduce_raw($self, @args);

}


sub map_reduce_raw {
  state $check = compile(Any, Str, Str, Optional[CodeRef]);
  my ( $self, $request, $content_type, $cb) = $check->(@_);

  my $body = RpbMapRedReq->encode(
    {
      request => $request,
      content_type => $content_type,
    }
  );

  # reset accumulator
  $self->_mapreduce_accumulator([]);

  $self->_parse_response( {
      request_code   => MAP_REDUCE_REQUEST_CODE,
      expected_code  => MAP_REDUCE_RESPONSE_CODE,
      operation => 'map_reduce',
      body_ref  => \$body,
      cb        => $cb,
      decode    => ($content_type eq 'application/json'),
      handle_response => \&_handle_map_reduce_response,
      lock_requests => 1,
  } );
}

sub _handle_map_reduce_response {
    my ( $self, $encoded_message, $args ) = @_;
    my $obj = RpbMapRedResp->decode( $encoded_message );

    # case 1 : no user callback
    my $cb = $args->{cb};
    if (! $cb ) {

        # all results were there, reset the accumulator and return the whole,
        if ($obj->done) {
            my $results = $self->_mapreduce_accumulator();
            $self->_mapreduce_accumulator([]);
            return \$results;
        }

        # accumulate results
        push @{$self->_mapreduce_accumulator},
          { phase => $obj->phase, response => ($args->{decode}) ? decode_json($obj->response // '[]') : $obj->response };

        # more stuff to come, say so
        return (undef, 1);

    }

    # case 2 : we have a user callback

    # means: nothing left to do, all results processed through callback
    $obj->done
      and return;

    $cb->($obj->response, $obj->phase, $obj->done);

    # more stuff to come, say so
    return (undef, 1);

}

sub _parse_response {
    my ( $self, $args ) = @_;

    my $socket = $self->_socket;
    _send_bytes($socket, $args->{request_code}, $args->{body_ref} // \'');

    while (1) {
        my $response;
        # get and check response
        my $raw_response_ref = _read_response($socket)
          or return _die_generic_error( $! || "Socket Closed", $args);

        my ( $response_code, $response_body ) = unpack( 'c a*', $$raw_response_ref );

        # in case of error msg
        if ($response_code == ERROR_RESPONSE_CODE) {
            my $decoded_message = RpbErrorResp->decode($response_body);
            my $errmsg  = $decoded_message->errmsg;
            my $errcode = $decoded_message->errcode;

            return _die_generic_error( "Riak Error (code: $errcode) '$errmsg'", $args);
        }


        # check if we have what we want
        $response_code != $args->{expected_code}
          and return _die_generic_error(
              "Unexpected Response Code in (got: $response_code, expected: $args->{expected_code})",
              $args );

        # default value if we don't need to handle the response.
        my ($ret, $more_to_come) = ( \1, undef);

        # handle the response.
        if (my $handle_response = $args->{handle_response}) {
            ($ret, $more_to_come) = $handle_response->( $self, $response_body, $args);
        }

        # it's a multiple response request, loop again
        $more_to_come
          and next;

        # there is a result, process or return it
        if ($ret) {
            $args->{cb} and return $args->{cb}->($$ret);
            return $$ret;
        }

        # ret was undef, means we have processed everything in the callback
        return;

    }
}

sub _die_generic_error {
    my ( $error, $args ) = @_;

    my ($operation_name, $bucket, $key) =
      map { $args->{$_} // "" } ( qw( operation_name bucket key) );

    my $extra = '';
    defined $bucket && defined $key
      and $extra = "(bucket: $bucket, key: $key) ";

    my $msg = "Error in '$operation_name' $extra: $error";
    if ( my $cb = $args->{cb} ) {
        $cb->((undef) x ($args->{cb_nb_args} // 0), $msg);
        return;
    } elsif (my $cv = $args->{cv}) {
        $cv->croak($msg);
    } else {
        croak $msg;
    }
    return;
}

sub _read_response {
    my ($socket) = @_;
    _read_bytes($socket, unpack( 'N', ${ _read_bytes($socket, 4) // return } ));
}

sub _read_bytes {
    my ( $socket, $length ) = @_;

    my $buffer;
    my $offset = 0;
    my $read = 0;

    while ($length > 0) {
        $read = $socket->sysread( $buffer, $length, $offset );
        if (! defined $read) {
            $! == EINTR
              and next;
            return;
        }

        $read > 0
          or return;

        $offset += $read;
        $length -= $read;
    }

    return \$buffer;
}


sub _send_bytes {
    my ( $socket, $request_code, $body_ref ) = @_;

    my $bytes = pack('N', my $length = (bytes::length($$body_ref) + 1)) . pack('c', $request_code) . $$body_ref;

    $length += 4;
    my $offset = 0;
    my $sent = 0;

    while ($length > 0) {
        $sent = $socket->syswrite( $bytes, $length, $offset );
        if (! defined $sent) {
            $! == EINTR
              and next;
            return;
        }

        $sent > 0
          or return;

        $offset += $sent;
        $length -= $sent;
    }

    return $offset;
}



1;

__END__