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());
}
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;
}
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());
}
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;
}
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;
}
Aggregations