com.basho.riak.client.http.util.ClientHelper.java Source code

Java tutorial

Introduction

Here is the source code for com.basho.riak.client.http.util.ClientHelper.java

Source

/*
 * This file is provided to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 * 
 * http://www.apache.org/licenses/LICENSE-2.0
 * 
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
 * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
 * License for the specific language governing permissions and limitations under
 * the License.
 */
package com.basho.riak.client.http.util;

import static com.basho.riak.client.util.CharsetUtils.utf8StringToBytes;
import java.io.IOException;
import java.io.InputStream;
import java.net.URI;
import java.net.URISyntaxException;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;

import org.apache.commons.codec.binary.Base64;
import org.apache.http.Header;
import org.apache.http.HttpEntity;
import org.apache.http.NameValuePair;
import org.apache.http.client.HttpClient;
import org.apache.http.client.methods.HttpDelete;
import org.apache.http.client.methods.HttpGet;
import org.apache.http.client.methods.HttpHead;
import org.apache.http.client.methods.HttpPost;
import org.apache.http.client.methods.HttpPut;
import org.apache.http.client.methods.HttpRequestBase;
import org.apache.http.client.utils.URIBuilder;
import org.apache.http.client.utils.URLEncodedUtils;
import org.apache.http.entity.ByteArrayEntity;
import org.apache.http.entity.ContentType;
import org.apache.http.entity.StringEntity;
import org.apache.http.message.BasicNameValuePair;
import org.apache.http.util.EntityUtils;
import org.json.JSONObject;

import com.basho.riak.client.http.RiakClient;
import com.basho.riak.client.http.RiakConfig;
import com.basho.riak.client.http.RiakObject;
import com.basho.riak.client.http.request.RequestMeta;
import com.basho.riak.client.http.response.BucketResponse;
import com.basho.riak.client.http.response.DefaultHttpResponse;
import com.basho.riak.client.http.response.HttpResponse;
import com.basho.riak.client.http.response.RiakExceptionHandler;
import com.basho.riak.client.http.response.RiakIORuntimeException;
import com.basho.riak.client.http.response.RiakResponseRuntimeException;
import com.basho.riak.client.http.response.StreamHandler;
import com.basho.riak.client.util.CharsetUtils;

/**
 * This class performs the actual HTTP requests underlying the operations in
 * RiakClient and returns the resulting HTTP responses. It is up to RiakClient
 * to interpret the responses and translate them into the appropriate format.
 */
public class ClientHelper {

    private RiakConfig config;
    private HttpClient httpClient;
    private String clientId = null;
    private RiakExceptionHandler exceptionHandler = null;

    public ClientHelper(RiakConfig config, String clientId) {
        this.config = config;
        httpClient = ClientUtils.newHttpClient(config);
        setClientId(clientId);
    }

    /** Used for testing -- inject an HttpClient */
    void setHttpClient(HttpClient httpClient) {
        this.httpClient = httpClient;
    }

    /**
     * See {@link RiakClient#getClientId()}
     */
    public byte[] getClientId() {
        return Base64.decodeBase64(utf8StringToBytes(clientId));
    }

    public void setClientId(String clientId) {
        if (clientId != null) {
            this.clientId = ClientUtils.encodeClientId(clientId);
        } else {
            this.clientId = ClientUtils.randomClientId();
        }
    }

    /**
     * See
     * {@link RiakClient#setBucketSchema(String, com.basho.riak.client.http.RiakBucketInfo, RequestMeta)}
     */
    public HttpResponse setBucketSchema(String bucket, JSONObject schema, RequestMeta meta) {
        if (schema == null) {
            schema = new JSONObject();
        }
        if (meta == null) {
            meta = new RequestMeta();
        }

        meta.setHeader(Constants.HDR_ACCEPT, Constants.CTYPE_JSON);

        HttpPut put = new HttpPut(ClientUtils.makeURI(config, bucket));
        ByteArrayEntity entity = new ByteArrayEntity(utf8StringToBytes(schema.toString()));
        entity.setContentType(Constants.CTYPE_JSON_UTF8);
        put.setEntity(entity);

        return executeMethod(bucket, null, put, meta);
    }

    /**
     * Same as {@link RiakClient#getBucketSchema(String, RequestMeta)}, except
     * only returning the HTTP response.
     */
    public HttpResponse getBucketSchema(String bucket, RequestMeta meta) {
        if (meta == null) {
            meta = new RequestMeta();
        }
        if (meta.getQueryParam(Constants.QP_KEYS) == null) {
            meta.setQueryParam(Constants.QP_KEYS, Constants.NO_KEYS);
        }
        return listBucket(bucket, meta, false);
    }

