All Downloads are FREE. Search and download functionalities are using the official Maven repository.

com.google.cloud.datastore.QueryResultsImpl Maven / Gradle / Ivy

There is a newer version: 2.24.2
Show newest version
/*
 * Copyright 2015 Google LLC
 *
 * Licensed 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.google.cloud.datastore;

import com.google.api.core.BetaApi;
import com.google.cloud.datastore.Query.ResultType;
import com.google.cloud.datastore.models.ExplainMetrics;
import com.google.common.base.Preconditions;
import com.google.common.collect.AbstractIterator;
import com.google.datastore.v1.ExplainOptions;
import com.google.datastore.v1.QueryResultBatch.MoreResultsType;
import com.google.datastore.v1.ReadOptions;
import com.google.protobuf.ByteString;
import java.util.Iterator;
import java.util.Objects;
import java.util.Optional;

class QueryResultsImpl extends AbstractIterator implements QueryResults {

  private final DatastoreImpl datastore;
  private final Optional readOptionsPb;
  private final com.google.datastore.v1.PartitionId partitionIdPb;
  private final ResultType queryResultType;
  private RecordQuery query;
  private ResultType actualResultType;
  private com.google.datastore.v1.RunQueryResponse runQueryResponsePb;
  private com.google.datastore.v1.Query mostRecentQueryPb;
  private boolean lastBatch;
  private Iterator entityResultPbIter;
  private ByteString cursor;
  private MoreResultsType moreResults;
  private final ExplainOptions explainOptions;
  private ExplainMetrics explainMetrics;

  QueryResultsImpl(
      DatastoreImpl datastore,
      Optional readOptionsPb,
      RecordQuery query,
      String namespace,
      ExplainOptions explainOptions) {
    this.datastore = datastore;
    this.readOptionsPb = readOptionsPb;
    this.query = query;
    queryResultType = query.getType();
    this.explainOptions = explainOptions;
    com.google.datastore.v1.PartitionId.Builder pbBuilder =
        com.google.datastore.v1.PartitionId.newBuilder();
    pbBuilder.setProjectId(datastore.getOptions().getProjectId());
    pbBuilder.setDatabaseId(datastore.getOptions().getDatabaseId());
    if (namespace != null) {
      pbBuilder.setNamespaceId(namespace);
    } else if (datastore.getOptions().getNamespace() != null) {
      pbBuilder.setNamespaceId(datastore.getOptions().getNamespace());
    }
    partitionIdPb = pbBuilder.build();
    sendRequest();
    if (runQueryResponsePb.getBatch().getSkippedResults() > 0) {
      cursor = runQueryResponsePb.getBatch().getSkippedCursor();
    } else {
      cursor = mostRecentQueryPb.getStartCursor();
    }
  }

  private void sendRequest() {
    com.google.datastore.v1.RunQueryRequest.Builder requestPb =
        com.google.datastore.v1.RunQueryRequest.newBuilder();
    readOptionsPb.ifPresent(requestPb::setReadOptions);
    requestPb.setPartitionId(partitionIdPb);
    requestPb.setProjectId(datastore.getOptions().getProjectId());
    requestPb.setDatabaseId(datastore.getOptions().getDatabaseId());
    if (explainOptions != null) {
      requestPb.setExplainOptions(explainOptions);
    }
    query.populatePb(requestPb);
    runQueryResponsePb = datastore.runQuery(requestPb.build());
    mostRecentQueryPb = requestPb.getQuery();
    moreResults = runQueryResponsePb.getBatch().getMoreResults();
    lastBatch = moreResults != MoreResultsType.NOT_FINISHED;
    entityResultPbIter = runQueryResponsePb.getBatch().getEntityResultsList().iterator();
    actualResultType = ResultType.fromPb(runQueryResponsePb.getBatch().getEntityResultType());
    if (Objects.equals(queryResultType, ResultType.PROJECTION_ENTITY)) {
      // projection entity can represent all type of results
      actualResultType = ResultType.PROJECTION_ENTITY;
    }
    boolean isExplain =
        explainOptions != null
            && actualResultType.getQueryType() == null
            && !explainOptions.getAnalyze();
    Preconditions.checkState(
        queryResultType.isAssignableFrom(actualResultType) || isExplain,
        "Unexpected result type or explain options set "
            + actualResultType
            + " vs "
            + queryResultType
            + ", explain options = false");
    if (runQueryResponsePb.hasExplainMetrics()) {
      this.explainMetrics = new ExplainMetrics(runQueryResponsePb.getExplainMetrics());
    }
  }

  @Override
  protected T computeNext() {
    while (!entityResultPbIter.hasNext() && !lastBatch) {
      query = query.nextQuery(runQueryResponsePb);
      sendRequest();
    }
    if (!entityResultPbIter.hasNext()) {
      cursor = runQueryResponsePb.getBatch().getEndCursor();
      return endOfData();
    }
    com.google.datastore.v1.EntityResult entityResultPb = entityResultPbIter.next();
    cursor = entityResultPb.getCursor();
    @SuppressWarnings("unchecked")
    T result = (T) actualResultType.convert(entityResultPb.getEntity());
    return result;
  }

  @Override
  public Class getResultClass() {
    return actualResultType.resultClass();
  }

  @Override
  public Cursor getCursorAfter() {
    return new Cursor(cursor);
  }

  @Override
  public int getSkippedResults() {
    return runQueryResponsePb.getBatch().getSkippedResults();
  }

  @Override
  public MoreResultsType getMoreResults() {
    return moreResults;
  }

  @Override
  @BetaApi
  public Optional getExplainMetrics() {
    return Optional.ofNullable(this.explainMetrics);
  }
}




© 2015 - 2024 Weber Informatics LLC | Privacy Policy