示例#1
0
int
main(int   argc,
     char *argv[])
{
    riak_args args;
    int operation = riak_parse_args(argc, argv, &args);

    event_enable_debug_mode();
//    event_use_pthreads();

    riak_context *ctx;
    riak_error err = riak_context_new_default(&ctx, args.host, args.portnum);
    if (err) {
        exit(1);
    }
    riak_event  *rev = NULL;
    riak_object *obj;
    riak_bucket_props *props;
    riak_put_options put_options;
    char output[10240];
    int it;

    riak_binary *bucket_bin = riak_binary_new_from_string(ctx, args.bucket); // Not copied
    riak_binary *key_bin    = riak_binary_new_from_string(ctx, args.key); // Not copied
    riak_binary *value_bin  = riak_binary_new_from_string(ctx, args.value); // Not copied
    if (bucket_bin == NULL ||
        key_bin    == NULL ||
        value_bin  == NULL) {
        fprintf(stderr, "Could not allocate bucket/key/value\n");
        exit(1);
    }

    for(it = 0; it < args.iterate; it++) {
        riak_log_context(ctx, RIAK_LOG_DEBUG, "Loop %d", it);

        if (args.async) {
            riak_event *rev;
            err = riak_event_new(ctx, &rev, NULL, NULL, NULL);
            if (err) {
                return err;
            }
            // For convenience have user callback know about its riak_event
            riak_event_set_cb_data(rev, rev);
        }
        switch (operation) {
        case MSG_RPBPINGREQ:
            if (args.async) {
                riak_event_set_response_cb(rev, (riak_response_callback)ping_cb);
                riak_encode_ping_request(rev, &(rev->pb_request));
            } else {
                err = riak_ping(ctx);
                if (err) {
                    fprintf(stderr, "No Ping [%s]\n", riak_strerror(err));
                    exit(1);
                }
            }
            break;
        case MSG_RPBGETSERVERINFOREQ:
            if (args.async) {
                riak_event_set_response_cb(rev, (riak_response_callback)serverinfo_cb);
                riak_encode_serverinfo_request(rev, &(rev->pb_request));
            } else {
                riak_serverinfo_response *serverinfo_response;
                err = riak_serverinfo(ctx, &serverinfo_response);
                if (err) {
                    fprintf(stderr, "Server Info Problems [%s]\n", riak_strerror(err));
                    exit(1);
                }
                riak_print_serverinfo_response(serverinfo_response, output, sizeof(output));
                printf("%s\n", output);
                riak_free_serverinfo_response(ctx, &serverinfo_response);
            }
            break;
        case MSG_RPBGETREQ:
            if (args.async) {
                riak_event_set_response_cb(rev, (riak_response_callback)get_cb);
                riak_encode_get_request(rev, bucket_bin, key_bin, NULL, &(rev->pb_request));
            } else {
                riak_get_response *get_response;
                err = riak_get(ctx, bucket_bin, key_bin, NULL, &get_response);
                if (err) {
                    fprintf(stderr, "Get Problems [%s]\n", riak_strerror(err));
                    exit(1);
                }
                riak_print_get_response(get_response, output, sizeof(output));
                printf("%s\n", output);
                riak_free_get_response(ctx, &get_response);
            }
            break;
        case MSG_RPBPUTREQ:
            obj = riak_object_new(ctx);
            if (obj == NULL) {
                riak_log(rev, RIAK_LOG_FATAL, "Could not allocate a Riak Object");
                return 1;
            }
            riak_object_set_bucket(obj, riak_binary_new_from_string(ctx, args.bucket)); // Not copied
            riak_object_set_value(obj, riak_binary_new_from_string(ctx, args.value)); // Not copied
            if (riak_object_get_bucket(obj) == NULL ||
                riak_object_get_value(obj) == NULL) {
                fprintf(stderr, "Could not allocate bucket/value\n");
                riak_free(ctx, &obj);
                exit(1);
            }
            memset(&put_options, '\0', sizeof(riak_put_options));
            put_options.has_return_head = RIAK_TRUE;
            put_options.return_head = RIAK_TRUE;
            put_options.has_return_body = RIAK_TRUE;
            put_options.return_body = RIAK_TRUE;
            if (args.async) {
                riak_event_set_response_cb(rev, (riak_response_callback)put_cb);
                riak_encode_put_request(rev, obj, &put_options, &(rev->pb_request));
            } else {
                riak_put_response *put_response;
                err = riak_put(ctx, obj, &put_options, &put_response);
                if (err) {
                    fprintf(stderr, "Put Problems [%s]\n", riak_strerror(err));
                    exit(1);
                }
                riak_print_put_response(put_response, output, sizeof(output));
                printf("%s\n", output);
                riak_free_put_response(ctx, &put_response);
            }
            riak_object_free(ctx, &obj);
            break;
        case MSG_RPBDELREQ:
            if (args.async) {
                riak_event_set_response_cb(rev, (riak_response_callback)delete_cb);
                riak_encode_delete_request(rev, bucket_bin, key_bin, NULL, &(rev->pb_request));
            } else {
                err = riak_delete(ctx, bucket_bin, key_bin, NULL);
                if (err) {
                    fprintf(stderr, "Delete Problems [%s]\n", riak_strerror(err));
                    exit(1);
                }
            }
            break;
        case MSG_RPBLISTBUCKETSREQ:
            if (args.async) {
                riak_event_set_response_cb(rev, (riak_response_callback)listbucket_cb);
                riak_encode_listbuckets_request(rev, &(rev->pb_request));
            } else {
                riak_listbuckets_response *bucket_response;
                err = riak_listbuckets(ctx, &bucket_response);
                if (err) {
                    fprintf(stderr, "List Buckets Problems [%s]\n", riak_strerror(err));
                    exit(1);
                }
                riak_print_listbuckets_response(bucket_response, output, sizeof(output));
                riak_log_context(ctx, RIAK_LOG_DEBUG, "%s", output);
                riak_free_listbuckets_response(ctx, &bucket_response);
            }
            break;
        case MSG_RPBLISTKEYSREQ:
            if (args.async) {
                riak_event_set_response_cb(rev, (riak_response_callback)listkey_cb);
                riak_encode_listkeys_request(rev, bucket_bin, args.timeout * 1000, &(rev->pb_request));
            } else {
                riak_listkeys_response *key_response;
                err = riak_listkeys(ctx, bucket_bin, args.timeout * 1000, &key_response);
                if (err) {
                    fprintf(stderr, "List Keys Problems [%s]\n", riak_strerror(err));
                    exit(1);
                }
                riak_print_listkeys_response(key_response, output, sizeof(output));
                riak_log_context(ctx, RIAK_LOG_DEBUG, "%s", output);
                riak_free_listkeys_response(ctx, &key_response);
            }
            break;
        case MSG_RPBGETCLIENTIDREQ:
            if (args.async) {
                riak_event_set_response_cb(rev, (riak_response_callback)getclientid_cb);
                riak_encode_get_clientid_request(rev, &(rev->pb_request));
            } else {
                riak_get_clientid_response *getcli_response;
                err = riak_get_clientid(ctx, &getcli_response);
                if (err) {
                    fprintf(stderr, "Get ClientID Problems [%s]\n", riak_strerror(err));
                    exit(1);
                }
                riak_print_get_clientid_response(getcli_response, output, sizeof(output));
                printf("%s\n", output);
                riak_free_get_clientid_response(ctx, &getcli_response);
            }
            break;
        case MSG_RPBSETCLIENTIDREQ:
            if (args.async) {
                riak_event_set_response_cb(rev, (riak_response_callback)setclientid_cb);
                riak_encode_set_clientid_request(rev, value_bin, &(rev->pb_request));
            } else {
                riak_set_clientid_response *setcli_response;
                err = riak_set_clientid(ctx, value_bin, &setcli_response);
                if (err) {
                    fprintf(stderr, "Set ClientID Problems [%s]\n", riak_strerror(err));
                    exit(1);
                }
                riak_free_set_clientid_response(ctx, &setcli_response);
            }
            break;
        case MSG_RPBGETBUCKETREQ:
            if (args.async) {
                riak_event_set_response_cb(rev, (riak_response_callback)getbucketprops_cb);
                riak_encode_get_bucketprops_request(rev, bucket_bin, &(rev->pb_request));
            } else {
                riak_get_bucketprops_response *bucket_response;
                err = riak_get_bucketprops(ctx, bucket_bin, &bucket_response);
                if (err) {
                    fprintf(stderr, "Get Bucket Properties Problems [%s]\n", riak_strerror(err));
                    exit(1);
                }
                riak_print_get_bucketprops_response(bucket_response, output, sizeof(output));
                riak_log_context(ctx, RIAK_LOG_DEBUG, "%s", output);
                riak_free_get_bucketprops_response(ctx, &bucket_response);
            }
            break;
        case MSG_RPBRESETBUCKETREQ:
            if (args.async) {
                riak_event_set_response_cb(rev, (riak_response_callback)resetbucketprops_cb);
                riak_encode_reset_bucketprops_request(rev, bucket_bin, &(rev->pb_request));
            } else {
                riak_reset_bucketprops_response *bucket_response;
                 err = riak_reset_bucketprops(ctx, bucket_bin, &bucket_response);
                if (err) {
                    fprintf(stderr, "Reset Bucket Properties Problems [%s]\n", riak_strerror(err));
                    exit(1);
                }
            }
            break;
        case MSG_RPBSETBUCKETREQ:
            props = riak_bucket_props_new(ctx);
            if (obj == NULL) {
                riak_log(rev, RIAK_LOG_FATAL, "Could not allocate a Riak Bucket Properties");
                return 1;
            }
            riak_bucket_props_set_last_write_wins(props, RIAK_FALSE);
            if (args.async) {
                riak_event_set_response_cb(rev, (riak_response_callback)setbucketprops_cb);
                riak_encode_set_bucketprops_request(rev, bucket_bin, props, &(rev->pb_request));
            } else {
                riak_set_bucketprops_response *bucket_response;
                 err = riak_set_bucketprops(ctx, bucket_bin, props, &bucket_response);
                if (err) {
                    fprintf(stderr, "Reset Bucket Properties Problems [%s]\n", riak_strerror(err));
                    exit(1);
                }
            }
            break;
        default:
            usage(stderr, argv[0]);
        }

        if (args.async) {
            err = riak_send_req(rev, rev->pb_request);
            if (err) {
                riak_log(rev, RIAK_LOG_FATAL, "Could not send request");
                exit(1);
            }
        }
    }
    // What has been queued up
    fflush(stdout);

    if (args.async) {
        // Terminates only on error or timeout
        riak_event_loop(ctx);
    }

    riak_free(ctx, &bucket_bin);
    riak_free(ctx, &key_bin);
    riak_free(ctx, &value_bin);
    riak_context_free(&ctx);

    return 0;
}
示例#2
0
// main function that retrieves data from riak kvs 
static int namak_handler(request_rec *r)
{
    // check handler name
    if (strcmp(r->handler, "namak")) {
        return DECLINED;
    }
    // if method is not GET do not process
    if (r->method_number != M_GET)
        return DECLINED;
    // a riak_config serves as your per-thread state to interact with Riak
    riak_config *cfg;
    // use the default configuration
    riak_error err = riak_config_new_default(&cfg);
    if (err) {
        ap_log_error(APLOG_MARK, APLOG_NOERRNO|APLOG_NOTICE, 0, r->server, "failed to initailize default config");
        return DECLINED;
    }
    // make a table to store GET requests
    apr_table_t *params = apr_table_make(r->pool, 5);
    // debug code
    // ap_log_error(APLOG_MARK, APLOG_NOERRNO|APLOG_NOTICE, 0, r->server, r->args);
    // parse request parameters
    if (r->args) {
        char *args = apr_pstrdup(r->pool, r->args);
        if (args == NULL) {
            ap_log_error(APLOG_MARK, APLOG_NOERRNO|APLOG_NOTICE, 0, r->server, "Can not allocate memory pool");
            return HTTP_BAD_REQUEST;
        }
        char *tok, *val;
        while (args && *args) {
            if ((val = ap_strchr(args, '='))) {
                *val++ = '\0';
                if ((tok = ap_strchr(val, '&')))
                    *tok++ = '\0';
                apr_table_setn(params, args, val);
                args = tok;
            }
            else
                return HTTP_BAD_REQUEST;
        }
    }
    // get params to send them to Riak
    const char *bucket = apr_table_get(params, "bucket");
    const char *key = apr_table_get(params, "key");
    if (bucket == NULL || key == NULL)
        return DECLINED;
    riak_binary *bucket_type_bin = riak_binary_copy_from_string(cfg, "default"); // only for Riak 2.0?
    riak_binary *bucket_bin   = riak_binary_copy_from_string(cfg, bucket);
    riak_binary *key_bin      = riak_binary_copy_from_string(cfg, key);
    // check for memory allocation problems
    if (bucket_bin == NULL || key_bin    == NULL) {
        riak_config_free(&cfg);
        return DECLINED;
    }
    // Supporting Options and outputs
    riak_get_options *get_options;
    // Every possible message response type
    riak_get_response *get_response = NULL;
    riak_connection  *cxn   = NULL;
    // get Riak host and port from httpd.conf
    namak_svr_cfg* svr = ap_get_module_config(r->per_dir_config, &namak_module);
    if (svr->riak_server == NULL || svr->riak_port == NULL)
    {
        riak_binary_free(cfg, &bucket_type_bin);
        riak_binary_free(cfg, &bucket_bin);
        riak_binary_free(cfg, &key_bin);
        riak_config_free(&cfg);
        return DECLINED;
    }
    // Create a connection with the default address resolver
    err = riak_connection_new(cfg, &cxn, svr->riak_server, svr->riak_port, NULL);
    if (err) {
        ap_log_error(APLOG_MARK, APLOG_NOERRNO|APLOG_NOTICE, 0, r->server,
                "failed to connect Riak(%s:%s)", svr->riak_server, svr->riak_port);
        riak_binary_free(cfg, &bucket_type_bin);
        riak_binary_free(cfg, &bucket_bin);
        riak_binary_free(cfg, &key_bin);
        riak_config_free(&cfg);
        return HTTP_SERVICE_UNAVAILABLE;
    }
    // handle possible operations from the command line
    get_options = riak_get_options_new(cfg);
    if (get_options == NULL) {
        ap_log_error(APLOG_MARK, APLOG_NOERRNO|APLOG_NOTICE, 0, r->server, "failed to initailize option config");
        riak_binary_free(cfg, &bucket_type_bin);
        riak_binary_free(cfg, &bucket_bin);
        riak_binary_free(cfg, &key_bin);
        riak_config_free(&cfg);
        return DECLINED;
    }
    riak_get_options_set_basic_quorum(get_options, RIAK_TRUE);
    riak_get_options_set_r(get_options, 1);
    // get data from Riak kvs
    err = riak_get(cxn, bucket_type_bin, bucket_bin, key_bin, get_options, &get_response);
    if (err == ERIAK_OK) {
        // if response does not have at least one object, just return 404
        if (riak_get_get_n_content(get_response) < 1)
        {
            riak_binary_free(cfg, &bucket_type_bin);
            riak_binary_free(cfg, &bucket_bin);
            riak_binary_free(cfg, &key_bin);
            riak_config_free(&cfg);
            return HTTP_NOT_FOUND;
        }
        // get the first object
        riak_object **objects = riak_get_get_content(get_response);
        int i = 0;
        for(i = 0; i < riak_get_get_n_content(get_response); i++) {
            riak_object *obj = objects[i];
            r->content_type = riak_binary_data(riak_object_get_content_type(obj));
            ap_rprintf(r,"%s\n",riak_binary_data(riak_object_get_value(obj)));
            break;
        }
    }
    riak_get_response_free(cfg, &get_response);
    riak_get_options_free(cfg, &get_options);
    if (err) {
        riak_binary_free(cfg, &bucket_type_bin);
        riak_binary_free(cfg, &bucket_bin);
        riak_binary_free(cfg, &key_bin);
        riak_config_free(&cfg);
        return HTTP_SERVICE_UNAVAILABLE;
    }
    // cleanup
    riak_binary_free(cfg, &bucket_type_bin);
    riak_binary_free(cfg, &bucket_bin);
    riak_binary_free(cfg, &key_bin);
    riak_config_free(&cfg);
    return OK;
}