Consider supporting Enum type for index

Closes gh-107
Restream · Dec 22, 2024 · 91cf29e · 91cf29e
1 parent b790ad6
commit 91cf29e
Show file tree

Hide file tree

Showing 6 changed files with 193 additions and 0 deletions.
diff --git a/src/main/java/ru/rt/restream/reindexer/EnumType.java b/src/main/java/ru/rt/restream/reindexer/EnumType.java
@@ -0,0 +1,33 @@
+/*
+ * Copyright 2020 Restream
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package ru.rt.restream.reindexer;
+
+/**
+ * Reindexer enum item field mapping type.
+ */
+public enum EnumType {
+
+    /**
+     * Persist enumerated type property or field as an integer. 
+     */
+    ORDINAL,
+
+    /**
+     * Persist enumerated type property or field as a string.
+     */
+    STRING
+
+}
diff --git a/src/main/java/ru/rt/restream/reindexer/annotations/Enumerated.java b/src/main/java/ru/rt/restream/reindexer/annotations/Enumerated.java
@@ -0,0 +1,40 @@
+/*
+ * Copyright 2020 Restream
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package ru.rt.restream.reindexer.annotations;
+
+import ru.rt.restream.reindexer.EnumType;
+
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+/**
+ * Specifies the mapping for an enum item field.
+ * Provides support to persist an enum as a String or Integer using {@link EnumType}.
+ */
+@Target(ElementType.FIELD)
+@Retention(RetentionPolicy.RUNTIME)
+public @interface Enumerated {
+
+    /**
+     * (Optional) The type used in mapping an enum type.
+     *
+     * @return The type used in mapping an enum type
+     */
+    EnumType value() default EnumType.ORDINAL;
+
+}
diff --git a/src/main/java/ru/rt/restream/reindexer/annotations/ReindexAnnotationScanner.java b/src/main/java/ru/rt/restream/reindexer/annotations/ReindexAnnotationScanner.java
@@ -16,6 +16,7 @@
 package ru.rt.restream.reindexer.annotations;
 
 import ru.rt.restream.reindexer.CollateMode;
+import ru.rt.restream.reindexer.EnumType;
 import ru.rt.restream.reindexer.FieldType;
 import ru.rt.restream.reindexer.IndexType;
 import ru.rt.restream.reindexer.ReindexScanner;
@@ -259,6 +260,9 @@ private FieldInfo getFieldInfo(Field field) {
                 fieldType = getFieldTypeByClass(componentType);
                 fieldInfo.componentType = componentType;
             }
+        } else if (Enum.class.isAssignableFrom(type)) {
+            Enumerated enumerated = field.getAnnotation(Enumerated.class);
+            fieldType = enumerated != null && enumerated.value() == EnumType.STRING ? STRING : INT;
         } else {
             fieldType = getFieldTypeByClass(type);
         }

diff --git a/src/main/java/ru/rt/restream/reindexer/binding/cproto/cjson/CJsonItemWriter.java b/src/main/java/ru/rt/restream/reindexer/binding/cproto/cjson/CJsonItemWriter.java
@@ -15,6 +15,8 @@
  */
 package ru.rt.restream.reindexer.binding.cproto.cjson;
 
+import ru.rt.restream.reindexer.EnumType;
+import ru.rt.restream.reindexer.annotations.Enumerated;
 import ru.rt.restream.reindexer.annotations.Json;
 import ru.rt.restream.reindexer.annotations.Reindex;
 import ru.rt.restream.reindexer.annotations.Transient;
@@ -90,6 +92,14 @@ private CjsonElement toCjson(Object source) {
                     if (field.getType() == String.class && field.isAnnotationPresent(Reindex.class)
                             && field.getAnnotation(Reindex.class).isUuid()) {
                         cjsonElement = new CjsonPrimitive(UUID.fromString((String) fieldValue));
+                    } else if (Enum.class.isAssignableFrom(field.getType())) {
+                        Enumerated enumerated = field.getAnnotation(Enumerated.class);
+                        if (enumerated != null && enumerated.value() == EnumType.STRING) {
+                            cjsonElement = new CjsonPrimitive(((Enum<?>) fieldValue).name());
+                        } else {
+                            int ordinal = ((Enum<?>) fieldValue).ordinal();
+                            cjsonElement = new CjsonPrimitive((long) ordinal);
+                        }
                     } else {
                         cjsonElement = toCjson(fieldValue);
                     }

diff --git a/src/main/java/ru/rt/restream/reindexer/binding/cproto/cjson/CjsonItemReader.java b/src/main/java/ru/rt/restream/reindexer/binding/cproto/cjson/CjsonItemReader.java
@@ -15,6 +15,8 @@
  */
 package ru.rt.restream.reindexer.binding.cproto.cjson;
 
+import ru.rt.restream.reindexer.EnumType;
+import ru.rt.restream.reindexer.annotations.Enumerated;
 import ru.rt.restream.reindexer.annotations.Json;
 import ru.rt.restream.reindexer.binding.cproto.ByteBuffer;
 import ru.rt.restream.reindexer.binding.cproto.ItemReader;
@@ -85,6 +87,12 @@ private Object getTargetValue(Field field, CjsonElement property) {
                 elements.add(convert(cjsonElement, (Class<?>) elementType));
             }
             return elements;
+        } else if (Enum.class.isAssignableFrom(fieldType)) {
+            Enumerated enumerated = field.getAnnotation(Enumerated.class);
+            if (enumerated != null && enumerated.value() == EnumType.STRING) {
+                return Enum.valueOf(fieldType.asSubclass(Enum.class), property.getAsString());
+            }
+            return fieldType.getEnumConstants()[property.getAsInteger()];
         } else {
             return convert(property, field.getType());
         }

