Allow Long as ID type in Entity for Cosmos DB

sdk/cosmos/azure-spring-data-cosmos/src/main/java/com/azure/spring/data/cosmos/common/CosmosUtils.java

@@ -6,8 +6,10 @@
 import com.azure.cosmos.models.FeedResponse;
 import com.azure.spring.data.cosmos.core.ResponseDiagnostics;
 import com.azure.spring.data.cosmos.core.ResponseDiagnosticsProcessor;
+import com.azure.spring.data.cosmos.exception.IllegalQueryException;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
+import org.springframework.util.Assert;
 
 /**
  * Util class to fill and process response diagnostics
@@ -45,4 +47,26 @@ public static <T> void fillAndProcessResponseDiagnostics(
         //  Process response diagnostics
         responseDiagnosticsProcessor.processResponseDiagnostics(responseDiagnostics);
     }
+
+    /**
+     * ID value should be string value, real id type will be String, Integer, Long,
+     * all of these must be converted to String type.
+     * @param idValue id value to find
+     * @throws IllegalArgumentException thrown if id value fail the validation.
+     * @throws IllegalQueryException thrown if id value fail the validation.
+     * @return String id value
+     */
+    public static Object getStringIDValue(Object idValue) {
+        Assert.notNull(idValue, "id should not be null");
+        if (idValue instanceof String) {
+            Assert.hasText(idValue.toString(), "id should not be empty or only whitespaces.");
+            return idValue;
+        } else if (idValue instanceof Integer) {
+            return Integer.toString((Integer) idValue);
+        } else if (idValue instanceof Long) {
+            return Long.toString((Long) idValue);
+        } else {
+            throw new IllegalQueryException("Type of id field must be String or Integer or Long");
+        }
+    }
 }

sdk/cosmos/azure-spring-data-cosmos/src/main/java/com/azure/spring/data/cosmos/core/CosmosTemplate.java

