Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Change DataSourceType from enum to class #2730

Merged
merged 9 commits into from
Jun 12, 2024
Merged
Show file tree
Hide file tree
Changes from 7 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -5,34 +5,34 @@

package org.opensearch.sql.datasource.model;

public enum DataSourceType {
PROMETHEUS("prometheus"),
OPENSEARCH("opensearch"),
SPARK("spark"),
S3GLUE("s3glue");
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;
}
private static DataSourceType[] KNOWN_VALUES = {PROMETHEUS, OPENSEARCH, SPARK, S3GLUE};
vmmusings marked this conversation as resolved.
Show resolved Hide resolved

private final String name;

public String getText() {
return this.text;
public String name() {
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)) {
public static DataSourceType fromString(String name) {
for (DataSourceType dataSourceType : KNOWN_VALUES) {
if (dataSourceType.name.equalsIgnoreCase(name)) {
return dataSourceType;
}
}
throw new IllegalArgumentException("No DataSourceType with text " + text + " found");
throw new IllegalArgumentException("No DataSourceType with name " + name + " found");
}

@Override
public String toString() {
return name;
}
}
51 changes: 51 additions & 0 deletions core/src/main/java/org/opensearch/sql/utils/SerializeUtils.java
Original file line number Diff line number Diff line change
@@ -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();
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,32 @@
/*
* 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"));
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,66 @@
/*
* 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);
System.out.println(json);
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

please remove

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for catching. removed.


// 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());
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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);
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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());
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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() {
Expand All @@ -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()
Expand All @@ -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 {
Expand All @@ -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() {
Expand All @@ -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());
}
};
}
}
Loading
Loading