Search in sources :

Example 16 with Key

use of com.google.recaptchaenterprise.v1beta1.Key in project beam by apache.

the class V1WriteIT method testDatastoreWriterFnWithDuplicatedEntities.

/**
 * Tests {@link DatastoreV1.DatastoreWriterFn} with duplicated entries. Once a duplicated entry is
 * found the batch gets flushed.
 */
@Test
public void testDatastoreWriterFnWithDuplicatedEntities() throws Exception {
    List<Mutation> mutations = new ArrayList<>(200);
    V1TestOptions options = TestPipeline.testingPipelineOptions().as(V1TestOptions.class);
    Pipeline pipeline = TestPipeline.create(options);
    for (int i = 1; i <= 200; i++) {
        Key key = makeKey("key" + i, i + 1).build();
        mutations.add(makeUpsert(Entity.newBuilder().setKey(key).build()).build());
        if (i % 30 == 0) {
            mutations.add(makeUpsert(Entity.newBuilder().setKey(key).build()).build());
        }
    }
    DatastoreV1.DatastoreWriterFn datastoreWriter = new DatastoreV1.DatastoreWriterFn(TestPipeline.testingPipelineOptions().as(GcpOptions.class).getProject(), null);
    PTransform<PCollection<? extends Mutation>, PCollection<Void>> datastoreWriterTransform = ParDo.of(datastoreWriter);
    /**
     * Following three lines turn the original arrayList into a member of the first PCollection
     */
    List<Mutation> newArrayList = new ArrayList<>(mutations);
    Create.Values<Iterable<Mutation>> mutationIterable = Create.of(Collections.singleton(newArrayList));
    PCollection<Iterable<Mutation>> input = pipeline.apply(mutationIterable);
    /**
     * Flatten divides the PCollection into several elements of the same bundle. By doing this we're
     * forcing the processing of the List of mutation in the same order the mutations were added to
     * the original List.
     */
    input.apply(Flatten.<Mutation>iterables()).apply(datastoreWriterTransform);
    PipelineResult pResult = pipeline.run();
    MetricQueryResults metricResults = pResult.metrics().queryMetrics(MetricsFilter.builder().addNameFilter(MetricNameFilter.named(DatastoreV1.DatastoreWriterFn.class, "batchSize")).build());
    AtomicLong timesCommitted = new AtomicLong();
    metricResults.getDistributions().forEach(distribution -> {
        if (distribution.getName().getName().equals("batchSize")) {
            timesCommitted.set(distribution.getCommitted().getCount());
        }
    });
    assertEquals(7, timesCommitted.get());
}
Also used : ArrayList(java.util.ArrayList) MetricQueryResults(org.apache.beam.sdk.metrics.MetricQueryResults) PipelineResult(org.apache.beam.sdk.PipelineResult) TestPipeline(org.apache.beam.sdk.testing.TestPipeline) Pipeline(org.apache.beam.sdk.Pipeline) PCollection(org.apache.beam.sdk.values.PCollection) AtomicLong(java.util.concurrent.atomic.AtomicLong) Create(org.apache.beam.sdk.transforms.Create) Mutation(com.google.datastore.v1.Mutation) Key(com.google.datastore.v1.Key) DatastoreHelper.makeKey(com.google.datastore.v1.client.DatastoreHelper.makeKey) Test(org.junit.Test)

Example 17 with Key

use of com.google.recaptchaenterprise.v1beta1.Key in project java-datastore by googleapis.

the class QuerySplitterImpl method getScatterKeys.

/**
 * Gets a list of split keys given a desired number of splits.
 *
 * <p>This list will contain multiple split keys for each split. Only a single split key will be
 * chosen as the split point, however providing multiple keys allows for more uniform sharding.
 *
 * @param numSplits the number of desired splits.
 * @param query the user query.
 * @param partition the partition to run the query in.
 * @param datastore the datastore containing the data.
 * @throws DatastoreException if there was an error when executing the datastore query.
 */