diff --git a/src/test/java/ru/rt/restream/reindexer/connector/ReindexerTest.java b/src/test/java/ru/rt/restream/reindexer/connector/ReindexerTest.java
@@ -22,11 +22,13 @@
 import lombok.NoArgsConstructor;
 import lombok.Setter;
 import org.junit.jupiter.api.Test;
+import ru.rt.restream.reindexer.EnumType;
 import ru.rt.restream.reindexer.Namespace;
 import ru.rt.restream.reindexer.NamespaceOptions;
 import ru.rt.restream.reindexer.QueryResultJsonIterator;
 import ru.rt.restream.reindexer.ResultIterator;
 import ru.rt.restream.reindexer.Transaction;
+import ru.rt.restream.reindexer.annotations.Enumerated;
 import ru.rt.restream.reindexer.annotations.Reindex;
 import ru.rt.restream.reindexer.annotations.Serial;
 import ru.rt.restream.reindexer.db.DbBaseTest;
@@ -2774,6 +2776,51 @@ public void testIsAppendableIndexItem() {
         assertThat(iterator.hasNext(), is(false));
     }
 
+    @Test
+    public void testTestItemEnumString() {
+        String namespaceName = "items";
+        db.openNamespace(namespaceName, NamespaceOptions.defaultOptions(), TestItemEnumString.class);
+        db.upsert(namespaceName, new TestItemEnumString(1, "TestName", TestEnum.TEST_CONSTANT));
+        Iterator<TestItemEnumString> iterator = db.query(namespaceName, TestItemEnumString.class)
+                .where("testEnumString", EQ, TestEnum.TEST_CONSTANT.name())
+                .execute();
+        assertThat(iterator.hasNext(), is(true));
+        TestItemEnumString foundByEnumString = iterator.next();
+        assertThat(foundByEnumString.id, is(1));
+        assertThat(foundByEnumString.name, is("TestName"));
+        assertThat(foundByEnumString.testEnum, is(TestEnum.TEST_CONSTANT));
+    }
+
+    @Test
+    public void testTestItemEnumOrdinal() {
+        String namespaceName = "items";
+        db.openNamespace(namespaceName, NamespaceOptions.defaultOptions(), TestItemEnumOrdinal.class);
+        db.upsert(namespaceName, new TestItemEnumOrdinal(1, "TestName", TestEnum.TEST_CONSTANT));
+        Iterator<TestItemEnumOrdinal> iterator = db.query(namespaceName, TestItemEnumOrdinal.class)
+                .where("testEnum", EQ, TestEnum.TEST_CONSTANT.ordinal())
+                .execute();
+        assertThat(iterator.hasNext(), is(true));
+        TestItemEnumOrdinal foundByEnumOrdinal = iterator.next();
+        assertThat(foundByEnumOrdinal.id, is(1));
+        assertThat(foundByEnumOrdinal.name, is("TestName"));
+        assertThat(foundByEnumOrdinal.testEnum, is(TestEnum.TEST_CONSTANT));
+    }
+
+    @Test
+    public void testTestItemEnumDefault() {
+        String namespaceName = "items";
+        db.openNamespace(namespaceName, NamespaceOptions.defaultOptions(), TestItemEnumDefault.class);
+        db.upsert(namespaceName, new TestItemEnumDefault(1, "TestName", TestEnum.TEST_CONSTANT));
+        Iterator<TestItemEnumDefault> iterator = db.query(namespaceName, TestItemEnumDefault.class)
+                .where("testEnum", EQ, TestEnum.TEST_CONSTANT.ordinal())
+                .execute();
+        assertThat(iterator.hasNext(), is(true));
+        TestItemEnumDefault foundByEnumOrdinal = iterator.next();
+        assertThat(foundByEnumOrdinal.id, is(1));
+        assertThat(foundByEnumOrdinal.name, is("TestName"));
+        assertThat(foundByEnumOrdinal.testEnum, is(TestEnum.TEST_CONSTANT));
+    }
+
     @Getter
     @Setter
     public static class SerialIdTestItem {
@@ -2877,4 +2924,55 @@ public static class ItemWithAppendableIndexes {
         @Reindex(name = "name", isAppendable = true, type = TEXT)
         private String description;
     }
+
+    @Getter
+    @Setter
+    @NoArgsConstructor
+    @AllArgsConstructor
+    public static class TestItemEnumString {
+        @Reindex(name = "id", isPrimaryKey = true)
+        private Integer id;
+
+        @Reindex(name = "name")
+        private String name;
+
+        @Enumerated(EnumType.STRING)
+        @Reindex(name = "testEnumString")
+        private TestEnum testEnum;
+    }
+
+    @Getter
+    @Setter
+    @NoArgsConstructor
+    @AllArgsConstructor
+    public static class TestItemEnumOrdinal {
+        @Reindex(name = "id", isPrimaryKey = true)
+        private Integer id;
+
+        @Reindex(name = "name")
+        private String name;
+
+        @Enumerated(EnumType.ORDINAL)
+        @Reindex(name = "testEnumOrdinal")
+        private TestEnum testEnum;
+    }
+
+    @Getter
+    @Setter
+    @NoArgsConstructor
+    @AllArgsConstructor
+    public static class TestItemEnumDefault {
+        @Reindex(name = "id", isPrimaryKey = true)
+        private Integer id;
+
+        @Reindex(name = "name")
+        private String name;
+
+        @Reindex(name = "testEnum")
+        private TestEnum testEnum;
+    }
+
+    public enum TestEnum {
+        TEST_CONSTANT
+    }
 }