    /**
     * List the buckets in Riak
     * 
     * @return an {@link HttpResponse} whose body should be the result of asking
     *         Riak to list buckets.
     */
    public HttpResponse listBuckets() {
        final RequestMeta meta = new RequestMeta();
        meta.setQueryParam(Constants.QP_BUCKETS, Constants.LIST_BUCKETS);
        HttpGet get = new HttpGet(config.getUrl());
        return executeMethod(null, null, get, meta);
    }

    /**
     * Same as {@link RiakClient}, except only returning the HTTP response, and
     * if streamResponse==true, the response will be streamed back, so the user
     * is responsible for calling {@link BucketResponse#close()}
     */
    public HttpResponse listBucket(String bucket, RequestMeta meta, boolean streamResponse) {
        if (meta == null) {
            meta = new RequestMeta();
        }
        if (meta.getQueryParam(Constants.QP_KEYS) == null) {
            if (streamResponse) {
                meta.setQueryParam(Constants.QP_KEYS, Constants.STREAM_KEYS);
            } else {
                meta.setQueryParam(Constants.QP_KEYS, Constants.INCLUDE_KEYS);
            }
        }
        if (meta.getHeader(Constants.HDR_CONTENT_TYPE) == null) {
            meta.setHeader(Constants.HDR_CONTENT_TYPE, Constants.CTYPE_JSON);
        }
        if (meta.getHeader(Constants.HDR_ACCEPT) == null) {
            meta.setHeader(Constants.HDR_ACCEPT, Constants.CTYPE_JSON);
        }

        HttpGet get = new HttpGet(ClientUtils.makeURI(config, bucket));
        return executeMethod(bucket, null, get, meta, streamResponse);
    }

    /**
     * Same as {@link RiakClient}, except only returning the HTTP response
     */
    public HttpResponse store(RiakObject object, RequestMeta meta) {
        if (meta == null) {
            meta = new RequestMeta();
        }
        if (meta.getClientId() == null) {
            meta.setClientId(clientId);
        }
        if (meta.getHeader(Constants.HDR_CONNECTION) == null) {
            meta.setHeader(Constants.HDR_CONNECTION, "keep-alive");
        }

        String bucket = object.getBucket();
        String key = object.getKey();
        String url = ClientUtils.makeURI(config, bucket, key);
        HttpRequestBase storeMethod = createStoreHttpMethod(key, url);

        object.writeToHttpMethod(storeMethod);
        return executeMethod(bucket, key, storeMethod, meta);
    }

    /**
     * Same as {@link RiakClient}, except only returning the HTTP response
     */
    public HttpResponse fetchMeta(String bucket, String key, RequestMeta meta) {
        if (meta == null) {
            meta = new RequestMeta();
        }

        HttpHead head = new HttpHead(ClientUtils.makeURI(config, bucket, key));
        return executeMethod(bucket, key, head, meta);
    }

    /**
     * Same as {@link RiakClient}, except only returning the HTTP response and
     * allows the response to be streamed.
     * 
     * @param bucket
     *            Same as {@link RiakClient}
     * @param key
     *            Same as {@link RiakClient}
     * @param meta
     *            Same as {@link RiakClient}
     * @param streamResponse
     *            If true, the connection will NOT be released. Use
     *            HttpResponse.getHttpMethod().getResponseBodyAsStream() to get
     *            the response stream; HttpResponse.getBody() will return null.
     * 
     * @return Same as {@link RiakClient}
     */
    public HttpResponse fetch(String bucket, String key, RequestMeta meta, boolean streamResponse) {
        if (meta == null) {
            meta = new RequestMeta();
        }

        HttpGet get = new HttpGet(ClientUtils.makeURI(config, bucket, key));
        return executeMethod(bucket, key, get, meta, streamResponse);
    }

    public HttpResponse fetch(String bucket, String key, RequestMeta meta) {
        return fetch(bucket, key, meta, false);
    }

