com.google.cloud.datastore.EntityQuery Java Examples

public List<Greeting> getGreetings() {
  // This query requires the index defined in index.yaml to work because of the orderBy on date.
  EntityQuery query =
      Query.newEntityQueryBuilder()
          .setKind("Greeting")
          .setFilter(hasAncestor(key))
          .setOrderBy(desc("date"))
          .setLimit(5)
          .build();

  QueryResults<Entity> results = getDatastore().run(query);

  Builder<Greeting> resultListBuilder = ImmutableList.builder();
  while (results.hasNext()) {
    resultListBuilder.add(new Greeting(results.next()));
  }

  return resultListBuilder.build();
}
 

private Slice executeSliceQuery(Object[] parameters) {
	StructuredQuery.Builder builder = getEntityOrProjectionQueryBuilder()
			.setKind(this.datastorePersistentEntity.kindName());
	StructuredQuery query = applyQueryBody(parameters, builder, false, false, null);
	DatastoreResultsIterable<?> resultList = this.datastoreOperations.queryKeysOrEntities(query, this.entityType);

	ParameterAccessor paramAccessor = new ParametersParameterAccessor(getQueryMethod().getParameters(), parameters);

	Pageable pageable = DatastorePageable.from(paramAccessor.getPageable(), resultList.getCursor(), null);

	EntityQuery.Builder builderNext = newEntityQueryBuilder().setKind(this.datastorePersistentEntity.kindName());
	StructuredQuery queryNext = applyQueryBody(parameters, builderNext, false, true, resultList.getCursor());
	Iterable nextResult = this.datastoreOperations.query(queryNext, x -> x);

	List<Object> result =
					StreamSupport.stream(resultList.spliterator(), false).collect(Collectors.toList());

	return (Slice) this.processRawObjectForProjection(
			new SliceImpl(result, pageable, nextResult.iterator().hasNext()));
}
 

@Test
public void findAllTestSort() {
	EntityQuery.Builder builder = Query.newEntityQueryBuilder().setKind("custom_test_kind");

	this.datastoreTemplate.findAll(TestEntity.class,
			new DatastoreQueryOptions.Builder().setSort(Sort.by("sortProperty"))
					.build());
	verify(this.datastore, times(1)).run(
			builder.setOrderBy(
					new StructuredQuery.OrderBy("prop", StructuredQuery.OrderBy.Direction.ASCENDING)).build());

	this.datastoreTemplate.findAll(TestEntity.class,
			new DatastoreQueryOptions.Builder()
					.setSort(Sort.by(Sort.Direction.DESC, "sortProperty")).build());
	verify(this.datastore, times(1)).run(
			builder.setOrderBy(
					new StructuredQuery.OrderBy("prop", StructuredQuery.OrderBy.Direction.DESCENDING)).build());
}
 

private EntityQuery.Builder backfillQueryBuilder(boolean showAll, Filter... filters) {
  final EntityQuery.Builder queryBuilder = Query.newEntityQueryBuilder().setKind(KIND_BACKFILL);

  final List<Filter> andedFilters = Lists.newArrayList(filters);

  if (!showAll) {
    andedFilters.add(PropertyFilter.eq(PROPERTY_ALL_TRIGGERED, false));
    andedFilters.add(PropertyFilter.eq(PROPERTY_HALTED, false));
  }

  if (!andedFilters.isEmpty()) {
    final Filter head = andedFilters.get(0);
    final Filter[] tail = andedFilters.stream().skip(1).toArray(Filter[]::new);
    queryBuilder.setFilter(CompositeFilter.and(head, tail));
  }

  return queryBuilder;
}
 

