Datastore reference query by example (backport) (#2614)

fixes #2611

Author: dmitry-s

spring-cloud-gcp-data-datastore/src/main/java/org/springframework/cloud/gcp/data/datastore/core/DatastoreTemplate.java

@@ -86,7 +86,6 @@
 import org.springframework.data.mapping.AssociationHandler;
 import org.springframework.data.mapping.PersistentProperty;
 import org.springframework.data.mapping.PersistentPropertyAccessor;
-import org.springframework.data.support.ExampleMatcherAccessor;
 import org.springframework.data.util.ClassTypeInformation;
 import org.springframework.lang.Nullable;
 import org.springframework.transaction.support.TransactionSynchronizationManager;
@@ -818,19 +817,30 @@ private <T> StructuredQuery exampleToQuery(Example<T> example, DatastoreQueryOpt
 		DatastorePersistentEntity<?> persistentEntity =
 				this.datastoreMappingContext.getPersistentEntity(example.getProbeType());
 
-		StructuredQuery.Builder builder = keyQuery ? Query.newKeyQueryBuilder() : Query.newEntityQueryBuilder();
-		builder.setKind(persistentEntity.kindName());
-
-		ExampleMatcherAccessor matcherAccessor = new ExampleMatcherAccessor(example.getMatcher());
-		matcherAccessor.getPropertySpecifiers();
 		LinkedList<StructuredQuery.Filter> filters = new LinkedList<>();
+		NullHandler nullHandler = example.getMatcher().getNullHandler();
 		persistentEntity.doWithColumnBackedProperties((persistentProperty) -> {
 			if (!ignoredProperty(example, persistentProperty)) {
 				Value<?> value = getValue(example, probeEntity, persistentEntity, persistentProperty);
-				NullHandler nullHandler = example.getMatcher().getNullHandler();
 				addFilter(nullHandler, filters, persistentProperty.getFieldName(), value);
 			}
 		});
+		persistentEntity.doWithAssociations((AssociationHandler<DatastorePersistentProperty>) association -> {
+			PersistentPropertyAccessor<?> accessor = persistentEntity.getPropertyAccessor(example.getProbe());
+			DatastorePersistentProperty property = association.getInverse();
+			Object value = accessor.getProperty(property);
+			Value<?> key = value == null
+					? NullValue.of()
+					: KeyValue.of(objectToKeyFactory.getKeyFromObject(value,
+						this.datastoreMappingContext.getPersistentEntity(value.getClass())));
+			addFilter(nullHandler, filters, property.getFieldName(), key);
+		});
+
+		StructuredQuery.Builder<?> builder =
+				keyQuery
+						? Query.newKeyQueryBuilder()
+						: Query.newEntityQueryBuilder();
+		builder.setKind(persistentEntity.kindName());
 
 		if (!filters.isEmpty()) {
 			builder.setFilter(

spring-cloud-gcp-data-datastore/src/test/java/org/springframework/cloud/gcp/data/datastore/entities/Product.java

@@ -0,0 +1,66 @@
+/*
+ * Copyright 2017-2019 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.cloud.gcp.data.datastore.entities;
+
+import com.google.cloud.datastore.Key;
+
+import org.springframework.data.annotation.Id;
+import org.springframework.data.annotation.Reference;
+
+public class Product {
+	@Id
+	private Key id;
+
+	@Reference
+	private Store store;
+
+	public Product(Store store) {
+		this.store = store;
+	}
+
+	@Override
+	public String toString() {
+		return "Product{" +
+				"id=" + id +
+				", store=" + store +
+				'}';
+	}
+
+	@Override
+	public boolean equals(Object o) {
+		if (this == o) {
+			return true;
+		}
+		if (!(o instanceof Product)) {
+			return false;
+		}
+
+		Product product = (Product) o;
+
+		if (id != null ? !id.equals(product.id) : product.id != null) {
+			return false;
+		}
+		return store != null ? store.equals(product.store) : product.store == null;
+	}
+
+	@Override
+	public int hashCode() {
+		int result = id != null ? id.hashCode() : 0;
+		result = 31 * result + (store != null ? store.hashCode() : 0);
+		return result;
+	}
+}

spring-cloud-gcp-data-datastore/src/test/java/org/springframework/cloud/gcp/data/datastore/entities/Store.java

@@ -0,0 +1,66 @@
+/*
+ * Copyright 2017-2019 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.cloud.gcp.data.datastore.entities;
+
+import com.google.cloud.datastore.Key;
+
+import org.springframework.cloud.gcp.data.datastore.core.mapping.Entity;
+import org.springframework.data.annotation.Id;
+
+@Entity(name = "store")
+public class Store {
+	@Id
+	private Key id;
+
+	private String name;
+
+	public Store(String name) {
+		this.name = name;
+	}
+
+	@Override
+	public String toString() {
+		return "Store{" +
+				"id=" + id +
+				", name='" + name + '\'' +
+				'}';
+	}
+
+	@Override
+	public boolean equals(Object o) {
+		if (this == o) {
+			return true;
+		}
+		if (!(o instanceof Store)) {
+			return false;
+		}
+
+		Store store = (Store) o;
+
+		if (id != null ? !id.equals(store.id) : store.id != null) {
+			return false;
+		}
+		return name != null ? name.equals(store.name) : store.name == null;
+	}
+
+	@Override
+	public int hashCode() {
+		int result = id != null ? id.hashCode() : 0;
+		result = 31 * result + (name != null ? name.hashCode() : 0);
+		return result;
+	}
+}

spring-cloud-gcp-data-datastore/src/test/java/org/springframework/cloud/gcp/data/datastore/it/DatastoreIntegrationTests.java

@@ -56,7 +56,9 @@
 import org.springframework.cloud.gcp.data.datastore.core.mapping.Entity;
 import org.springframework.cloud.gcp.data.datastore.core.mapping.Unindexed;
 import org.springframework.cloud.gcp.data.datastore.entities.CustomMap;
+import org.springframework.cloud.gcp.data.datastore.entities.Product;
 import org.springframework.cloud.gcp.data.datastore.entities.ServiceConfiguration;
+import org.springframework.cloud.gcp.data.datastore.entities.Store;
 import org.springframework.cloud.gcp.data.datastore.it.TestEntity.Shape;
 import org.springframework.cloud.gcp.data.datastore.repository.DatastoreRepository;
 import org.springframework.cloud.gcp.data.datastore.repository.query.Query;
@@ -67,6 +69,7 @@
 import org.springframework.data.domain.ExampleMatcher;
 import org.springframework.data.domain.Page;
 import org.springframework.data.domain.PageRequest;
+import org.springframework.data.domain.Pageable;
 import org.springframework.data.domain.Slice;
 import org.springframework.data.domain.Sort;
 import org.springframework.test.context.ContextConfiguration;
@@ -97,6 +100,9 @@ public class DatastoreIntegrationTests extends AbstractDatastoreIntegrationTests
 	// consistency.
 	private static final int WAIT_FOR_EVENTUAL_CONSISTENCY_SAFETY_MULTIPLE = 3;
 
+	@Autowired
+	private ProductRepository productRepository;
+
 	@Autowired
 	private TestEntityRepository testEntityRepository;
 
@@ -162,6 +168,8 @@ public void deleteAll() {
 		this.datastoreTemplate.deleteAll(PetOwner.class);
 		this.datastoreTemplate.deleteAll(Event.class);
 		this.datastoreTemplate.deleteAll(LazyEntity.class);
+		this.datastoreTemplate.deleteAll(Product.class);
+		this.datastoreTemplate.deleteAll(Store.class);
 		this.testEntityRepository.deleteAll();
 		if (this.keyForMap != null) {
 			this.datastore.delete(this.keyForMap);
@@ -177,6 +185,40 @@ public void saveEntities() {
 
 	}
 
+	@Test
+	public void testFindByExampleReference() {
+		Store store1 = new Store("store1");
+		Product product1 = new Product(store1);
+
+		productRepository.save(product1);
+
+		Store store2 = new Store("store2");
+		Product product2 = new Product(store2);
+
+		productRepository.save(product2);
+
+		Pageable pageable = PageRequest.of(0, 3);
+		Product product = new Product(store1);
+		Example<Product> example = Example.of(product);
+		Page<Product> pagedProduct = this.productRepository.findAll(example, pageable);
+
+		assertThat(pagedProduct).containsOnly(product1);
+
+		product = new Product(null);
+		example = Example.of(product);
+		pagedProduct = this.productRepository.findAll(example, pageable);
+
+		assertThat(pagedProduct).containsExactlyInAnyOrder(product1, product2);
+
+		product = new Product(null);
+		example = Example.of(product, ExampleMatcher.matching()
+				.withIgnorePaths("id")
+				.withIncludeNullValues());
+		pagedProduct = this.productRepository.findAll(example, pageable);
+
+		assertThat(pagedProduct).isEmpty();
+	}
+
 	@Test
 	public void testFindByExample() {
 		assertThat(this.testEntityRepository
@@ -216,7 +258,7 @@ public void testFindByExample() {
 		assertThat(this.testEntityRepository
 				.findAll(
 						Example.of(new TestEntity(null, null, null, null, null)),
-						Sort.by(Sort.Direction.ASC, "size")))
+						Sort.by(Sort.Direction.ASC, "id")))
 				.containsExactly(this.testEntityA, this.testEntityB, this.testEntityC, this.testEntityD);
 
 		assertThat(this.testEntityRepository

spring-cloud-gcp-data-datastore/src/test/java/org/springframework/cloud/gcp/data/datastore/it/ProductRepository.java

@@ -0,0 +1,28 @@
+/*
+ * Copyright 2017-2019 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.cloud.gcp.data.datastore.it;
+
+import com.google.cloud.datastore.Key;
+
+import org.springframework.cloud.gcp.data.datastore.entities.Product;
+import org.springframework.cloud.gcp.data.datastore.repository.DatastoreRepository;
+import org.springframework.stereotype.Repository;
+
+@Repository
+public interface ProductRepository extends DatastoreRepository<Product, Key> {
+
+}