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

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

import java.util.*;

import com.microsoft.azure.documentdb.*;
import com.microsoft.azure.documentdb.internal.*;
import com.microsoft.azure.documentdb.internal.routing.CollectionCache;

final class ProxyQueryExecutionContext implements QueryExecutionContext {
    private static final int DEFAULT_PAGE_SIZE = 1000;
    private final DocumentQueryClientInternal client;
    private final ResourceType resourceType;
    private final Class classT;
    private final SqlQuerySpec querySpec;
    private final FeedOptionsBase options;
    private final String resourceLink;
    private final QueryExecutionContext queryExecutionContext;
    private T prefetchedResource;
    private volatile boolean hasPrefetchedResource;

    @SuppressWarnings("unchecked")
    public ProxyQueryExecutionContext(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;

        QueryExecutionContext currentQueryExecutionContext = null;

        FeedOptions feedOptions = this.options instanceof FeedOptions ? (FeedOptions) this.options : null;

        // Prefer to get query execution information from ServiceJNI if it's available.
        // Otherwise, fallback to Gateway as usual.
        // Keep the query execution info to use it in DefaultQueryExecutionContext
        // in case query is not executed with ParallelQueryExecutionContext
        PartitionedQueryExecutionInfo partitionedQueryExecutionInfo = null;
        if (ServiceJNIWrapper.isServiceJNIAvailable()
                && Utils.isCollectionChild(resourceType) && resourceType.isPartitioned()
                && feedOptions != null
                && feedOptions.getEnableCrossPartitionQuery() != null
                && feedOptions.getEnableCrossPartitionQuery()) {

            DocumentServiceRequest request = DocumentServiceRequest.create(
                    OperationType.Query, resourceType, resourceLink, null);
            CollectionCache collectionCache = this.client.getCollectionCache();
            DocumentCollection collection = collectionCache.resolveCollection(request);

            QueryPartitionProvider queryPartitionProvider = this.client.getQueryPartitionProvider();
            partitionedQueryExecutionInfo = queryPartitionProvider.getPartitionQueryExcecutionInfo(
                    querySpec, collection.getPartitionKey());

            if (shouldCreateSpecializedDocumentQueryExecutionContext(resourceType, feedOptions, partitionedQueryExecutionInfo)) {
                currentQueryExecutionContext = (QueryExecutionContext) new PipelinedQueryExecutionContext(this.client,
                        collection.getSelfLink(), this.querySpec, feedOptions, this.resourceLink, partitionedQueryExecutionInfo);
            }
        }

        if (currentQueryExecutionContext == null) {
            currentQueryExecutionContext = new DefaultQueryExecutionContext(
                    client, resourceType, classT, querySpec, partitionedQueryExecutionInfo, options, resourceLink);

            // To be able to answer hasNext reliably, we have to call next now.
            try {
                if (currentQueryExecutionContext.hasNext()) {
                    this.prefetchedResource = currentQueryExecutionContext.next();

                    if (this.prefetchedResource != null) {
                        this.hasPrefetchedResource = true;
                    }
                }
            } catch (IllegalStateException e) {
                DocumentClientException dce;
                if (!(e.getCause() instanceof DocumentClientException) || !this
                        .shouldCreatePipelinedQueryExecutionContext((dce = (DocumentClientException) e.getCause()))) {
                    throw e;
                }

                DocumentServiceRequest request = DocumentServiceRequest.create(
                        OperationType.Query, resourceType, resourceLink, null);
                CollectionCache collectionCache = this.client.getCollectionCache();
                DocumentCollection collection = collectionCache.resolveCollection(request);

                currentQueryExecutionContext = (QueryExecutionContext) new PipelinedQueryExecutionContext(this.client,
                        collection.getSelfLink(),
                        this.querySpec, feedOptions, this.resourceLink,
                        new PartitionedQueryExecutionInfo(dce.getError().getPartitionedQueryExecutionInfo()));
            }
        }

        this.queryExecutionContext = currentQueryExecutionContext;
    }