@Test
public void nonCollectionReturnType() throws NoSuchMethodException {
	Trade trade = new Trade();
	queryWithMockResult("findByAction", null,
			getClass().getMethod("findByAction", String.class), true, null);

	Object[] params = new Object[] { "BUY", };

	when(this.datastoreTemplate.queryKeysOrEntities(any(), any())).thenAnswer((invocation) -> {
		EntityQuery statement = invocation.getArgument(0);

		EntityQuery expected = StructuredQuery.newEntityQueryBuilder()
				.setFilter(PropertyFilter.eq("action", "BUY"))
				.setKind("trades")
				.setLimit(1).build();

		assertThat(statement).isEqualTo(expected);

		List<Trade> results = Collections.singletonList(trade);
		return new DatastoreResultsIterable(results.iterator(), null);
	});

	assertThat(this.partTreeDatastoreQuery.execute(params)).isEqualTo(trade);
}
 

@Test
public void sliceQuery() throws NoSuchMethodException {
	queryWithMockResult("findByActionAndSymbolAndPriceLessThanAndPriceGreater"
					+ "ThanEqualAndIdIsNull", null,
			getClass().getMethod("tradeMethod", String.class, String.class, double.class, double.class,
					Pageable.class));

	this.partTreeDatastoreQuery = createQuery(false, true, null);

	Object[] params = new Object[] { "BUY", "abcd", 8.88, 3.33, PageRequest.of(0, 2, Sort.Direction.DESC, "id") };

	prepareSliceResults(0, 2, false);

	when(this.queryMethod.getCollectionReturnType()).thenReturn(List.class);

	Slice result = (Slice) this.partTreeDatastoreQuery.execute(params);
	assertThat(result.hasNext()).isEqualTo(false);

	verify(this.datastoreTemplate, times(1))
			.query(isA(EntityQuery.class), (Function) any());

	verify(this.datastoreTemplate, times(0))
			.queryKeysOrEntities(isA(KeyQuery.class), any());
}
 

@Test
public void pageableQueryMissingPageableParamReturnsAllResults() throws NoSuchMethodException {
	queryWithMockResult("findByActionAndSymbolAndPriceLessThanAndPriceGreater"
					+ "ThanEqualAndIdIsNullOrderByIdDesc", null,
			getClass().getMethod("tradeMethod", String.class, String.class, double.class, double.class));

	this.partTreeDatastoreQuery = createQuery(true, false, null);

	Object[] params = new Object[] { "BUY", "abcd", 8.88, 3.33 };

	preparePageResults(0, null, null, Arrays.asList(1, 2, 3, 4), Arrays.asList(1, 2, 3, 4));

	when(this.queryMethod.getCollectionReturnType()).thenReturn(List.class);

	Page result = (Page) this.partTreeDatastoreQuery.execute(params);
	assertThat(result.getTotalElements()).isEqualTo(4);
	assertThat(result.getTotalPages()).isEqualTo(1);

	verify(this.datastoreTemplate, times(1))
			.queryKeysOrEntities(isA(EntityQuery.class), any());

	verify(this.datastoreTemplate, times(1))
			.queryKeysOrEntities(isA(KeyQuery.class), any());
}
 

@Test
public void pageableQuery() throws NoSuchMethodException {
	queryWithMockResult("findByActionAndSymbolAndPriceLessThanAndPriceGreater"
					+ "ThanEqualAndIdIsNull", null,
			getClass().getMethod("tradeMethod", String.class, String.class, double.class, double.class,
					Pageable.class));

	this.partTreeDatastoreQuery = createQuery(true, false, null);

	Object[] params = new Object[] { "BUY", "abcd", 8.88, 3.33, PageRequest.of(1, 2, Sort.Direction.DESC, "id") };

	preparePageResults(2, 2, null, Arrays.asList(3, 4), Arrays.asList(1, 2, 3, 4));

	when(this.queryMethod.getCollectionReturnType()).thenReturn(List.class);

	Page result = (Page) this.partTreeDatastoreQuery.execute(params);
	assertThat(result.getTotalElements()).isEqualTo(4);
	assertThat(result.getTotalPages()).isEqualTo(2);
	assertThat(result.getNumberOfElements()).isEqualTo(2);

	verify(this.datastoreTemplate, times(1))
			.queryKeysOrEntities(isA(EntityQuery.class), any());

	verify(this.datastoreTemplate, times(1))
			.queryKeysOrEntities(isA(KeyQuery.class), any());
}
 

