Support and test different id types through findUsingQuery

src/main/java/org/springframework/data/aerospike/query/Qualifier.java

@@ -361,10 +361,40 @@ private IdQualifierBuilder setId(String id) {
             return this;
         }
 
+        private IdQualifierBuilder setId(Short id) {
+            this.map.put(SINGLE_ID_FIELD, id);
+            return this;
+        }
+
+        private IdQualifierBuilder setId(Integer id) {
+            this.map.put(SINGLE_ID_FIELD, id);
+            return this;
+        }
+
+        private IdQualifierBuilder setId(Long id) {
+            this.map.put(SINGLE_ID_FIELD, id);
+            return this;
+        }
+
         private IdQualifierBuilder setIds(String... ids) {
             this.map.put(MULTIPLE_IDS_FIELD, ids);
             return this;
         }
+
+        private IdQualifierBuilder setIds(Short... ids) {
+            this.map.put(MULTIPLE_IDS_FIELD, ids);
+            return this;
+        }
+
+        private IdQualifierBuilder setIds(Integer... ids) {
+            this.map.put(MULTIPLE_IDS_FIELD, ids);
+            return this;
+        }
+
+        private IdQualifierBuilder setIds(Long... ids) {
+            this.map.put(MULTIPLE_IDS_FIELD, ids);
+            return this;
+        }
     }
 
     private static class ConjunctionQualifierBuilder extends BaseQualifierBuilder<ConjunctionQualifierBuilder> {
@@ -439,26 +469,96 @@ protected void validate() {
      * @param id String value
      * @return Single id qualifier
      */
-    // TODO: Support other than string (based on flag)
     public static Qualifier idEquals(String id) {
         return new Qualifier(new IdQualifierBuilder()
             .setId(id)
             .setFilterOperation(FilterOperation.EQ));
     }
 
+    /**
+     * Create a qualifier for the condition when the primary key equals the given short
+     *
+     * @param id Short value
+     * @return Single id qualifier
+     */
+    public static Qualifier idEquals(Short id) {
+        return new Qualifier(new IdQualifierBuilder()
+            .setId(id)
+            .setFilterOperation(FilterOperation.EQ));
+    }
+
+    /**
+     * Create a qualifier for the condition when the primary key equals the given integer
+     *
+     * @param id Integer value
+     * @return Single id qualifier
+     */
+    public static Qualifier idEquals(Integer id) {
+        return new Qualifier(new IdQualifierBuilder()
+            .setId(id)
+            .setFilterOperation(FilterOperation.EQ));
+    }
+
+    /**
+     * Create a qualifier for the condition when the primary key equals the given long
+     *
+     * @param id Long value
+     * @return Single id qualifier
+     */
+    public static Qualifier idEquals(Long id) {
+        return new Qualifier(new IdQualifierBuilder()
+            .setId(id)
+            .setFilterOperation(FilterOperation.EQ));
+    }
+
     /**
      * Create a qualifier for the condition when the primary key equals one of the given strings (logical OR)
      *
      * @param ids String values
      * @return Multiple ids qualifier with OR condition
      */
-    // TODO: Support other than string (based on flag)
     public static Qualifier idIn(String... ids) {
         return new Qualifier(new IdQualifierBuilder()
             .setIds(ids)
             .setFilterOperation(FilterOperation.EQ));
     }
 
+    /**
+     * Create a qualifier for the condition when the primary key equals one of the given shorts (logical OR)
+     *
+     * @param ids Short values
+     * @return Multiple ids qualifier with OR condition
+     */
+    public static Qualifier idIn(Short... ids) {
+        return new Qualifier(new IdQualifierBuilder()
+            .setIds(ids)
+            .setFilterOperation(FilterOperation.EQ));
+    }
+
+    /**
+     * Create a qualifier for the condition when the primary key equals one of the given integers (logical OR)
+     *
+     * @param ids Integer values
+     * @return Multiple ids qualifier with OR condition
+     */
+    public static Qualifier idIn(Integer... ids) {
+        return new Qualifier(new IdQualifierBuilder()
+            .setIds(ids)
+            .setFilterOperation(FilterOperation.EQ));
+    }
+
+    /**
+     * Create a qualifier for the condition when the primary key equals one of the given longs (logical OR)
+     *
+     * @param ids Long values
+     * @return Multiple ids qualifier with OR condition
+     */
+    public static Qualifier idIn(Long... ids) {
+        return new Qualifier(new IdQualifierBuilder()
+            .setIds(ids)
+            .setFilterOperation(FilterOperation.EQ));
+    }
+
     /**
      * Create a parent qualifier that contains the given qualifiers combined using logical OR
      *

src/test/java/org/springframework/data/aerospike/convert/MappingAerospikeConverterTypesTests.java

@@ -94,7 +94,7 @@ void primitiveByteId(int converterOption) {
     @ParameterizedTest()
     @ValueSource(ints = {0, 1})
     void shortId(int converterOption) {
-        DocumentWithShortId object = new DocumentWithShortId((short) 5);
+        DocumentWithShortId object = DocumentWithShortId.builder().id((short) 5).build();
 
         assertWriteAndRead(converterOption, object, "DocumentWithShortId", (short) 5,
             new Bin("@_class", DocumentWithShortId.class.getName())
@@ -104,7 +104,7 @@ void shortId(int converterOption) {
     @ParameterizedTest()
     @ValueSource(ints = {0, 1})
     void integerId(int converterOption) {
-        DocumentWithIntegerId object = new DocumentWithIntegerId(5);
+        DocumentWithIntegerId object = DocumentWithIntegerId.builder().id(5).build();
 
         assertWriteAndRead(converterOption, object, "DocumentWithIntegerId", 5,
             new Bin("@_class", DocumentWithIntegerId.class.getName())
@@ -114,7 +114,7 @@ void integerId(int converterOption) {
     @ParameterizedTest()
     @ValueSource(ints = {0, 1})
     void longId(int converterOption) {
-        DocumentWithLongId object = new DocumentWithLongId(5L);
+        DocumentWithLongId object = DocumentWithLongId.builder().id(5L).build();
 
         assertWriteAndRead(converterOption, object, "DocumentWithLongId", 5L,
             new Bin("@_class", DocumentWithLongId.class.getName())
@@ -144,7 +144,7 @@ void byteId(int converterOption) {
     @ParameterizedTest()
     @ValueSource(ints = {0, 1})
     void stringId(int converterOption) {
-        DocumentWithStringId object = new DocumentWithStringId("my-amazing-string-id");
+        DocumentWithStringId object = DocumentWithStringId.builder().id("my-amazing-string-id").build();
 
         assertWriteAndRead(converterOption, object, "DocumentWithStringId",
             "my-amazing-string-id", new Bin("@_class", DocumentWithStringId.class.getName())

src/test/java/org/springframework/data/aerospike/core/AerospikeTemplateFindByIdTests.java

@@ -193,7 +193,7 @@ public void findById_shouldReadClassWithNumericKeyMap() {
     public void findById_shouldReadClassWithNonStringId() {
         if (template.getAerospikeConverter().getAerospikeDataSettings().isKeepOriginalKeyTypes()) {
             long longId = 10L;
-            SampleClasses.DocumentWithLongId document = new SampleClasses.DocumentWithLongId(longId);
+            SampleClasses.DocumentWithLongId document = SampleClasses.DocumentWithLongId.builder().id(longId).build();
             template.save(document);
             SampleClasses.DocumentWithLongId result = template.findById(longId, SampleClasses.DocumentWithLongId.class);
             assertThat(result.getId().equals(longId)).isTrue();

src/test/java/org/springframework/data/aerospike/core/reactive/ReactiveAerospikeTemplateFindByIdTests.java

@@ -108,7 +108,7 @@ public void findByIds_shouldFindExisting() {
     public void findById_shouldReadClassWithNonStringId() {
         if (reactiveTemplate.getAerospikeConverter().getAerospikeDataSettings().isKeepOriginalKeyTypes()) {
             long longId = 10L;
-            SampleClasses.DocumentWithLongId document = new SampleClasses.DocumentWithLongId(longId);
+            SampleClasses.DocumentWithLongId document = SampleClasses.DocumentWithLongId.builder().id(longId).build();
             reactiveTemplate.save(document).block();
             SampleClasses.DocumentWithLongId result = reactiveTemplate.findById(longId,
                     SampleClasses.DocumentWithLongId.class)

src/test/java/org/springframework/data/aerospike/repository/IdTypesRepositoryQueryTests.java

@@ -0,0 +1,120 @@
+package org.springframework.data.aerospike.repository;
+
+import org.junit.jupiter.api.AfterAll;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.TestInstance;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.data.aerospike.BaseBlockingIntegrationTests;
+import org.springframework.data.aerospike.query.Qualifier;
+import org.springframework.data.aerospike.repository.query.Query;
+import org.springframework.data.aerospike.sample.DocumentIntegerIdRepository;
+import org.springframework.data.aerospike.sample.DocumentLongIdRepository;
+import org.springframework.data.aerospike.sample.DocumentShortIdRepository;
+import org.springframework.data.aerospike.sample.DocumentStringIdRepository;
+import org.springframework.data.aerospike.sample.SampleClasses.DocumentWithIntegerId;
+import org.springframework.data.aerospike.sample.SampleClasses.DocumentWithLongId;
+import org.springframework.data.aerospike.sample.SampleClasses.DocumentWithShortId;
+import org.springframework.data.aerospike.sample.SampleClasses.DocumentWithStringId;
+
+import java.util.List;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+@TestInstance(TestInstance.Lifecycle.PER_CLASS)
+public class IdTypesRepositoryQueryTests extends BaseBlockingIntegrationTests {
+
+    @Autowired
+    DocumentStringIdRepository documentStringIdRepository;
+    @Autowired
+    DocumentShortIdRepository documentShortIdRepository;
+    @Autowired
+    DocumentIntegerIdRepository documentIntegerIdRepository;
+    @Autowired
+    DocumentLongIdRepository documentLongIdRepository;
+
+    DocumentWithStringId docStringId1 = DocumentWithStringId.builder().id("id1").content("contentString1").build();
+    DocumentWithStringId docStringId2 = DocumentWithStringId.builder().id("id2").content("contentString2").build();
+    DocumentWithStringId docStringId3 = DocumentWithStringId.builder().id("id3").content("contentString3").build();
+    DocumentWithShortId docShortId1 = DocumentWithShortId.builder().id((short) 1).content("contentShort1").build();
+    DocumentWithShortId docShortId2 = DocumentWithShortId.builder().id((short) 2).content("contentShort2").build();
+    DocumentWithShortId docShortId3 = DocumentWithShortId.builder().id((short) 3).content("contentShort3").build();
+    DocumentWithIntegerId docIntegerId1 = DocumentWithIntegerId.builder().id(1).content("contentInteger1").build();
+    DocumentWithIntegerId docIntegerId2 = DocumentWithIntegerId.builder().id(2).content("contentInteger2").build();
+    DocumentWithIntegerId docIntegerId3 = DocumentWithIntegerId.builder().id(3).content("contentInteger3").build();
+    DocumentWithLongId docLongId1 = DocumentWithLongId.builder().id(1L).content("contentLong1").build();
+    DocumentWithLongId docLongId2 = DocumentWithLongId.builder().id(2L).content("contentLong2").build();
+    DocumentWithLongId docLongId3 = DocumentWithLongId.builder().id(3L).content("contentLong3").build();
+
+    @BeforeAll
+    void beforeAll() {
+        documentStringIdRepository.saveAll(List.of(docStringId1, docStringId2, docStringId3));
+        documentShortIdRepository.saveAll(List.of(docShortId1, docShortId2, docShortId3));
+        documentIntegerIdRepository.saveAll(List.of(docIntegerId1, docIntegerId2, docIntegerId3));
+        documentLongIdRepository.saveAll(List.of(docLongId1, docLongId2, docLongId3));
+    }
+
+    @AfterAll
+    void afterAll() {
+        documentStringIdRepository.deleteAll();
+        documentShortIdRepository.deleteAll();
+        documentIntegerIdRepository.deleteAll();
+        documentLongIdRepository.deleteAll();
+    }
+
+    @Test
+    void findUsingQueryStringIdEquals() {
+        Qualifier idEquals = Qualifier.idEquals(docStringId1.getId());
+        Iterable<DocumentWithStringId> result = documentStringIdRepository.findUsingQuery(new Query(idEquals));
+        assertThat(result).containsOnly(docStringId1);
+    }
+
+    @Test
+    void findUsingQueryShortIdEquals() {
+        Qualifier idEquals = Qualifier.idEquals(docShortId1.getId());
+        Iterable<DocumentWithShortId> result = documentShortIdRepository.findUsingQuery(new Query(idEquals));
+        assertThat(result).containsOnly(docShortId1);
+    }
+
+    @Test
+    void findUsingQueryIntegerIdEquals() {
+        Qualifier idEquals = Qualifier.idEquals(docIntegerId1.getId());
+        Iterable<DocumentWithIntegerId> result = documentIntegerIdRepository.findUsingQuery(new Query(idEquals));
+        assertThat(result).containsOnly(docIntegerId1);
+    }
+
+    @Test
+    void findUsingQueryLongIdEquals() {
+        Qualifier idEquals = Qualifier.idEquals(docLongId1.getId());
+        Iterable<DocumentWithLongId> result = documentLongIdRepository.findUsingQuery(new Query(idEquals));
+        assertThat(result).containsOnly(docLongId1);
+    }
+
+    @Test
+    void findUsingQueryStringIdIn() {
+        Qualifier idIn = Qualifier.idIn(docStringId1.getId(), docStringId3.getId());
+        Iterable<DocumentWithStringId> result = documentStringIdRepository.findUsingQuery(new Query(idIn));
+        assertThat(result).containsOnly(docStringId1, docStringId3);
+    }
+
+    @Test
+    void findUsingQueryShortIdIn() {
+        Qualifier idIn = Qualifier.idIn(docShortId1.getId(), docShortId3.getId());
+        Iterable<DocumentWithShortId> result = documentShortIdRepository.findUsingQuery(new Query(idIn));
+        assertThat(result).containsOnly(docShortId1, docShortId3);
+    }
+
+    @Test
+    void findUsingQueryIntegerIdIn() {
+        Qualifier idIn = Qualifier.idIn(docIntegerId2.getId(), docIntegerId3.getId());
+        Iterable<DocumentWithIntegerId> result = documentIntegerIdRepository.findUsingQuery(new Query(idIn));
+        assertThat(result).containsOnly(docIntegerId2, docIntegerId3);
+    }
+
+    @Test
+    void findUsingQueryLongIdIn() {
+        Qualifier idIn = Qualifier.idIn(docLongId1.getId(), docLongId2.getId());
+        Iterable<DocumentWithLongId> result = documentLongIdRepository.findUsingQuery(new Query(idIn));
+        assertThat(result).containsOnly(docLongId1, docLongId2);
+    }
+}

src/test/java/org/springframework/data/aerospike/sample/DocumentIntegerIdRepository.java

@@ -0,0 +1,7 @@
+package org.springframework.data.aerospike.sample;
+
+import org.springframework.data.aerospike.repository.AerospikeRepository;
+
+public interface DocumentIntegerIdRepository extends AerospikeRepository<SampleClasses.DocumentWithIntegerId, Integer> {
+
+}

src/test/java/org/springframework/data/aerospike/sample/DocumentLongIdRepository.java

@@ -0,0 +1,7 @@
+package org.springframework.data.aerospike.sample;
+
+import org.springframework.data.aerospike.repository.AerospikeRepository;
+
+public interface DocumentLongIdRepository extends AerospikeRepository<SampleClasses.DocumentWithLongId, Long> {
+
+}

src/test/java/org/springframework/data/aerospike/sample/DocumentShortIdRepository.java

@@ -0,0 +1,7 @@
+package org.springframework.data.aerospike.sample;
+
+import org.springframework.data.aerospike.repository.AerospikeRepository;
+
+public interface DocumentShortIdRepository extends AerospikeRepository<SampleClasses.DocumentWithShortId, Short> {
+
+}

src/test/java/org/springframework/data/aerospike/sample/DocumentStringIdRepository.java

@@ -0,0 +1,7 @@
+package org.springframework.data.aerospike.sample;
+
+import org.springframework.data.aerospike.repository.AerospikeRepository;
+
+public interface DocumentStringIdRepository extends AerospikeRepository<SampleClasses.DocumentWithStringId, String> {
+
+}