| /* |
| * Licensed to the Apache Software Foundation (ASF) under one |
| * or more contributor license agreements. See the NOTICE file |
| * distributed with this work for additional information |
| * regarding copyright ownership. The ASF licenses this file |
| * 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 org.apache.beam.sdk.io.gcp.datastore; |
| |
| import static com.google.datastore.v1.PropertyFilter.Operator.EQUAL; |
| import static com.google.datastore.v1.PropertyOrder.Direction.DESCENDING; |
| import static com.google.datastore.v1.QueryResultBatch.MoreResultsType.NOT_FINISHED; |
| import static com.google.datastore.v1.client.DatastoreHelper.makeAndFilter; |
| import static com.google.datastore.v1.client.DatastoreHelper.makeDelete; |
| import static com.google.datastore.v1.client.DatastoreHelper.makeFilter; |
| import static com.google.datastore.v1.client.DatastoreHelper.makeOrder; |
| import static com.google.datastore.v1.client.DatastoreHelper.makeUpsert; |
| import static com.google.datastore.v1.client.DatastoreHelper.makeValue; |
| import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument; |
| import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkNotNull; |
| import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Verify.verify; |
| |
| import com.google.api.client.http.HttpRequestInitializer; |
| import com.google.auth.Credentials; |
| import com.google.auth.http.HttpCredentialsAdapter; |
| import com.google.auto.value.AutoValue; |
| import com.google.cloud.hadoop.util.ChainingHttpRequestInitializer; |
| import com.google.datastore.v1.CommitRequest; |
| import com.google.datastore.v1.Entity; |
| import com.google.datastore.v1.EntityResult; |
| import com.google.datastore.v1.GqlQuery; |
| import com.google.datastore.v1.Key; |
| import com.google.datastore.v1.Key.PathElement; |
| import com.google.datastore.v1.Mutation; |
| import com.google.datastore.v1.PartitionId; |
| import com.google.datastore.v1.Query; |
| import com.google.datastore.v1.QueryResultBatch; |
| import com.google.datastore.v1.RunQueryRequest; |
| import com.google.datastore.v1.RunQueryResponse; |
| import com.google.datastore.v1.client.Datastore; |
| import com.google.datastore.v1.client.DatastoreException; |
| import com.google.datastore.v1.client.DatastoreFactory; |
| import com.google.datastore.v1.client.DatastoreHelper; |
| import com.google.datastore.v1.client.DatastoreOptions; |
| import com.google.datastore.v1.client.QuerySplitter; |
| import com.google.protobuf.Int32Value; |
| import com.google.rpc.Code; |
| import java.io.IOException; |
| import java.io.Serializable; |
| import java.util.ArrayList; |
| import java.util.List; |
| import java.util.NoSuchElementException; |
| import java.util.Set; |
| import javax.annotation.Nullable; |
| import org.apache.beam.sdk.PipelineRunner; |
| import org.apache.beam.sdk.annotations.Experimental; |
| import org.apache.beam.sdk.annotations.Experimental.Kind; |
| import org.apache.beam.sdk.coders.StringUtf8Coder; |
| import org.apache.beam.sdk.extensions.gcp.options.GcpOptions; |
| import org.apache.beam.sdk.extensions.gcp.util.RetryHttpRequestInitializer; |
| import org.apache.beam.sdk.metrics.Counter; |
| import org.apache.beam.sdk.metrics.Metrics; |
| import org.apache.beam.sdk.options.PipelineOptions; |
| import org.apache.beam.sdk.options.ValueProvider; |
| import org.apache.beam.sdk.options.ValueProvider.StaticValueProvider; |
| import org.apache.beam.sdk.transforms.Create; |
| import org.apache.beam.sdk.transforms.DoFn; |
| import org.apache.beam.sdk.transforms.MapElements; |
| import org.apache.beam.sdk.transforms.PTransform; |
| import org.apache.beam.sdk.transforms.ParDo; |
| import org.apache.beam.sdk.transforms.Reshuffle; |
| import org.apache.beam.sdk.transforms.SimpleFunction; |
| import org.apache.beam.sdk.transforms.display.DisplayData; |
| import org.apache.beam.sdk.transforms.display.DisplayData.Builder; |
| import org.apache.beam.sdk.transforms.display.HasDisplayData; |
| import org.apache.beam.sdk.util.BackOff; |
| import org.apache.beam.sdk.util.BackOffUtils; |
| import org.apache.beam.sdk.util.FluentBackoff; |
| import org.apache.beam.sdk.util.Sleeper; |
| import org.apache.beam.sdk.values.KV; |
| import org.apache.beam.sdk.values.PBegin; |
| import org.apache.beam.sdk.values.PCollection; |
| import org.apache.beam.sdk.values.PDone; |
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting; |
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.MoreObjects; |
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Strings; |
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList; |
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableSet; |
| import org.joda.time.Duration; |
| import org.slf4j.Logger; |
| import org.slf4j.LoggerFactory; |
| |
| /** |
| * {@link DatastoreV1} provides an API to Read, Write and Delete {@link PCollection PCollections} of |
| * <a href="https://developers.google.com/datastore/">Google Cloud Datastore</a> version v1 {@link |
| * Entity} objects. Read is only supported for Bounded PCollections while Write and Delete are |
| * supported for both Bounded and Unbounded PCollections. |
| * |
| * <p>This API currently requires an authentication workaround. To use {@link DatastoreV1}, users |
| * must use the {@code gcloud} command line tool to get credentials for Cloud Datastore: |
| * |
| * <pre> |
| * $ gcloud auth login |
| * </pre> |
| * |
| * <p>To read a {@link PCollection} from a query to Cloud Datastore, use {@link DatastoreV1#read} |
| * and its methods {@link DatastoreV1.Read#withProjectId} and {@link DatastoreV1.Read#withQuery} to |
| * specify the project to query and the query to read from. You can optionally provide a namespace |
| * to query within using {@link DatastoreV1.Read#withNamespace}. You could also optionally specify |
| * how many splits you want for the query using {@link DatastoreV1.Read#withNumQuerySplits}. |
| * |
| * <p>For example: |
| * |
| * <pre>{@code |
| * // Read a query from Datastore |
| * PipelineOptions options = PipelineOptionsFactory.fromArgs(args).create(); |
| * Query query = ...; |
| * String projectId = "..."; |
| * |
| * Pipeline p = Pipeline.create(options); |
| * PCollection<Entity> entities = p.apply( |
| * DatastoreIO.v1().read() |
| * .withProjectId(projectId) |
| * .withQuery(query)); |
| * }</pre> |
| * |
| * <p><b>Note:</b> A runner may read from Cloud Datastore in parallel across many workers. However, |
| * when the {@link Query} is configured with a limit using {@link |
| * com.google.datastore.v1.Query.Builder#setLimit(Int32Value)} or if the Query contains inequality |
| * filters like {@code GREATER_THAN, LESS_THAN} etc., then all returned results will be read by a |
| * single worker in order to ensure correct data. Since data is read from a single worker, this |
| * could have a significant impact on the performance of the job. |
| * |
| * <p>To write a {@link PCollection} to a Cloud Datastore, use {@link DatastoreV1#write}, specifying |
| * the Cloud Datastore project to write to: |
| * |
| * <pre>{@code |
| * PCollection<Entity> entities = ...; |
| * entities.apply(DatastoreIO.v1().write().withProjectId(projectId)); |
| * p.run(); |
| * }</pre> |
| * |
| * <p>To delete a {@link PCollection} of {@link Entity Entities} from Cloud Datastore, use {@link |
| * DatastoreV1#deleteEntity()}, specifying the Cloud Datastore project to write to: |
| * |
| * <pre>{@code |
| * PCollection<Entity> entities = ...; |
| * entities.apply(DatastoreIO.v1().deleteEntity().withProjectId(projectId)); |
| * p.run(); |
| * }</pre> |
| * |
| * <p>To delete entities associated with a {@link PCollection} of {@link Key Keys} from Cloud |
| * Datastore, use {@link DatastoreV1#deleteKey}, specifying the Cloud Datastore project to write to: |
| * |
| * <pre>{@code |
| * PCollection<Entity> entities = ...; |
| * entities.apply(DatastoreIO.v1().deleteKey().withProjectId(projectId)); |
| * p.run(); |
| * }</pre> |
| * |
| * <p>{@link Entity Entities} in the {@code PCollection} to be written or deleted must have complete |
| * {@link Key Keys}. Complete {@code Keys} specify the {@code name} and {@code id} of the {@code |
| * Entity}, where incomplete {@code Keys} do not. A {@code namespace} other than {@code projectId} |
| * default may be used by specifying it in the {@code Entity} {@code Keys}. |
| * |
| * <pre>{@code |
| * Key.Builder keyBuilder = DatastoreHelper.makeKey(...); |
| * keyBuilder.getPartitionIdBuilder().setNamespace(namespace); |
| * }</pre> |
| * |
| * <p>{@code Entities} will be committed as upsert (update or insert) or delete mutations. Please |
| * read <a href="https://cloud.google.com/datastore/docs/concepts/entities">Entities, Properties, |
| * and Keys</a> for more information about {@code Entity} keys. |
| * |
| * <h3>Permissions</h3> |
| * |
| * Permission requirements depend on the {@code PipelineRunner} that is used to execute the |
| * pipeline. Please refer to the documentation of corresponding {@code PipelineRunner}s for more |
| * details. |
| * |
| * <p>Please see <a href="https://cloud.google.com/datastore/docs/activate">Cloud Datastore Sign Up |
| * </a>for security and permission related information specific to Cloud Datastore. |
| * |
| * <p>Optionally, Cloud Datastore V1 Emulator, running locally, could be used for testing purposes |
| * by providing the host port information through {@code withLocalhost("host:port"} for all the |
| * above transforms. In such a case, all the Cloud Datastore API calls are directed to the Emulator. |
| * |
| * @see PipelineRunner |
| */ |
| public class DatastoreV1 { |
| |
| // A package-private constructor to prevent direct instantiation from outside of this package |
| DatastoreV1() {} |
| |
| /** |
| * The number of entity updates written per RPC, initially. We buffer updates in the connector and |
| * write a batch to Datastore once we have collected a certain number. This is the initial batch |
| * size; it is adjusted at runtime based on the performance of previous writes (see {@link |
| * DatastoreV1.WriteBatcher}). |
| * |
| * <p>Testing has found that a batch of 200 entities will generally finish within the timeout even |
| * in adverse conditions. |
| */ |
| @VisibleForTesting static final int DATASTORE_BATCH_UPDATE_ENTITIES_START = 200; |
| |
| /** |
| * When choosing the number of updates in a single RPC, never exceed the maximum allowed by the |
| * API. |
| */ |
| @VisibleForTesting static final int DATASTORE_BATCH_UPDATE_ENTITIES_LIMIT = 500; |
| |
| /** |
| * When choosing the number of updates in a single RPC, do not go below this value. The actual |
| * number of entities per request may be lower when we flush for the end of a bundle or if we hit |
| * {@link DatastoreV1.DATASTORE_BATCH_UPDATE_BYTES_LIMIT}. |
| */ |
| @VisibleForTesting static final int DATASTORE_BATCH_UPDATE_ENTITIES_MIN = 10; |
| |
| /** |
| * Cloud Datastore has a limit of 10MB per RPC, so we also flush if the total size of mutations |
| * exceeds this limit. This is set lower than the 10MB limit on the RPC, as this only accounts for |
| * the mutations themselves and not the CommitRequest wrapper around them. |
| */ |
| @VisibleForTesting static final int DATASTORE_BATCH_UPDATE_BYTES_LIMIT = 9_000_000; |
| |
| /** |
| * Non-retryable errors. See https://cloud.google.com/datastore/docs/concepts/errors#Error_Codes . |
| */ |
| private static final Set<Code> NON_RETRYABLE_ERRORS = |
| ImmutableSet.of( |
| Code.FAILED_PRECONDITION, |
| Code.INVALID_ARGUMENT, |
| Code.PERMISSION_DENIED, |
| Code.UNAUTHENTICATED); |
| |
| /** |
| * Returns an empty {@link DatastoreV1.Read} builder. Configure the source {@code projectId}, |
| * {@code query}, and optionally {@code namespace} and {@code numQuerySplits} using {@link |
| * DatastoreV1.Read#withProjectId}, {@link DatastoreV1.Read#withQuery}, {@link |
| * DatastoreV1.Read#withNamespace}, {@link DatastoreV1.Read#withNumQuerySplits}. |
| */ |
| public DatastoreV1.Read read() { |
| return new AutoValue_DatastoreV1_Read.Builder().setNumQuerySplits(0).build(); |
| } |
| |
| /** |
| * A {@link PTransform} that reads the result rows of a Cloud Datastore query as {@code Entity} |
| * objects. |
| * |
| * @see DatastoreIO |
| */ |
| @AutoValue |
| public abstract static class Read extends PTransform<PBegin, PCollection<Entity>> { |
| private static final Logger LOG = LoggerFactory.getLogger(Read.class); |
| |
| /** An upper bound on the number of splits for a query. */ |
| public static final int NUM_QUERY_SPLITS_MAX = 50000; |
| |
| /** A lower bound on the number of splits for a query. */ |
| static final int NUM_QUERY_SPLITS_MIN = 12; |
| |
| /** Default bundle size of 64MB. */ |
| static final long DEFAULT_BUNDLE_SIZE_BYTES = 64L * 1024L * 1024L; |
| |
| /** |
| * Maximum number of results to request per query. |
| * |
| * <p>Must be set, or it may result in an I/O error when querying Cloud Datastore. |
| */ |
| static final int QUERY_BATCH_LIMIT = 500; |
| |
| @Nullable |
| public abstract ValueProvider<String> getProjectId(); |
| |
| @Nullable |
| public abstract Query getQuery(); |
| |
| @Nullable |
| public abstract ValueProvider<String> getLiteralGqlQuery(); |
| |
| @Nullable |
| public abstract ValueProvider<String> getNamespace(); |
| |
| public abstract int getNumQuerySplits(); |
| |
| @Nullable |
| public abstract String getLocalhost(); |
| |
| @Override |
| public abstract String toString(); |
| |
| abstract Builder toBuilder(); |
| |
| @AutoValue.Builder |
| abstract static class Builder { |
| abstract Builder setProjectId(ValueProvider<String> projectId); |
| |
| abstract Builder setQuery(Query query); |
| |
| abstract Builder setLiteralGqlQuery(ValueProvider<String> literalGqlQuery); |
| |
| abstract Builder setNamespace(ValueProvider<String> namespace); |
| |
| abstract Builder setNumQuerySplits(int numQuerySplits); |
| |
| abstract Builder setLocalhost(String localhost); |
| |
| abstract Read build(); |
| } |
| |
| /** |
| * Computes the number of splits to be performed on the given query by querying the estimated |
| * size from Cloud Datastore. |
| */ |
| static int getEstimatedNumSplits(Datastore datastore, Query query, @Nullable String namespace) { |
| int numSplits; |
| try { |
| long estimatedSizeBytes = getEstimatedSizeBytes(datastore, query, namespace); |
| LOG.info("Estimated size bytes for the query is: {}", estimatedSizeBytes); |
| numSplits = |
| (int) |
| Math.min( |
| NUM_QUERY_SPLITS_MAX, |
| Math.round(((double) estimatedSizeBytes) / DEFAULT_BUNDLE_SIZE_BYTES)); |
| } catch (Exception e) { |
| LOG.warn("Failed the fetch estimatedSizeBytes for query: {}", query, e); |
| // Fallback in case estimated size is unavailable. |
| numSplits = NUM_QUERY_SPLITS_MIN; |
| } |
| return Math.max(numSplits, NUM_QUERY_SPLITS_MIN); |
| } |
| |
| /** |
| * Cloud Datastore system tables with statistics are periodically updated. This method fetches |
| * the latest timestamp (in microseconds) of statistics update using the {@code __Stat_Total__} |
| * table. |
| */ |
| private static long queryLatestStatisticsTimestamp( |
| Datastore datastore, @Nullable String namespace) throws DatastoreException { |
| Query.Builder query = Query.newBuilder(); |
| // Note: namespace either being null or empty represents the default namespace, in which |
| // case we treat it as not provided by the user. |
| if (Strings.isNullOrEmpty(namespace)) { |
| query.addKindBuilder().setName("__Stat_Total__"); |
| } else { |
| query.addKindBuilder().setName("__Stat_Ns_Total__"); |
| } |
| query.addOrder(makeOrder("timestamp", DESCENDING)); |
| query.setLimit(Int32Value.newBuilder().setValue(1)); |
| RunQueryRequest request = makeRequest(query.build(), namespace); |
| |
| RunQueryResponse response = datastore.runQuery(request); |
| QueryResultBatch batch = response.getBatch(); |
| if (batch.getEntityResultsCount() == 0) { |
| throw new NoSuchElementException("Datastore total statistics unavailable"); |
| } |
| Entity entity = batch.getEntityResults(0).getEntity(); |
| return entity.getProperties().get("timestamp").getTimestampValue().getSeconds() * 1000000; |
| } |
| |
| /** |
| * Get the estimated size of the data returned by the given query. |
| * |
| * <p>Cloud Datastore provides no way to get a good estimate of how large the result of a query |
| * entity kind being queried, using the __Stat_Kind__ system table, assuming exactly 1 kind is |
| * specified in the query. |
| * |
| * <p>See https://cloud.google.com/datastore/docs/concepts/stats. |
| */ |
| static long getEstimatedSizeBytes(Datastore datastore, Query query, @Nullable String namespace) |
| throws DatastoreException { |
| String ourKind = query.getKind(0).getName(); |
| long latestTimestamp = queryLatestStatisticsTimestamp(datastore, namespace); |
| LOG.info("Latest stats timestamp for kind {} is {}", ourKind, latestTimestamp); |
| |
| Query.Builder queryBuilder = Query.newBuilder(); |
| if (Strings.isNullOrEmpty(namespace)) { |
| queryBuilder.addKindBuilder().setName("__Stat_Kind__"); |
| } else { |
| queryBuilder.addKindBuilder().setName("__Stat_Ns_Kind__"); |
| } |
| |
| queryBuilder.setFilter( |
| makeAndFilter( |
| makeFilter("kind_name", EQUAL, makeValue(ourKind).build()).build(), |
| makeFilter("timestamp", EQUAL, makeValue(latestTimestamp).build()).build())); |
| |
| RunQueryRequest request = makeRequest(queryBuilder.build(), namespace); |
| |
| long now = System.currentTimeMillis(); |
| RunQueryResponse response = datastore.runQuery(request); |
| LOG.debug("Query for per-kind statistics took {}ms", System.currentTimeMillis() - now); |
| |
| QueryResultBatch batch = response.getBatch(); |
| if (batch.getEntityResultsCount() == 0) { |
| throw new NoSuchElementException( |
| "Datastore statistics for kind " + ourKind + " unavailable"); |
| } |
| Entity entity = batch.getEntityResults(0).getEntity(); |
| return entity.getProperties().get("entity_bytes").getIntegerValue(); |
| } |
| |
| private static PartitionId.Builder forNamespace(@Nullable String namespace) { |
| PartitionId.Builder partitionBuilder = PartitionId.newBuilder(); |
| // Namespace either being null or empty represents the default namespace. |
| // Datastore Client libraries expect users to not set the namespace proto field in |
| // either of these cases. |
| if (!Strings.isNullOrEmpty(namespace)) { |
| partitionBuilder.setNamespaceId(namespace); |
| } |
| return partitionBuilder; |
| } |
| |
| /** Builds a {@link RunQueryRequest} from the {@code query} and {@code namespace}. */ |
| static RunQueryRequest makeRequest(Query query, @Nullable String namespace) { |
| return RunQueryRequest.newBuilder() |
| .setQuery(query) |
| .setPartitionId(forNamespace(namespace)) |
| .build(); |
| } |
| |
| @VisibleForTesting |
| /** Builds a {@link RunQueryRequest} from the {@code GqlQuery} and {@code namespace}. */ |
| static RunQueryRequest makeRequest(GqlQuery gqlQuery, @Nullable String namespace) { |
| return RunQueryRequest.newBuilder() |
| .setGqlQuery(gqlQuery) |
| .setPartitionId(forNamespace(namespace)) |
| .build(); |
| } |
| |
| /** |
| * A helper function to get the split queries, taking into account the optional {@code |
| * namespace}. |
| */ |
| private static List<Query> splitQuery( |
| Query query, |
| @Nullable String namespace, |
| Datastore datastore, |
| QuerySplitter querySplitter, |
| int numSplits) |
| throws DatastoreException { |
| // If namespace is set, include it in the split request so splits are calculated accordingly. |
| return querySplitter.getSplits(query, forNamespace(namespace).build(), numSplits, datastore); |
| } |
| |
| /** |
| * Translates a Cloud Datastore gql query string to {@link Query}. |
| * |
| * <p>Currently, the only way to translate a gql query string to a Query is to run the query |
| * against Cloud Datastore and extract the {@code Query} from the response. To prevent reading |
| * any data, we set the {@code LIMIT} to 0 but if the gql query already has a limit set, we |
| * catch the exception with {@code INVALID_ARGUMENT} error code and retry the translation |
| * without the zero limit. |
| * |
| * <p>Note: This may result in reading actual data from Cloud Datastore but the service has a |
| * cap on the number of entities returned for a single rpc request, so this should not be a |
| * problem in practice. |
| */ |
| @VisibleForTesting |
| static Query translateGqlQueryWithLimitCheck(String gql, Datastore datastore, String namespace) |
| throws DatastoreException { |
| String gqlQueryWithZeroLimit = gql + " LIMIT 0"; |
| try { |
| Query translatedQuery = translateGqlQuery(gqlQueryWithZeroLimit, datastore, namespace); |
| // Clear the limit that we set. |
| return translatedQuery.toBuilder().clearLimit().build(); |
| } catch (DatastoreException e) { |
| // Note: There is no specific error code or message to detect if the query already has a |
| // limit, so we just check for INVALID_ARGUMENT and assume that that the query might have |
| // a limit already set. |
| if (e.getCode() == Code.INVALID_ARGUMENT) { |
| LOG.warn("Failed to translate Gql query '{}': {}", gqlQueryWithZeroLimit, e.getMessage()); |
| LOG.warn("User query might have a limit already set, so trying without zero limit"); |
| // Retry without the zero limit. |
| return translateGqlQuery(gql, datastore, namespace); |
| } else { |
| throw e; |
| } |
| } |
| } |
| |
| /** Translates a gql query string to {@link Query}. */ |
| private static Query translateGqlQuery(String gql, Datastore datastore, String namespace) |
| throws DatastoreException { |
| GqlQuery gqlQuery = GqlQuery.newBuilder().setQueryString(gql).setAllowLiterals(true).build(); |
| RunQueryRequest req = makeRequest(gqlQuery, namespace); |
| return datastore.runQuery(req).getQuery(); |
| } |
| |
| /** |
| * Returns a new {@link DatastoreV1.Read} that reads from the Cloud Datastore for the specified |
| * project. |
| */ |
| public DatastoreV1.Read withProjectId(String projectId) { |
| checkArgument(projectId != null, "projectId can not be null"); |
| return toBuilder().setProjectId(StaticValueProvider.of(projectId)).build(); |
| } |
| |
| /** Same as {@link Read#withProjectId(String)} but with a {@link ValueProvider}. */ |
| public DatastoreV1.Read withProjectId(ValueProvider<String> projectId) { |
| checkArgument(projectId != null, "projectId can not be null"); |
| return toBuilder().setProjectId(projectId).build(); |
| } |
| |
| /** |
| * Returns a new {@link DatastoreV1.Read} that reads the results of the specified query. |
| * |
| * <p><b>Note:</b> Normally, {@code DatastoreIO} will read from Cloud Datastore in parallel |
| * across many workers. However, when the {@link Query} is configured with a limit using {@link |
| * Query.Builder#setLimit}, then all results will be read by a single worker in order to ensure |
| * correct results. |
| */ |
| public DatastoreV1.Read withQuery(Query query) { |
| checkArgument(query != null, "query can not be null"); |
| checkArgument( |
| !query.hasLimit() || query.getLimit().getValue() > 0, |
| "Invalid query limit %s: must be positive", |
| query.getLimit().getValue()); |
| return toBuilder().setQuery(query).build(); |
| } |
| |
| /** |
| * Returns a new {@link DatastoreV1.Read} that reads the results of the specified GQL query. See |
| * <a href="https://cloud.google.com/datastore/docs/reference/gql_reference">GQL Reference </a> |
| * to know more about GQL grammar. |
| * |
| * <p><b><i>Note:</i></b> This query is executed with literals allowed, so the users should |
| * ensure that the query is originated from trusted sources to avoid any security |
| * vulnerabilities via SQL Injection. |
| * |
| * <p><b><i>Experimental</i></b>: Cloud Datastore does not a provide a clean way to translate a |
| * gql query string to {@link Query}, so we end up making a query to the service for translation |
| * but this may read the actual data, although it will be a small amount. It needs more |
| * validation through production use cases before marking it as stable. |
| */ |
| @Experimental(Kind.SOURCE_SINK) |
| public DatastoreV1.Read withLiteralGqlQuery(String gqlQuery) { |
| checkArgument(gqlQuery != null, "gqlQuery can not be null"); |
| return toBuilder().setLiteralGqlQuery(StaticValueProvider.of(gqlQuery)).build(); |
| } |
| |
| /** Same as {@link Read#withLiteralGqlQuery(String)} but with a {@link ValueProvider}. */ |
| @Experimental(Kind.SOURCE_SINK) |
| public DatastoreV1.Read withLiteralGqlQuery(ValueProvider<String> gqlQuery) { |
| checkArgument(gqlQuery != null, "gqlQuery can not be null"); |
| if (gqlQuery.isAccessible()) { |
| checkArgument(gqlQuery.get() != null, "gqlQuery can not be null"); |
| } |
| return toBuilder().setLiteralGqlQuery(gqlQuery).build(); |
| } |
| |
| /** Returns a new {@link DatastoreV1.Read} that reads from the given namespace. */ |
| public DatastoreV1.Read withNamespace(String namespace) { |
| return toBuilder().setNamespace(StaticValueProvider.of(namespace)).build(); |
| } |
| |
| /** Same as {@link Read#withNamespace(String)} but with a {@link ValueProvider}. */ |
| public DatastoreV1.Read withNamespace(ValueProvider<String> namespace) { |
| return toBuilder().setNamespace(namespace).build(); |
| } |
| |
| /** |
| * Returns a new {@link DatastoreV1.Read} that reads by splitting the given {@code query} into |
| * {@code numQuerySplits}. |
| * |
| * <p>The semantics for the query splitting is defined below: |
| * |
| * <ul> |
| * <li>Any value less than or equal to 0 will be ignored, and the number of splits will be |
| * chosen dynamically at runtime based on the query data size. |
| * <li>Any value greater than {@link Read#NUM_QUERY_SPLITS_MAX} will be capped at {@code |
| * NUM_QUERY_SPLITS_MAX}. |
| * <li>If the {@code query} has a user limit set, then {@code numQuerySplits} will be ignored |
| * and no split will be performed. |
| * <li>Under certain cases Cloud Datastore is unable to split query to the requested number of |
| * splits. In such cases we just use whatever the Cloud Datastore returns. |
| * </ul> |
| */ |
| public DatastoreV1.Read withNumQuerySplits(int numQuerySplits) { |
| return toBuilder() |
| .setNumQuerySplits(Math.min(Math.max(numQuerySplits, 0), NUM_QUERY_SPLITS_MAX)) |
| .build(); |
| } |
| |
| /** |
| * Returns a new {@link DatastoreV1.Read} that reads from a Datastore Emulator running at the |
| * given localhost address. |
| */ |
| public DatastoreV1.Read withLocalhost(String localhost) { |
| return toBuilder().setLocalhost(localhost).build(); |
| } |
| |
| @Override |
| public PCollection<Entity> expand(PBegin input) { |
| checkArgument(getProjectId() != null, "projectId provider cannot be null"); |
| if (getProjectId().isAccessible()) { |
| checkArgument(getProjectId().get() != null, "projectId cannot be null"); |
| } |
| |
| checkArgument( |
| getQuery() != null || getLiteralGqlQuery() != null, |
| "Either withQuery() or withLiteralGqlQuery() is required"); |
| checkArgument( |
| getQuery() == null || getLiteralGqlQuery() == null, |
| "withQuery() and withLiteralGqlQuery() are exclusive"); |
| |
| V1Options v1Options = V1Options.from(getProjectId(), getNamespace(), getLocalhost()); |
| |
| /* |
| * This composite transform involves the following steps: |
| * 1. Create a singleton of the user provided {@code query} or if {@code gqlQuery} is |
| * provided apply a {@link ParDo} that translates the {@code gqlQuery} into a {@code query}. |
| * |
| * 2. A {@link ParDo} splits the resulting query into {@code numQuerySplits} and |
| * assign each split query a unique {@code Integer} as the key. The resulting output is |
| * of the type {@code PCollection<KV<Integer, Query>>}. |
| * |
| * If the value of {@code numQuerySplits} is less than or equal to 0, then the number of |
| * splits will be computed dynamically based on the size of the data for the {@code query}. |
| * |
| * 3. The resulting {@code PCollection} is sharded using a {@link GroupByKey} operation. The |
| * queries are extracted from they {@code KV<Integer, Iterable<Query>>} and flattened to |
| * output a {@code PCollection<Query>}. |
| * |
| * 4. In the third step, a {@code ParDo} reads entities for each query and outputs |
| * a {@code PCollection<Entity>}. |
| */ |
| |
| PCollection<Query> inputQuery; |
| if (getQuery() != null) { |
| inputQuery = input.apply(Create.of(getQuery())); |
| } else { |
| inputQuery = |
| input |
| .apply(Create.ofProvider(getLiteralGqlQuery(), StringUtf8Coder.of())) |
| .apply(ParDo.of(new GqlQueryTranslateFn(v1Options))); |
| } |
| |
| return inputQuery |
| .apply("Split", ParDo.of(new SplitQueryFn(v1Options, getNumQuerySplits()))) |
| .apply("Reshuffle", Reshuffle.viaRandomKey()) |
| .apply("Read", ParDo.of(new ReadFn(v1Options))); |
| } |
| |
| @Override |
| public void populateDisplayData(DisplayData.Builder builder) { |
| super.populateDisplayData(builder); |
| String query = getQuery() == null ? null : getQuery().toString(); |
| builder |
| .addIfNotNull(DisplayData.item("projectId", getProjectId()).withLabel("ProjectId")) |
| .addIfNotNull(DisplayData.item("namespace", getNamespace()).withLabel("Namespace")) |
| .addIfNotNull(DisplayData.item("query", query).withLabel("Query")) |
| .addIfNotNull(DisplayData.item("gqlQuery", getLiteralGqlQuery()).withLabel("GqlQuery")); |
| } |
| |
| @VisibleForTesting |
| static class V1Options implements HasDisplayData, Serializable { |
| private final ValueProvider<String> project; |
| @Nullable private final ValueProvider<String> namespace; |
| @Nullable private final String localhost; |
| |
| private V1Options( |
| ValueProvider<String> project, ValueProvider<String> namespace, String localhost) { |
| this.project = project; |
| this.namespace = namespace; |
| this.localhost = localhost; |
| } |
| |
| public static V1Options from(String projectId, String namespace, String localhost) { |
| return from( |
| StaticValueProvider.of(projectId), StaticValueProvider.of(namespace), localhost); |
| } |
| |
| public static V1Options from( |
| ValueProvider<String> project, ValueProvider<String> namespace, String localhost) { |
| return new V1Options(project, namespace, localhost); |
| } |
| |
| public String getProjectId() { |
| return project.get(); |
| } |
| |
| @Nullable |
| public String getNamespace() { |
| return namespace == null ? null : namespace.get(); |
| } |
| |
| public ValueProvider<String> getProjectValueProvider() { |
| return project; |
| } |
| |
| @Nullable |
| public ValueProvider<String> getNamespaceValueProvider() { |
| return namespace; |
| } |
| |
| @Nullable |
| public String getLocalhost() { |
| return localhost; |
| } |
| |
| @Override |
| public void populateDisplayData(DisplayData.Builder builder) { |
| builder |
| .addIfNotNull( |
| DisplayData.item("projectId", getProjectValueProvider()).withLabel("ProjectId")) |
| .addIfNotNull( |
| DisplayData.item("namespace", getNamespaceValueProvider()).withLabel("Namespace")); |
| } |
| } |
| |
| /** A DoFn that translates a Cloud Datastore gql query string to {@code Query}. */ |
| static class GqlQueryTranslateFn extends DoFn<String, Query> { |
| private final V1Options v1Options; |
| private transient Datastore datastore; |
| private final V1DatastoreFactory datastoreFactory; |
| |
| GqlQueryTranslateFn(V1Options options) { |
| this(options, new V1DatastoreFactory()); |
| } |
| |
| GqlQueryTranslateFn(V1Options options, V1DatastoreFactory datastoreFactory) { |
| this.v1Options = options; |
| this.datastoreFactory = datastoreFactory; |
| } |
| |
| @StartBundle |
| public void startBundle(StartBundleContext c) throws Exception { |
| datastore = |
| datastoreFactory.getDatastore( |
| c.getPipelineOptions(), v1Options.getProjectId(), v1Options.getLocalhost()); |
| } |
| |
| @ProcessElement |
| public void processElement(ProcessContext c) throws Exception { |
| String gqlQuery = c.element(); |
| LOG.info("User query: '{}'", gqlQuery); |
| Query query = |
| translateGqlQueryWithLimitCheck(gqlQuery, datastore, v1Options.getNamespace()); |
| LOG.info("User gql query translated to Query({})", query); |
| c.output(query); |
| } |
| } |
| |
| /** |
| * A {@link DoFn} that splits a given query into multiple sub-queries, assigns them unique keys |
| * and outputs them as {@link KV}. |
| */ |
| @VisibleForTesting |
| static class SplitQueryFn extends DoFn<Query, Query> { |
| private final V1Options options; |
| // number of splits to make for a given query |
| private final int numSplits; |
| |
| private final V1DatastoreFactory datastoreFactory; |
| // Datastore client |
| private transient Datastore datastore; |
| // Query splitter |
| private transient QuerySplitter querySplitter; |
| |
| public SplitQueryFn(V1Options options, int numSplits) { |
| this(options, numSplits, new V1DatastoreFactory()); |
| } |
| |
| @VisibleForTesting |
| SplitQueryFn(V1Options options, int numSplits, V1DatastoreFactory datastoreFactory) { |
| this.options = options; |
| this.numSplits = numSplits; |
| this.datastoreFactory = datastoreFactory; |
| } |
| |
| @StartBundle |
| public void startBundle(StartBundleContext c) throws Exception { |
| datastore = |
| datastoreFactory.getDatastore( |
| c.getPipelineOptions(), options.getProjectId(), options.getLocalhost()); |
| querySplitter = datastoreFactory.getQuerySplitter(); |
| } |
| |
| @ProcessElement |
| public void processElement(ProcessContext c) throws Exception { |
| Query query = c.element(); |
| |
| // If query has a user set limit, then do not split. |
| if (query.hasLimit()) { |
| c.output(query); |
| return; |
| } |
| |
| int estimatedNumSplits; |
| // Compute the estimated numSplits if numSplits is not specified by the user. |
| if (numSplits <= 0) { |
| estimatedNumSplits = getEstimatedNumSplits(datastore, query, options.getNamespace()); |
| } else { |
| estimatedNumSplits = numSplits; |
| } |
| |
| LOG.info("Splitting the query into {} splits", estimatedNumSplits); |
| List<Query> querySplits; |
| try { |
| querySplits = |
| splitQuery( |
| query, options.getNamespace(), datastore, querySplitter, estimatedNumSplits); |
| } catch (Exception e) { |
| LOG.warn("Unable to parallelize the given query: {}", query, e); |
| querySplits = ImmutableList.of(query); |
| } |
| |
| // assign unique keys to query splits. |
| for (Query subquery : querySplits) { |
| c.output(subquery); |
| } |
| } |
| |
| @Override |
| public void populateDisplayData(DisplayData.Builder builder) { |
| super.populateDisplayData(builder); |
| builder.include("options", options); |
| if (numSplits > 0) { |
| builder.add( |
| DisplayData.item("numQuerySplits", numSplits) |
| .withLabel("Requested number of Query splits")); |
| } |
| } |
| } |
| |
| /** A {@link DoFn} that reads entities from Cloud Datastore for each query. */ |
| @VisibleForTesting |
| static class ReadFn extends DoFn<Query, Entity> { |
| private final V1Options options; |
| private final V1DatastoreFactory datastoreFactory; |
| // Datastore client |
| private transient Datastore datastore; |
| private final Counter rpcErrors = |
| Metrics.counter(DatastoreWriterFn.class, "datastoreRpcErrors"); |
| private final Counter rpcSuccesses = |
| Metrics.counter(DatastoreWriterFn.class, "datastoreRpcSuccesses"); |
| private static final int MAX_RETRIES = 5; |
| private static final FluentBackoff RUNQUERY_BACKOFF = |
| FluentBackoff.DEFAULT |
| .withMaxRetries(MAX_RETRIES) |
| .withInitialBackoff(Duration.standardSeconds(5)); |
| |
| public ReadFn(V1Options options) { |
| this(options, new V1DatastoreFactory()); |
| } |
| |
| @VisibleForTesting |
| ReadFn(V1Options options, V1DatastoreFactory datastoreFactory) { |
| this.options = options; |
| this.datastoreFactory = datastoreFactory; |
| } |
| |
| @StartBundle |
| public void startBundle(StartBundleContext c) throws Exception { |
| datastore = |
| datastoreFactory.getDatastore( |
| c.getPipelineOptions(), options.getProjectId(), options.getLocalhost()); |
| } |
| |
| private RunQueryResponse runQueryWithRetries(RunQueryRequest request) throws Exception { |
| Sleeper sleeper = Sleeper.DEFAULT; |
| BackOff backoff = RUNQUERY_BACKOFF.backoff(); |
| while (true) { |
| try { |
| RunQueryResponse response = datastore.runQuery(request); |
| rpcSuccesses.inc(); |
| return response; |
| } catch (DatastoreException exception) { |
| rpcErrors.inc(); |
| |
| if (NON_RETRYABLE_ERRORS.contains(exception.getCode())) { |
| throw exception; |
| } |
| if (!BackOffUtils.next(sleeper, backoff)) { |
| LOG.error("Aborting after {} retries.", MAX_RETRIES); |
| throw exception; |
| } |
| } |
| } |
| } |
| |
| /** Read and output entities for the given query. */ |
| @ProcessElement |
| public void processElement(ProcessContext context) throws Exception { |
| Query query = context.element(); |
| String namespace = options.getNamespace(); |
| int userLimit = query.hasLimit() ? query.getLimit().getValue() : Integer.MAX_VALUE; |
| |
| boolean moreResults = true; |
| QueryResultBatch currentBatch = null; |
| |
| while (moreResults) { |
| Query.Builder queryBuilder = query.toBuilder(); |
| queryBuilder.setLimit( |
| Int32Value.newBuilder().setValue(Math.min(userLimit, QUERY_BATCH_LIMIT))); |
| |
| if (currentBatch != null && !currentBatch.getEndCursor().isEmpty()) { |
| queryBuilder.setStartCursor(currentBatch.getEndCursor()); |
| } |
| |
| RunQueryRequest request = makeRequest(queryBuilder.build(), namespace); |
| RunQueryResponse response = runQueryWithRetries(request); |
| |
| currentBatch = response.getBatch(); |
| |
| // MORE_RESULTS_AFTER_LIMIT is not implemented yet: |
| // https://groups.google.com/forum/#!topic/gcd-discuss/iNs6M1jA2Vw, so |
| // use result count to determine if more results might exist. |
| int numFetch = currentBatch.getEntityResultsCount(); |
| if (query.hasLimit()) { |
| verify( |
| userLimit >= numFetch, |
| "Expected userLimit %s >= numFetch %s, because query limit %s must be <= userLimit", |
| userLimit, |
| numFetch, |
| query.getLimit()); |
| userLimit -= numFetch; |
| } |
| |
| // output all the entities from the current batch. |
| for (EntityResult entityResult : currentBatch.getEntityResultsList()) { |
| context.output(entityResult.getEntity()); |
| } |
| |
| // Check if we have more entities to be read. |
| moreResults = |
| // User-limit does not exist (so userLimit == MAX_VALUE) and/or has not been satisfied |
| (userLimit > 0) |
| // All indications from the API are that there are/may be more results. |
| && ((numFetch == QUERY_BATCH_LIMIT) |
| || (currentBatch.getMoreResults() == NOT_FINISHED)); |
| } |
| } |
| |
| @Override |
| public void populateDisplayData(DisplayData.Builder builder) { |
| super.populateDisplayData(builder); |
| builder.include("options", options); |
| } |
| } |
| } |
| |
| /** |
| * Returns an empty {@link DatastoreV1.Write} builder. Configure the destination {@code projectId} |
| * using {@link DatastoreV1.Write#withProjectId}. |
| */ |
| public Write write() { |
| return new Write(null, null); |
| } |
| |
| /** |
| * Returns an empty {@link DeleteEntity} builder. Configure the destination {@code projectId} |
| * using {@link DeleteEntity#withProjectId}. |
| */ |
| public DeleteEntity deleteEntity() { |
| return new DeleteEntity(null, null); |
| } |
| |
| /** |
| * Returns an empty {@link DeleteKey} builder. Configure the destination {@code projectId} using |
| * {@link DeleteKey#withProjectId}. |
| */ |
| public DeleteKey deleteKey() { |
| return new DeleteKey(null, null); |
| } |
| |
| /** |
| * A {@link PTransform} that writes {@link Entity} objects to Cloud Datastore. |
| * |
| * @see DatastoreIO |
| */ |
| public static class Write extends Mutate<Entity> { |
| /** |
| * Note that {@code projectId} is only {@code @Nullable} as a matter of build order, but if it |
| * is {@code null} at instantiation time, an error will be thrown. |
| */ |
| Write(@Nullable ValueProvider<String> projectId, @Nullable String localhost) { |
| super(projectId, localhost, new UpsertFn()); |
| } |
| |
| /** Returns a new {@link Write} that writes to the Cloud Datastore for the specified project. */ |
| public Write withProjectId(String projectId) { |
| checkArgument(projectId != null, "projectId can not be null"); |
| return withProjectId(StaticValueProvider.of(projectId)); |
| } |
| |
| /** Same as {@link Write#withProjectId(String)} but with a {@link ValueProvider}. */ |
| public Write withProjectId(ValueProvider<String> projectId) { |
| checkArgument(projectId != null, "projectId can not be null"); |
| return new Write(projectId, localhost); |
| } |
| |
| /** |
| * Returns a new {@link Write} that writes to the Cloud Datastore Emulator running locally on |
| * the specified host port. |
| */ |
| public Write withLocalhost(String localhost) { |
| checkArgument(localhost != null, "localhost can not be null"); |
| return new Write(projectId, localhost); |
| } |
| } |
| |
| /** |
| * A {@link PTransform} that deletes {@link Entity Entities} from Cloud Datastore. |
| * |
| * @see DatastoreIO |
| */ |
| public static class DeleteEntity extends Mutate<Entity> { |
| /** |
| * Note that {@code projectId} is only {@code @Nullable} as a matter of build order, but if it |
| * is {@code null} at instantiation time, an error will be thrown. |
| */ |
| DeleteEntity(@Nullable ValueProvider<String> projectId, @Nullable String localhost) { |
| super(projectId, localhost, new DeleteEntityFn()); |
| } |
| |
| /** |
| * Returns a new {@link DeleteEntity} that deletes entities from the Cloud Datastore for the |
| * specified project. |
| */ |
| public DeleteEntity withProjectId(String projectId) { |
| checkArgument(projectId != null, "projectId can not be null"); |
| return withProjectId(StaticValueProvider.of(projectId)); |
| } |
| |
| /** Same as {@link DeleteEntity#withProjectId(String)} but with a {@link ValueProvider}. */ |
| public DeleteEntity withProjectId(ValueProvider<String> projectId) { |
| checkArgument(projectId != null, "projectId can not be null"); |
| return new DeleteEntity(projectId, localhost); |
| } |
| |
| /** |
| * Returns a new {@link DeleteEntity} that deletes entities from the Cloud Datastore Emulator |
| * running locally on the specified host port. |
| */ |
| public DeleteEntity withLocalhost(String localhost) { |
| checkArgument(localhost != null, "localhost can not be null"); |
| return new DeleteEntity(projectId, localhost); |
| } |
| } |
| |
| /** |
| * A {@link PTransform} that deletes {@link Entity Entities} associated with the given {@link Key |
| * Keys} from Cloud Datastore. |
| * |
| * @see DatastoreIO |
| */ |
| public static class DeleteKey extends Mutate<Key> { |
| /** |
| * Note that {@code projectId} is only {@code @Nullable} as a matter of build order, but if it |
| * is {@code null} at instantiation time, an error will be thrown. |
| */ |
| DeleteKey(@Nullable ValueProvider<String> projectId, @Nullable String localhost) { |
| super(projectId, localhost, new DeleteKeyFn()); |
| } |
| |
| /** |
| * Returns a new {@link DeleteKey} that deletes entities from the Cloud Datastore for the |
| * specified project. |
| */ |
| public DeleteKey withProjectId(String projectId) { |
| checkArgument(projectId != null, "projectId can not be null"); |
| return withProjectId(StaticValueProvider.of(projectId)); |
| } |
| |
| /** |
| * Returns a new {@link DeleteKey} that deletes entities from the Cloud Datastore Emulator |
| * running locally on the specified host port. |
| */ |
| public DeleteKey withLocalhost(String localhost) { |
| checkArgument(localhost != null, "localhost can not be null"); |
| return new DeleteKey(projectId, localhost); |
| } |
| |
| /** Same as {@link DeleteKey#withProjectId(String)} but with a {@link ValueProvider}. */ |
| public DeleteKey withProjectId(ValueProvider<String> projectId) { |
| checkArgument(projectId != null, "projectId can not be null"); |
| return new DeleteKey(projectId, localhost); |
| } |
| } |
| |
| /** |
| * A {@link PTransform} that writes mutations to Cloud Datastore. |
| * |
| * <p>It requires a {@link DoFn} that tranforms an object of type {@code T} to a {@link Mutation}. |
| * {@code T} is usually either an {@link Entity} or a {@link Key} <b>Note:</b> Only idempotent |
| * Cloud Datastore mutation operations (upsert and delete) should be used by the {@code DoFn} |
| * provided, as the commits are retried when failures occur. |
| */ |
| private abstract static class Mutate<T> extends PTransform<PCollection<T>, PDone> { |
| protected ValueProvider<String> projectId; |
| @Nullable protected String localhost; |
| /** A function that transforms each {@code T} into a mutation. */ |
| private final SimpleFunction<T, Mutation> mutationFn; |
| |
| /** |
| * Note that {@code projectId} is only {@code @Nullable} as a matter of build order, but if it |
| * is {@code null} at instantiation time, an error will be thrown. |
| */ |
| Mutate( |
| @Nullable ValueProvider<String> projectId, |
| @Nullable String localhost, |
| SimpleFunction<T, Mutation> mutationFn) { |
| this.projectId = projectId; |
| this.localhost = localhost; |
| this.mutationFn = checkNotNull(mutationFn); |
| } |
| |
| @Override |
| public PDone expand(PCollection<T> input) { |
| checkArgument(projectId != null, "withProjectId() is required"); |
| if (projectId.isAccessible()) { |
| checkArgument(projectId.get() != null, "projectId can not be null"); |
| } |
| checkArgument(mutationFn != null, "mutationFn can not be null"); |
| |
| input |
| .apply("Convert to Mutation", MapElements.via(mutationFn)) |
| .apply( |
| "Write Mutation to Datastore", ParDo.of(new DatastoreWriterFn(projectId, localhost))); |
| |
| return PDone.in(input.getPipeline()); |
| } |
| |
| @Override |
| public String toString() { |
| return MoreObjects.toStringHelper(getClass()) |
| .add("projectId", projectId) |
| .add("mutationFn", mutationFn.getClass().getName()) |
| .toString(); |
| } |
| |
| @Override |
| public void populateDisplayData(DisplayData.Builder builder) { |
| super.populateDisplayData(builder); |
| builder |
| .addIfNotNull(DisplayData.item("projectId", projectId).withLabel("Output Project")) |
| .include("mutationFn", mutationFn); |
| } |
| |
| public String getProjectId() { |
| return projectId.get(); |
| } |
| } |
| |
| /** Determines batch sizes for commit RPCs. */ |
| @VisibleForTesting |
| interface WriteBatcher { |
| /** Call before using this WriteBatcher. */ |
| void start(); |
| |
| /** |
| * Reports the latency of a previous commit RPC, and the number of mutations that it contained. |
| */ |
| void addRequestLatency(long timeSinceEpochMillis, long latencyMillis, int numMutations); |
| |
| /** Returns the number of entities to include in the next CommitRequest. */ |
| int nextBatchSize(long timeSinceEpochMillis); |
| } |
| |
| /** |
| * Determines batch sizes for commit RPCs based on past performance. |
| * |
| * <p>It aims for a target response time per RPC: it uses the response times for previous RPCs and |
| * the number of entities contained in them, calculates a rolling average time-per-entity, and |
| * chooses the number of entities for future writes to hit the target time. |
| * |
| * <p>This enables us to send large batches without sending over-large requests in the case of |
| * expensive entity writes that may timeout before the server can apply them all. |
| */ |
| @VisibleForTesting |
| static class WriteBatcherImpl implements WriteBatcher, Serializable { |
| /** Target time per RPC for writes. */ |
| static final int DATASTORE_BATCH_TARGET_LATENCY_MS = 5000; |
| |
| @Override |
| public void start() { |
| meanLatencyPerEntityMs = |
| new MovingAverage( |
| 120000 /* sample period 2 minutes */, 10000 /* sample interval 10s */, |
| 1 /* numSignificantBuckets */, 1 /* numSignificantSamples */); |
| } |
| |
| @Override |
| public void addRequestLatency(long timeSinceEpochMillis, long latencyMillis, int numMutations) { |
| meanLatencyPerEntityMs.add(timeSinceEpochMillis, latencyMillis / numMutations); |
| } |
| |
| @Override |
| public int nextBatchSize(long timeSinceEpochMillis) { |
| if (!meanLatencyPerEntityMs.hasValue(timeSinceEpochMillis)) { |
| return DATASTORE_BATCH_UPDATE_ENTITIES_START; |
| } |
| long recentMeanLatency = Math.max(meanLatencyPerEntityMs.get(timeSinceEpochMillis), 1); |
| return (int) |
| Math.max( |
| DATASTORE_BATCH_UPDATE_ENTITIES_MIN, |
| Math.min( |
| DATASTORE_BATCH_UPDATE_ENTITIES_LIMIT, |
| DATASTORE_BATCH_TARGET_LATENCY_MS / recentMeanLatency)); |
| } |
| |
| private transient MovingAverage meanLatencyPerEntityMs; |
| } |
| |
| /** |
| * {@link DoFn} that writes {@link Mutation}s to Cloud Datastore. Mutations are written in |
| * batches; see {@link DatastoreV1.WriteBatcherImpl}. |
| * |
| * <p>See <a href="https://cloud.google.com/datastore/docs/concepts/entities">Datastore: Entities, |
| * Properties, and Keys</a> for information about entity keys and mutations. |
| * |
| * <p>Commits are non-transactional. If a commit fails because of a conflict over an entity group, |
| * the commit will be retried (up to {@link DatastoreV1.DatastoreWriterFn#MAX_RETRIES} times). |
| * This means that the mutation operation should be idempotent. Thus, the writer should only be |
| * used for {@code upsert} and {@code delete} mutation operations, as these are the only two Cloud |
| * Datastore mutations that are idempotent. |
| */ |
| @VisibleForTesting |
| static class DatastoreWriterFn extends DoFn<Mutation, Void> { |
| private static final Logger LOG = LoggerFactory.getLogger(DatastoreWriterFn.class); |
| private final ValueProvider<String> projectId; |
| @Nullable private final String localhost; |
| private transient Datastore datastore; |
| private final V1DatastoreFactory datastoreFactory; |
| // Current batch of mutations to be written. |
| private final List<Mutation> mutations = new ArrayList<>(); |
| private int mutationsSize = 0; // Accumulated size of protos in mutations. |
| private WriteBatcher writeBatcher; |
| private transient AdaptiveThrottler throttler; |
| private final Counter throttledSeconds = |
| Metrics.counter(DatastoreWriterFn.class, "cumulativeThrottlingSeconds"); |
| private final Counter rpcErrors = |
| Metrics.counter(DatastoreWriterFn.class, "datastoreRpcErrors"); |
| private final Counter rpcSuccesses = |
| Metrics.counter(DatastoreWriterFn.class, "datastoreRpcSuccesses"); |
| |
| private static final int MAX_RETRIES = 5; |
| private static final FluentBackoff BUNDLE_WRITE_BACKOFF = |
| FluentBackoff.DEFAULT |
| .withMaxRetries(MAX_RETRIES) |
| .withInitialBackoff(Duration.standardSeconds(5)); |
| |
| DatastoreWriterFn(String projectId, @Nullable String localhost) { |
| this( |
| StaticValueProvider.of(projectId), |
| localhost, |
| new V1DatastoreFactory(), |
| new WriteBatcherImpl()); |
| } |
| |
| DatastoreWriterFn(ValueProvider<String> projectId, @Nullable String localhost) { |
| this(projectId, localhost, new V1DatastoreFactory(), new WriteBatcherImpl()); |
| } |
| |
| @VisibleForTesting |
| DatastoreWriterFn( |
| ValueProvider<String> projectId, |
| @Nullable String localhost, |
| V1DatastoreFactory datastoreFactory, |
| WriteBatcher writeBatcher) { |
| this.projectId = checkNotNull(projectId, "projectId"); |
| this.localhost = localhost; |
| this.datastoreFactory = datastoreFactory; |
| this.writeBatcher = writeBatcher; |
| } |
| |
| @StartBundle |
| public void startBundle(StartBundleContext c) { |
| datastore = datastoreFactory.getDatastore(c.getPipelineOptions(), projectId.get(), localhost); |
| writeBatcher.start(); |
| if (throttler == null) { |
| // Initialize throttler at first use, because it is not serializable. |
| throttler = new AdaptiveThrottler(120000, 10000, 1.25); |
| } |
| } |
| |
| @ProcessElement |
| public void processElement(ProcessContext c) throws Exception { |
| Mutation write = c.element(); |
| int size = write.getSerializedSize(); |
| if (mutations.size() > 0 |
| && mutationsSize + size >= DatastoreV1.DATASTORE_BATCH_UPDATE_BYTES_LIMIT) { |
| flushBatch(); |
| } |
| mutations.add(c.element()); |
| mutationsSize += size; |
| if (mutations.size() >= writeBatcher.nextBatchSize(System.currentTimeMillis())) { |
| flushBatch(); |
| } |
| } |
| |
| @FinishBundle |
| public void finishBundle() throws Exception { |
| if (!mutations.isEmpty()) { |
| flushBatch(); |
| } |
| } |
| |
| /** |
| * Writes a batch of mutations to Cloud Datastore. |
| * |
| * <p>If a commit fails, it will be retried up to {@link #MAX_RETRIES} times. All mutations in |
| * the batch will be committed again, even if the commit was partially successful. If the retry |
| * limit is exceeded, the last exception from Cloud Datastore will be thrown. |
| * |
| * @throws DatastoreException if the commit fails or IOException or InterruptedException if |
| * backing off between retries fails. |
| */ |
| private void flushBatch() throws DatastoreException, IOException, InterruptedException { |
| LOG.debug("Writing batch of {} mutations", mutations.size()); |
| Sleeper sleeper = Sleeper.DEFAULT; |
| BackOff backoff = BUNDLE_WRITE_BACKOFF.backoff(); |
| |
| while (true) { |
| // Batch upsert entities. |
| CommitRequest.Builder commitRequest = CommitRequest.newBuilder(); |
| commitRequest.addAllMutations(mutations); |
| commitRequest.setMode(CommitRequest.Mode.NON_TRANSACTIONAL); |
| long startTime = System.currentTimeMillis(), endTime; |
| |
| if (throttler.throttleRequest(startTime)) { |
| LOG.info("Delaying request due to previous failures"); |
| throttledSeconds.inc(WriteBatcherImpl.DATASTORE_BATCH_TARGET_LATENCY_MS / 1000); |
| sleeper.sleep(WriteBatcherImpl.DATASTORE_BATCH_TARGET_LATENCY_MS); |
| continue; |
| } |
| |
| try { |
| datastore.commit(commitRequest.build()); |
| endTime = System.currentTimeMillis(); |
| |
| writeBatcher.addRequestLatency(endTime, endTime - startTime, mutations.size()); |
| throttler.successfulRequest(startTime); |
| rpcSuccesses.inc(); |
| |
| // Break if the commit threw no exception. |
| break; |
| } catch (DatastoreException exception) { |
| if (exception.getCode() == Code.DEADLINE_EXCEEDED) { |
| /* Most errors are not related to request size, and should not change our expectation of |
| * the latency of successful requests. DEADLINE_EXCEEDED can be taken into |
| * consideration, though. */ |
| endTime = System.currentTimeMillis(); |
| writeBatcher.addRequestLatency(endTime, endTime - startTime, mutations.size()); |
| } |
| // Only log the code and message for potentially-transient errors. The entire exception |
| // will be propagated upon the last retry. |
| LOG.error( |
| "Error writing batch of {} mutations to Datastore ({}): {}", |
| mutations.size(), |
| exception.getCode(), |
| exception.getMessage()); |
| rpcErrors.inc(); |
| |
| if (NON_RETRYABLE_ERRORS.contains(exception.getCode())) { |
| throw exception; |
| } |
| if (!BackOffUtils.next(sleeper, backoff)) { |
| LOG.error("Aborting after {} retries.", MAX_RETRIES); |
| throw exception; |
| } |
| } |
| } |
| LOG.debug("Successfully wrote {} mutations", mutations.size()); |
| mutations.clear(); |
| mutationsSize = 0; |
| } |
| |
| @Override |
| public void populateDisplayData(Builder builder) { |
| super.populateDisplayData(builder); |
| builder.addIfNotNull(DisplayData.item("projectId", projectId).withLabel("Output Project")); |
| } |
| } |
| |
| /** |
| * Returns true if a Cloud Datastore key is complete. A key is complete if its last element has |
| * either an id or a name. |
| */ |
| static boolean isValidKey(Key key) { |
| List<PathElement> elementList = key.getPathList(); |
| if (elementList.isEmpty()) { |
| return false; |
| } |
| PathElement lastElement = elementList.get(elementList.size() - 1); |
| return (lastElement.getId() != 0 || !lastElement.getName().isEmpty()); |
| } |
| |
| /** A function that constructs an upsert {@link Mutation} from an {@link Entity}. */ |
| @VisibleForTesting |
| static class UpsertFn extends SimpleFunction<Entity, Mutation> { |
| @Override |
| public Mutation apply(Entity entity) { |
| // Verify that the entity to write has a complete key. |
| checkArgument( |
| isValidKey(entity.getKey()), |
| "Entities to be written to the Cloud Datastore must have complete keys:\n%s", |
| entity); |
| |
| return makeUpsert(entity).build(); |
| } |
| |
| @Override |
| public void populateDisplayData(Builder builder) { |
| builder.add( |
| DisplayData.item("upsertFn", this.getClass()).withLabel("Create Upsert Mutation")); |
| } |
| } |
| |
| /** A function that constructs a delete {@link Mutation} from an {@link Entity}. */ |
| @VisibleForTesting |
| static class DeleteEntityFn extends SimpleFunction<Entity, Mutation> { |
| @Override |
| public Mutation apply(Entity entity) { |
| // Verify that the entity to delete has a complete key. |
| checkArgument( |
| isValidKey(entity.getKey()), |
| "Entities to be deleted from the Cloud Datastore must have complete keys:\n%s", |
| entity); |
| |
| return makeDelete(entity.getKey()).build(); |
| } |
| |
| @Override |
| public void populateDisplayData(Builder builder) { |
| builder.add( |
| DisplayData.item("deleteEntityFn", this.getClass()).withLabel("Create Delete Mutation")); |
| } |
| } |
| |
| /** A function that constructs a delete {@link Mutation} from a {@link Key}. */ |
| @VisibleForTesting |
| static class DeleteKeyFn extends SimpleFunction<Key, Mutation> { |
| @Override |
| public Mutation apply(Key key) { |
| // Verify that the entity to delete has a complete key. |
| checkArgument( |
| isValidKey(key), |
| "Keys to be deleted from the Cloud Datastore must be complete:\n%s", |
| key); |
| |
| return makeDelete(key).build(); |
| } |
| |
| @Override |
| public void populateDisplayData(Builder builder) { |
| builder.add( |
| DisplayData.item("deleteKeyFn", this.getClass()).withLabel("Create Delete Mutation")); |
| } |
| } |
| |
| /** |
| * A wrapper factory class for Cloud Datastore singleton classes {@link DatastoreFactory} and |
| * {@link QuerySplitter} |
| * |
| * <p>{@link DatastoreFactory} and {@link QuerySplitter} are not java serializable, hence wrapping |
| * them under this class, which implements {@link Serializable}. |
| */ |
| @VisibleForTesting |
| static class V1DatastoreFactory implements Serializable { |
| |
| /** Builds a Cloud Datastore client for the given pipeline options and project. */ |
| public Datastore getDatastore(PipelineOptions pipelineOptions, String projectId) { |
| return getDatastore(pipelineOptions, projectId, null); |
| } |
| |
| /** |
| * Builds a Cloud Datastore client for the given pipeline options, project and an optional |
| * locahost. |
| */ |
| public Datastore getDatastore( |
| PipelineOptions pipelineOptions, String projectId, @Nullable String localhost) { |
| Credentials credential = pipelineOptions.as(GcpOptions.class).getGcpCredential(); |
| HttpRequestInitializer initializer; |
| if (credential != null) { |
| initializer = |
| new ChainingHttpRequestInitializer( |
| new HttpCredentialsAdapter(credential), new RetryHttpRequestInitializer()); |
| } else { |
| initializer = new RetryHttpRequestInitializer(); |
| } |
| |
| DatastoreOptions.Builder builder = |
| new DatastoreOptions.Builder().projectId(projectId).initializer(initializer); |
| |
| if (localhost != null) { |
| builder.localHost(localhost); |
| } else { |
| builder.host("batch-datastore.googleapis.com"); |
| } |
| |
| return DatastoreFactory.get().create(builder.build()); |
| } |
| |
| /** Builds a Cloud Datastore {@link QuerySplitter}. */ |
| public QuerySplitter getQuerySplitter() { |
| return DatastoreHelper.getQuerySplitter(); |
| } |
| } |
| } |
