com.microsoft.azure.documentdb.internal.query.AbstractQueryExecutionContext Maven / Gradle / Ivy

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

import static com.microsoft.azure.documentdb.internal.Constants.QueryExecutionContext.INCREMENTAL_FEED_HEADER_VALUE;

import java.util.HashMap;
import java.util.Map;

import org.apache.commons.lang3.StringUtils;
import org.apache.http.HttpStatus;

import com.microsoft.azure.documentdb.ChangeFeedOptions;
import com.microsoft.azure.documentdb.DocumentClientException;
import com.microsoft.azure.documentdb.DocumentCollection;
import com.microsoft.azure.documentdb.DocumentQueryClientInternal;
import com.microsoft.azure.documentdb.FeedOptions;
import com.microsoft.azure.documentdb.FeedOptionsBase;
import com.microsoft.azure.documentdb.Resource;
import com.microsoft.azure.documentdb.SqlQuerySpec;
import com.microsoft.azure.documentdb.internal.DocumentServiceRequest;
import com.microsoft.azure.documentdb.internal.DocumentServiceResponse;
import com.microsoft.azure.documentdb.internal.HttpConstants;
import com.microsoft.azure.documentdb.internal.OperationType;
import com.microsoft.azure.documentdb.internal.ResourceType;
import com.microsoft.azure.documentdb.internal.Utils;
import com.microsoft.azure.documentdb.internal.routing.PartitionKeyInternal;
import com.microsoft.azure.documentdb.internal.routing.PartitionKeyRangeIdentity;

abstract class AbstractQueryExecutionContext implements QueryExecutionContext {
    private static final String IF_NONE_MATCH_ALL_HEADER_VALUE = "*";
    protected final DocumentQueryClientInternal client;
    protected final ResourceType resourceType;
    protected final Class classT;
    protected final SqlQuerySpec querySpec;
    protected final FeedOptionsBase options;
    protected final String resourceLink;
    private boolean isChangeFeedContext;
    protected int prefetchedStatusCode;

    protected Map responseHeaders;

    protected AbstractQueryExecutionContext(DocumentQueryClientInternal client, ResourceType resourceType, Class classT,
            SqlQuerySpec querySpec, FeedOptionsBase options, String resourceLink) {
        this.client = client;
        this.resourceType = resourceType;
        this.classT = classT;
        this.querySpec = querySpec;
        this.options = options;
        this.resourceLink = resourceLink;
        this.responseHeaders = null;
        this.isChangeFeedContext = options instanceof ChangeFeedOptions;
    }

    @Override
    public void remove() {
        throw new UnsupportedOperationException("remove");
    }

    @Override
    public Map getResponseHeaders() {
        return this.responseHeaders;
    }

    protected boolean hasNextInternal() {
        return this.responseHeaders == null
                || !this.isChangeFeedContext && !StringUtils.isEmpty(this.getContinuationToken())

                // For ChangeFeed, we check whether the next batch contains any documents or not
                || this.isChangeFeedContext && this.prefetchedStatusCode != HttpStatus.SC_NOT_MODIFIED;
    }

    public boolean shouldExecuteQuery() {
        return this.querySpec != null;
    }

    protected String getContinuationToken() {
        return this.responseHeaders == null ?
            null :
            this.responseHeaders.get(
                !this.isChangeFeedContext ?
                    HttpConstants.HttpHeaders.CONTINUATION :
                    HttpConstants.HttpHeaders.E_TAG);
    }

