Change DataSourceType from enum to class (opensearch-project#2730)

* Change DataSourceType from enum to class Signed-off-by: Tomoyuki Morita <[email protected]> * Fix test failure Signed-off-by: Tomoyuki Morita <[email protected]> * Fix serialization issue Signed-off-by: Tomoyuki Morita <[email protected]> * Fix format Signed-off-by: Tomoyuki Morita <[email protected]> * Fix integTest Signed-off-by: Tomoyuki Morita <[email protected]> * Fix style Signed-off-by: Tomoyuki Morita <[email protected]> * Fix failing test Signed-off-by: Tomoyuki Morita <[email protected]> * Address comment Signed-off-by: Tomoyuki Morita <[email protected]> * Fix DataSourceType to allow registering new type Signed-off-by: Tomoyuki Morita <[email protected]> --------- Signed-off-by: Tomoyuki Morita <[email protected]>
jzonthemtn · Aug 28, 2024 · c86fd87 · c86fd87
1 parent 5c02a66
commit c86fd87
Show file tree

Hide file tree

Showing 14 changed files with 302 additions and 65 deletions.
diff --git a/core/src/main/java/org/opensearch/sql/datasource/model/DataSourceType.java b/core/src/main/java/org/opensearch/sql/datasource/model/DataSourceType.java
@@ -5,34 +5,53 @@
 
 package org.opensearch.sql.datasource.model;
 
-public enum DataSourceType {
-  PROMETHEUS("prometheus"),
-  OPENSEARCH("opensearch"),
-  SPARK("spark"),
-  S3GLUE("s3glue");
+import java.util.HashMap;
+import java.util.Map;
+import lombok.RequiredArgsConstructor;
 
-  private String text;
+@RequiredArgsConstructor
+public class DataSourceType {
+  public static DataSourceType PROMETHEUS = new DataSourceType("PROMETHEUS");
+  public static DataSourceType OPENSEARCH = new DataSourceType("OPENSEARCH");
+  public static DataSourceType SPARK = new DataSourceType("SPARK");
+  public static DataSourceType S3GLUE = new DataSourceType("S3GLUE");
 
-  DataSourceType(String text) {
-    this.text = text;
+  // Map from uppercase DataSourceType name to DataSourceType object
+  private static Map<String, DataSourceType> knownValues = new HashMap<>();
+
+  static {
+    register(PROMETHEUS, OPENSEARCH, SPARK, S3GLUE);
+  }
+
+  private final String name;
+
+  public String name() {
+    return name;
   }
 
-  public String getText() {
-    return this.text;
+  @Override
+  public String toString() {
+    return name;
   }
 
-  /**
-   * Get DataSourceType from text.
-   *
-   * @param text text.
-   * @return DataSourceType {@link DataSourceType}.
-   */
-  public static DataSourceType fromString(String text) {
-    for (DataSourceType dataSourceType : DataSourceType.values()) {
-      if (dataSourceType.text.equalsIgnoreCase(text)) {
-        return dataSourceType;
+  /** Register DataSourceType to be used in fromString method */
+  public static void register(DataSourceType ... dataSourceTypes) {
+    for (DataSourceType type : dataSourceTypes) {
+      String upperCaseName = type.name().toUpperCase();
+      if (!knownValues.containsKey(upperCaseName)) {
+        knownValues.put(type.name().toUpperCase(), type);
+      } else {
+        throw new IllegalArgumentException("DataSourceType with name " + type.name() + " already exists");
       }
     }
-    throw new IllegalArgumentException("No DataSourceType with text " + text + " found");
+  }
+
+  public static DataSourceType fromString(String name) {
+    String upperCaseName = name.toUpperCase();
+    if (knownValues.containsKey(upperCaseName)) {
+      return knownValues.get(upperCaseName);
+    } else {
+      throw new IllegalArgumentException("No DataSourceType with name " + name + " found");
+    }
   }
 }
diff --git a/core/src/main/java/org/opensearch/sql/utils/SerializeUtils.java b/core/src/main/java/org/opensearch/sql/utils/SerializeUtils.java
@@ -0,0 +1,51 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.utils;
+
+import com.google.gson.Gson;
+import com.google.gson.GsonBuilder;
+import com.google.gson.JsonDeserializationContext;
+import com.google.gson.JsonDeserializer;
+import com.google.gson.JsonElement;
+import com.google.gson.JsonParseException;
+import com.google.gson.JsonPrimitive;
+import com.google.gson.JsonSerializationContext;
+import com.google.gson.JsonSerializer;
+import java.lang.reflect.Type;
+import lombok.experimental.UtilityClass;
+import org.opensearch.sql.datasource.model.DataSourceType;
+
+@UtilityClass
+public class SerializeUtils {
+  private static class DataSourceTypeSerializer implements JsonSerializer<DataSourceType> {
+    @Override
+    public JsonElement serialize(
+        DataSourceType dataSourceType,
+        Type type,
+        JsonSerializationContext jsonSerializationContext) {
+      return new JsonPrimitive(dataSourceType.name());
+    }
+  }
+
+  private static class DataSourceTypeDeserializer implements JsonDeserializer<DataSourceType> {
+    @Override
+    public DataSourceType deserialize(
+        JsonElement jsonElement, Type type, JsonDeserializationContext jsonDeserializationContext)
+        throws JsonParseException {
+      return DataSourceType.fromString(jsonElement.getAsString());
+    }
+  }
+
+  public static GsonBuilder getGsonBuilder() {
+    return new GsonBuilder()
+        .registerTypeAdapter(DataSourceType.class, new DataSourceTypeSerializer())
+        .registerTypeAdapter(DataSourceType.class, new DataSourceTypeDeserializer());
+  }
+
+  public static Gson buildGson() {
+    return getGsonBuilder().create();
+  }
+}
diff --git a/core/src/test/java/org/opensearch/sql/datasource/model/DataSourceTypeTest.java b/core/src/test/java/org/opensearch/sql/datasource/model/DataSourceTypeTest.java
@@ -0,0 +1,39 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.datasource.model;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+import org.junit.jupiter.api.Test;
+
+class DataSourceTypeTest {
+  @Test
+  public void fromString_succeed() {
+    testFromString("PROMETHEUS", DataSourceType.PROMETHEUS);
+    testFromString("OPENSEARCH", DataSourceType.OPENSEARCH);
+    testFromString("SPARK", DataSourceType.SPARK);
+    testFromString("S3GLUE", DataSourceType.S3GLUE);
+
+    testFromString("prometheus", DataSourceType.PROMETHEUS);
+  }
+
+  private void testFromString(String name, DataSourceType expectedType) {
+    assertEquals(expectedType, DataSourceType.fromString(name));
+  }
+
+  @Test
+  public void fromStringWithUnknownName_throws() {
+    assertThrows(IllegalArgumentException.class, () -> DataSourceType.fromString("UnknownName"));
+  }
+
+  @Test
+  public void registerExistingType_throwsException() {
+    assertThrows(
+        IllegalArgumentException.class,
+        () -> DataSourceType.register(new DataSourceType("s3glue")));
+  }
+}
diff --git a/core/src/test/java/org/opensearch/sql/utils/SerializeUtilsTest.java b/core/src/test/java/org/opensearch/sql/utils/SerializeUtilsTest.java
@@ -0,0 +1,65 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.utils;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import com.google.common.collect.ImmutableList;
+import com.google.gson.Gson;
+import org.junit.jupiter.api.Test;
+import org.opensearch.sql.datasource.model.DataSourceMetadata;
+import org.opensearch.sql.datasource.model.DataSourceStatus;
+import org.opensearch.sql.datasource.model.DataSourceType;
+
+class SerializeUtilsTest {
+  @Test
+  public void buildGson_serializeDataSourceTypeAsString() {
+    DataSourceMetadata dataSourceMetadata =
+        new DataSourceMetadata.Builder()
+            .setName("DATASOURCE_NAME")
+            .setDescription("DESCRIPTION")
+            .setConnector(DataSourceType.S3GLUE)
+            .setAllowedRoles(ImmutableList.of("ROLE"))
+            .setResultIndex("query_execution_result_test")
+            .setDataSourceStatus(DataSourceStatus.ACTIVE)
+            .build();
+
+    Gson gson = SerializeUtils.buildGson();
+    String json = gson.toJson(dataSourceMetadata);
+
+    // connector should be serialized as string (not as object)
+    assertJsonAttribute(json, "connector", "S3GLUE");
+    // other attribute should be serialized as normal
+    assertJsonAttribute(json, "name", "DATASOURCE_NAME");
+    assertJsonAttribute(json, "description", "DESCRIPTION");
+    assertJsonAttribute(json, "resultIndex", "query_execution_result_test");
+    assertJsonAttribute(json, "status", "ACTIVE");
+    assertTrue(json.contains("\"allowedRoles\":[\"ROLE\"]"));
+  }
+
+  private void assertJsonAttribute(String json, String attribute, String value) {
+    assertTrue(json.contains("\"" + attribute + "\":\"" + value + "\""));
+  }
+
+  @Test
+  public void buildGson_deserializeDataSourceTypeFromString() {
+    String json =
+        "{\"name\":\"DATASOURCE_NAME\","
+            + "\"description\":\"DESCRIPTION\","
+            + "\"connector\":\"S3GLUE\","
+            + "\"allowedRoles\":[\"ROLE\"],"
+            + "\"properties\":{},"
+            + "\"resultIndex\":\"query_execution_result_test\","
+            + "\"status\":\"ACTIVE\""
+            + "}";
+
+    Gson gson = SerializeUtils.buildGson();
+    DataSourceMetadata metadata = gson.fromJson(json, DataSourceMetadata.class);
+
+    assertEquals(DataSourceType.S3GLUE, metadata.getConnector());
+  }
+}
diff --git a/datasources/src/main/java/org/opensearch/sql/datasources/exceptions/ErrorMessage.java b/datasources/src/main/java/org/opensearch/sql/datasources/exceptions/ErrorMessage.java
@@ -6,10 +6,10 @@
 package org.opensearch.sql.datasources.exceptions;
 
 import com.google.gson.Gson;
-import com.google.gson.GsonBuilder;
 import com.google.gson.JsonObject;
 import lombok.Getter;
 import org.opensearch.core.rest.RestStatus;
+import org.opensearch.sql.utils.SerializeUtils;
 
 /** Error Message. */
 public class ErrorMessage {
@@ -61,7 +61,7 @@ public String toString() {
     JsonObject jsonObject = new JsonObject();
     jsonObject.addProperty("status", status);
     jsonObject.add("error", getErrorAsJson());
-    Gson gson = new GsonBuilder().setPrettyPrinting().create();
+    Gson gson = SerializeUtils.getGsonBuilder().setPrettyPrinting().create();
     return gson.toJson(jsonObject);
   }
 

diff --git a/datasources/src/test/java/org/opensearch/sql/datasources/exceptions/ErrorMessageTest.java b/datasources/src/test/java/org/opensearch/sql/datasources/exceptions/ErrorMessageTest.java
@@ -13,9 +13,33 @@
 class ErrorMessageTest {
 
   @Test
-  void fetchReason() {
+  void toString_returnPrettyPrintedJson() {
     ErrorMessage errorMessage =
         new ErrorMessage(new RuntimeException(), RestStatus.TOO_MANY_REQUESTS.getStatus());
-    assertEquals("Too Many Requests", errorMessage.getReason());
+
+    assertEquals(
+        "{\n"
+            + "  \"status\": 429,\n"
+            + "  \"error\": {\n"
+            + "    \"type\": \"RuntimeException\",\n"
+            + "    \"reason\": \"Too Many Requests\",\n"
+            + "    \"details\": \"\"\n"
+            + "  }\n"
+            + "}",
+        errorMessage.toString());
+  }
+
+  @Test
+  void getReason() {
+    testGetReason(RestStatus.TOO_MANY_REQUESTS, "Too Many Requests");
+    testGetReason(RestStatus.BAD_REQUEST, "Invalid Request");
+    // other status
+    testGetReason(RestStatus.BAD_GATEWAY, "There was internal problem at backend");
+  }
+
+  void testGetReason(RestStatus status, String expectedReason) {
+    ErrorMessage errorMessage = new ErrorMessage(new RuntimeException(), status.getStatus());
+
+    assertEquals(expectedReason, errorMessage.getReason());
   }
 }
diff --git a/.../java/org/opensearch/sql/datasources/storage/OpenSearchDataSourceMetadataStorageTest.java b/.../java/org/opensearch/sql/datasources/storage/OpenSearchDataSourceMetadataStorageTest.java
@@ -8,8 +8,13 @@
 import static org.opensearch.sql.datasource.model.DataSourceStatus.ACTIVE;
 import static org.opensearch.sql.datasources.storage.OpenSearchDataSourceMetadataStorage.DATASOURCE_INDEX_NAME;
 
+import com.fasterxml.jackson.core.JsonGenerator;
 import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.databind.SerializerProvider;
+import com.fasterxml.jackson.databind.module.SimpleModule;
+import com.fasterxml.jackson.databind.ser.std.StdSerializer;
+import java.io.IOException;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
@@ -645,8 +650,7 @@ private String getBasicDataSourceMetadataString() throws JsonProcessingException
             .setConnector(DataSourceType.PROMETHEUS)
             .setAllowedRoles(Collections.singletonList("prometheus_access"))
             .build();
-    ObjectMapper objectMapper = new ObjectMapper();
-    return objectMapper.writeValueAsString(dataSourceMetadata);
+    return serialize(dataSourceMetadata);
   }
 
   private String getOldDataSourceMetadataStringWithOutStatusEnum() {
@@ -666,8 +670,7 @@ private String getAWSSigv4DataSourceMetadataString() throws JsonProcessingExcept
             .setConnector(DataSourceType.PROMETHEUS)
             .setAllowedRoles(Collections.singletonList("prometheus_access"))
             .build();
-    ObjectMapper objectMapper = new ObjectMapper();
-    return objectMapper.writeValueAsString(dataSourceMetadata);
+    return serialize(dataSourceMetadata);
   }
 
   private String getDataSourceMetadataStringWithBasicAuthentication()
@@ -684,8 +687,7 @@ private String getDataSourceMetadataStringWithBasicAuthentication()
             .setConnector(DataSourceType.PROMETHEUS)
             .setAllowedRoles(Collections.singletonList("prometheus_access"))
             .build();
-    ObjectMapper objectMapper = new ObjectMapper();
-    return objectMapper.writeValueAsString(dataSourceMetadata);
+    return serialize(dataSourceMetadata);
   }
 
   private String getDataSourceMetadataStringWithNoAuthentication() throws JsonProcessingException {
@@ -698,8 +700,7 @@ private String getDataSourceMetadataStringWithNoAuthentication() throws JsonProc
             .setConnector(DataSourceType.PROMETHEUS)
             .setAllowedRoles(Collections.singletonList("prometheus_access"))
             .build();
-    ObjectMapper objectMapper = new ObjectMapper();
-    return objectMapper.writeValueAsString(dataSourceMetadata);
+    return serialize(dataSourceMetadata);
   }
 
   private DataSourceMetadata getDataSourceMetadata() {
@@ -715,4 +716,32 @@ private DataSourceMetadata getDataSourceMetadata() {
         .setAllowedRoles(Collections.singletonList("prometheus_access"))
         .build();
   }
+
+  private String serialize(DataSourceMetadata dataSourceMetadata) throws JsonProcessingException {
+    return getObjectMapper().writeValueAsString(dataSourceMetadata);
+  }
+
+  private ObjectMapper getObjectMapper() {
+    ObjectMapper mapper = new ObjectMapper();
+    addSerializerForDataSourceType(mapper);
+    return mapper;
+  }
+
+  /** It is needed to serialize DataSourceType as string. */
+  private void addSerializerForDataSourceType(ObjectMapper mapper) {
+    SimpleModule module = new SimpleModule();
+    module.addSerializer(DataSourceType.class, getDataSourceTypeSerializer());
+    mapper.registerModule(module);
+  }
+
+  private StdSerializer<DataSourceType> getDataSourceTypeSerializer() {
+    return new StdSerializer<>(DataSourceType.class) {
+      @Override
+      public void serialize(
+          DataSourceType dsType, JsonGenerator jsonGen, SerializerProvider provider)
+          throws IOException {
+        jsonGen.writeString(dsType.name());
+      }
+    };
+  }
 }