@@ -166,13 +166,12 @@ public <T> T findById(Object id, Class<T> domainType) {
     public <T> T findById(Object id, Class<T> domainType, PartitionKey partitionKey) {
         Assert.notNull(domainType, "domainType should not be null");
         Assert.notNull(partitionKey, "partitionKey should not be null");
-        assertValidId(id);
-
+        String idToQuery = (String) CosmosUtils.getStringIDValue(id);
         final String containerName = getContainerName(domainType);
         return cosmosAsyncClient
             .getDatabase(databaseName)
             .getContainer(containerName)
-            .readItem(id.toString(), partitionKey, JsonNode.class)
+            .readItem(idToQuery, partitionKey, JsonNode.class)
             .flatMap(cosmosItemResponse -> {
                 CosmosUtils.fillAndProcessResponseDiagnostics(responseDiagnosticsProcessor,
                     cosmosItemResponse.getDiagnostics(), null);
@@ -195,10 +194,9 @@ public <T> T findById(Object id, Class<T> domainType, PartitionKey partitionKey)
     public <T> T findById(String containerName, Object id, Class<T> domainType) {
         Assert.hasText(containerName, "containerName should not be null, empty or only whitespaces");
         Assert.notNull(domainType, "domainType should not be null");
-        assertValidId(id);
 
         final String query = String.format("select * from root where root.id = '%s'",
-            id.toString());
+            CosmosUtils.getStringIDValue(id));
         final CosmosQueryRequestOptions options = new CosmosQueryRequestOptions();
         options.setQueryMetricsEnabled(enableQueryMetrics);
         return cosmosAsyncClient
@@ -429,8 +427,7 @@ public CosmosContainerProperties createContainerIfNotExists(CosmosEntityInformat
      */
     public void deleteById(String containerName, Object id, PartitionKey partitionKey) {
         Assert.hasText(containerName, "containerName should not be null, empty or only whitespaces");
-        assertValidId(id);
-
+        String idToDelete = (String) CosmosUtils.getStringIDValue(id);
         LOGGER.debug("execute deleteById in database {} container {}", this.databaseName,
             containerName);
 
@@ -439,7 +436,7 @@ public void deleteById(String containerName, Object id, PartitionKey partitionKe
         }
         cosmosAsyncClient.getDatabase(this.databaseName)
                          .getContainer(containerName)
-                         .deleteItem(id.toString(), partitionKey)
+                         .deleteItem(idToDelete, partitionKey)
                          .doOnNext(response ->
                              CosmosUtils.fillAndProcessResponseDiagnostics(responseDiagnosticsProcessor,
                                  response.getDiagnostics(), null))
@@ -454,9 +451,12 @@ public <T, ID> List<T> findByIds(Iterable<ID> ids, Class<T> domainType, String c
         Assert.notNull(ids, "Id list should not be null");
         Assert.notNull(domainType, "domainType should not be null.");
         Assert.hasText(containerName, "container should not be null, empty or only whitespaces");
-
+        final List<Object> idList = new ArrayList<>();
+        for (ID id : ids) {
+            idList.add(CosmosUtils.getStringIDValue(id));
+        }
         final DocumentQuery query = new DocumentQuery(Criteria.getInstance(CriteriaType.IN, "id",
-            Collections.singletonList(ids), Part.IgnoreCaseType.NEVER));
+            Collections.singletonList(idList), Part.IgnoreCaseType.NEVER));
         return find(query, domainType, containerName);
     }
 
@@ -670,13 +670,6 @@ private List<String> getPartitionKeyNames(Class<?> domainType) {
         return Collections.singletonList(entityInfo.getPartitionKeyFieldName());
     }
 
-    private void assertValidId(Object id) {
-        Assert.notNull(id, "id should not be null");
-        if (id instanceof String) {
-            Assert.hasText(id.toString(), "id should not be empty or only whitespaces.");
-        }
-    }
-
     private <T> List<JsonNode> findItems(@NonNull DocumentQuery query,
                                          @NonNull String containerName) {
         final SqlQuerySpec sqlQuerySpec = new FindQuerySpecGenerator().generateCosmos(query);

sdk/cosmos/azure-spring-data-cosmos/src/main/java/com/azure/spring/data/cosmos/core/ReactiveCosmosTemplate.java

@@ -212,10 +212,9 @@ public <T> Mono<T> findById(Object id, Class<T> domainType) {
     public <T> Mono<T> findById(String containerName, Object id, Class<T> domainType) {
         Assert.hasText(containerName, "containerName should not be null, empty or only whitespaces");
         Assert.notNull(domainType, "domainType should not be null");
-        assertValidId(id);
 
         final String query = String.format("select * from root where root.id = '%s'",
-            id.toString());
+            CosmosUtils.getStringIDValue(id));
         final CosmosQueryRequestOptions options = new CosmosQueryRequestOptions();
         options.setQueryMetricsEnabled(isPopulateQueryMetrics);
 
@@ -249,12 +248,12 @@ public <T> Mono<T> findById(String containerName, Object id, Class<T> domainType
     @Override
     public <T> Mono<T> findById(Object id, Class<T> domainType, PartitionKey partitionKey) {
         Assert.notNull(domainType, "domainType should not be null");
-        assertValidId(id);
+        String idToFind = (String) CosmosUtils.getStringIDValue(id);
 
         final String containerName = getContainerName(domainType);
         return cosmosAsyncClient.getDatabase(databaseName)
                                 .getContainer(containerName)
-                                .readItem(id.toString(), partitionKey, JsonNode.class)
+                                .readItem(idToFind, partitionKey, JsonNode.class)
                                 .flatMap(cosmosItemResponse -> {
                                     CosmosUtils.fillAndProcessResponseDiagnostics(responseDiagnosticsProcessor,
                                         cosmosItemResponse.getDiagnostics(), null);
@@ -390,15 +389,15 @@ public <T> Mono<T> upsert(String containerName, T object) {
     @Override
     public Mono<Void> deleteById(String containerName, Object id, PartitionKey partitionKey) {
         Assert.hasText(containerName, "container name should not be null, empty or only whitespaces");
-        assertValidId(id);
+        String idToDelete = (String) CosmosUtils.getStringIDValue(id);
 
         if (partitionKey == null) {
             partitionKey = PartitionKey.NONE;
         }
 
         return cosmosAsyncClient.getDatabase(this.databaseName)
                                 .getContainer(containerName)
-                                .deleteItem(id.toString(), partitionKey)
+                                .deleteItem(idToDelete, partitionKey)
                                 .doOnNext(cosmosItemResponse ->
                                     CosmosUtils.fillAndProcessResponseDiagnostics(responseDiagnosticsProcessor,
                                         cosmosItemResponse.getDiagnostics(), null))
@@ -589,13 +588,6 @@ private Flux<JsonNode> findItems(@NonNull DocumentQuery query,
                 CosmosExceptionUtils.exceptionHandler("Failed to query items", throwable));
     }
 
-    private void assertValidId(Object id) {
-        Assert.notNull(id, "id should not be null");
-        if (id instanceof String) {
-            Assert.hasText(id.toString(), "id should not be empty or only whitespaces.");
-        }
-    }
-
     private List<String> getPartitionKeyNames(Class<?> domainType) {
         final CosmosEntityInformation<?, ?> entityInfo = entityInfoCreator.apply(domainType);
 

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

@@ -201,8 +201,10 @@ private Field getIdField(Class<?> domainType) {
             throw new IllegalArgumentException("domain should contain @Id field or field named id");
         } else if (idField.getType() != String.class
                 && idField.getType() != Integer.class
-                && idField.getType() != int.class) {
-            throw new IllegalArgumentException("type of id field must be String or Integer");
+                && idField.getType() != int.class
+                && idField.getType() != Long.class
+                && idField.getType() != long.class) {
+            throw new IllegalArgumentException("type of id field must be String, Integer or Long");
         }
 
         return idField;

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

@@ -0,0 +1,70 @@
+// 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.Document;
+import org.springframework.data.annotation.Id;
+
+import java.util.Objects;
+
+@Document
+public class LongIdDomain {
+
+    @Id
+    private Long number;
+
+    private String name;
+
+    public LongIdDomain(Long number, String name) {
+        this.number = number;
+        this.name = name;
+    }
+
+    public LongIdDomain() {
+    }
+
+    public Long getNumber() {
+        return number;
+    }
+
+    public void setNumber(Long number) {
+        this.number = number;
+    }
+
+    public String getName() {
+        return name;
+    }
+
+    public void setName(String name) {
+        this.name = name;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) {
+            return true;
+        }
+        if (o == null || getClass() != o.getClass()) {
+            return false;
+        }
+        LongIdDomain that = (LongIdDomain) o;
+        return Objects.equals(number, that.number)
+            && Objects.equals(name, that.name);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(number, name);
+    }
+
+    @Override
+    public String toString() {
+        return "LongIdDomain{"
+            + "number="
+            + number
+            + ", name='"
+            + name
+            + '\''
+            + '}';
+    }
+}

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

@@ -0,0 +1,72 @@
+// 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.Document;
+import com.azure.spring.data.cosmos.core.mapping.PartitionKey;
+import org.springframework.data.annotation.Id;
+
+import java.util.Objects;
+
+@Document
+public class LongIdDomainPartition {
+
+    @Id
+    private Long number;
+
+    @PartitionKey
+    private String name;
+
+    public LongIdDomainPartition(Long number, String name) {
+        this.number = number;
+        this.name = name;
+    }
+
+    public LongIdDomainPartition() {
+    }
+
+    public Long getNumber() {
+        return number;
+    }
+
+    public void setNumber(Long number) {
+        this.number = number;
+    }
+
+    public String getName() {
+        return name;
+    }
+
+    public void setName(String name) {
+        this.name = name;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) {
+            return true;
+        }
+        if (o == null || getClass() != o.getClass()) {
+            return false;
+        }
+        LongIdDomainPartition that = (LongIdDomainPartition) o;
+        return Objects.equals(number, that.number)
+            && Objects.equals(name, that.name);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(number, name);
+    }
+
+    @Override
+    public String toString() {
+        return "LongIdDomain{"
+            + "number="
+            + number
+            + ", name='"
+            + name
+            + '\''
+            + '}';
+    }
+}

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

@@ -13,7 +13,6 @@
 import org.junit.AfterClass;
 import org.junit.Assert;
 import org.junit.Before;
-import org.junit.Ignore;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.springframework.beans.factory.annotation.Autowired;
@@ -122,7 +121,6 @@ public void testSaveAllAndFindAll() {
     }
 
     @Test
-    @Ignore //  TODO(kuthapar): findById IN clause not working in case of Integer
     public void testFindAllById() {
         final Iterable<IntegerIdDomain> allById =
             this.repository.findAllById(Collections.singleton(DOMAIN.getNumber()));