    /**
     * Same as {@link RiakClient}, except only returning the HTTP response
     */
    public boolean stream(String bucket, String key, StreamHandler handler, RequestMeta meta) throws IOException {
        if (meta == null) {
            meta = new RequestMeta();
        }

        HttpGet get = new HttpGet(ClientUtils.makeURI(config, bucket, key));
        try {
            org.apache.http.HttpResponse response = httpClient.execute(get);
            HttpEntity entity = response.getEntity();

            boolean result = true;
            if (handler != null) {

                result = handler.process(bucket, key, response.getStatusLine().getStatusCode(),
                        ClientUtils.asHeaderMap(response.getAllHeaders()), entity.getContent(), response);
            }
            EntityUtils.consume(entity);

            return result;
        } catch (IOException e) {
            get.abort();
            throw e;
        }
    }

    /**
     * Same as {@link RiakClient}, except only returning the HTTP response
     */
    public HttpResponse delete(String bucket, String key, RequestMeta meta) {
        if (meta == null) {
            meta = new RequestMeta();
        }
        String url = ClientUtils.makeURI(config, bucket, key);
        HttpDelete delete = new HttpDelete(url);
        return executeMethod(bucket, key, delete, meta);
    }

    /**
     * Same as {@link RiakClient}, except only returning the HTTP response
     */
    public HttpResponse walk(String bucket, String key, String walkSpec, RequestMeta meta) {
        HttpGet get = new HttpGet(ClientUtils.makeURI(config, bucket, key, walkSpec));
        return executeMethod(bucket, key, get, meta);
    }

    /**
     * Same as {@link RiakClient}, except only returning the HTTP response
     */
    public HttpResponse mapReduce(String job, RequestMeta meta) {
        HttpPost post = new HttpPost(config.getMapReduceUrl());
        StringEntity entity = new StringEntity(job, ContentType.APPLICATION_JSON);
        post.setEntity(entity);
        return executeMethod(null, null, post, meta);
    }

    /**
     * Perform the fetch/query on an index
     * 
     * @param bucket
     *            the bucket
     * @param indexName
     *            the name of the index
     * @param values
     *            an array of Strings, if only 1 long, then a value, if longer
     *            the 1st 2 elements are treated as the bounds for a range
     * @return an {@link HttpResponse}
     */
    public HttpResponse fetchIndex(String bucket, String indexName, String[] values) {
        HttpGet get = new HttpGet(ClientUtils.makeURI(config, bucket, indexName, values));
        return executeMethod(bucket, null, get, null);
    }

    /**
     * Perform the fetch/query on an index
     * 
     * @param bucket
     *            the bucket
     * @param indexName
     *            the name of the index
     * @param values
     *            an array of longs, if only 1 element, then a value, if longer
     *            the 1st 2 elements are treated as the bounds for a range
     * @return an {@link HttpResponse}
     */
    public HttpResponse fetchIndex(String bucket, String indexName, long[] values) {
        HttpGet get = new HttpGet(ClientUtils.makeURI(config, bucket, indexName, values));
        return executeMethod(bucket, null, get, null);
    }

    /**
     * Same as {@link RiakClient#ping}
     * @return the ping HttpResponse
     */
    public HttpResponse ping() {
        HttpGet get = new HttpGet(config.getPingUrl());
        return executeMethod(null, null, get, null);
    }

    /** @return the installed exception handler or null if not installed */
    public RiakExceptionHandler getExceptionHandler() {
        return exceptionHandler;
    }

    /**
     * Same as {@link RiakClient#stats}  
     * 
     * @return an {@link HttpResponse} whose body should be the result of asking
     *         Riak for its stats (status).
     */
    public HttpResponse stats() {
        HttpGet get = new HttpGet(config.getStatsUrl());
        return executeMethod(null, null, get, null);
    }

    /**
     * Install an exception handler. If an exception handler is provided, then
     * the Riak client will hand exceptions to the handler rather than throwing
     * them.
     */
    public void setExceptionHandler(RiakExceptionHandler exceptionHandler) {
        this.exceptionHandler = exceptionHandler;
    }

    /**
     * Hands exception <code>e</code> to installed exception handler if there is
     * one or throw it.
     * 
     * @return A 0-status {@link HttpResponse}.
     */
    public HttpResponse toss(RiakIORuntimeException e) {
        if (exceptionHandler != null) {
            exceptionHandler.handle(e);
            return new DefaultHttpResponse(null, null, 0, null, null, null, null, null);
        } else
            throw e;
    }

    public HttpResponse toss(RiakResponseRuntimeException e) {
        if (exceptionHandler != null) {
            exceptionHandler.handle(e);
            return new DefaultHttpResponse(null, null, 0, null, null, null, null, null);
        } else
            throw e;
    }

    /**
     * Return the {@link HttpClient} used to make requests, which can be
     * configured.
     */
    public HttpClient getHttpClient() {
        return httpClient;
    }