@Test
public void shouldIncrementCounter() throws IOException {
  // init counter
  assertEquals(0, shardedCounter.getCounter(COUNTER_ID1));

  //increment counter by 1
  updateCounterInTransaction(COUNTER_ID1, 1L);

  QueryResults<Entity> results = datastore.run(EntityQuery.newEntityQueryBuilder()
      .setKind(KIND_COUNTER_SHARD)
      .setFilter(CompositeFilter.and(PropertyFilter.eq(PROPERTY_COUNTER_ID, COUNTER_ID1),
          PropertyFilter.eq(PROPERTY_SHARD_VALUE,1)))
      .build());
  // assert there's one and only one shard with the value set to 1
  assertEquals(1L, results.next().getLong(PROPERTY_SHARD_VALUE));
  assertFalse(results.hasNext());

  // assert the correct value is fetched after cache expiry
  shardedCounter.inMemSnapshot.invalidate(COUNTER_ID1);
  assertEquals(1L, shardedCounter.getCounter(COUNTER_ID1));
}
 

private Cursor cursorPaging(int pageSize, Cursor pageCursor) {
  // [START datastore_cursor_paging]
  EntityQuery.Builder queryBuilder = Query.newEntityQueryBuilder().setKind("Task")
      .setLimit(pageSize);
  if (pageCursor != null) {
    queryBuilder.setStartCursor(pageCursor);
  }
  QueryResults<Entity> tasks = datastore.run(queryBuilder.build());
  while (tasks.hasNext()) {
    Entity task = tasks.next();
    // do something with the task
  }
  Cursor nextPageCursor = tasks.getCursorAfter();
  // [END datastore_cursor_paging]
  return nextPageCursor;
}
 

List<Backfill> getBackfillsForWorkflow(boolean showAll, String workflow) throws IOException {
  final EntityQuery query = backfillQueryBuilder(showAll,
                                                 PropertyFilter.eq(PROPERTY_WORKFLOW, workflow))
      .build();

  return backfillsForQuery(query);
}
 

List<Backfill> getBackfillsForWorkflowId(boolean showAll, WorkflowId workflowId) throws IOException {
  final EntityQuery query = backfillQueryBuilder(
      showAll,
      PropertyFilter.eq(PROPERTY_COMPONENT, workflowId.componentId()),
      PropertyFilter.eq(PROPERTY_WORKFLOW, workflowId.id()))
      .build();

  return backfillsForQuery(query);
}
 

List<Backfill> getBackfillsForComponent(boolean showAll, String component) throws IOException {
  final EntityQuery query = backfillQueryBuilder(showAll,
                                                 PropertyFilter.eq(PROPERTY_COMPONENT, component))
      .build();

  return backfillsForQuery(query);
}
 

@Test
public void queryAndConsumeShouldPropagateIOException() throws IOException {
  final IOException cause = new IOException("foo");
  when(queryResults.results()).thenReturn(Stream.of(entity1, entity2).iterator());
  doReturn(queryResults).when(rw).run(any(EntityQuery.class));
  exception.expect(is(cause));
  sut.query(query, entity -> { throw cause; });
}
 

@Test
public void queryAndConsumeShouldHandleForEachRemainingThrowing() throws IOException {
  doThrow(CAUSE).when(queryResults).forEachRemaining(any());
  doReturn(queryResults).when(rw).run(any(EntityQuery.class));
  exception.expect(IOException.class);
  exception.expectCause(is((Throwable) CAUSE));
  sut.query(query, entity -> fail());
}
 

private List<Entity> entitiesOfKind(String kind) {
  EntityQuery query = Query.newEntityQueryBuilder().setKind(kind).build();
  QueryResults<Entity> keys = datastoreClient.run(query);
  List<Entity> entities = new ArrayList<>();
  while (keys.hasNext()) {
    entities.add(keys.next());
  }
  return entities;
}
 