    private static boolean shouldCreateSpecializedDocumentQueryExecutionContext(
            ResourceType resourceType,
            FeedOptions feedOptions,
            PartitionedQueryExecutionInfo partitionedQueryExecutionInfo) {
        return isCrossPartitionQuery(resourceType, feedOptions, partitionedQueryExecutionInfo)
                && (isTopOrderByQuery(partitionedQueryExecutionInfo)
                        || isAggregateQuery(partitionedQueryExecutionInfo)
                        || isParallelQuery(feedOptions));
    }

    private static boolean isParallelQuery(FeedOptions feedOptions) {
        return (feedOptions.getMaxDegreeOfParallelism() != 0);
    }

    private static boolean isTopOrderByQuery(PartitionedQueryExecutionInfo partitionedQueryExecutionInfo) {
        return partitionedQueryExecutionInfo.getQueryInfo() != null
                && (partitionedQueryExecutionInfo.getQueryInfo().hasOrderBy()
                        || partitionedQueryExecutionInfo.getQueryInfo().hasTop());
    }

    private static boolean isAggregateQuery(PartitionedQueryExecutionInfo partitionedQueryExecutionInfo) {
        return partitionedQueryExecutionInfo.getQueryInfo() != null
                && partitionedQueryExecutionInfo.getQueryInfo().hasAggregates();
    }

    private static boolean isCrossPartitionQuery(
            ResourceType resourceType,
            FeedOptions feedOptions,
            PartitionedQueryExecutionInfo partitionedQueryExecutionInfo) {
        return resourceType.isPartitioned()
                && (feedOptions.getPartitionKey() == null
                && feedOptions.getEnableCrossPartitionQuery() != null && feedOptions.getEnableCrossPartitionQuery().booleanValue())
                && !(partitionedQueryExecutionInfo.getQueryRanges().size() == 1
                && partitionedQueryExecutionInfo.getQueryRanges().get(0).isSingleValue());
    }

    @Override
    public boolean hasNext() {
        return this.hasPrefetchedResource || this.queryExecutionContext.hasNext();
    }

    @Override
    public T next() {
        T item = null;
        if (this.hasPrefetchedResource) {
            synchronized (this) {
                if (this.hasPrefetchedResource) {
                    T result = this.prefetchedResource;
                    this.hasPrefetchedResource = false;
                    item = result;
                }
            }
        }

        while (item == null && this.queryExecutionContext.hasNext()) {
            item = this.queryExecutionContext.next();
        }

        return item;
    }

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

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

    @Override
    public List fetchNextBlock() throws DocumentClientException {
        if (this.hasPrefetchedResource) {
            synchronized (this) {
                if (this.hasPrefetchedResource) {
                    int pageSize = options.getPageSize() == null || options.getPageSize() < 1
                            ? ProxyQueryExecutionContext.DEFAULT_PAGE_SIZE
                                    : options.getPageSize();

                    List result = new ArrayList(pageSize);
                    result.add(this.prefetchedResource);
                    while (this.queryExecutionContext.hasNext() && result.size() < pageSize) {
                        result.add(this.queryExecutionContext.next());
                    }
                    this.hasPrefetchedResource = false;
                    return result;
                }
            }
        }

        return this.queryExecutionContext.fetchNextBlock();
    }

    @Override
    public void onNotifyStop() {
        this.queryExecutionContext.onNotifyStop();
    }

    private boolean shouldCreatePipelinedQueryExecutionContext(DocumentClientException e) {
        return !(this.queryExecutionContext instanceof PipelinedQueryExecutionContext)
                && this.resourceType == ResourceType.Document && Document.class.equals(this.classT)
                && (e.getStatusCode() == HttpConstants.StatusCodes.BADREQUEST && e.getSubStatusCode() != null
                && e.getSubStatusCode() == HttpConstants.SubStatusCodes.CROSS_PARTITION_QUERY_NOT_SERVABLE);
    }
}