use of com.google.datastore.v1.Key in project YCSB by brianfrankcooper.
the class GoogleDatastoreClient method read.
@Override
public Status read(String table, String key, Set<String> fields, HashMap<String, ByteIterator> result) {
LookupRequest.Builder lookupRequest = LookupRequest.newBuilder();
lookupRequest.addKeys(buildPrimaryKey(table, key));
lookupRequest.getReadOptionsBuilder().setReadConsistency(this.readConsistency);
// Note above, datastore lookupRequest always reads the entire entity, it
// does not support reading a subset of "fields" (properties) of an entity.
logger.debug("Built lookup request as: " + lookupRequest.toString());
LookupResponse response = null;
try {
response = datastore.lookup(lookupRequest.build());
} catch (DatastoreException exception) {
logger.error(String.format("Datastore Exception when reading (%s): %s %s", exception.getMessage(), exception.getMethodName(), exception.getCode()));
// will bubble up to the user as part of the YCSB Status "name".
return new Status("ERROR-" + exception.getCode(), exception.getMessage());
}
if (response.getFoundCount() == 0) {
return new Status("ERROR-404", "Not Found, key is: " + key);
} else if (response.getFoundCount() > 1) {
// entity back. Unexpected State.
return Status.UNEXPECTED_STATE;
}
Entity entity = response.getFound(0).getEntity();
logger.debug("Read entity: " + entity.toString());
Map<String, Value> properties = entity.getProperties();
Set<String> propertiesToReturn = (fields == null ? properties.keySet() : fields);
for (String name : propertiesToReturn) {
if (properties.containsKey(name)) {
result.put(name, new StringByteIterator(properties.get(name).getStringValue()));
}
}
return Status.OK;
}
use of com.google.datastore.v1.Key in project YCSB by brianfrankcooper.
the class GoogleDatastoreClient method doSingleItemMutation.
private Status doSingleItemMutation(String table, String key, @Nullable HashMap<String, ByteIterator> values, MutationType mutationType) {
// First build the key.
Key.Builder datastoreKey = buildPrimaryKey(table, key);
// Build a commit request in non-transactional mode.
// Single item mutation to google datastore
// is always atomic and strongly consistent. Transaction is only necessary
// for multi-item mutation, or Read-modify-write operation.
CommitRequest.Builder commitRequest = CommitRequest.newBuilder();
commitRequest.setMode(Mode.NON_TRANSACTIONAL);
if (mutationType == MutationType.DELETE) {
commitRequest.addMutationsBuilder().setDelete(datastoreKey);
} else {
// If this is not for delete, build the entity.
Entity.Builder entityBuilder = Entity.newBuilder();
entityBuilder.setKey(datastoreKey);
for (Entry<String, ByteIterator> val : values.entrySet()) {
entityBuilder.getMutableProperties().put(val.getKey(), Value.newBuilder().setStringValue(val.getValue().toString()).build());
}
Entity entity = entityBuilder.build();
logger.debug("entity built as: " + entity.toString());
if (mutationType == MutationType.UPSERT) {
commitRequest.addMutationsBuilder().setUpsert(entity);
} else if (mutationType == MutationType.UPDATE) {
commitRequest.addMutationsBuilder().setUpdate(entity);
} else {
throw new RuntimeException("Impossible MutationType, code bug.");
}
}
try {
datastore.commit(commitRequest.build());
logger.debug("successfully committed.");
} catch (DatastoreException exception) {
// Catch all Datastore rpc errors.
// Log the exception, the name of the method called and the error code.
logger.error(String.format("Datastore Exception when committing (%s): %s %s", exception.getMessage(), exception.getMethodName(), exception.getCode()));
// will bubble up to the user as part of the YCSB Status "name".
return new Status("ERROR-" + exception.getCode(), exception.getMessage());
}
return Status.OK;
}
use of com.google.datastore.v1.Key in project beam by apache.
the class DatastoreV1Test method testDeleteKeys.
/**
* Test that valid keys are transformed to delete mutations.
*/
@Test
public void testDeleteKeys() throws Exception {
Key key = makeKey("bird", "finch").build();
DeleteKeyFn deleteKeyFn = new DeleteKeyFn();
Mutation exceptedMutation = makeDelete(key).build();
assertEquals(deleteKeyFn.apply(key), exceptedMutation);
}
use of com.google.datastore.v1.Key in project beam by apache.
the class DatastoreV1Test method testAddEntities.
@Test
public /**
* Test that entities with valid keys are transformed to upsert mutations.
*/
void testAddEntities() throws Exception {
Key key = makeKey("bird", "finch").build();
Entity entity = Entity.newBuilder().setKey(key).build();
UpsertFn upsertFn = new UpsertFn();
Mutation exceptedMutation = makeUpsert(entity).build();
assertEquals(upsertFn.apply(entity), exceptedMutation);
}
use of com.google.datastore.v1.Key in project beam by apache.
the class DatastoreV1Test method datastoreWriterFnTest.
// A helper method to test DatastoreWriterFn for various batch sizes.
private void datastoreWriterFnTest(int numMutations) throws Exception {
// Create the requested number of mutations.
List<Mutation> mutations = new ArrayList<>(numMutations);
for (int i = 0; i < numMutations; ++i) {
mutations.add(makeUpsert(Entity.newBuilder().setKey(makeKey("key" + i, i + 1)).build()).build());
}
DatastoreWriterFn datastoreWriter = new DatastoreWriterFn(StaticValueProvider.of(PROJECT_ID), null, mockDatastoreFactory);
DoFnTester<Mutation, Void> doFnTester = DoFnTester.of(datastoreWriter);
doFnTester.setCloningBehavior(CloningBehavior.DO_NOT_CLONE);
doFnTester.processBundle(mutations);
int start = 0;
while (start < numMutations) {
int end = Math.min(numMutations, start + DATASTORE_BATCH_UPDATE_LIMIT);
CommitRequest.Builder commitRequest = CommitRequest.newBuilder();
commitRequest.setMode(CommitRequest.Mode.NON_TRANSACTIONAL);
commitRequest.addAllMutations(mutations.subList(start, end));
// Verify all the batch requests were made with the expected mutations.
verify(mockDatastore, times(1)).commit(commitRequest.build());
start = end;
}
}
Aggregations