private void deleteShardsForCounter(String counterId) throws IOException {
  final List<Key> shards = new ArrayList<>();
  datastore.query(EntityQuery.newEntityQueryBuilder()
      .setKind(KIND_COUNTER_SHARD)
      .setFilter(PropertyFilter.eq(PROPERTY_COUNTER_ID, counterId))
      .build(), entity -> shards.add(entity.getKey()));

  // this is a safe guard to not to exceed max number of entities in one batch write
  // because in practice number of shards is much smaller
  for (List<Key> batch : Lists.partition(shards, MAX_NUMBER_OF_ENTITIES_IN_ONE_BATCH_WRITE)) {
      datastore.delete(batch.toArray(new Key[0]));
  }
}
 

List<Resource> getResources() throws IOException {
  final EntityQuery query = Query.newEntityQueryBuilder().setKind(KIND_COUNTER_LIMIT).build();
  final List<Resource> resources = Lists.newArrayList();
  datastore.query(query, entity ->
      resources.add(entityToResource(entity)));
  return resources;
}
 

private Map<WorkflowInstance, RunState> queryActiveStates(EntityQuery activeStatesQuery)
    throws IOException {
  final ImmutableMap.Builder<WorkflowInstance, RunState> mapBuilder = ImmutableMap.builder();
  datastore.query(activeStatesQuery, entity -> {
    final WorkflowInstance instance = parseWorkflowInstance(entity);
    mapBuilder.put(instance, entityToRunState(entity, instance));
  });

  return mapBuilder.build();
}
 

Map<WorkflowInstance, RunState> activeStatesByTriggerId(
    String triggerId) throws IOException {
  final EntityQuery query =
      Query.newEntityQueryBuilder().setKind(KIND_ACTIVE_WORKFLOW_INSTANCE)
          .setFilter(PropertyFilter.eq(PROPERTY_STATE_TRIGGER_ID, triggerId))
          .build();

  return queryActiveStates(query);
}
 

Map<WorkflowInstance, RunState> readActiveStates(String componentId, String workflowId) throws IOException {
  final EntityQuery query =
      Query.newEntityQueryBuilder().setKind(KIND_ACTIVE_WORKFLOW_INSTANCE)
          .setFilter(CompositeFilter.and(PropertyFilter.eq(PROPERTY_COMPONENT, componentId),
          PropertyFilter.eq(PROPERTY_WORKFLOW, workflowId)))
          .build();

  return queryActiveStates(query);
}
 

Map<WorkflowInstance, RunState> readActiveStates(String componentId) throws IOException {
  final EntityQuery query =
      Query.newEntityQueryBuilder().setKind(KIND_ACTIVE_WORKFLOW_INSTANCE)
          .setFilter(PropertyFilter.eq(PROPERTY_COMPONENT, componentId))
          .build();

  return queryActiveStates(query);
}
 

@Test
public void queryByExampleIncludeNullValuesTest() {
	EntityQuery.Builder builder = Query.newEntityQueryBuilder().setKind("test_kind");
	this.datastoreTemplate.queryByExample(
			Example.of(this.simpleTestEntityNullVallues,
					ExampleMatcher.matching().withIgnorePaths("id").withIncludeNullValues()),
			null);

	StructuredQuery.CompositeFilter filter = StructuredQuery.CompositeFilter
			.and(PropertyFilter.eq("color", NullValue.of()),
					PropertyFilter.eq("int_field", NullValue.of()));
	verify(this.datastore, times(1)).run(builder.setFilter(filter).build());
}
 

