use of com.google.datastore.v1.Entity in project google-cloud-java by GoogleCloudPlatform.
the class DatastoreTest method testQueryPaginationWithLimit.
@Test
public void testQueryPaginationWithLimit() throws DatastoreException {
List<RunQueryResponse> responses = buildResponsesForQueryPaginationWithLimit();
List<ByteString> endCursors = Lists.newArrayListWithCapacity(responses.size());
for (RunQueryResponse response : responses) {
EasyMock.expect(rpcMock.runQuery(EasyMock.anyObject(RunQueryRequest.class))).andReturn(response);
if (response.getBatch().getMoreResults() != QueryResultBatch.MoreResultsType.NOT_FINISHED) {
endCursors.add(response.getBatch().getEndCursor());
}
}
EasyMock.replay(rpcFactoryMock, rpcMock);
Datastore datastore = rpcMockOptions.getService();
int limit = 2;
int totalCount = 0;
Iterator<ByteString> cursorIter = endCursors.iterator();
StructuredQuery<Entity> query = Query.newEntityQueryBuilder().setLimit(limit).build();
while (true) {
QueryResults<Entity> results = datastore.run(query);
int resultCount = 0;
while (results.hasNext()) {
results.next();
resultCount++;
totalCount++;
}
assertTrue(cursorIter.hasNext());
Cursor expectedEndCursor = Cursor.copyFrom(cursorIter.next().toByteArray());
assertEquals(expectedEndCursor, results.getCursorAfter());
if (resultCount < limit) {
break;
}
query = query.toBuilder().setStartCursor(results.getCursorAfter()).build();
}
assertEquals(5, totalCount);
EasyMock.verify(rpcFactoryMock, rpcMock);
}
use of com.google.datastore.v1.Entity in project google-cloud-java by GoogleCloudPlatform.
the class DatastoreTest method buildResponsesForQueryPaginationWithLimit.
private List<RunQueryResponse> buildResponsesForQueryPaginationWithLimit() {
Entity entity4 = Entity.newBuilder(KEY4).set("value", StringValue.of("value")).build();
Entity entity5 = Entity.newBuilder(KEY5).set("value", "value").build();
datastore.add(ENTITY3, entity4, entity5);
DatastoreRpc datastoreRpc = datastore.getOptions().getDatastoreRpcV1();
List<RunQueryResponse> responses = new ArrayList<>();
Query<Entity> query = Query.newEntityQueryBuilder().build();
RunQueryRequest.Builder requestPb = RunQueryRequest.newBuilder();
query.populatePb(requestPb);
QueryResultBatch queryResultBatchPb = RunQueryResponse.newBuilder().mergeFrom(datastoreRpc.runQuery(requestPb.build())).getBatch();
QueryResultBatch queryResultBatchPb1 = QueryResultBatch.newBuilder().mergeFrom(queryResultBatchPb).setMoreResults(QueryResultBatch.MoreResultsType.NOT_FINISHED).clearEntityResults().addAllEntityResults(queryResultBatchPb.getEntityResultsList().subList(0, 1)).setEndCursor(ByteString.copyFromUtf8("a")).build();
responses.add(RunQueryResponse.newBuilder().setBatch(queryResultBatchPb1).build());
QueryResultBatch queryResultBatchPb2 = QueryResultBatch.newBuilder().mergeFrom(queryResultBatchPb).setMoreResults(QueryResultBatch.MoreResultsType.MORE_RESULTS_AFTER_LIMIT).clearEntityResults().addAllEntityResults(queryResultBatchPb.getEntityResultsList().subList(1, 2)).setEndCursor(// test invalid UTF-8 string
ByteString.copyFrom(new byte[] { (byte) 0x80 })).build();
responses.add(RunQueryResponse.newBuilder().setBatch(queryResultBatchPb2).build());
QueryResultBatch queryResultBatchPb3 = QueryResultBatch.newBuilder().mergeFrom(queryResultBatchPb).setMoreResults(QueryResultBatch.MoreResultsType.MORE_RESULTS_AFTER_LIMIT).clearEntityResults().addAllEntityResults(queryResultBatchPb.getEntityResultsList().subList(2, 4)).setEndCursor(ByteString.copyFromUtf8("b")).build();
responses.add(RunQueryResponse.newBuilder().setBatch(queryResultBatchPb3).build());
QueryResultBatch queryResultBatchPb4 = QueryResultBatch.newBuilder().mergeFrom(queryResultBatchPb).setMoreResults(QueryResultBatch.MoreResultsType.NO_MORE_RESULTS).clearEntityResults().addAllEntityResults(queryResultBatchPb.getEntityResultsList().subList(4, 5)).setEndCursor(ByteString.copyFromUtf8("c")).build();
responses.add(RunQueryResponse.newBuilder().setBatch(queryResultBatchPb4).build());
return responses;
}
use of com.google.datastore.v1.Entity in project GeoGig by boundlessgeo.
the class OSMExport method getFeatures.
private Iterator<EntityContainer> getFeatures(String ref) {
Optional<ObjectId> id = geogig.command(RevParse.class).setRefSpec(ref).call();
if (!id.isPresent()) {
return Iterators.emptyIterator();
}
LsTreeOp op = geogig.command(LsTreeOp.class).setStrategy(Strategy.DEPTHFIRST_ONLY_FEATURES).setReference(ref);
if (bbox != null) {
final Envelope env;
try {
env = new Envelope(Double.parseDouble(bbox.get(0)), Double.parseDouble(bbox.get(2)), Double.parseDouble(bbox.get(1)), Double.parseDouble(bbox.get(3)));
} catch (NumberFormatException e) {
throw new IllegalArgumentException("Wrong bbox definition");
}
Predicate<Bounded> filter = new Predicate<Bounded>() {
@Override
public boolean apply(final Bounded bounded) {
boolean intersects = bounded.intersects(env);
return intersects;
}
};
op.setBoundsFilter(filter);
}
Iterator<NodeRef> iterator = op.call();
final EntityConverter converter = new EntityConverter();
Function<NodeRef, EntityContainer> function = new Function<NodeRef, EntityContainer>() {
@Override
@Nullable
public EntityContainer apply(@Nullable NodeRef ref) {
RevFeature revFeature = geogig.command(RevObjectParse.class).setObjectId(ref.objectId()).call(RevFeature.class).get();
SimpleFeatureType featureType;
if (ref.path().startsWith(OSMUtils.NODE_TYPE_NAME)) {
featureType = OSMUtils.nodeType();
} else {
featureType = OSMUtils.wayType();
}
SimpleFeatureBuilder featureBuilder = new SimpleFeatureBuilder(featureType);
RevFeatureType revFeatureType = RevFeatureTypeImpl.build(featureType);
List<PropertyDescriptor> descriptors = revFeatureType.sortedDescriptors();
ImmutableList<Optional<Object>> values = revFeature.getValues();
for (int i = 0; i < descriptors.size(); i++) {
PropertyDescriptor descriptor = descriptors.get(i);
Optional<Object> value = values.get(i);
featureBuilder.set(descriptor.getName(), value.orNull());
}
SimpleFeature feature = featureBuilder.buildFeature(ref.name());
Entity entity = converter.toEntity(feature, null);
EntityContainer container;
if (entity instanceof Node) {
container = new NodeContainer((Node) entity);
} else {
container = new WayContainer((Way) entity);
}
return container;
}
};
return Iterators.transform(iterator, function);
}
use of com.google.datastore.v1.Entity in project YCSB by brianfrankcooper.
the class GoogleDatastoreClient method init.
/**
* Initialize any state for this DB. Called once per DB instance; there is
* one DB instance per client thread.
*/
@Override
public void init() throws DBException {
String debug = getProperties().getProperty("googledatastore.debug", null);
if (null != debug && "true".equalsIgnoreCase(debug)) {
logger.setLevel(Level.DEBUG);
}
// We need the following 3 essential properties to initialize datastore:
//
// - DatasetId,
// - Path to private key file,
// - Service account email address.
String datasetId = getProperties().getProperty("googledatastore.datasetId", null);
if (datasetId == null) {
throw new DBException("Required property \"datasetId\" missing.");
}
String privateKeyFile = getProperties().getProperty("googledatastore.privateKeyFile", null);
if (privateKeyFile == null) {
throw new DBException("Required property \"privateKeyFile\" missing.");
}
String serviceAccountEmail = getProperties().getProperty("googledatastore.serviceAccountEmail", null);
if (serviceAccountEmail == null) {
throw new DBException("Required property \"serviceAccountEmail\" missing.");
}
// Below are properties related to benchmarking.
String readConsistencyConfig = getProperties().getProperty("googledatastore.readConsistency", null);
if (readConsistencyConfig != null) {
try {
this.readConsistency = ReadConsistency.valueOf(readConsistencyConfig.trim().toUpperCase());
} catch (IllegalArgumentException e) {
throw new DBException("Invalid read consistency specified: " + readConsistencyConfig + ". Expecting STRONG or EVENTUAL.");
}
}
//
// Entity Grouping Mode (googledatastore.entitygroupingmode), see
// documentation in conf/googledatastore.properties.
//
String entityGroupingConfig = getProperties().getProperty("googledatastore.entityGroupingMode", null);
if (entityGroupingConfig != null) {
try {
this.entityGroupingMode = EntityGroupingMode.valueOf(entityGroupingConfig.trim().toUpperCase());
} catch (IllegalArgumentException e) {
throw new DBException("Invalid entity grouping mode specified: " + entityGroupingConfig + ". Expecting ONE_ENTITY_PER_GROUP or " + "MULTI_ENTITY_PER_GROUP.");
}
}
this.rootEntityName = getProperties().getProperty("googledatastore.rootEntityName", "YCSB_ROOT_ENTITY");
try {
// Setup the connection to Google Cloud Datastore with the credentials
// obtained from the configure.
DatastoreOptions.Builder options = new DatastoreOptions.Builder();
Credential credential = DatastoreHelper.getServiceAccountCredential(serviceAccountEmail, privateKeyFile);
logger.info("Using JWT Service Account credential.");
logger.info("DatasetID: " + datasetId + ", Service Account Email: " + serviceAccountEmail + ", Private Key File Path: " + privateKeyFile);
datastore = DatastoreFactory.get().create(options.credential(credential).projectId(datasetId).build());
} catch (GeneralSecurityException exception) {
throw new DBException("Security error connecting to the datastore: " + exception.getMessage(), exception);
} catch (IOException exception) {
throw new DBException("I/O error connecting to the datastore: " + exception.getMessage(), exception);
}
logger.info("Datastore client instance created: " + datastore.toString());
}
use of com.google.datastore.v1.Entity in project beam by apache.
the class V1ReadIT method writeEntitiesToDatastore.
// Creates entities and write them to datastore
private static void writeEntitiesToDatastore(V1TestOptions options, String project, String ancestor, long numEntities) throws Exception {
Datastore datastore = getDatastore(options, project);
// Write test entities to datastore
V1TestWriter writer = new V1TestWriter(datastore, new UpsertMutationBuilder());
Key ancestorKey = makeAncestorKey(options.getNamespace(), options.getKind(), ancestor);
for (long i = 0; i < numEntities; i++) {
Entity entity = makeEntity(i, ancestorKey, options.getKind(), options.getNamespace(), 0);
writer.write(entity);
}
writer.close();
}
Aggregations