    /**
     * @return The config used to construct the HttpClient connecting to Riak.
     */
    public RiakConfig getConfig() {
        return config;
    }

    /**
     * Perform and HTTP request and return the resulting response using the
     * internal HttpClient.
     * 
     * @param bucket
     *            Bucket of the object receiving the request.
     * @param key
     *            Key of the object receiving the request or null if the request
     *            is for a bucket.
     * @param httpMethod
     *            The HTTP request to perform; must not be null.
     * @param meta
     *            Extra HTTP headers to attach to the request. Query parameters
     *            are ignored; they should have already been used to construct
     *            <code>httpMethod</code> and query parameters.
     * @param streamResponse
     *            If true, the connection will NOT be released. Use
     *            HttpResponse.getHttpMethod().getResponseBodyAsStream() to get
     *            the response stream; HttpResponse.getBody() will return null.
     * 
     * @return The HTTP response returned by Riak from executing
     *         <code>httpMethod</code>.
     * 
     * @throws RiakIORuntimeException
     *             If an error occurs during communication with the Riak server
     *             (i.e. HttpClient threw an IOException)
     */
    HttpResponse executeMethod(String bucket, String key, HttpRequestBase httpMethod, RequestMeta meta,
            boolean streamResponse) {

        if (meta != null) {
            Map<String, String> headers = meta.getHeaders();
            for (String header : headers.keySet()) {
                httpMethod.addHeader(header, headers.get(header));
            }

            Map<String, String> queryParams = meta.getQueryParamMap();
            if (!queryParams.isEmpty()) {
                URI originalURI = httpMethod.getURI();
                List<NameValuePair> currentQuery = URLEncodedUtils.parse(originalURI, CharsetUtils.UTF_8.name());
                List<NameValuePair> newQuery = new LinkedList<NameValuePair>(currentQuery);

                for (Map.Entry<String, String> qp : queryParams.entrySet()) {
                    newQuery.add(new BasicNameValuePair(qp.getKey(), qp.getValue()));
                }

                // For this, HC4.1 authors, I hate you
                URI newURI;
                try {
                    newURI = new URIBuilder(originalURI).setQuery(URLEncodedUtils.format(newQuery, "UTF-8"))
                            .build();
                } catch (URISyntaxException e) {
                    e.printStackTrace();
                    throw new RiakIORuntimeException(e);
                }
                httpMethod.setURI(newURI);
            }
        }
        HttpEntity entity = null;
        try {
            org.apache.http.HttpResponse response = httpClient.execute(httpMethod);

            int status = 0;
            if (response.getStatusLine() != null) {
                status = response.getStatusLine().getStatusCode();
            }

            Map<String, String> headers = ClientUtils.asHeaderMap(response.getAllHeaders());
            byte[] body = null;
            InputStream stream = null;
            entity = response.getEntity();

            if (streamResponse) {
                stream = entity.getContent();
            } else {
                if (null != entity) {
                    body = EntityUtils.toByteArray(entity);
                }
            }

            key = extractKeyFromResponseIfItWasNotAlreadyProvided(key, response);

            return new DefaultHttpResponse(bucket, key, status, headers, body, stream, response, httpMethod);
        } catch (IOException e) {
            httpMethod.abort();
            return toss(new RiakIORuntimeException(e));
        } finally {
            if (!streamResponse && entity != null) {
                try {
                    EntityUtils.consume(entity);
                } catch (IOException e) {
                    // NO-OP
                }
            }
        }
    }

    private String extractKeyFromResponseIfItWasNotAlreadyProvided(String key,
            org.apache.http.HttpResponse response) {
        if (key == null) {
            Header locationHeader = response.getFirstHeader("Location");
            if (locationHeader != null) {
                String location = locationHeader.getValue();
                if (location != null) {
                    int indexOfLastSlash = location.lastIndexOf("/");
                    key = location.substring(indexOfLastSlash + 1);
                }
            }
        }
        return key;
    }

    HttpResponse executeMethod(String bucket, String key, HttpRequestBase httpMethod, RequestMeta meta) {
        return executeMethod(bucket, key, httpMethod, meta, false);
    }

    private HttpRequestBase createStoreHttpMethod(String key, String url) {
        HttpRequestBase storeMethod = null;
        if (key == null) {
            storeMethod = new HttpPost(url);
        } else {
            storeMethod = new HttpPut(url);
        }
        return storeMethod;
    }
}