private List<Key> getScatterKeys(int numSplits, Query query, PartitionId partition, Datastore datastore) throws DatastoreException {
    Query.Builder scatterPointQuery = createScatterQuery(query, numSplits);
    List<Key> keySplits = new ArrayList<Key>();
    QueryResultBatch batch;
    do {
        RunQueryRequest scatterRequest = RunQueryRequest.newBuilder().setPartitionId(partition).setQuery(scatterPointQuery).build();
        batch = datastore.runQuery(scatterRequest).getBatch();
        for (EntityResult result : batch.getEntityResultsList()) {
            keySplits.add(result.getEntity().getKey());
        }
        scatterPointQuery.setStartCursor(batch.getEndCursor());
        scatterPointQuery.getLimitBuilder().setValue(scatterPointQuery.getLimit().getValue() - batch.getEntityResultsCount());
    } while (batch.getMoreResults() == MoreResultsType.NOT_FINISHED);
    Collections.sort(keySplits, DatastoreHelper.getKeyComparator());
    return keySplits;
}
Also used : QueryResultBatch(com.google.datastore.v1.QueryResultBatch) Query(com.google.datastore.v1.Query) RunQueryRequest(com.google.datastore.v1.RunQueryRequest) ArrayList(java.util.ArrayList) EntityResult(com.google.datastore.v1.EntityResult) Key(com.google.datastore.v1.Key)

Example 18 with Key

use of com.google.recaptchaenterprise.v1beta1.Key in project java-datastore by googleapis.

the class DatastoreHelperTest method testMakeKey_PartitionId.

@Test
public void testMakeKey_PartitionId() {
    PartitionId partitionId = PartitionId.newBuilder().setNamespaceId("namespace-id").build();
    Key parent = PARENT.toBuilder().setPartitionId(partitionId).build();
    assertEquals(Key.newBuilder().setPartitionId(partitionId).addPath(PARENT.getPath(0)).addPath(Key.PathElement.newBuilder().setKind("Child")).build(), makeKey(parent, "Child").build());
}
Also used : PartitionId(com.google.datastore.v1.PartitionId) Key(com.google.datastore.v1.Key) DatastoreHelper.makeKey(com.google.datastore.v1.client.DatastoreHelper.makeKey) Test(org.junit.Test)

Example 19 with Key

use of com.google.recaptchaenterprise.v1beta1.Key in project java-datastore by googleapis.

the class DatastoreHelper method makeKey.

/**
 * Make a key from the specified path of kind/id-or-name pairs and/or Keys.
 *
 * <p>The id-or-name values must be either String, Long, Integer or Short.
 *
 * <p>The last id-or-name value may be omitted, in which case an entity without an id is created
 * (for use with automatic id allocation).
 *
 * <p>The PartitionIds of all Keys in the path must be equal. The returned Key.Builder will use
 * this PartitionId.
 */
