From 3ae374f0f8f698d4e6391c60c4ae49b330b32219 Mon Sep 17 00:00:00 2001
From: kxmas <kxmas@users.noreply.github.com>
Date: Tue, 24 Sep 2019 12:57:54 -0500
Subject: [PATCH] Custom partition key name (#383)

---
 .../cosmosdb/core/mapping/PartitionKey.java   |  7 ++
 .../support/DocumentDbEntityInformation.java  |  7 +-
 .../DocumentDbEntityInformationUnitTest.java  | 74 +++++++++++++++++++
 3 files changed, 87 insertions(+), 1 deletion(-)

diff --git a/src/main/java/com/microsoft/azure/spring/data/cosmosdb/core/mapping/PartitionKey.java b/src/main/java/com/microsoft/azure/spring/data/cosmosdb/core/mapping/PartitionKey.java
index 29b2c90..6b052f0 100644
--- a/src/main/java/com/microsoft/azure/spring/data/cosmosdb/core/mapping/PartitionKey.java
+++ b/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 "";
 }
diff --git a/src/main/java/com/microsoft/azure/spring/data/cosmosdb/repository/support/DocumentDbEntityInformation.java b/src/main/java/com/microsoft/azure/spring/data/cosmosdb/repository/support/DocumentDbEntityInformation.java
index 7c6967f..80361e9 100644
--- a/src/main/java/com/microsoft/azure/spring/data/cosmosdb/repository/support/DocumentDbEntityInformation.java
+++ b/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) {
diff --git a/src/test/java/com/microsoft/azure/spring/data/cosmosdb/repository/support/DocumentDbEntityInformationUnitTest.java b/src/test/java/com/microsoft/azure/spring/data/cosmosdb/repository/support/DocumentDbEntityInformationUnitTest.java
index 20b071c..632d2e7 100644
--- a/src/test/java/com/microsoft/azure/spring/data/cosmosdb/repository/support/DocumentDbEntityInformationUnitTest.java
+++ b/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;
+        }
+    }
 }