Skip to content

Commit

Permalink
Resolve references and descendants for subclasses (#377)
Browse files Browse the repository at this point in the history
Fixes: #356.

Co-authored-by: Mike Eltsufin <meltsufin@google.com>
  • Loading branch information
THCoulon and meltsufin authored Mar 18, 2021
1 parent 7442572 commit cc4a44e
Show file tree
Hide file tree
Showing 6 changed files with 290 additions and 12 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -611,19 +611,13 @@ private <T> List<T> convertEntitiesForRead(Collection<? extends BaseKey> keys, C
return Collections.emptyList();
}

DatastorePersistentEntity datastorePersistentEntity = this.datastoreMappingContext
.getPersistentEntity(entityClass);

return keys.stream()
.map(key -> convertEntityResolveDescendantsAndReferences(entityClass,
datastorePersistentEntity,
key,
context)).filter(Objects::nonNull)
key, context)).filter(Objects::nonNull)
.collect(Collectors.toList());
}

private <T> T convertEntityResolveDescendantsAndReferences(Class<T> entityClass,
DatastorePersistentEntity datastorePersistentEntity,
BaseKey key, ReadContext context) {
T convertedObject;
if (context.converted(key)) {
Expand All @@ -640,8 +634,10 @@ private <T> T convertEntityResolveDescendantsAndReferences(Class<T> entityClass,
//raw Datastore entity is no longer needed
context.removeReadEntity(key);
if (convertedObject != null) {
resolveDescendantProperties(datastorePersistentEntity, readEntity, convertedObject, context);
resolveReferenceProperties(datastorePersistentEntity, readEntity, convertedObject, context);
DatastorePersistentEntity<T> discriminatedEntity = this.datastoreEntityConverter
.getDiscriminationPersistentEntity(entityClass, readEntity);
resolveDescendantProperties(discriminatedEntity, readEntity, convertedObject, context);
resolveReferenceProperties(discriminatedEntity, readEntity, convertedObject, context);
}
}

Expand Down Expand Up @@ -737,10 +733,19 @@ private <T> void resolveDescendantProperties(DatastorePersistentEntity datastore
Key entityKey = (Key) entity.getKey();
Key ancestorKey = KeyUtil.getKeyWithoutAncestors(entityKey);

DatastorePersistentEntity descendantEntityType = this.datastoreMappingContext
.getPersistentEntity(descendantType);

Filter ancestorFilter = descendantEntityType.getDiscriminationFieldName() != null ?
StructuredQuery.CompositeFilter.and(
PropertyFilter.eq(descendantEntityType.getDiscriminationFieldName(),
descendantEntityType.getDiscriminatorValue()),
PropertyFilter.hasAncestor(ancestorKey)
) : PropertyFilter.hasAncestor(ancestorKey);

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

List entities = convertEntitiesForRead(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@
import java.util.Map;

import com.google.cloud.datastore.BaseEntity;
import com.google.cloud.spring.data.datastore.core.mapping.DatastorePersistentEntity;

import org.springframework.data.convert.EntityReader;
import org.springframework.data.convert.EntityWriter;
Expand All @@ -40,6 +41,15 @@ public interface DatastoreEntityConverter extends
*/
ReadWriteConversions getConversions();

/**
* Provide a {@link DatastorePersistentEntity} with support for discriminator fields.
* @param entityClass the entity class
* @param entity the Datastore entity
* @param <T> the type of the entity
* @return {@link DatastorePersistentEntity} for the entity type with support for discriminator fields.
*/
<T> DatastorePersistentEntity<T> getDiscriminationPersistentEntity(Class<T> entityClass, BaseEntity<?> entity);

/**
* Read the entity as a {@link Map}.
* @param <T> the type of the key in the map
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -122,6 +122,20 @@ public <T, R> Map<T, R> readAsMap(Class<T> keyType, TypeInformation<R> component
return readAsMap(entity, ClassTypeInformation.from(HashMap.class));
}

public <T> DatastorePersistentEntity<T> getDiscriminationPersistentEntity(Class<T> entityClass, BaseEntity<?> entity) {
DatastorePersistentEntity ostensiblePersistentEntity = this.mappingContext
.getPersistentEntity(entityClass);

if (ostensiblePersistentEntity == null) {
throw new DatastoreDataException("Unable to convert Datastore Entity to " + entityClass);
}

EntityPropertyValueProvider propertyValueProvider = new EntityPropertyValueProvider(entity, this.conversions);

return getDiscriminationPersistentEntity(ostensiblePersistentEntity,
propertyValueProvider);
}

@Override
@SuppressWarnings("unchecked")
public <R> R read(Class<R> aClass, BaseEntity entity) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,7 @@
import com.google.cloud.spring.data.datastore.core.convert.ReadWriteConversions;
import com.google.cloud.spring.data.datastore.core.mapping.DatastoreDataException;
import com.google.cloud.spring.data.datastore.core.mapping.DatastoreMappingContext;
import com.google.cloud.spring.data.datastore.core.mapping.DatastorePersistentEntity;
import com.google.cloud.spring.data.datastore.core.mapping.Descendants;
import com.google.cloud.spring.data.datastore.core.mapping.DiscriminatorField;
import com.google.cloud.spring.data.datastore.core.mapping.DiscriminatorValue;
Expand Down Expand Up @@ -237,12 +238,20 @@ private void setUpConverters(Entity ce1, Query childTestEntityQuery,
QueryResults childTestEntityQueryResults, QueryResults testEntityQueryResults) {
// mocking the converter to return the final objects corresponding to their
// specific entities.
DatastorePersistentEntity testEntityPE = new DatastoreMappingContext().getDatastorePersistentEntity(TestEntity.class);
DatastorePersistentEntity childEntityPE = new DatastoreMappingContext().getDatastorePersistentEntity(ChildEntity.class);
when(this.datastoreEntityConverter.read(TestEntity.class, this.e1))
.thenReturn(this.ob1);
when(this.datastoreEntityConverter.getDiscriminationPersistentEntity(TestEntity.class, this.e1))
.thenReturn(testEntityPE);
when(this.datastoreEntityConverter.read(TestEntity.class, this.e2))
.thenReturn(this.ob2);
when(this.datastoreEntityConverter.getDiscriminationPersistentEntity(TestEntity.class, this.e2))
.thenReturn(testEntityPE);
when(this.datastoreEntityConverter.read(eq(ChildEntity.class), same(ce1)))
.thenAnswer(invocationOnMock -> createChildEntity());
when(this.datastoreEntityConverter.getDiscriminationPersistentEntity(eq(ChildEntity.class), same(ce1)))
.thenReturn(childEntityPE);

doAnswer(invocation -> {
FullEntity.Builder builder = invocation.getArgument(1);
Expand Down Expand Up @@ -483,13 +492,21 @@ public void findAllReferenceLoopTest() {
ReferenceTestEntity childEntity = new ReferenceTestEntity();
ReferenceTestEntity childEntity2 = new ReferenceTestEntity();

DatastorePersistentEntity referenceTestEntityPE = new DatastoreMappingContext().getDatastorePersistentEntity(ReferenceTestEntity.class);

when(this.datastoreEntityConverter.read(eq(ReferenceTestEntity.class), same(referenceTestDatastoreEntity)))
.thenAnswer(invocationOnMock -> referenceTestEntity);
when(this.datastoreEntityConverter.getDiscriminationPersistentEntity(eq(ReferenceTestEntity.class), same(referenceTestDatastoreEntity)))
.thenReturn(referenceTestEntityPE);

when(this.datastoreEntityConverter.read(eq(ReferenceTestEntity.class), same(child)))
.thenAnswer(invocationOnMock -> childEntity);
when(this.datastoreEntityConverter.getDiscriminationPersistentEntity(eq(ReferenceTestEntity.class), same(child)))
.thenReturn(referenceTestEntityPE);
when(this.datastoreEntityConverter.read(eq(ReferenceTestEntity.class), same(child2)))
.thenAnswer(invocationOnMock -> childEntity2);
when(this.datastoreEntityConverter.getDiscriminationPersistentEntity(eq(ReferenceTestEntity.class), same(child2)))
.thenReturn(referenceTestEntityPE);

verifyBeforeAndAfterEvents(null,
new AfterFindByKeyEvent(Collections.singletonList(referenceTestEntity),
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,118 @@
/*
* Copyright 2017-2018 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 com.google.cloud.spring.data.datastore.it.subclasses.descendants;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

import com.google.cloud.datastore.Key;
import com.google.cloud.spring.data.datastore.core.DatastoreTemplate;
import com.google.cloud.spring.data.datastore.core.mapping.Descendants;
import com.google.cloud.spring.data.datastore.core.mapping.DiscriminatorField;
import com.google.cloud.spring.data.datastore.core.mapping.DiscriminatorValue;
import com.google.cloud.spring.data.datastore.core.mapping.Entity;
import com.google.cloud.spring.data.datastore.it.AbstractDatastoreIntegrationTests;
import com.google.cloud.spring.data.datastore.it.DatastoreIntegrationTestConfiguration;
import com.google.cloud.spring.data.datastore.repository.DatastoreRepository;
import org.junit.After;
import org.junit.BeforeClass;
import org.junit.Test;
import org.junit.runner.RunWith;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.annotation.Id;
import org.springframework.stereotype.Repository;
import org.springframework.test.context.ContextConfiguration;
import org.springframework.test.context.junit4.SpringRunner;

import static org.assertj.core.api.Assertions.assertThat;
import static org.hamcrest.Matchers.is;
import static org.junit.Assume.assumeThat;

@Repository
interface SubclassesDescendantsEntityARepository extends DatastoreRepository<EntityA, Key> {
}

@RunWith(SpringRunner.class)
@ContextConfiguration(classes = { DatastoreIntegrationTestConfiguration.class })
public class SubclassesDescendantsIntegrationTests extends AbstractDatastoreIntegrationTests {

@Autowired
SubclassesDescendantsEntityARepository entityARepository;

@Autowired
private DatastoreTemplate datastoreTemplate;

@BeforeClass
public static void checkToRun() {
assumeThat(
"Datastore integration tests are disabled. Please use '-Dit.datastore=true' "
+ "to enable them. ",
System.getProperty("it.datastore"), is("true"));
}

@After
public void deleteAll() {
datastoreTemplate.deleteAll(EntityA.class);
datastoreTemplate.deleteAll(EntityB.class);
datastoreTemplate.deleteAll(EntityC.class);
}

@Test
public void testEntityCContainsReferenceToEntityB() {
EntityB entityB_1 = new EntityB();
EntityC entityC_1 = new EntityC();
entityB_1.addEntityC(entityC_1);
entityARepository.saveAll(Arrays.asList(entityB_1, entityC_1));
EntityB fetchedB = (EntityB) entityARepository.findById(entityB_1.getId()).get();
List<EntityC> entitiesCOfB = fetchedB.getEntitiesC();
assertThat(entitiesCOfB).hasSize(1);
}

}

@Entity(name = "A")
@DiscriminatorField(field = "type")
abstract class EntityA {
@Id
private Key id;

public Key getId() {
return id;
}
}

@Entity(name = "A")
@DiscriminatorValue("B")
class EntityB extends EntityA {
@Descendants
private List<EntityC> entitiesC = new ArrayList<>();

public void addEntityC(EntityC entityCDescendants) {
this.entitiesC.add(entityCDescendants);
}

public List<EntityC> getEntitiesC() {
return entitiesC;
}
}

@Entity(name = "A")
@DiscriminatorValue("C")
class EntityC extends EntityA {
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,114 @@
/*
* Copyright 2017-2018 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 com.google.cloud.spring.data.datastore.it.subclasses.references;

import java.util.Arrays;

import com.google.cloud.datastore.Key;
import com.google.cloud.spring.data.datastore.core.DatastoreTemplate;
import com.google.cloud.spring.data.datastore.core.mapping.DiscriminatorField;
import com.google.cloud.spring.data.datastore.core.mapping.DiscriminatorValue;
import com.google.cloud.spring.data.datastore.core.mapping.Entity;
import com.google.cloud.spring.data.datastore.it.AbstractDatastoreIntegrationTests;
import com.google.cloud.spring.data.datastore.it.DatastoreIntegrationTestConfiguration;
import com.google.cloud.spring.data.datastore.repository.DatastoreRepository;
import org.junit.After;
import org.junit.BeforeClass;
import org.junit.Test;
import org.junit.runner.RunWith;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.annotation.Id;
import org.springframework.data.annotation.Reference;
import org.springframework.stereotype.Repository;
import org.springframework.test.context.ContextConfiguration;
import org.springframework.test.context.junit4.SpringRunner;

import static org.assertj.core.api.Assertions.assertThat;
import static org.hamcrest.Matchers.is;
import static org.junit.Assume.assumeThat;

@Repository
interface SubclassesReferencesEntityARepository extends DatastoreRepository<EntityA, Key> {
}

@RunWith(SpringRunner.class)
@ContextConfiguration(classes = { DatastoreIntegrationTestConfiguration.class })
public class SubclassesReferencesIntegrationTests extends AbstractDatastoreIntegrationTests {

@Autowired
SubclassesReferencesEntityARepository entityARepository;

@Autowired
private DatastoreTemplate datastoreTemplate;

@BeforeClass
public static void checkToRun() {
assumeThat(
"Datastore integration tests are disabled. Please use '-Dit.datastore=true' "
+ "to enable them. ",
System.getProperty("it.datastore"), is("true"));
}

@After
public void deleteAll() {
datastoreTemplate.deleteAll(EntityA.class);
datastoreTemplate.deleteAll(EntityB.class);
datastoreTemplate.deleteAll(EntityC.class);
}

@Test
public void testEntityCContainsReferenceToEntityB() {
EntityB entityB_1 = new EntityB();
EntityC entityC_1 = new EntityC(entityB_1);
entityARepository.saveAll(Arrays.asList(entityB_1, entityC_1));
EntityC fetchedC = (EntityC) entityARepository.findById(entityC_1.getId()).get();
assertThat(fetchedC.getEntityB()).isNotNull();
}

}

@Entity(name = "A")
@DiscriminatorField(field = "type")
abstract class EntityA {
@Id
private Key id;

public Key getId() {
return id;
}
}

@Entity(name = "A")
@DiscriminatorValue("B")
class EntityB extends EntityA {
}

@Entity(name = "A")
@DiscriminatorValue("C")
class EntityC extends EntityA {
@Reference
private EntityB entityB;

EntityC(EntityB entityB) {
this.entityB = entityB;
}

public EntityB getEntityB() {
return entityB;
}
}

0 comments on commit cc4a44e

Please sign in to comment.