Fixing bug when PK is inhertied from a parent entity.

sdk/spring/azure-spring-data-cosmos/CHANGELOG.md

@@ -9,6 +9,7 @@
 #### Breaking Changes
 
 #### Bugs Fixed
+* Fixing bug when the PartitionKey is on a parent entity that was introduced when bulk support was added - See [PR 43801](https://github.com/Azure/azure-sdk-for-java/pull/43801).
 
 #### Other Changes
 

sdk/spring/azure-spring-data-cosmos/src/main/java/com/azure/spring/data/cosmos/repository/support/CosmosEntityInformation.java

@@ -306,11 +306,23 @@ public Object getPartitionKeyFieldValue(T entity) {
             parts.forEach(part -> {
                 if (!part.isEmpty()) {
                     Field f = null;
-                    try {
-                        f = currentObject[0].getClass().getDeclaredField(part);
-                    } catch (NoSuchFieldException e) {
-                        throw new RuntimeException(e);
+                    NoSuchFieldException noSuchFieldException = null;
+                    Class<?> currentClass = currentObject[0].getClass();
+                    while (currentClass != null) {
+                        try {
+                            f = currentClass.getDeclaredField(part);
+                        } catch (NoSuchFieldException e) {
+                            currentClass = currentClass.getSuperclass();
+                            noSuchFieldException = e;
+                        }
+                        if (f != null) {
+                            break;
+                        }
                     }
+                    if (f == null && noSuchFieldException != null) {
+                        throw new RuntimeException(noSuchFieldException);
+                    }
+
                     ReflectionUtils.makeAccessible(f);
                     currentObject[0] = ReflectionUtils.getField(f, currentObject[0]);
                 }

sdk/spring/azure-spring-data-cosmos/src/test/java/com/azure/spring/data/cosmos/domain/AbstractEntity.java

@@ -0,0 +1,18 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+package com.azure.spring.data.cosmos.domain;
+
+import org.springframework.data.annotation.Id;
+
+public abstract class AbstractEntity {
+    @Id
+    private String id;
+
+    public String getId() {
+        return id;
+    }
+
+    public void setId(String id) {
+        this.id = id;
+    }
+}

sdk/spring/azure-spring-data-cosmos/src/test/java/com/azure/spring/data/cosmos/domain/EntityImpl.java

@@ -0,0 +1,8 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+package com.azure.spring.data.cosmos.domain;
+import com.azure.spring.data.cosmos.core.mapping.Container;
+
+@Container(partitionKeyPath = "/id")
+public class EntityImpl extends AbstractEntity {
+}

sdk/spring/azure-spring-data-cosmos/src/test/java/com/azure/spring/data/cosmos/repository/EntityImplRepositoryIT.java

@@ -0,0 +1,75 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+package com.azure.spring.data.cosmos.repository;
+
+import com.azure.spring.data.cosmos.IntegrationTestCollectionManager;
+import com.azure.spring.data.cosmos.common.ResponseDiagnosticsTestUtils;
+import com.azure.spring.data.cosmos.config.CosmosConfig;
+import com.azure.spring.data.cosmos.core.CosmosTemplate;
+import com.azure.spring.data.cosmos.domain.Address;
+import com.azure.spring.data.cosmos.domain.EntityImpl;
+import com.azure.spring.data.cosmos.repository.repository.EntityImplRepository;
+import org.assertj.core.util.Lists;
+import org.junit.Before;
+import org.junit.ClassRule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.test.context.ContextConfiguration;
+import org.springframework.test.context.junit4.SpringJUnit4ClassRunner;
+
+import java.util.Iterator;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+@RunWith(SpringJUnit4ClassRunner.class)
+@ContextConfiguration(classes = TestRepositoryConfig.class)
+public class EntityImplRepositoryIT {
+
+    @ClassRule
+    public static final IntegrationTestCollectionManager collectionManager = new IntegrationTestCollectionManager();
+
+    @Autowired
+    EntityImplRepository repository;
+
+    @Autowired
+    CosmosConfig cosmosConfig;
+
+    @Autowired
+    private CosmosTemplate template;
+
+    @Autowired
+    private ResponseDiagnosticsTestUtils responseDiagnosticsTestUtils;
+
+    private final EntityImpl ENTITYIMPL_1 = new EntityImpl();
+    private final EntityImpl ENTITYIMPL_2 = new EntityImpl();
+
+    @Before
+    public void setUp() {
+        collectionManager.ensureContainersCreatedAndEmpty(template, Address.class);
+        ENTITYIMPL_1.setId("entityImpl1");
+        ENTITYIMPL_2.setId("entityImpl2");
+    }
+
+    @Test
+    public void testSave() {
+        EntityImpl savedEntityImpl = repository.save(ENTITYIMPL_1);
+        assertThat(savedEntityImpl.getId()).isEqualTo(ENTITYIMPL_1.getId());
+
+        Iterable<EntityImpl> savedEI = repository.findAll();
+        assertThat(savedEI.iterator().next().getId()).isEqualTo(ENTITYIMPL_1.getId());
+    }
+
+    @Test
+    public void testSaveAll() {
+        Iterable<EntityImpl> savedEntityImpl = repository.saveAll(Lists.newArrayList(ENTITYIMPL_1, ENTITYIMPL_2));
+        Iterator<EntityImpl> iter = savedEntityImpl.iterator();
+        assertThat(iter.next().getId()).isEqualTo(ENTITYIMPL_1.getId());
+        assertThat(iter.next().getId()).isEqualTo(ENTITYIMPL_2.getId());
+
+        Iterable<EntityImpl> savedEI = repository.findAll();
+        Iterator<EntityImpl> iter2 = savedEI.iterator();
+        assertThat(iter2.next().getId()).isEqualTo(ENTITYIMPL_1.getId());
+        assertThat(iter2.next().getId()).isEqualTo(ENTITYIMPL_2.getId());
+    }
+}

sdk/spring/azure-spring-data-cosmos/src/test/java/com/azure/spring/data/cosmos/repository/repository/EntityImplRepository.java

@@ -0,0 +1,11 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+package com.azure.spring.data.cosmos.repository.repository;
+
+import com.azure.spring.data.cosmos.domain.EntityImpl;
+import com.azure.spring.data.cosmos.repository.CosmosRepository;
+import org.springframework.stereotype.Repository;
+
+@Repository
+public interface EntityImplRepository extends CosmosRepository<EntityImpl, String> {
+}