From b90ed852dd095e4421d801ca5c84448f29671b08 Mon Sep 17 00:00:00 2001 From: Min Zhu Date: Wed, 28 Jul 2021 17:26:41 -0400 Subject: [PATCH] Support Stream as return type to datastore queries (#551) fixes #452. Fixed code in GqlDatastoreQuery and PartTreeDatastoreQuery to support for Stream as a return type and added corresponding unit tests and integration tests. --- docs/src/main/asciidoc/datastore.adoc | 6 ++++ .../repository/query/GqlDatastoreQuery.java | 5 +++- .../query/PartTreeDatastoreQuery.java | 17 +++++++++-- .../it/DatastoreIntegrationTests.java | 15 ++++++++++ .../datastore/it/TestEntityRepository.java | 6 ++++ .../query/GqlDatastoreQueryTests.java | 30 +++++++++++++++++++ .../query/PartTreeDatastoreQueryTests.java | 20 +++++++++++++ .../example/DatastoreRepositoryExample.java | 5 ++++ .../java/com/example/SingerRepository.java | 3 ++ ...toreSampleApplicationIntegrationTests.java | 8 +++++ 10 files changed, 111 insertions(+), 4 deletions(-) diff --git a/docs/src/main/asciidoc/datastore.adoc b/docs/src/main/asciidoc/datastore.adoc index 0ee4606f69..80e4bbf7aa 100644 --- a/docs/src/main/asciidoc/datastore.adoc +++ b/docs/src/main/asciidoc/datastore.adoc @@ -739,6 +739,7 @@ In addition to retrieving entities by their IDs, you can also submit queries. ---- These methods, respectively, allow querying for: + * entities mapped by a given entity class using all the same mapping and converting features * arbitrary types produced by a given mapping function * only the Cloud Datastore keys of the entities found by the query @@ -949,6 +950,8 @@ public interface TradeRepository extends DatastoreRepository { Slice findBySymbol(String symbol, Pageable pageable); List findBySymbol(String symbol, Sort sort); + + Stream findBySymbol(String symbol); } ---- @@ -1039,6 +1042,9 @@ public interface TraderRepository extends DatastoreRepository { @Query("SELECT * FROM traders WHERE name = @trader_name") List tradersByName(@Param("trader_name") String traderName); + @Query("SELECT * FROM traders WHERE name = @trader_name") + Stream tradersStreamByName(@Param("trader_name") String traderName); + @Query("SELECT * FROM test_entities_ci WHERE name = @trader_name") TestEntity getOneTestEntity(@Param("trader_name") String traderName); diff --git a/spring-cloud-gcp-data-datastore/src/main/java/com/google/cloud/spring/data/datastore/repository/query/GqlDatastoreQuery.java b/spring-cloud-gcp-data-datastore/src/main/java/com/google/cloud/spring/data/datastore/repository/query/GqlDatastoreQuery.java index 3dccc8de19..7546c333e8 100644 --- a/spring-cloud-gcp-data-datastore/src/main/java/com/google/cloud/spring/data/datastore/repository/query/GqlDatastoreQuery.java +++ b/spring-cloud-gcp-data-datastore/src/main/java/com/google/cloud/spring/data/datastore/repository/query/GqlDatastoreQuery.java @@ -149,7 +149,7 @@ public Object execute(Object[] parameters) { if (isPageQuery() || isSliceQuery()) { result = buildPageOrSlice(parameters, parsedQueryWithTagsAndValues, found); } - else if (this.queryMethod.isCollectionQuery()) { + else if (this.queryMethod.isCollectionQuery() || this.queryMethod.isStreamQuery()) { result = convertCollectionResult(returnedItemType, found); } else { @@ -199,6 +199,9 @@ private Page buildPage(Pageable pageableParam, ParsedQueryWithTagsAndValues pars } private Object convertCollectionResult(Class returnedItemType, Iterable rawResult) { + if (this.queryMethod.isStreamQuery()) { + return StreamSupport.stream(rawResult.spliterator(), false); + } Object result = this.datastoreOperations.getDatastoreEntityConverter() .getConversions().convertOnRead( rawResult, this.queryMethod.getCollectionReturnType(), returnedItemType); diff --git a/spring-cloud-gcp-data-datastore/src/main/java/com/google/cloud/spring/data/datastore/repository/query/PartTreeDatastoreQuery.java b/spring-cloud-gcp-data-datastore/src/main/java/com/google/cloud/spring/data/datastore/repository/query/PartTreeDatastoreQuery.java index f3b8b9a099..1bb977c4a3 100644 --- a/spring-cloud-gcp-data-datastore/src/main/java/com/google/cloud/spring/data/datastore/repository/query/PartTreeDatastoreQuery.java +++ b/spring-cloud-gcp-data-datastore/src/main/java/com/google/cloud/spring/data/datastore/repository/query/PartTreeDatastoreQuery.java @@ -186,7 +186,8 @@ public Object execute(Object[] parameters) { } private Object runQuery(Object[] parameters, Class returnedElementType, Class collectionType, boolean requiresCount) { - ExecutionOptions options = new ExecutionOptions(returnedElementType, collectionType, requiresCount); + ExecutionOptions options = new ExecutionOptions(returnedElementType, collectionType, requiresCount, + getQueryMethod().isStreamQuery()); DatastoreResultsIterable rawResults = getDatastoreOperations() .queryKeysOrEntities( @@ -194,6 +195,10 @@ private Object runQuery(Object[] parameters, Class returnedElementType, Class requiresCount, options.isSingularResult(), null), this.entityType); + if (getQueryMethod().isStreamQuery()) { + return StreamSupport.stream(rawResults.spliterator(), false); + } + Object result = StreamSupport.stream(rawResults.spliterator(), false) .map(options.isReturnedTypeIsNumber() ? Function.identity() : this::processRawObjectForProjection) .collect(options.getResultsCollector()); @@ -422,7 +427,7 @@ private class ExecutionOptions { private boolean singularResult; - ExecutionOptions(Class returnedElementType, Class collectionType, boolean requiresCount) { + ExecutionOptions(Class returnedElementType, Class collectionType, boolean requiresCount, boolean isStreamQuery) { returnedTypeIsNumber = Number.class.isAssignableFrom(returnedElementType) || returnedElementType == int.class || returnedElementType == long.class; @@ -439,7 +444,13 @@ private class ExecutionOptions { structuredQueryBuilder.setKind(PartTreeDatastoreQuery.this.datastorePersistentEntity.kindName()); - singularResult = (!isCountingQuery && collectionType == null) && !PartTreeDatastoreQuery.this.tree.isDelete(); + if (isCountingQuery || collectionType != null || isStreamQuery + || PartTreeDatastoreQuery.this.tree.isDelete()) { + singularResult = false; + } + else { + singularResult = true; + } } boolean isReturnedTypeIsNumber() { diff --git a/spring-cloud-gcp-data-datastore/src/test/java/com/google/cloud/spring/data/datastore/it/DatastoreIntegrationTests.java b/spring-cloud-gcp-data-datastore/src/test/java/com/google/cloud/spring/data/datastore/it/DatastoreIntegrationTests.java index a9f8ab4f65..d22eb20736 100644 --- a/spring-cloud-gcp-data-datastore/src/test/java/com/google/cloud/spring/data/datastore/it/DatastoreIntegrationTests.java +++ b/spring-cloud-gcp-data-datastore/src/test/java/com/google/cloud/spring/data/datastore/it/DatastoreIntegrationTests.java @@ -29,6 +29,7 @@ import java.util.concurrent.TimeUnit; import java.util.stream.Collectors; import java.util.stream.IntStream; +import java.util.stream.Stream; import com.google.cloud.datastore.Blob; import com.google.cloud.datastore.DatastoreException; @@ -1037,6 +1038,20 @@ public void newFieldTest() { assertThat(companyWithBooleanPrimitive.name).isEqualTo(company.name); assertThat(companyWithBooleanPrimitive.active).isFalse(); } + + @Test + public void returnStreamPartTreeTest() { + this.testEntityRepository.saveAll(this.allTestEntities); + Stream resultStream = this.testEntityRepository.findPartTreeStreamByColor("red"); + assertThat(resultStream).hasSize(3).contains(testEntityA, testEntityC, testEntityD); + } + + @Test + public void returnStreamGqlTest() { + this.testEntityRepository.saveAll(this.allTestEntities); + Stream resultStream = this.testEntityRepository.findGqlStreamByColor("red"); + assertThat(resultStream).hasSize(3).contains(testEntityA, testEntityC, testEntityD); + } } /** diff --git a/spring-cloud-gcp-data-datastore/src/test/java/com/google/cloud/spring/data/datastore/it/TestEntityRepository.java b/spring-cloud-gcp-data-datastore/src/test/java/com/google/cloud/spring/data/datastore/it/TestEntityRepository.java index 4c800909ff..d6f08516c2 100644 --- a/spring-cloud-gcp-data-datastore/src/test/java/com/google/cloud/spring/data/datastore/it/TestEntityRepository.java +++ b/spring-cloud-gcp-data-datastore/src/test/java/com/google/cloud/spring/data/datastore/it/TestEntityRepository.java @@ -20,6 +20,7 @@ import java.util.List; import java.util.Optional; import java.util.Set; +import java.util.stream.Stream; import javax.annotation.Nonnull; @@ -134,6 +135,11 @@ public interface TestEntityRepository extends DatastoreRepository findFirstByColor(String color); + Stream findPartTreeStreamByColor(String color); + + @Query("select * from test_entities_ci where color = @color") + Stream findGqlStreamByColor(@Param("color") String color); + @Nullable TestEntity getByColor(String color); diff --git a/spring-cloud-gcp-data-datastore/src/test/java/com/google/cloud/spring/data/datastore/repository/query/GqlDatastoreQueryTests.java b/spring-cloud-gcp-data-datastore/src/test/java/com/google/cloud/spring/data/datastore/repository/query/GqlDatastoreQueryTests.java index a11bcc7d74..b1ae5de617 100644 --- a/spring-cloud-gcp-data-datastore/src/test/java/com/google/cloud/spring/data/datastore/repository/query/GqlDatastoreQueryTests.java +++ b/spring-cloud-gcp-data-datastore/src/test/java/com/google/cloud/spring/data/datastore/repository/query/GqlDatastoreQueryTests.java @@ -22,6 +22,7 @@ import java.util.List; import java.util.Map; import java.util.Optional; +import java.util.stream.Stream; import com.google.cloud.datastore.Cursor; import com.google.cloud.datastore.DoubleValue; @@ -447,6 +448,35 @@ public void pageableTestPageCursor() { .queryKeysOrEntities(any(), eq(Trade.class)); } + @Test + public void streamResultTest() { + Mockito.when(this.queryMethod.getReturnedObjectType()).thenReturn(Trade.class); + Parameters parameters = mock(Parameters.class); + when(this.queryMethod.getParameters()).thenReturn(parameters); + when(parameters.getNumberOfParameters()).thenReturn(0); + when(this.queryMethod.isStreamQuery()).thenReturn(true); + + Trade tradeA = new Trade(); + tradeA.id = "a"; + Trade tradeB = new Trade(); + tradeB.id = "b"; + doAnswer(invocation -> { + GqlQuery statement = invocation.getArgument(0); + assertThat(statement.getQueryString()).isEqualTo("unusedGqlString"); + + Cursor cursor = Cursor.copyFrom("abc".getBytes()); + DatastoreResultsIterable datastoreResultsIterable = new DatastoreResultsIterable( + Arrays.asList(tradeA, tradeB), cursor); + return datastoreResultsIterable; + }).when(this.datastoreTemplate).queryKeysOrEntities(any(), eq(Trade.class)); + + GqlDatastoreQuery gqlDatastoreQuery = createQuery("unusedGqlString", false, false); + + Object result = gqlDatastoreQuery.execute(new Parameters[0]); + assertThat(result).isInstanceOf(Stream.class); + assertThat((Stream) result).hasSize(2).containsExactly(tradeA, tradeB); + } + private Parameters buildParameters(Object[] params, String[] paramNames) { Parameters parameters = mock(Parameters.class); diff --git a/spring-cloud-gcp-data-datastore/src/test/java/com/google/cloud/spring/data/datastore/repository/query/PartTreeDatastoreQueryTests.java b/spring-cloud-gcp-data-datastore/src/test/java/com/google/cloud/spring/data/datastore/repository/query/PartTreeDatastoreQueryTests.java index 3839170993..6cad19cb0e 100644 --- a/spring-cloud-gcp-data-datastore/src/test/java/com/google/cloud/spring/data/datastore/repository/query/PartTreeDatastoreQueryTests.java +++ b/spring-cloud-gcp-data-datastore/src/test/java/com/google/cloud/spring/data/datastore/repository/query/PartTreeDatastoreQueryTests.java @@ -25,6 +25,7 @@ import java.util.List; import java.util.Optional; import java.util.function.Function; +import java.util.stream.Stream; import com.google.cloud.datastore.Cursor; import com.google.cloud.datastore.EntityQuery; @@ -846,6 +847,21 @@ public void nonCollectionReturnTypeNoResultsOptional() throws NoSuchMethodExcept assertThat((Optional) this.partTreeDatastoreQuery.execute(params)).isNotPresent(); } + @Test + public void streamResultTest() throws NoSuchMethodException { + Trade tradeA = new Trade(); + tradeA.id = "a"; + Trade tradeB = new Trade(); + tradeB.id = "b"; + queryWithMockResult("findStreamByAction", Arrays.asList(tradeA, tradeB), + getClass().getMethod("findStreamByAction", String.class)); + when(this.queryMethod.isStreamQuery()).thenReturn(true); + Object[] params = new Object[] { "BUY", }; + Object result = this.partTreeDatastoreQuery.execute(params); + assertThat(result).isInstanceOf(Stream.class); + assertThat((Stream) result).hasSize(2).contains(tradeA, tradeB); + } + private void queryWithMockResult(String queryName, List results, Method m, ProjectionInformation projectionInformation) { queryWithMockResult(queryName, results, m, false, projectionInformation); @@ -884,6 +900,10 @@ public Trade findByAction(String action) { return null; } + public Stream findStreamByAction(String action) { + return null; + } + @Nullable public Trade findByActionNullable(String action) { return null; diff --git a/spring-cloud-gcp-samples/spring-cloud-gcp-data-datastore-sample/src/main/java/com/example/DatastoreRepositoryExample.java b/spring-cloud-gcp-samples/spring-cloud-gcp-data-datastore-sample/src/main/java/com/example/DatastoreRepositoryExample.java index 56129a0b8b..2c0e6fce6a 100644 --- a/spring-cloud-gcp-samples/spring-cloud-gcp-data-datastore-sample/src/main/java/com/example/DatastoreRepositoryExample.java +++ b/spring-cloud-gcp-samples/spring-cloud-gcp-data-datastore-sample/src/main/java/com/example/DatastoreRepositoryExample.java @@ -23,6 +23,7 @@ import java.util.List; import java.util.TreeSet; import java.util.stream.Collectors; +import java.util.stream.Stream; import org.springframework.beans.factory.annotation.Autowired; import org.springframework.boot.CommandLineRunner; @@ -99,6 +100,10 @@ private void retrieveAndPrintSingers() { .findAllById(Arrays.asList("singer1", "singer2", "singer3")) .forEach(x -> System.out.println("retrieved singer: " + x)); + System.out.println("Query results can also be returned as Stream: "); + Stream streamResult = singerRepository.findSingersByLastName("Doe"); + streamResult.forEach(System.out::println); + //Query by example: find all singers with the last name "Doe" Iterable singers = this.singerRepository.findAll( Example.of(new Singer(null, null, "Doe", null), ExampleMatcher.matching().withIgnorePaths("singerId"))); diff --git a/spring-cloud-gcp-samples/spring-cloud-gcp-data-datastore-sample/src/main/java/com/example/SingerRepository.java b/spring-cloud-gcp-samples/spring-cloud-gcp-data-datastore-sample/src/main/java/com/example/SingerRepository.java index 5512817974..cd2019eefd 100644 --- a/spring-cloud-gcp-samples/spring-cloud-gcp-data-datastore-sample/src/main/java/com/example/SingerRepository.java +++ b/spring-cloud-gcp-samples/spring-cloud-gcp-data-datastore-sample/src/main/java/com/example/SingerRepository.java @@ -17,6 +17,7 @@ package com.example; import java.util.List; +import java.util.stream.Stream; import com.google.cloud.spring.data.datastore.repository.DatastoreRepository; import com.google.cloud.spring.data.datastore.repository.query.Query; @@ -39,4 +40,6 @@ public interface SingerRepository extends DatastoreRepository { List findSingersByFirstBand(@Param("band") Band band); List findByFirstBand(Band band); + + Stream findSingersByLastName(String name); } diff --git a/spring-cloud-gcp-samples/spring-cloud-gcp-data-datastore-sample/src/test/java/com/example/DatastoreSampleApplicationIntegrationTests.java b/spring-cloud-gcp-samples/spring-cloud-gcp-data-datastore-sample/src/test/java/com/example/DatastoreSampleApplicationIntegrationTests.java index 3c58e26b62..21dd4075a6 100644 --- a/spring-cloud-gcp-samples/spring-cloud-gcp-data-datastore-sample/src/test/java/com/example/DatastoreSampleApplicationIntegrationTests.java +++ b/spring-cloud-gcp-samples/spring-cloud-gcp-data-datastore-sample/src/test/java/com/example/DatastoreSampleApplicationIntegrationTests.java @@ -25,6 +25,7 @@ import java.util.Map; import java.util.concurrent.TimeUnit; import java.util.stream.Collectors; +import java.util.stream.Stream; import com.fasterxml.jackson.core.type.TypeReference; import com.fasterxml.jackson.databind.ObjectMapper; @@ -240,4 +241,11 @@ private List getSingers(String url) throws java.io.IOException { (String) som.get("firstName"), (String) som.get("lastName"), null)) .collect(Collectors.toList()); } + + @Test + public void testQueryReturnStream() { + Stream streamResult = singerRepository.findSingersByLastName("Doe"); + assertThat(streamResult).isInstanceOf(Stream.class); + streamResult.map(Singer::getLastName).forEach(x -> assertThat(x).isEqualTo("Doe")); + } }