public static Key.Builder makeKey(Object... elements) {
    Key.Builder key = Key.newBuilder();
    PartitionId partitionId = null;
    for (int pathIndex = 0; pathIndex < elements.length; pathIndex += 2) {
        PathElement.Builder pathElement = PathElement.newBuilder();
        Object element = elements[pathIndex];
        if (element instanceof Key) {
            Key subKey = (Key) element;
            if (partitionId == null) {
                partitionId = subKey.getPartitionId();
            } else if (!partitionId.equals(subKey.getPartitionId())) {
                throw new IllegalArgumentException("Partition IDs did not match, found: " + partitionId + " and " + subKey.getPartitionId());
            }
            key.addAllPath(((Key) element).getPathList());
            // We increment by 2, but since we got a Key argument we're only consuming 1 element in this
            // iteration of the loop. Decrement the index so that when we jump by 2 we end up in the
            // right spot.
            pathIndex--;
        } else {
            String kind;
            try {
                kind = (String) element;
            } catch (ClassCastException e) {
                throw new IllegalArgumentException("Expected string or Key, got: " + element.getClass());
            }
            pathElement.setKind(kind);
            if (pathIndex + 1 < elements.length) {
                Object value = elements[pathIndex + 1];
                if (value instanceof String) {
                    pathElement.setName((String) value);
                } else if (value instanceof Long) {
                    pathElement.setId((Long) value);
                } else if (value instanceof Integer) {
                    pathElement.setId((Integer) value);
                } else if (value instanceof Short) {
                    pathElement.setId((Short) value);
                } else {
                    throw new IllegalArgumentException("Expected string or integer, got: " + value.getClass());
                }
            }
            key.addPath(pathElement);
        }
    }
    if (partitionId != null && !partitionId.equals(PartitionId.getDefaultInstance())) {
        key.setPartitionId(partitionId);
    }
    return key;
}
Also used : ByteString(com.google.protobuf.ByteString) PartitionId(com.google.datastore.v1.PartitionId) PathElement(com.google.datastore.v1.Key.PathElement) Key(com.google.datastore.v1.Key) PrivateKey(java.security.PrivateKey)

Example 20 with Key

use of com.google.recaptchaenterprise.v1beta1.Key in project java-datastore by googleapis.

the class QuerySplitterImpl method getSplits.

@Override
public List<Query> getSplits(Query query, PartitionId partition, int numSplits, Datastore datastore) throws DatastoreException, IllegalArgumentException {
    List<Query> splits = new ArrayList<Query>(numSplits);
    if (numSplits == 1) {
        splits.add(query);
        return splits;
    }
    validateQuery(query);
    validateSplitSize(numSplits);
    List<Key> scatterKeys = getScatterKeys(numSplits, query, partition, datastore);
    Key lastKey = null;
    for (Key nextKey : getSplitKey(scatterKeys, numSplits)) {
        splits.add(createSplit(lastKey, nextKey, query));
        lastKey = nextKey;
    }
    splits.add(createSplit(lastKey, null, query));
    return splits;
}
Also used : Query(com.google.datastore.v1.Query) ArrayList(java.util.ArrayList) Key(com.google.datastore.v1.Key)

Aggregations

Test (org.junit.Test)23 Key (com.google.datastore.v1.Key)17 DatastoreHelper.makeKey (com.google.datastore.v1.client.DatastoreHelper.makeKey)11 AbstractMessage (com.google.protobuf.AbstractMessage)10 Key (com.google.recaptchaenterprise.v1.Key)10 Entity (com.google.datastore.v1.Entity)7 DeleteKey (org.apache.beam.sdk.io.gcp.datastore.DatastoreV1.DeleteKey)7 DatastoreV1.isValidKey (org.apache.beam.sdk.io.gcp.datastore.DatastoreV1.isValidKey)7 RecaptchaEnterpriseServiceClient (com.google.cloud.recaptchaenterprise.v1.RecaptchaEnterpriseServiceClient)5 Mutation (com.google.datastore.v1.Mutation)5 Key (com.google.recaptchaenterprise.v1beta1.Key)4 HashMap (java.util.HashMap)4 DeleteEntity (org.apache.beam.sdk.io.gcp.datastore.DatastoreV1.DeleteEntity)4 Query (com.google.datastore.v1.Query)3 GetKeyRequest (com.google.recaptchaenterprise.v1.GetKeyRequest)3 ArrayList (java.util.ArrayList)3 InvalidArgumentException (com.google.api.gax.rpc.InvalidArgumentException)2 ListKeysPagedResponse (com.google.cloud.recaptchaenterprise.v1.RecaptchaEnterpriseServiceClient.ListKeysPagedResponse)2 PathElement (com.google.datastore.v1.Key.PathElement)2 PartitionId (com.google.datastore.v1.PartitionId)2