    protected Map getFeedHeaders(FeedOptionsBase options) {
        if (options == null)
            return new HashMap<>();

        Map headers = new HashMap<>();

        if (options.getPageSize() != null) {
            headers.put(HttpConstants.HttpHeaders.PAGE_SIZE, options.getPageSize().toString());
        }

        if (this.isChangeFeedContext) {
            ChangeFeedOptions changeFeedOptions = (ChangeFeedOptions) options;

            String ifNoneMatchValue = null;
            if (changeFeedOptions.getRequestContinuation() != null) {
                ifNoneMatchValue = changeFeedOptions.getRequestContinuation();
            } else if (changeFeedOptions.getStartDateTime() != null) {
                String dateTimeInHttpFormat = Utils.dateTimeAsUTCRFC1123(changeFeedOptions.getStartDateTime());
                headers.put(HttpConstants.HttpHeaders.IF_MODIFIED_SINCE, dateTimeInHttpFormat);
            } else if (!changeFeedOptions.isStartFromBeginning()) {
                ifNoneMatchValue = IF_NONE_MATCH_ALL_HEADER_VALUE;
            }
            // On REST level, change feed is using IfNoneMatch/ETag instead of continuation.
            if (ifNoneMatchValue != null) {
                headers.put(HttpConstants.HttpHeaders.IF_NONE_MATCH, ifNoneMatchValue);
            }

            headers.put(HttpConstants.HttpHeaders.A_IM, INCREMENTAL_FEED_HEADER_VALUE);
        } else if (options.getRequestContinuation() != null) {
            headers.put(HttpConstants.HttpHeaders.CONTINUATION, options.getRequestContinuation());
        }

        FeedOptions feedOptions = options instanceof FeedOptions ? (FeedOptions) options : null;
        if (feedOptions != null) {
            if (feedOptions.getSessionToken() != null) {
                headers.put(HttpConstants.HttpHeaders.SESSION_TOKEN, feedOptions.getSessionToken());
            }

            if (feedOptions.getEnableScanInQuery() != null) {
                headers.put(HttpConstants.HttpHeaders.ENABLE_SCAN_IN_QUERY, feedOptions.getEnableScanInQuery().toString());
            }

            if (feedOptions.getEmitVerboseTracesInQuery() != null) {
                headers.put(HttpConstants.HttpHeaders.EMIT_VERBOSE_TRACES_IN_QUERY,
                        feedOptions.getEmitVerboseTracesInQuery().toString());
            }

            if (feedOptions.getEnableCrossPartitionQuery() != null) {
                headers.put(HttpConstants.HttpHeaders.ENABLE_CROSS_PARTITION_QUERY,
                        feedOptions.getEnableCrossPartitionQuery().toString());
            }

            if (feedOptions.getMaxDegreeOfParallelism() != 0) {
                headers.put(HttpConstants.HttpHeaders.PARALLELIZE_CROSS_PARTITION_QUERY, Boolean.TRUE.toString());
            }

            if (feedOptions.getDisableRUPerMinuteUsage() != null) {
                headers.put(HttpConstants.HttpHeaders.DISABLE_RU_PER_MINUTE_USAGE, feedOptions.getDisableRUPerMinuteUsage().toString());
            }

            if (feedOptions.getResponseContinuationTokenLimitInKb() > 0) {
                headers.put(HttpConstants.HttpHeaders.RESPONSE_CONTINUATION_TOKEN_LIMIT_IN_KB,
                        Integer.toString(feedOptions.getResponseContinuationTokenLimitInKb()));
            }
        }

        return headers;
    }

    private void populatePartitionKeyInfo(DocumentServiceRequest request, PartitionKeyInternal partitionKey) {
        if (request == null) {
            throw new IllegalArgumentException("request");
        }

        if (request.getResourceType().isPartitioned()) {
            if (partitionKey != null) {
                request.getHeaders().put(HttpConstants.HttpHeaders.PARTITION_KEY, partitionKey.toJson());
            }
        }
    }

    private void populatePartitionKeyRangeInfo(DocumentServiceRequest request, String partitionKeyRangeId) {
        if (request == null) {
            throw new IllegalArgumentException("request");
        }

        if (partitionKeyRangeId == null) {
            throw new IllegalArgumentException("partitionKeyRangeId");
        }

        DocumentCollection collection = this.client.getCollectionCache().resolveCollection(request);

        if (request.getResourceType().isPartitioned()) {
            request.routeTo(new PartitionKeyRangeIdentity(collection.getResourceId(), partitionKeyRangeId));
        }
    }

    DocumentServiceRequest createRequest(Map requestHeaders, SqlQuerySpec querySpec, PartitionKeyInternal partitionKeyInternal) {
        DocumentServiceRequest request;

        if (querySpec == null) {
            request = DocumentServiceRequest.create(
                    OperationType.ReadFeed, this.resourceType, this.resourceLink, requestHeaders);
        } else {
            request = DocumentServiceRequest.create(this.resourceType, this.resourceLink, querySpec,
                    this.client.getQueryCompatiblityMode(), requestHeaders);
        }

        this.populatePartitionKeyInfo(request, partitionKeyInternal);

        return request;
    }

    DocumentServiceRequest createRequest(Map requestHeaders, SqlQuerySpec querySpec, String partitionKeyRangeId) {
        DocumentServiceRequest request;

        if (querySpec == null) {
            request = DocumentServiceRequest.create(
                    OperationType.ReadFeed, this.resourceType, this.resourceLink, requestHeaders);
        } else {
            request = DocumentServiceRequest.create(this.resourceType, this.resourceLink, querySpec,
                    this.client.getQueryCompatiblityMode(), requestHeaders);
        }

        this.populatePartitionKeyRangeInfo(request, partitionKeyRangeId);

        return request;
    }

    PartitionKeyInternal getPartitionKeyInternal() {
        FeedOptions feedOptions = this.options instanceof FeedOptions ? (FeedOptions) this.options : null;
        if (feedOptions == null) {
            return null;
        }
        return feedOptions.getPartitionKey() == null ? null : feedOptions.getPartitionKey().getInternalPartitionKey();
    }

    DocumentServiceResponse executeRequest(DocumentServiceRequest request) throws DocumentClientException {
        return this.shouldExecuteQuery()
                ? this.client.doQuery(request)
                        : this.client.doReadFeed(request);
    }
}