Custom partition key name (#383)

src/main/java/com/microsoft/azure/spring/data/cosmosdb/core/mapping/PartitionKey.java

@@ -11,4 +11,11 @@ import java.lang.annotation.*;
 @Retention(RetentionPolicy.RUNTIME)
 @Target({ElementType.FIELD})
 public @interface PartitionKey {
+
+    /**
+     * The name of the partition key if the serialized attribute name differs from the field name
+     *
+     * @return
+     */
+    String value() default "";
 }

src/main/java/com/microsoft/azure/spring/data/cosmosdb/repository/support/DocumentDbEntityInformation.java

@@ -84,7 +84,12 @@ public class DocumentDbEntityInformation<T, ID> extends AbstractEntityInformatio
     }
 
     public String getPartitionKeyFieldName() {
-        return partitionKeyField == null ? null : partitionKeyField.getName();
+        if (partitionKeyField == null) {
+            return null;
+        } else {
+            final PartitionKey partitionKey = partitionKeyField.getAnnotation(PartitionKey.class);
+            return partitionKey.value().equals("") ? partitionKeyField.getName() : partitionKey.value();
+        }
     }
 
     public String getPartitionKeyFieldValue(T entity) {

src/test/java/com/microsoft/azure/spring/data/cosmosdb/repository/support/DocumentDbEntityInformationUnitTest.java

@@ -7,6 +7,7 @@ package com.microsoft.azure.spring.data.cosmosdb.repository.support;
 
 import com.microsoft.azure.spring.data.cosmosdb.common.TestConstants;
 import com.microsoft.azure.spring.data.cosmosdb.core.mapping.Document;
+import com.microsoft.azure.spring.data.cosmosdb.core.mapping.PartitionKey;
 import com.microsoft.azure.spring.data.cosmosdb.domain.Address;
 import com.microsoft.azure.spring.data.cosmosdb.domain.Person;
 import org.junit.Test;
@@ -60,9 +61,82 @@ public class DocumentDbEntityInformationUnitTest {
         assertThat(collectionName).isEqualTo("testCollection");
     }
 
+    @Test
+    public void testGetPartitionKeyName() {
+        final DocumentDbEntityInformation<VolunteerWithPartitionKey, String> entityInformation =
+                new DocumentDbEntityInformation<>(VolunteerWithPartitionKey.class);
+
+        final String partitionKeyName = entityInformation.getPartitionKeyFieldName();
+        assertThat(partitionKeyName).isEqualTo("name");
+    }
+
+    @Test
+    public void testNullPartitionKeyName() {
+        final DocumentDbEntityInformation<Volunteer, String> entityInformation =
+                new DocumentDbEntityInformation<>(Volunteer.class);
+
+        final String partitionKeyName = entityInformation.getPartitionKeyFieldName();
+        assertThat(partitionKeyName).isEqualTo(null);
+    }
+
+    @Test
+    public void testCustomPartitionKeyName() {
+        final DocumentDbEntityInformation<VolunteerWithCustomPartitionKey, String> entityInformation =
+                new DocumentDbEntityInformation<>(VolunteerWithCustomPartitionKey.class);
+
+        final String partitionKeyName = entityInformation.getPartitionKeyFieldName();
+        assertThat(partitionKeyName).isEqualTo("vol_name");
+    }
+
     @Document(collection = "testCollection")
     class Volunteer {
         String id;
         String name;
     }
+
+    @Document
+    private class VolunteerWithCustomPartitionKey {
+        private String id;
+        @PartitionKey("vol_name")
+        private String name;
+
+        public String getId() {
+            return id;
+        }
+
+        public void setId(String id) {
+            this.id = id;
+        }
+
+        public String getName() {
+            return name;
+        }
+
+        public void setName(String name) {
+            this.name = name;
+        }
+    }
+
+    @Document
+    private class VolunteerWithPartitionKey {
+        private String id;
+        @PartitionKey
+        private String name;
+
+        public String getId() {
+            return id;
+        }
+
+        public void setId(String id) {
+            this.id = id;
+        }
+
+        public String getName() {
+            return name;
+        }
+
+        public void setName(String name) {
+            this.name = name;
+        }
+    }
 }