Map<Workflow, TriggerInstantSpec> workflowsWithNextNaturalTrigger() throws IOException {
  final Map<Workflow, TriggerInstantSpec> map = Maps.newHashMap();
  final EntityQuery query =
      Query.newEntityQueryBuilder().setKind(KIND_WORKFLOW).build();
  datastore.query(query, entity -> {
    final Workflow workflow;
    try {
      workflow = OBJECT_MAPPER.readValue(entity.getString(PROPERTY_WORKFLOW_JSON), Workflow.class);
    } catch (IOException e) {
      log.warn("Failed to read workflow {}.", entity.getKey(), e);
      return;
    }

    if (entity.contains(PROPERTY_NEXT_NATURAL_TRIGGER)) {
      Instant instant = timestampToInstant(entity.getTimestamp(PROPERTY_NEXT_NATURAL_TRIGGER));
      final Instant triggerInstant;

      // todo: this check is only needed during a transition period
      if (!entity.contains(PROPERTY_NEXT_NATURAL_OFFSET_TRIGGER)) {
        // instant has to be moved one schedule interval back
        final Schedule schedule = workflow.configuration().schedule();
        if (TimeUtil.isAligned(instant, schedule)) {
          instant = TimeUtil.previousInstant(instant, schedule);
        }
        triggerInstant = workflow.configuration().addOffset(instant);
      } else {
        triggerInstant = timestampToInstant(entity.getTimestamp(PROPERTY_NEXT_NATURAL_OFFSET_TRIGGER));
      }

      map.put(workflow, TriggerInstantSpec.create(instant, triggerInstant));
    }
  });
  return map;
}
 

@Test
public void queryByExampleNoNullValuesTest() {
	EntityQuery.Builder builder = Query.newEntityQueryBuilder().setKind("test_kind");
	this.datastoreTemplate.queryByExample(
			Example.of(this.simpleTestEntityNullVallues, ExampleMatcher.matching().withIgnorePaths("id")), null);

	verify(this.datastore, times(1)).run(builder.build());
}
 

@Override
public <T> DatastoreResultsCollection<T> findAll(Class<T> entityClass, DatastoreQueryOptions queryOptions) {
	DatastorePersistentEntity<?> persistentEntity = this.datastoreMappingContext.getPersistentEntity(entityClass);
	EntityQuery.Builder builder = Query.newEntityQueryBuilder()
			.setKind(persistentEntity.kindName());
	applyQueryOptions(builder, queryOptions, persistentEntity);
	Query query = builder.build();
	QueryResults queryResults = getDatastoreReadWriter().run(query);
	Collection<T> convertedResults = convertEntitiesForRead(queryResults, entityClass);
	maybeEmitEvent(new AfterQueryEvent(convertedResults, query));
	return new DatastoreResultsCollection<>(convertedResults,
			queryResults != null ? queryResults.getCursorAfter() : null);
}
 

private <T> void resolveDescendantProperties(DatastorePersistentEntity datastorePersistentEntity,
		BaseEntity entity, T convertedObject, ReadContext context) {
	datastorePersistentEntity
			.doWithDescendantProperties((descendantPersistentProperty) -> {

				Class descendantType = descendantPersistentProperty
						.getComponentType();

				Key entityKey = (Key) entity.getKey();
				Key ancestorKey = KeyUtil.getKeyWithoutAncestors(entityKey);

				EntityQuery descendantQuery = Query.newEntityQueryBuilder()
						.setKind(this.datastoreMappingContext
								.getPersistentEntity(descendantType).kindName())
						.setFilter(PropertyFilter.hasAncestor(ancestorKey))
						.build();

				List entities = convertEntitiesForRead(
						getDatastoreReadWriter().run(descendantQuery), descendantType, context);

				datastorePersistentEntity.getPropertyAccessor(convertedObject)
						.setProperty(descendantPersistentProperty,
								// Converting the collection type.
								this.datastoreEntityConverter.getConversions()
										.convertOnRead(
												entities,
												descendantPersistentProperty.getType(),
												descendantType));
			});
}
 

