com.microsoft.azure.documentdb.internal.directconnectivity.StoreResponse Maven / Gradle / Ivy

Go to download
package com.microsoft.azure.documentdb.internal.directconnectivity;

import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.UnsupportedEncodingException;
import java.net.URLDecoder;
import java.util.Scanner;

import org.apache.commons.codec.Charsets;
import org.apache.commons.lang3.StringUtils;
import org.apache.http.Header;
import org.apache.http.HttpEntity;
import org.apache.http.HttpResponse;
import org.apache.http.util.EntityUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.microsoft.azure.documentdb.internal.HttpConstants;
import com.microsoft.azure.documentdb.ClientSideRequestStatistics;

/**
 * Used internally to represents a response from the store.
 */
public class StoreResponse {
    final static Logger LOGGER = LoggerFactory.getLogger(StoreResponse.class);
    final private int status;
    final private String[] responseHeaderNames;
    final private String[] responseHeaderValues;
    final private InputStream httpEntityStream;
    private volatile String httpEntityContent;

    private ClientSideRequestStatistics clientSideRequestStatistics;

    public StoreResponse(
            String[] responseHeaderNames,
            String[] responseHeaderValues,
            int status,
            HttpEntity responseBody,
            boolean isMedia,
            ClientSideRequestStatistics clientSideRequestStatistics) {
        
        this(responseHeaderNames, responseHeaderValues, status,
                responseBody != null ? asInputStream(responseBody, isMedia) : null,
                clientSideRequestStatistics);
    }

    public StoreResponse(
            String[] responseHeaderNames,
            String[] responseHeaderValues,
            int status,
            InputStream inputStream,
            ClientSideRequestStatistics clientSideRequestStatistics) {
        this.responseHeaderNames = responseHeaderNames;
        this.responseHeaderValues = responseHeaderValues;
        this.status = status;

        this.httpEntityStream = inputStream;
        this.clientSideRequestStatistics = clientSideRequestStatistics;
    }

    public int getStatus() {
        return status;
    }

    public String[] getResponseHeaderNames() {
        return responseHeaderNames;
    }

    public String[] getResponseHeaderValues() {
        return responseHeaderValues;
    }

    public String getResponseBody() {
        if (this.httpEntityContent == null) {
            if (this.httpEntityStream == null) {
                // Response from some operations such as Delete doesn't have a response body
                return null;
            }

            synchronized (this) {
                if (this.httpEntityContent != null) {
                    return this.httpEntityContent;
                }
                // Read httpEntityStream into string. We can use standard Java library only by using Scanner with \\A delimiter.
                // The scanner iterates over tokens in the stream. When we use \\A, which is beginning of the input boundary,
                // the scanner will return one token which is the content of the whole stream.
                Scanner s = new Scanner(this.httpEntityStream, Charsets.UTF_8.toString()).useDelimiter("\\A");
                this.httpEntityContent = s.hasNext() ? s.next() : "";
            }
        }

        return this.httpEntityContent;
    }

    public InputStream getResponseStream() {
        // Some operation type doesn't have a response stream so this can be null
        return this.httpEntityStream;
    }

    public ClientSideRequestStatistics getClientSideRequestStatistics() {
        return clientSideRequestStatistics;
    }

    public void setClientSideRequestStatistics(ClientSideRequestStatistics clientSideRequestStatistics) {
        this.clientSideRequestStatistics = clientSideRequestStatistics;
    }

    private static InputStream asInputStream(HttpEntity responseBody, boolean isMedia) {
        try {
            if (isMedia && responseBody.isStreaming()) {
                return  responseBody.getContent();
            }
            else {
                byte[] ba = EntityUtils.toByteArray(responseBody);
                return new ByteArrayInputStream(ba);
            }
        } catch (IOException e) {
            throw new IllegalStateException(e);
        }
    }

    public long getLSN() {
        String lsnString = this.getHeaderValue(WFConstants.BackendHeaders.LSN);
        if (StringUtils.isNotEmpty(lsnString)) {
            return Long.parseLong(lsnString);
        }

        return -1;
    }

    public String getPartitionKeyRangeId() {
        return this.getHeaderValue(WFConstants.BackendHeaders.PARTITION_KEY_RANGE_ID);
    }

    public String getContinuation() {
        return this.getHeaderValue(HttpConstants.HttpHeaders.CONTINUATION);
    }

    public String getHeaderValue(String attribute) {
        if (this.responseHeaderValues == null || this.responseHeaderNames.length != this.responseHeaderValues.length) {
            return null;
        }

        for (int i = 0; i < responseHeaderNames.length; i++) {
            if (responseHeaderNames[i].equalsIgnoreCase(attribute)) {
                return responseHeaderValues[i];
            }
        }

        return null;
    }

    /**
     * Construct an instance of StoreResponse from HttpResponse and close the connection associated with it.
     *
     * @param httpResponse the HttpResponse instance
     * @param isMedia      a boolean indicating whether this response is from a media request or not
     * @return             a new instance of StoreResponse
     */
    public static StoreResponse fromHttpResponse(HttpResponse httpResponse, boolean isMedia, ClientSideRequestStatistics clientSideRequestStatistics) {
        Header[] allHeaders = httpResponse.getAllHeaders();
        String[] headers = new String[allHeaders.length];
        String[] values = new String[allHeaders.length];

        for (int i = 0; i < allHeaders.length; i++) {
            headers[i] = allHeaders[i].getName();
            if (headers[i].equals(HttpConstants.HttpHeaders.OWNER_FULL_NAME)) {
                try {
                    values[i] = URLDecoder.decode(allHeaders[i].getValue(),  "UTF-8");
                } catch (UnsupportedEncodingException e) {
                    throw new IllegalStateException("Unable to decode exception." + e.getMessage());
                }
            } else {
                values[i] = allHeaders[i].getValue();
            }
        }

        StoreResponse storeResponse = new StoreResponse(
                headers,
                values,
                httpResponse.getStatusLine().getStatusCode(),
                httpResponse.getEntity(), isMedia, clientSideRequestStatistics);

        // Close the connection associated with the httpEntity instance.
        // For media requests, the user is responsible for connection handling.
        if (!isMedia) {
            try {
                EntityUtils.consume(httpResponse.getEntity());
            } catch (IOException e) {
                throw new IllegalStateException("Failed to consume HttpResponse.", e);
            }
        }

        return storeResponse;
    }
}