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

[Backport 2.x] Change DataSourceType from enum to class #2742

Closed
wants to merge 1 commit into from
Closed
Show file tree
Hide file tree
Changes from all 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,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");
}
}
}
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,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")));
}
}
Original file line number Diff line number Diff line change
@@ -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());
}
}
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