@Test
public void compoundNameConventionTest() throws NoSuchMethodException {
	queryWithMockResult("findTop333ByActionAndSymbolAndPriceLessThan"
					+ "AndPriceGreaterThanEqual"
					+ "AndEmbeddedEntityStringFieldEquals"
					+ "AndIdIsNullOrderByIdDesc", null,
			getClass().getMethod("tradeMethod", String.class, String.class, double.class, double.class, String.class));

	Object[] params = new Object[] { "BUY", "abcd",
			// this int param requires custom conversion
			8, 3.33, "abc" };

	when(this.datastoreTemplate.queryKeysOrEntities(any(), any())).thenAnswer((invocation) -> {
		EntityQuery statement = invocation.getArgument(0);

		EntityQuery expected = StructuredQuery.newEntityQueryBuilder()
				.setFilter(CompositeFilter.and(PropertyFilter.eq("action", "BUY"),
						PropertyFilter.eq("ticker", "abcd"),
						PropertyFilter.lt("price", 8L),
						PropertyFilter.ge("price", 3.33),
						PropertyFilter.eq("embeddedEntity.stringField", "abc"),
						PropertyFilter.isNull("__key__")))
				.setKind("trades")
				.setOrderBy(OrderBy.desc("__key__")).setLimit(333).build();

		assertThat(statement).isEqualTo(expected);

		return EMPTY_RESPONSE;
	});

	when(this.queryMethod.getCollectionReturnType()).thenReturn(List.class);

	this.partTreeDatastoreQuery.execute(params);
	verify(this.datastoreTemplate, times(1))
			.queryKeysOrEntities(any(), any());
}
 

@Test
public void ambiguousSortPageableParam() throws NoSuchMethodException {
	queryWithMockResult("findTop333ByActionAndSymbolAndPriceLessThanAndPriceGreater"
					+ "ThanEqualAndIdIsNullOrderByIdDesc", null,
			getClass().getMethod("tradeMethod", String.class, String.class, double.class, double.class,
					Pageable.class));

	Object[] params = new Object[] { "BUY", "abcd", 8.88, 3.33, PageRequest.of(1, 444, Sort.Direction.ASC, "price") };

	when(this.datastoreTemplate.queryKeysOrEntities(any(), any())).thenAnswer((invocation) -> {
		EntityQuery statement = invocation.getArgument(0);

		EntityQuery expected = StructuredQuery.newEntityQueryBuilder()
				.setFilter(FILTER)
				.setKind("trades")
				.setOffset(444)
				.setLimit(444)
				.setOrderBy(OrderBy.desc("__key__"), OrderBy.asc("price")).build();

		assertThat(statement).isEqualTo(expected);

		return EMPTY_RESPONSE;
	});

	when(this.queryMethod.getCollectionReturnType()).thenReturn(List.class);

	this.partTreeDatastoreQuery.execute(params);
	verify(this.datastoreTemplate, times(1))
			.queryKeysOrEntities(any(), any());
}
 

@Test
public void nullPageable() throws NoSuchMethodException {
	queryWithMockResult("findTop333ByActionAndSymbolAndPriceLessThanAndPriceGreater"
					+ "ThanEqualAndIdIsNullOrderByIdDesc", null,
			getClass().getMethod("tradeMethod", String.class, String.class, double.class, double.class,
					Pageable.class));

	Object[] params = new Object[] { "BUY", "abcd", 8.88, 3.33, null};

	when(this.datastoreTemplate.queryKeysOrEntities(any(), any())).thenAnswer((invocation) -> {
		EntityQuery statement = invocation.getArgument(0);

		EntityQuery expected = StructuredQuery.newEntityQueryBuilder()
				.setFilter(FILTER)
				.setKind("trades")
				.setLimit(333)
				.setOrderBy(OrderBy.desc("__key__")).build();

		assertThat(statement).isEqualTo(expected);

		return EMPTY_RESPONSE;
	});

	when(this.queryMethod.getCollectionReturnType()).thenReturn(List.class);

	this.partTreeDatastoreQuery.execute(params);
	verify(this.datastoreTemplate, times(1))
			.queryKeysOrEntities(any(), any());
}