diff --git a/core/src/main/java/org/opensearch/sql/datasource/DataSourceService.java b/core/src/main/java/org/opensearch/sql/datasource/DataSourceService.java index 162fe9e8f8..ffa818f02a 100644 --- a/core/src/main/java/org/opensearch/sql/datasource/DataSourceService.java +++ b/core/src/main/java/org/opensearch/sql/datasource/DataSourceService.java @@ -14,7 +14,8 @@ public interface DataSourceService { /** - * Returns {@link DataSource} corresponding to the DataSource name. + * Returns {@link DataSource} corresponding to the DataSource name only if the datasource is + * active and authorized. * * @param dataSourceName Name of the {@link DataSource}. * @return {@link DataSource}. @@ -84,4 +85,15 @@ public interface DataSourceService { * @param dataSourceName name of the {@link DataSource}. */ Boolean dataSourceExists(String dataSourceName); + + /** + * Performs authorization and datasource status check. We could have exposed the api using only + * datasource name, In order to avoid multiple calls form SparkQueryDispatcher, we are exposing + * API using {@link DataSourceMetadata}. + * + *

+ * + * @param dataSourceMetadata name of the {@link DataSourceMetadata} + */ + void verifyDataSourceAccess(DataSourceMetadata dataSourceMetadata); } diff --git a/core/src/main/java/org/opensearch/sql/datasource/model/DataSourceMetadata.java b/core/src/main/java/org/opensearch/sql/datasource/model/DataSourceMetadata.java index 9e47f9b37e..61918d1330 100644 --- a/core/src/main/java/org/opensearch/sql/datasource/model/DataSourceMetadata.java +++ b/core/src/main/java/org/opensearch/sql/datasource/model/DataSourceMetadata.java @@ -10,6 +10,7 @@ import com.fasterxml.jackson.annotation.JsonFormat; import com.fasterxml.jackson.annotation.JsonIgnoreProperties; import com.fasterxml.jackson.annotation.JsonProperty; +import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import java.util.ArrayList; import java.util.Collections; @@ -57,96 +58,147 @@ public class DataSourceMetadata { @JsonProperty private String resultIndex; + @JsonProperty private DataSourceStatus status; + public static Function DATASOURCE_TO_RESULT_INDEX = datasourceName -> String.format("%s_%s", DEFAULT_RESULT_INDEX, datasourceName); - public DataSourceMetadata( - String name, - String description, - DataSourceType connector, - List allowedRoles, - Map properties, - String resultIndex) { - this.name = name; - String errorMessage = validateCustomResultIndex(resultIndex); - if (errorMessage != null) { - throw new IllegalArgumentException(errorMessage); + private DataSourceMetadata(Builder builder) { + this.name = builder.name; + this.description = builder.description; + this.connector = builder.connector; + this.allowedRoles = builder.allowedRoles; + this.properties = builder.properties; + this.resultIndex = builder.resultIndex; + this.status = builder.status; + } + + public static class Builder { + private String name; + private String description; + private DataSourceType connector; + private List allowedRoles; + private Map properties; + private String resultIndex; // Optional + private DataSourceStatus status = + DataSourceStatus.ACTIVE; // Optional, default to false if not provided + + public Builder setName(String name) { + this.name = name; + return this; } - if (resultIndex == null) { - this.resultIndex = fromNameToCustomResultIndex(); - } else { - this.resultIndex = resultIndex; + + public Builder setDescription(String description) { + this.description = description; + return this; } - this.connector = connector; - this.description = description; - this.properties = properties; - this.allowedRoles = allowedRoles; - } + public Builder setConnector(DataSourceType connector) { + this.connector = connector; + return this; + } - public DataSourceMetadata() { - this.description = StringUtils.EMPTY; - this.allowedRoles = new ArrayList<>(); - this.properties = new HashMap<>(); - } + public Builder setAllowedRoles(List allowedRoles) { + this.allowedRoles = new ArrayList<>(allowedRoles); + return this; + } - /** - * Default OpenSearch {@link DataSourceMetadata}. Which is used to register default OpenSearch - * {@link DataSource} to {@link DataSourceService}. - */ - public static DataSourceMetadata defaultOpenSearchDataSourceMetadata() { - return new DataSourceMetadata( - DEFAULT_DATASOURCE_NAME, - StringUtils.EMPTY, - DataSourceType.OPENSEARCH, - Collections.emptyList(), - ImmutableMap.of(), - null); - } + public Builder setProperties(Map properties) { + this.properties = new HashMap<>(properties); + return this; + } - public String validateCustomResultIndex(String resultIndex) { - if (resultIndex == null) { - return null; + public Builder setResultIndex(String resultIndex) { + this.resultIndex = resultIndex; + return this; } - if (resultIndex.length() > MAX_RESULT_INDEX_NAME_SIZE) { - return INVALID_RESULT_INDEX_NAME_SIZE; + + public Builder setDataSourceStatus(DataSourceStatus status) { + this.status = status; + return this; } - if (!resultIndex.matches(RESULT_INDEX_NAME_PATTERN)) { - return INVALID_CHAR_IN_RESULT_INDEX_NAME; + + public DataSourceMetadata build() { + String errorMessage = validateCustomResultIndex(resultIndex); + if (errorMessage != null) { + throw new IllegalArgumentException(errorMessage); + } + if (resultIndex == null) { + this.resultIndex = fromNameToCustomResultIndex(); + } + if (status == null) { + this.status = DataSourceStatus.ACTIVE; + } + if (description == null) { + this.description = StringUtils.EMPTY; + } + if (properties == null) { + this.properties = ImmutableMap.of(); + } + if (allowedRoles == null) { + this.allowedRoles = ImmutableList.of(); + } + return new DataSourceMetadata(this); } - if (resultIndex != null && !resultIndex.startsWith(DEFAULT_RESULT_INDEX)) { - return INVALID_RESULT_INDEX_PREFIX; + + public String validateCustomResultIndex(String resultIndex) { + if (resultIndex == null) { + return null; + } + if (resultIndex.length() > MAX_RESULT_INDEX_NAME_SIZE) { + return INVALID_RESULT_INDEX_NAME_SIZE; + } + if (!resultIndex.matches(RESULT_INDEX_NAME_PATTERN)) { + return INVALID_CHAR_IN_RESULT_INDEX_NAME; + } + if (resultIndex != null && !resultIndex.startsWith(DEFAULT_RESULT_INDEX)) { + return INVALID_RESULT_INDEX_PREFIX; + } + return null; } - return null; - } - /** - * Since we are using datasource name to create result index, we need to make sure that the final - * name is valid - * - * @param resultIndex result index name - * @return valid result index name - */ - private String convertToValidResultIndex(String resultIndex) { - // Limit Length - if (resultIndex.length() > MAX_RESULT_INDEX_NAME_SIZE) { - resultIndex = resultIndex.substring(0, MAX_RESULT_INDEX_NAME_SIZE); + /** + * Since we are using datasource name to create result index, we need to make sure that the + * final name is valid + * + * @param resultIndex result index name + * @return valid result index name + */ + private String convertToValidResultIndex(String resultIndex) { + // Limit Length + if (resultIndex.length() > MAX_RESULT_INDEX_NAME_SIZE) { + resultIndex = resultIndex.substring(0, MAX_RESULT_INDEX_NAME_SIZE); + } + + // Pattern Matching: Remove characters that don't match the pattern + StringBuilder validChars = new StringBuilder(); + for (char c : resultIndex.toCharArray()) { + if (String.valueOf(c).matches(RESULT_INDEX_NAME_PATTERN)) { + validChars.append(c); + } + } + return validChars.toString(); } - // Pattern Matching: Remove characters that don't match the pattern - StringBuilder validChars = new StringBuilder(); - for (char c : resultIndex.toCharArray()) { - if (String.valueOf(c).matches(RESULT_INDEX_NAME_PATTERN)) { - validChars.append(c); + public String fromNameToCustomResultIndex() { + if (name == null) { + throw new IllegalArgumentException("Datasource name cannot be null"); } + return convertToValidResultIndex(DATASOURCE_TO_RESULT_INDEX.apply(name.toLowerCase())); } - return validChars.toString(); } - public String fromNameToCustomResultIndex() { - if (name == null) { - throw new IllegalArgumentException("Datasource name cannot be null"); - } - return convertToValidResultIndex(DATASOURCE_TO_RESULT_INDEX.apply(name.toLowerCase())); + /** + * Default OpenSearch {@link DataSourceMetadata}. Which is used to register default OpenSearch + * {@link DataSource} to {@link DataSourceService}. + */ + public static DataSourceMetadata defaultOpenSearchDataSourceMetadata() { + return new DataSourceMetadata.Builder() + .setName(DEFAULT_DATASOURCE_NAME) + .setDescription(StringUtils.EMPTY) + .setConnector(DataSourceType.OPENSEARCH) + .setAllowedRoles(Collections.emptyList()) + .setProperties(ImmutableMap.of()) + .build(); } } diff --git a/core/src/main/java/org/opensearch/sql/datasource/model/DataSourceStatus.java b/core/src/main/java/org/opensearch/sql/datasource/model/DataSourceStatus.java new file mode 100644 index 0000000000..e382d6b55f --- /dev/null +++ b/core/src/main/java/org/opensearch/sql/datasource/model/DataSourceStatus.java @@ -0,0 +1,36 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.datasource.model; + +public enum DataSourceStatus { + ACTIVE("active"), + DISABLED("disabled"); + + private String text; + + DataSourceStatus(String text) { + this.text = text; + } + + public String getText() { + return this.text; + } + + /** + * Get DataSourceStatus from text. + * + * @param text text. + * @return DataSourceStatus {@link DataSourceStatus}. + */ + public static DataSourceStatus fromString(String text) { + for (DataSourceStatus dataSourceStatus : DataSourceStatus.values()) { + if (dataSourceStatus.text.equalsIgnoreCase(text)) { + return dataSourceStatus; + } + } + throw new IllegalArgumentException("No DataSourceStatus with text " + text + " found"); + } +} diff --git a/core/src/test/java/org/opensearch/sql/analysis/AnalyzerTestBase.java b/core/src/test/java/org/opensearch/sql/analysis/AnalyzerTestBase.java index bfd68ee53a..2c2ddbdd88 100644 --- a/core/src/test/java/org/opensearch/sql/analysis/AnalyzerTestBase.java +++ b/core/src/test/java/org/opensearch/sql/analysis/AnalyzerTestBase.java @@ -19,7 +19,6 @@ import java.util.Set; import java.util.stream.Collectors; import java.util.stream.Stream; -import org.apache.commons.lang3.StringUtils; import org.apache.commons.lang3.tuple.Pair; import org.opensearch.sql.DataSourceSchemaName; import org.opensearch.sql.analysis.symbol.Namespace; @@ -196,13 +195,10 @@ public Set getDataSourceMetadata(boolean isDefaultDataSource return Stream.of(opensearchDataSource, prometheusDataSource) .map( ds -> - new DataSourceMetadata( - ds.getName(), - StringUtils.EMPTY, - ds.getConnectorType(), - Collections.emptyList(), - ImmutableMap.of(), - null)) + new DataSourceMetadata.Builder() + .setName(ds.getName()) + .setConnector(ds.getConnectorType()) + .build()) .collect(Collectors.toSet()); } @@ -243,6 +239,9 @@ public void deleteDataSource(String dataSourceName) {} public Boolean dataSourceExists(String dataSourceName) { return dataSourceName.equals(DEFAULT_DATASOURCE_NAME) || dataSourceName.equals("prometheus"); } + + @Override + public void verifyDataSourceAccess(DataSourceMetadata dataSourceMetadata) {} } private class TestTableFunctionImplementation implements TableFunctionImplementation { diff --git a/core/src/test/java/org/opensearch/sql/planner/physical/datasource/DataSourceTableScanTest.java b/core/src/test/java/org/opensearch/sql/planner/physical/datasource/DataSourceTableScanTest.java index 5c7182a752..222b9da6c9 100644 --- a/core/src/test/java/org/opensearch/sql/planner/physical/datasource/DataSourceTableScanTest.java +++ b/core/src/test/java/org/opensearch/sql/planner/physical/datasource/DataSourceTableScanTest.java @@ -13,12 +13,10 @@ import static org.mockito.Mockito.when; import com.google.common.collect.ImmutableMap; -import java.util.Collections; import java.util.HashSet; import java.util.LinkedHashMap; import java.util.Set; import java.util.stream.Collectors; -import org.apache.commons.lang3.StringUtils; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; @@ -61,13 +59,10 @@ void testIterator() { dataSourceSet.stream() .map( dataSource -> - new DataSourceMetadata( - dataSource.getName(), - StringUtils.EMPTY, - dataSource.getConnectorType(), - Collections.emptyList(), - ImmutableMap.of(), - null)) + new DataSourceMetadata.Builder() + .setName(dataSource.getName()) + .setConnector(dataSource.getConnectorType()) + .build()) .collect(Collectors.toSet()); when(dataSourceService.getDataSourceMetadata(false)).thenReturn(dataSourceMetadata); diff --git a/datasources/src/main/java/org/opensearch/sql/datasources/exceptions/DatasourceDisabledException.java b/datasources/src/main/java/org/opensearch/sql/datasources/exceptions/DatasourceDisabledException.java new file mode 100644 index 0000000000..9eebc9ddc1 --- /dev/null +++ b/datasources/src/main/java/org/opensearch/sql/datasources/exceptions/DatasourceDisabledException.java @@ -0,0 +1,7 @@ +package org.opensearch.sql.datasources.exceptions; + +public class DatasourceDisabledException extends RuntimeException { + public DatasourceDisabledException(String message) { + super(message); + } +} diff --git a/datasources/src/main/java/org/opensearch/sql/datasources/service/DataSourceServiceImpl.java b/datasources/src/main/java/org/opensearch/sql/datasources/service/DataSourceServiceImpl.java index 8ba618fb44..c77121e327 100644 --- a/datasources/src/main/java/org/opensearch/sql/datasources/service/DataSourceServiceImpl.java +++ b/datasources/src/main/java/org/opensearch/sql/datasources/service/DataSourceServiceImpl.java @@ -15,8 +15,10 @@ import org.opensearch.sql.datasource.DataSourceService; import org.opensearch.sql.datasource.model.DataSource; import org.opensearch.sql.datasource.model.DataSourceMetadata; +import org.opensearch.sql.datasource.model.DataSourceStatus; import org.opensearch.sql.datasources.auth.DataSourceUserAuthorizationHelper; import org.opensearch.sql.datasources.exceptions.DataSourceNotFoundException; +import org.opensearch.sql.datasources.exceptions.DatasourceDisabledException; import org.opensearch.sql.storage.DataSourceFactory; /** @@ -71,7 +73,7 @@ public DataSourceMetadata getDataSourceMetadata(String dataSourceName) { @Override public DataSource getDataSource(String dataSourceName) { DataSourceMetadata dataSourceMetadata = getRawDataSourceMetadata(dataSourceName); - this.dataSourceUserAuthorizationHelper.authorizeDataSource(dataSourceMetadata); + verifyDataSourceAccess(dataSourceMetadata); return dataSourceLoaderCache.getOrLoadDataSource(dataSourceMetadata); } @@ -125,6 +127,15 @@ public Boolean dataSourceExists(String dataSourceName) { || this.dataSourceMetadataStorage.getDataSourceMetadata(dataSourceName).isPresent(); } + @Override + public void verifyDataSourceAccess(DataSourceMetadata dataSourceMetadata) { + if (dataSourceMetadata.getStatus().equals(DataSourceStatus.DISABLED)) { + throw new DatasourceDisabledException( + String.format("Datasource %s is disabled.", dataSourceMetadata)); + } + this.dataSourceUserAuthorizationHelper.authorizeDataSource(dataSourceMetadata); + } + /** * This can be moved to a different validator class when we introduce more connectors. * diff --git a/datasources/src/main/java/org/opensearch/sql/datasources/utils/XContentParserUtils.java b/datasources/src/main/java/org/opensearch/sql/datasources/utils/XContentParserUtils.java index 6af2a5a761..7e7ae458ed 100644 --- a/datasources/src/main/java/org/opensearch/sql/datasources/utils/XContentParserUtils.java +++ b/datasources/src/main/java/org/opensearch/sql/datasources/utils/XContentParserUtils.java @@ -21,6 +21,7 @@ import org.opensearch.core.xcontent.XContentBuilder; import org.opensearch.core.xcontent.XContentParser; import org.opensearch.sql.datasource.model.DataSourceMetadata; +import org.opensearch.sql.datasource.model.DataSourceStatus; import org.opensearch.sql.datasource.model.DataSourceType; /** Utitlity class to serialize and deserialize objects in XContent. */ @@ -33,6 +34,7 @@ public class XContentParserUtils { public static final String ALLOWED_ROLES_FIELD = "allowedRoles"; public static final String RESULT_INDEX_FIELD = "resultIndex"; + public static final String STATUS_FIELD = "status"; /** * Convert xcontent parser to DataSourceMetadata. @@ -48,6 +50,7 @@ public static DataSourceMetadata toDataSourceMetadata(XContentParser parser) thr List allowedRoles = new ArrayList<>(); Map properties = new HashMap<>(); String resultIndex = null; + DataSourceStatus status = null; ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.nextToken(), parser); while (parser.nextToken() != XContentParser.Token.END_OBJECT) { String fieldName = parser.currentName(); @@ -79,6 +82,9 @@ public static DataSourceMetadata toDataSourceMetadata(XContentParser parser) thr case RESULT_INDEX_FIELD: resultIndex = parser.textOrNull(); break; + case STATUS_FIELD: + status = DataSourceStatus.fromString(parser.textOrNull()); + break; default: throw new IllegalArgumentException("Unknown field: " + fieldName); } @@ -86,8 +92,15 @@ public static DataSourceMetadata toDataSourceMetadata(XContentParser parser) thr if (name == null || connector == null) { throw new IllegalArgumentException("name and connector are required fields."); } - return new DataSourceMetadata( - name, description, connector, allowedRoles, properties, resultIndex); + return new DataSourceMetadata.Builder() + .setName(name) + .setDescription(description) + .setConnector(connector) + .setProperties(properties) + .setAllowedRoles(allowedRoles) + .setResultIndex(resultIndex) + .setDataSourceStatus(status) + .build(); } public static Map toMap(XContentParser parser) throws IOException { @@ -97,6 +110,7 @@ public static Map toMap(XContentParser parser) throws IOExceptio List allowedRoles = new ArrayList<>(); Map properties = new HashMap<>(); String resultIndex; + DataSourceStatus status; ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.nextToken(), parser); while (parser.nextToken() != XContentParser.Token.END_OBJECT) { String fieldName = parser.currentName(); @@ -133,6 +147,10 @@ public static Map toMap(XContentParser parser) throws IOExceptio resultIndex = parser.textOrNull(); resultMap.put(RESULT_INDEX_FIELD, resultIndex); break; + case STATUS_FIELD: + status = DataSourceStatus.fromString(parser.textOrNull()); + resultMap.put(STATUS_FIELD, status); + break; default: throw new IllegalArgumentException("Unknown field: " + fieldName); } diff --git a/datasources/src/test/java/org/opensearch/sql/datasources/auth/DataSourceUserAuthorizationHelperImplTest.java b/datasources/src/test/java/org/opensearch/sql/datasources/auth/DataSourceUserAuthorizationHelperImplTest.java index 6ee3c12edd..abf1cb9cfe 100644 --- a/datasources/src/test/java/org/opensearch/sql/datasources/auth/DataSourceUserAuthorizationHelperImplTest.java +++ b/datasources/src/test/java/org/opensearch/sql/datasources/auth/DataSourceUserAuthorizationHelperImplTest.java @@ -7,7 +7,6 @@ import static org.opensearch.commons.ConfigConstants.OPENSEARCH_SECURITY_USER_INFO_THREAD_CONTEXT; -import java.util.HashMap; import java.util.List; import org.junit.Assert; import org.junit.jupiter.api.Test; @@ -102,11 +101,10 @@ public void testAuthorizeDataSourceWithException() { } private DataSourceMetadata dataSourceMetadata() { - DataSourceMetadata dataSourceMetadata = new DataSourceMetadata(); - dataSourceMetadata.setName("test"); - dataSourceMetadata.setConnector(DataSourceType.PROMETHEUS); - dataSourceMetadata.setAllowedRoles(List.of("prometheus_access")); - dataSourceMetadata.setProperties(new HashMap<>()); - return dataSourceMetadata; + return new DataSourceMetadata.Builder() + .setName("test_datasource") + .setAllowedRoles(List.of("prometheus_access")) + .setConnector(DataSourceType.PROMETHEUS) + .build(); } } diff --git a/datasources/src/test/java/org/opensearch/sql/datasources/glue/GlueDataSourceFactoryTest.java b/datasources/src/test/java/org/opensearch/sql/datasources/glue/GlueDataSourceFactoryTest.java index 4dd054de70..52f8ec9cd1 100644 --- a/datasources/src/test/java/org/opensearch/sql/datasources/glue/GlueDataSourceFactoryTest.java +++ b/datasources/src/test/java/org/opensearch/sql/datasources/glue/GlueDataSourceFactoryTest.java @@ -35,17 +35,18 @@ void testCreateGLueDatSource() { .thenReturn(Collections.emptyList()); GlueDataSourceFactory glueDatasourceFactory = new GlueDataSourceFactory(settings); - DataSourceMetadata metadata = new DataSourceMetadata(); HashMap properties = new HashMap<>(); properties.put("glue.auth.type", "iam_role"); properties.put("glue.auth.role_arn", "role_arn"); properties.put("glue.indexstore.opensearch.uri", "http://localhost:9200"); properties.put("glue.indexstore.opensearch.auth", "noauth"); properties.put("glue.indexstore.opensearch.region", "us-west-2"); - - metadata.setName("my_glue"); - metadata.setConnector(DataSourceType.S3GLUE); - metadata.setProperties(properties); + DataSourceMetadata metadata = + new DataSourceMetadata.Builder() + .setName("my_glue") + .setConnector(DataSourceType.S3GLUE) + .setProperties(properties) + .build(); DataSource dataSource = glueDatasourceFactory.createDataSource(metadata); Assertions.assertEquals(DataSourceType.S3GLUE, dataSource.getConnectorType()); UnsupportedOperationException unsupportedOperationException = @@ -66,7 +67,6 @@ void testCreateGLueDatSourceWithBasicAuthForIndexStore() { .thenReturn(Collections.emptyList()); GlueDataSourceFactory glueDatasourceFactory = new GlueDataSourceFactory(settings); - DataSourceMetadata metadata = new DataSourceMetadata(); HashMap properties = new HashMap<>(); properties.put("glue.auth.type", "iam_role"); properties.put("glue.auth.role_arn", "role_arn"); @@ -75,10 +75,12 @@ void testCreateGLueDatSourceWithBasicAuthForIndexStore() { properties.put("glue.indexstore.opensearch.auth.username", "username"); properties.put("glue.indexstore.opensearch.auth.password", "password"); properties.put("glue.indexstore.opensearch.region", "us-west-2"); - - metadata.setName("my_glue"); - metadata.setConnector(DataSourceType.S3GLUE); - metadata.setProperties(properties); + DataSourceMetadata metadata = + new DataSourceMetadata.Builder() + .setName("my_glue") + .setConnector(DataSourceType.S3GLUE) + .setProperties(properties) + .build(); DataSource dataSource = glueDatasourceFactory.createDataSource(metadata); Assertions.assertEquals(DataSourceType.S3GLUE, dataSource.getConnectorType()); UnsupportedOperationException unsupportedOperationException = @@ -99,17 +101,18 @@ void testCreateGLueDatSourceWithAwsSigV4AuthForIndexStore() { .thenReturn(Collections.emptyList()); GlueDataSourceFactory glueDatasourceFactory = new GlueDataSourceFactory(settings); - DataSourceMetadata metadata = new DataSourceMetadata(); HashMap properties = new HashMap<>(); properties.put("glue.auth.type", "iam_role"); properties.put("glue.auth.role_arn", "role_arn"); properties.put("glue.indexstore.opensearch.uri", "http://localhost:9200"); properties.put("glue.indexstore.opensearch.auth", "awssigv4"); properties.put("glue.indexstore.opensearch.region", "us-west-2"); - - metadata.setName("my_glue"); - metadata.setConnector(DataSourceType.S3GLUE); - metadata.setProperties(properties); + DataSourceMetadata metadata = + new DataSourceMetadata.Builder() + .setName("my_glue") + .setConnector(DataSourceType.S3GLUE) + .setProperties(properties) + .build(); DataSource dataSource = glueDatasourceFactory.createDataSource(metadata); Assertions.assertEquals(DataSourceType.S3GLUE, dataSource.getConnectorType()); UnsupportedOperationException unsupportedOperationException = @@ -128,16 +131,19 @@ void testCreateGLueDatSourceWithAwsSigV4AuthForIndexStore() { void testCreateGLueDatSourceWithBasicAuthForIndexStoreAndMissingFields() { GlueDataSourceFactory glueDatasourceFactory = new GlueDataSourceFactory(settings); - DataSourceMetadata metadata = new DataSourceMetadata(); HashMap properties = new HashMap<>(); properties.put("glue.auth.type", "iam_role"); properties.put("glue.auth.role_arn", "role_arn"); properties.put("glue.indexstore.opensearch.uri", "http://localhost:9200"); properties.put("glue.indexstore.opensearch.auth", "basicauth"); - metadata.setName("my_glue"); - metadata.setConnector(DataSourceType.S3GLUE); - metadata.setProperties(properties); + DataSourceMetadata metadata = + new DataSourceMetadata.Builder() + .setName("my_glue") + .setConnector(DataSourceType.S3GLUE) + .setProperties(properties) + .build(); + IllegalArgumentException illegalArgumentException = Assertions.assertThrows( IllegalArgumentException.class, () -> glueDatasourceFactory.createDataSource(metadata)); @@ -154,7 +160,6 @@ void testCreateGLueDatSourceWithInvalidFlintHost() { .thenReturn(List.of("127.0.0.0/8")); GlueDataSourceFactory glueDatasourceFactory = new GlueDataSourceFactory(settings); - DataSourceMetadata metadata = new DataSourceMetadata(); HashMap properties = new HashMap<>(); properties.put("glue.auth.type", "iam_role"); properties.put("glue.auth.role_arn", "role_arn"); @@ -162,9 +167,12 @@ void testCreateGLueDatSourceWithInvalidFlintHost() { properties.put("glue.indexstore.opensearch.auth", "noauth"); properties.put("glue.indexstore.opensearch.region", "us-west-2"); - metadata.setName("my_glue"); - metadata.setConnector(DataSourceType.S3GLUE); - metadata.setProperties(properties); + DataSourceMetadata metadata = + new DataSourceMetadata.Builder() + .setName("my_glue") + .setConnector(DataSourceType.S3GLUE) + .setProperties(properties) + .build(); IllegalArgumentException illegalArgumentException = Assertions.assertThrows( IllegalArgumentException.class, () -> glueDatasourceFactory.createDataSource(metadata)); @@ -181,7 +189,6 @@ void testCreateGLueDatSourceWithInvalidFlintHostSyntax() { .thenReturn(List.of("127.0.0.0/8")); GlueDataSourceFactory glueDatasourceFactory = new GlueDataSourceFactory(settings); - DataSourceMetadata metadata = new DataSourceMetadata(); HashMap properties = new HashMap<>(); properties.put("glue.auth.type", "iam_role"); properties.put("glue.auth.role_arn", "role_arn"); @@ -191,9 +198,12 @@ void testCreateGLueDatSourceWithInvalidFlintHostSyntax() { properties.put("glue.indexstore.opensearch.auth", "noauth"); properties.put("glue.indexstore.opensearch.region", "us-west-2"); - metadata.setName("my_glue"); - metadata.setConnector(DataSourceType.S3GLUE); - metadata.setProperties(properties); + DataSourceMetadata metadata = + new DataSourceMetadata.Builder() + .setName("my_glue") + .setConnector(DataSourceType.S3GLUE) + .setProperties(properties) + .build(); IllegalArgumentException illegalArgumentException = Assertions.assertThrows( IllegalArgumentException.class, () -> glueDatasourceFactory.createDataSource(metadata)); diff --git a/datasources/src/test/java/org/opensearch/sql/datasources/service/DataSourceLoaderCacheImplTest.java b/datasources/src/test/java/org/opensearch/sql/datasources/service/DataSourceLoaderCacheImplTest.java index b2ea221eb7..1df92d29e9 100644 --- a/datasources/src/test/java/org/opensearch/sql/datasources/service/DataSourceLoaderCacheImplTest.java +++ b/datasources/src/test/java/org/opensearch/sql/datasources/service/DataSourceLoaderCacheImplTest.java @@ -7,7 +7,6 @@ import static org.mockito.Mockito.verifyNoMoreInteractions; import static org.mockito.Mockito.when; -import com.google.common.collect.ImmutableMap; import java.util.Collections; import java.util.List; import org.junit.jupiter.api.Assertions; @@ -35,7 +34,8 @@ public void setup() { .doAnswer( invocation -> { DataSourceMetadata metadata = invocation.getArgument(0); - return new DataSource(metadata.getName(), metadata.getConnector(), storageEngine); + return new DataSource( + metadata.getName(), metadata.getConnector(), storageEngine); }) .when(dataSourceFactory) .createDataSource(any()); @@ -46,11 +46,7 @@ public void setup() { void testGetOrLoadDataSource() { DataSourceLoaderCache dataSourceLoaderCache = new DataSourceLoaderCacheImpl(Collections.singleton(dataSourceFactory)); - DataSourceMetadata dataSourceMetadata = new DataSourceMetadata(); - dataSourceMetadata.setName("testDS"); - dataSourceMetadata.setConnector(DataSourceType.OPENSEARCH); - dataSourceMetadata.setAllowedRoles(Collections.emptyList()); - dataSourceMetadata.setProperties(ImmutableMap.of()); + DataSourceMetadata dataSourceMetadata = getMetadata(); DataSource dataSource = dataSourceLoaderCache.getOrLoadDataSource(dataSourceMetadata); verify(dataSourceFactory, times(1)).createDataSource(dataSourceMetadata); Assertions.assertEquals( @@ -72,11 +68,9 @@ void testGetOrLoadDataSourceWithMetadataUpdate() { } private DataSourceMetadata getMetadata() { - DataSourceMetadata dataSourceMetadata = new DataSourceMetadata(); - dataSourceMetadata.setName("testDS"); - dataSourceMetadata.setConnector(DataSourceType.OPENSEARCH); - dataSourceMetadata.setAllowedRoles(Collections.emptyList()); - dataSourceMetadata.setProperties(ImmutableMap.of()); - return dataSourceMetadata; + return new DataSourceMetadata.Builder() + .setName("testDS") + .setConnector(DataSourceType.OPENSEARCH) + .build(); } } diff --git a/datasources/src/test/java/org/opensearch/sql/datasources/service/DataSourceServiceImplTest.java b/datasources/src/test/java/org/opensearch/sql/datasources/service/DataSourceServiceImplTest.java index bf88302833..012959b959 100644 --- a/datasources/src/test/java/org/opensearch/sql/datasources/service/DataSourceServiceImplTest.java +++ b/datasources/src/test/java/org/opensearch/sql/datasources/service/DataSourceServiceImplTest.java @@ -29,7 +29,6 @@ import java.util.Map; import java.util.Optional; import java.util.Set; -import org.apache.commons.lang3.StringUtils; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -365,12 +364,12 @@ void testDataSourceExistsForDefaultDataSource() { DataSourceMetadata metadata( String name, DataSourceType type, List allowedRoles, Map properties) { - DataSourceMetadata dataSourceMetadata = new DataSourceMetadata(); - dataSourceMetadata.setName(name); - dataSourceMetadata.setConnector(type); - dataSourceMetadata.setAllowedRoles(allowedRoles); - dataSourceMetadata.setProperties(properties); - return dataSourceMetadata; + return new DataSourceMetadata.Builder() + .setName(name) + .setConnector(type) + .setAllowedRoles(allowedRoles) + .setProperties(properties) + .build(); } @Test @@ -381,13 +380,12 @@ void testRemovalOfAuthorizationInfo() { properties.put("prometheus.auth.username", "username"); properties.put("prometheus.auth.password", "password"); DataSourceMetadata dataSourceMetadata = - new DataSourceMetadata( - "testDS", - StringUtils.EMPTY, - DataSourceType.PROMETHEUS, - Collections.singletonList("prometheus_access"), - properties, - null); + new DataSourceMetadata.Builder() + .setName("testDS") + .setProperties(properties) + .setConnector(DataSourceType.PROMETHEUS) + .setAllowedRoles(Collections.singletonList("prometheus_access")) + .build(); when(dataSourceMetadataStorage.getDataSourceMetadata("testDS")) .thenReturn(Optional.of(dataSourceMetadata)); @@ -407,13 +405,12 @@ void testRemovalOfAuthorizationInfoForAccessKeyAndSecretKye() { properties.put("prometheus.auth.access_key", "access_key"); properties.put("prometheus.auth.secret_key", "secret_key"); DataSourceMetadata dataSourceMetadata = - new DataSourceMetadata( - "testDS", - StringUtils.EMPTY, - DataSourceType.PROMETHEUS, - Collections.singletonList("prometheus_access"), - properties, - null); + new DataSourceMetadata.Builder() + .setName("testDS") + .setProperties(properties) + .setConnector(DataSourceType.PROMETHEUS) + .setAllowedRoles(Collections.singletonList("prometheus_access")) + .build(); when(dataSourceMetadataStorage.getDataSourceMetadata("testDS")) .thenReturn(Optional.of(dataSourceMetadata)); @@ -435,13 +432,12 @@ void testRemovalOfAuthorizationInfoForGlueWithRoleARN() { properties.put("glue.indexstore.opensearch.auth.username", "username"); properties.put("glue.indexstore.opensearch.auth.password", "password"); DataSourceMetadata dataSourceMetadata = - new DataSourceMetadata( - "testGlue", - StringUtils.EMPTY, - DataSourceType.S3GLUE, - Collections.singletonList("glue_access"), - properties, - null); + new DataSourceMetadata.Builder() + .setName("testGlue") + .setProperties(properties) + .setConnector(DataSourceType.S3GLUE) + .setAllowedRoles(Collections.singletonList("glue_access")) + .build(); when(dataSourceMetadataStorage.getDataSourceMetadata("testGlue")) .thenReturn(Optional.of(dataSourceMetadata)); @@ -500,13 +496,12 @@ void testGetRawDataSourceMetadata() { properties.put("prometheus.auth.username", "username"); properties.put("prometheus.auth.password", "password"); DataSourceMetadata dataSourceMetadata = - new DataSourceMetadata( - "testDS", - StringUtils.EMPTY, - DataSourceType.PROMETHEUS, - Collections.singletonList("prometheus_access"), - properties, - null); + new DataSourceMetadata.Builder() + .setName("testDS") + .setProperties(properties) + .setConnector(DataSourceType.PROMETHEUS) + .setAllowedRoles(Collections.singletonList("prometheus_access")) + .build(); when(dataSourceMetadataStorage.getDataSourceMetadata("testDS")) .thenReturn(Optional.of(dataSourceMetadata)); diff --git a/datasources/src/test/java/org/opensearch/sql/datasources/storage/OpenSearchDataSourceMetadataStorageTest.java b/datasources/src/test/java/org/opensearch/sql/datasources/storage/OpenSearchDataSourceMetadataStorageTest.java index 7d41737b2d..593ed8781c 100644 --- a/datasources/src/test/java/org/opensearch/sql/datasources/storage/OpenSearchDataSourceMetadataStorageTest.java +++ b/datasources/src/test/java/org/opensearch/sql/datasources/storage/OpenSearchDataSourceMetadataStorageTest.java @@ -599,74 +599,82 @@ public void testDeleteDataSourceMetadataWithUnexpectedResult() { } private String getBasicDataSourceMetadataString() throws JsonProcessingException { - DataSourceMetadata dataSourceMetadata = new DataSourceMetadata(); - dataSourceMetadata.setName("testDS"); - dataSourceMetadata.setConnector(DataSourceType.PROMETHEUS); - dataSourceMetadata.setAllowedRoles(Collections.singletonList("prometheus_access")); Map properties = new HashMap<>(); properties.put("prometheus.auth.type", "basicauth"); properties.put("prometheus.auth.username", "username"); properties.put("prometheus.auth.uri", "https://localhost:9090"); properties.put("prometheus.auth.password", "password"); - dataSourceMetadata.setProperties(properties); + DataSourceMetadata dataSourceMetadata = + new DataSourceMetadata.Builder() + .setName("testDS") + .setProperties(properties) + .setConnector(DataSourceType.PROMETHEUS) + .setAllowedRoles(Collections.singletonList("prometheus_access")) + .build(); ObjectMapper objectMapper = new ObjectMapper(); return objectMapper.writeValueAsString(dataSourceMetadata); } private String getAWSSigv4DataSourceMetadataString() throws JsonProcessingException { - DataSourceMetadata dataSourceMetadata = new DataSourceMetadata(); - dataSourceMetadata.setName("testDS"); - dataSourceMetadata.setConnector(DataSourceType.PROMETHEUS); - dataSourceMetadata.setAllowedRoles(Collections.singletonList("prometheus_access")); Map properties = new HashMap<>(); properties.put("prometheus.auth.type", "awssigv4"); properties.put("prometheus.auth.secret_key", "secret_key"); properties.put("prometheus.auth.uri", "https://localhost:9090"); properties.put("prometheus.auth.access_key", "access_key"); - dataSourceMetadata.setProperties(properties); + DataSourceMetadata dataSourceMetadata = + new DataSourceMetadata.Builder() + .setName("testDS") + .setProperties(properties) + .setConnector(DataSourceType.PROMETHEUS) + .setAllowedRoles(Collections.singletonList("prometheus_access")) + .build(); ObjectMapper objectMapper = new ObjectMapper(); return objectMapper.writeValueAsString(dataSourceMetadata); } private String getDataSourceMetadataStringWithBasicAuthentication() throws JsonProcessingException { - DataSourceMetadata dataSourceMetadata = new DataSourceMetadata(); - dataSourceMetadata.setName("testDS"); - dataSourceMetadata.setConnector(DataSourceType.PROMETHEUS); - dataSourceMetadata.setAllowedRoles(Collections.singletonList("prometheus_access")); Map properties = new HashMap<>(); properties.put("prometheus.auth.uri", "https://localhost:9090"); properties.put("prometheus.auth.type", "basicauth"); properties.put("prometheus.auth.username", "username"); properties.put("prometheus.auth.password", "password"); - dataSourceMetadata.setProperties(properties); + DataSourceMetadata dataSourceMetadata = + new DataSourceMetadata.Builder() + .setName("testDS") + .setProperties(properties) + .setConnector(DataSourceType.PROMETHEUS) + .setAllowedRoles(Collections.singletonList("prometheus_access")) + .build(); ObjectMapper objectMapper = new ObjectMapper(); return objectMapper.writeValueAsString(dataSourceMetadata); } private String getDataSourceMetadataStringWithNoAuthentication() throws JsonProcessingException { - DataSourceMetadata dataSourceMetadata = new DataSourceMetadata(); - dataSourceMetadata.setName("testDS"); - dataSourceMetadata.setConnector(DataSourceType.PROMETHEUS); - dataSourceMetadata.setAllowedRoles(Collections.singletonList("prometheus_access")); Map properties = new HashMap<>(); properties.put("prometheus.auth.uri", "https://localhost:9090"); - dataSourceMetadata.setProperties(properties); + DataSourceMetadata dataSourceMetadata = + new DataSourceMetadata.Builder() + .setName("testDS") + .setProperties(properties) + .setConnector(DataSourceType.PROMETHEUS) + .setAllowedRoles(Collections.singletonList("prometheus_access")) + .build(); ObjectMapper objectMapper = new ObjectMapper(); return objectMapper.writeValueAsString(dataSourceMetadata); } private DataSourceMetadata getDataSourceMetadata() { - DataSourceMetadata dataSourceMetadata = new DataSourceMetadata(); - dataSourceMetadata.setName("testDS"); - dataSourceMetadata.setConnector(DataSourceType.PROMETHEUS); - dataSourceMetadata.setAllowedRoles(Collections.singletonList("prometheus_access")); Map properties = new HashMap<>(); properties.put("prometheus.auth.type", "awssigv4"); properties.put("prometheus.auth.secret_key", "secret_key"); properties.put("prometheus.auth.uri", "https://localhost:9090"); properties.put("prometheus.auth.access_key", "access_key"); - dataSourceMetadata.setProperties(properties); - return dataSourceMetadata; + return new DataSourceMetadata.Builder() + .setName("testDS") + .setProperties(properties) + .setConnector(DataSourceType.PROMETHEUS) + .setAllowedRoles(Collections.singletonList("prometheus_access")) + .build(); } } diff --git a/datasources/src/test/java/org/opensearch/sql/datasources/transport/TransportCreateDataSourceActionTest.java b/datasources/src/test/java/org/opensearch/sql/datasources/transport/TransportCreateDataSourceActionTest.java index 9088d3c4ad..ba93890883 100644 --- a/datasources/src/test/java/org/opensearch/sql/datasources/transport/TransportCreateDataSourceActionTest.java +++ b/datasources/src/test/java/org/opensearch/sql/datasources/transport/TransportCreateDataSourceActionTest.java @@ -71,9 +71,11 @@ public void setUp() { @Test public void testDoExecute() { - DataSourceMetadata dataSourceMetadata = new DataSourceMetadata(); - dataSourceMetadata.setName("test_datasource"); - dataSourceMetadata.setConnector(DataSourceType.PROMETHEUS); + DataSourceMetadata dataSourceMetadata = + new DataSourceMetadata.Builder() + .setName("test_datasource") + .setConnector(DataSourceType.PROMETHEUS) + .build(); CreateDataSourceActionRequest request = new CreateDataSourceActionRequest(dataSourceMetadata); action.doExecute(task, request, actionListener); @@ -88,9 +90,11 @@ public void testDoExecute() { @Test public void testDoExecuteWithException() { - DataSourceMetadata dataSourceMetadata = new DataSourceMetadata(); - dataSourceMetadata.setName("test_datasource"); - dataSourceMetadata.setConnector(DataSourceType.PROMETHEUS); + DataSourceMetadata dataSourceMetadata = + new DataSourceMetadata.Builder() + .setName("test_datasource") + .setConnector(DataSourceType.PROMETHEUS) + .build(); doThrow(new RuntimeException("Error")) .when(dataSourceService) .createDataSource(dataSourceMetadata); @@ -105,9 +109,11 @@ public void testDoExecuteWithException() { @Test public void testDataSourcesLimit() { - DataSourceMetadata dataSourceMetadata = new DataSourceMetadata(); - dataSourceMetadata.setName("test_datasource"); - dataSourceMetadata.setConnector(DataSourceType.PROMETHEUS); + DataSourceMetadata dataSourceMetadata = + new DataSourceMetadata.Builder() + .setName("test_datasource") + .setConnector(DataSourceType.PROMETHEUS) + .build(); CreateDataSourceActionRequest request = new CreateDataSourceActionRequest(dataSourceMetadata); when(dataSourceService.getDataSourceMetadata(false).size()).thenReturn(1); when(settings.getSettingValue(DATASOURCES_LIMIT)).thenReturn(1); diff --git a/datasources/src/test/java/org/opensearch/sql/datasources/transport/TransportGetDataSourceActionTest.java b/datasources/src/test/java/org/opensearch/sql/datasources/transport/TransportGetDataSourceActionTest.java index 286f308402..90bd7bb025 100644 --- a/datasources/src/test/java/org/opensearch/sql/datasources/transport/TransportGetDataSourceActionTest.java +++ b/datasources/src/test/java/org/opensearch/sql/datasources/transport/TransportGetDataSourceActionTest.java @@ -68,9 +68,11 @@ public void setUp() { @Test public void testDoExecute() { - DataSourceMetadata dataSourceMetadata = new DataSourceMetadata(); - dataSourceMetadata.setName("test_datasource"); - dataSourceMetadata.setConnector(DataSourceType.PROMETHEUS); + DataSourceMetadata dataSourceMetadata = + new DataSourceMetadata.Builder() + .setName("test_datasource") + .setConnector(DataSourceType.PROMETHEUS) + .build(); GetDataSourceActionRequest request = new GetDataSourceActionRequest("test_datasource"); when(dataSourceService.getDataSourceMetadata("test_datasource")).thenReturn(dataSourceMetadata); @@ -97,10 +99,11 @@ protected Object buildJsonObject(DataSourceMetadata response) { @Test public void testDoExecuteForGetAllDataSources() { - DataSourceMetadata dataSourceMetadata = new DataSourceMetadata(); - dataSourceMetadata.setName("test_datasource"); - dataSourceMetadata.setConnector(DataSourceType.PROMETHEUS); - + DataSourceMetadata dataSourceMetadata = + new DataSourceMetadata.Builder() + .setName("test_datasource") + .setConnector(DataSourceType.PROMETHEUS) + .build(); GetDataSourceActionRequest request = new GetDataSourceActionRequest(); when(dataSourceService.getDataSourceMetadata(false)) .thenReturn(Collections.singleton(dataSourceMetadata)); diff --git a/datasources/src/test/java/org/opensearch/sql/datasources/transport/TransportUpdateDataSourceActionTest.java b/datasources/src/test/java/org/opensearch/sql/datasources/transport/TransportUpdateDataSourceActionTest.java index 4d42cdb2fa..e086813938 100644 --- a/datasources/src/test/java/org/opensearch/sql/datasources/transport/TransportUpdateDataSourceActionTest.java +++ b/datasources/src/test/java/org/opensearch/sql/datasources/transport/TransportUpdateDataSourceActionTest.java @@ -62,9 +62,11 @@ public void setUp() { @Test public void testDoExecute() { - DataSourceMetadata dataSourceMetadata = new DataSourceMetadata(); - dataSourceMetadata.setName("test_datasource"); - dataSourceMetadata.setConnector(DataSourceType.PROMETHEUS); + DataSourceMetadata dataSourceMetadata = + new DataSourceMetadata.Builder() + .setName("test_datasource") + .setConnector(DataSourceType.PROMETHEUS) + .build(); UpdateDataSourceActionRequest request = new UpdateDataSourceActionRequest(dataSourceMetadata); action.doExecute(task, request, actionListener); @@ -80,9 +82,12 @@ public void testDoExecute() { @Test public void testDoExecuteWithException() { - DataSourceMetadata dataSourceMetadata = new DataSourceMetadata(); - dataSourceMetadata.setName("test_datasource"); - dataSourceMetadata.setConnector(DataSourceType.PROMETHEUS); + DataSourceMetadata dataSourceMetadata = + new DataSourceMetadata.Builder() + .setName("test_datasource") + .setConnector(DataSourceType.PROMETHEUS) + .build(); + doThrow(new RuntimeException("Error")) .when(dataSourceService) .updateDataSource(dataSourceMetadata); diff --git a/datasources/src/test/java/org/opensearch/sql/datasources/utils/XContentParserUtilsTest.java b/datasources/src/test/java/org/opensearch/sql/datasources/utils/XContentParserUtilsTest.java index 5a1f5e155f..a739169978 100644 --- a/datasources/src/test/java/org/opensearch/sql/datasources/utils/XContentParserUtilsTest.java +++ b/datasources/src/test/java/org/opensearch/sql/datasources/utils/XContentParserUtilsTest.java @@ -23,11 +23,13 @@ public class XContentParserUtilsTest { @SneakyThrows @Test public void testConvertToXContent() { - DataSourceMetadata dataSourceMetadata = new DataSourceMetadata(); - dataSourceMetadata.setName("testDS"); - dataSourceMetadata.setConnector(DataSourceType.PROMETHEUS); - dataSourceMetadata.setAllowedRoles(List.of("prometheus_access")); - dataSourceMetadata.setProperties(Map.of("prometheus.uri", "https://localhost:9090")); + DataSourceMetadata dataSourceMetadata = + new DataSourceMetadata.Builder() + .setName("testDS") + .setConnector(DataSourceType.PROMETHEUS) + .setAllowedRoles(List.of("prometheus_access")) + .setProperties(Map.of("prometheus.uri", "https://localhost:9090")) + .build(); XContentBuilder contentBuilder = XContentParserUtils.convertToXContent(dataSourceMetadata); String contentString = BytesReference.bytes(contentBuilder).utf8ToString(); @@ -39,12 +41,14 @@ public void testConvertToXContent() { @SneakyThrows @Test public void testToDataSourceMetadataFromJson() { - DataSourceMetadata dataSourceMetadata = new DataSourceMetadata(); - dataSourceMetadata.setName("testDS"); - dataSourceMetadata.setConnector(DataSourceType.PROMETHEUS); - dataSourceMetadata.setAllowedRoles(List.of("prometheus_access")); - dataSourceMetadata.setProperties(Map.of("prometheus.uri", "https://localhost:9090")); - dataSourceMetadata.setResultIndex("query_execution_result2"); + DataSourceMetadata dataSourceMetadata = + new DataSourceMetadata.Builder() + .setName("testDS") + .setConnector(DataSourceType.PROMETHEUS) + .setAllowedRoles(List.of("prometheus_access")) + .setProperties(Map.of("prometheus.uri", "https://localhost:9090")) + .setResultIndex("query_execution_result2") + .build(); Gson gson = new Gson(); String json = gson.toJson(dataSourceMetadata); @@ -97,10 +101,12 @@ public void testToMapFromJson() { @SneakyThrows @Test public void testToDataSourceMetadataFromJsonWithoutName() { - DataSourceMetadata dataSourceMetadata = new DataSourceMetadata(); - dataSourceMetadata.setConnector(DataSourceType.PROMETHEUS); - dataSourceMetadata.setAllowedRoles(List.of("prometheus_access")); - dataSourceMetadata.setProperties(Map.of("prometheus.uri", "https://localhost:9090")); + DataSourceMetadata dataSourceMetadata = + new DataSourceMetadata.Builder() + .setConnector(DataSourceType.PROMETHEUS) + .setAllowedRoles(List.of("prometheus_access")) + .setProperties(Map.of("prometheus.uri", "https://localhost:9090")) + .build(); Gson gson = new Gson(); String json = gson.toJson(dataSourceMetadata); @@ -132,10 +138,12 @@ public void testToMapFromJsonWithoutName() { @SneakyThrows @Test public void testToDataSourceMetadataFromJsonWithoutConnector() { - DataSourceMetadata dataSourceMetadata = new DataSourceMetadata(); - dataSourceMetadata.setName("name"); - dataSourceMetadata.setAllowedRoles(List.of("prometheus_access")); - dataSourceMetadata.setProperties(Map.of("prometheus.uri", "https://localhost:9090")); + DataSourceMetadata dataSourceMetadata = + new DataSourceMetadata.Builder() + .setName("name") + .setAllowedRoles(List.of("prometheus_access")) + .setProperties(Map.of("prometheus.uri", "https://localhost:9090")) + .build(); Gson gson = new Gson(); String json = gson.toJson(dataSourceMetadata); diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/SQLPlugin.java b/plugin/src/main/java/org/opensearch/sql/plugin/SQLPlugin.java index 2b75a8b2c9..ff56ef8634 100644 --- a/plugin/src/main/java/org/opensearch/sql/plugin/SQLPlugin.java +++ b/plugin/src/main/java/org/opensearch/sql/plugin/SQLPlugin.java @@ -194,6 +194,7 @@ public Collection createComponents( NamedWriteableRegistry namedWriteableRegistry, IndexNameExpressionResolver indexNameResolver, Supplier repositoriesServiceSupplier) { + this.clusterService.getClusterManagerService(); this.clusterService = clusterService; this.pluginSettings = new OpenSearchSettings(clusterService.getClusterSettings()); this.client = (NodeClient) client; diff --git a/spark/src/main/antlr/FlintSparkSqlExtensions.g4 b/spark/src/main/antlr/FlintSparkSqlExtensions.g4 index 4de5bfaa66..219bbe782b 100644 --- a/spark/src/main/antlr/FlintSparkSqlExtensions.g4 +++ b/spark/src/main/antlr/FlintSparkSqlExtensions.g4 @@ -18,6 +18,7 @@ statement : skippingIndexStatement | coveringIndexStatement | materializedViewStatement + | indexManagementStatement | indexJobManagementStatement ; @@ -125,6 +126,14 @@ vacuumMaterializedViewStatement : VACUUM MATERIALIZED VIEW mvName=multipartIdentifier ; +indexManagementStatement + : showFlintIndexStatement + ; + +showFlintIndexStatement + : SHOW FLINT (INDEX | INDEXES) IN catalogDb=multipartIdentifier + ; + indexJobManagementStatement : recoverIndexJobStatement ; diff --git a/spark/src/main/antlr/SparkSqlBase.g4 b/spark/src/main/antlr/SparkSqlBase.g4 index 82c890a618..01f45016d6 100644 --- a/spark/src/main/antlr/SparkSqlBase.g4 +++ b/spark/src/main/antlr/SparkSqlBase.g4 @@ -161,6 +161,7 @@ DESCRIBE: 'DESCRIBE'; DROP: 'DROP'; EXISTS: 'EXISTS'; FALSE: 'FALSE'; +FLINT: 'FLINT'; IF: 'IF'; IN: 'IN'; INDEX: 'INDEX'; diff --git a/spark/src/main/antlr/SqlBaseParser.g4 b/spark/src/main/antlr/SqlBaseParser.g4 index 737d5196e7..ca01de4ffd 100644 --- a/spark/src/main/antlr/SqlBaseParser.g4 +++ b/spark/src/main/antlr/SqlBaseParser.g4 @@ -989,6 +989,7 @@ primaryExpression | CASE whenClause+ (ELSE elseExpression=expression)? END #searchedCase | CASE value=expression whenClause+ (ELSE elseExpression=expression)? END #simpleCase | name=(CAST | TRY_CAST) LEFT_PAREN expression AS dataType RIGHT_PAREN #cast + | primaryExpression collateClause #collate | primaryExpression DOUBLE_COLON dataType #castByColon | STRUCT LEFT_PAREN (argument+=namedExpression (COMMA argument+=namedExpression)*)? RIGHT_PAREN #struct | FIRST LEFT_PAREN expression (IGNORE NULLS)? RIGHT_PAREN #first @@ -1094,6 +1095,10 @@ colPosition : position=FIRST | position=AFTER afterCol=errorCapturingIdentifier ; +collateClause + : COLLATE collationName=stringLit + ; + type : BOOLEAN | TINYINT | BYTE @@ -1104,7 +1109,7 @@ type | DOUBLE | DATE | TIMESTAMP | TIMESTAMP_NTZ | TIMESTAMP_LTZ - | STRING + | STRING collateClause? | CHARACTER | CHAR | VARCHAR | BINARY diff --git a/spark/src/main/java/org/opensearch/sql/spark/dispatcher/SparkQueryDispatcher.java b/spark/src/main/java/org/opensearch/sql/spark/dispatcher/SparkQueryDispatcher.java index 498a3b9af5..29e183a45c 100644 --- a/spark/src/main/java/org/opensearch/sql/spark/dispatcher/SparkQueryDispatcher.java +++ b/spark/src/main/java/org/opensearch/sql/spark/dispatcher/SparkQueryDispatcher.java @@ -61,7 +61,7 @@ public DispatchQueryResponse dispatch(DispatchQueryRequest dispatchQueryRequest) EMRServerlessClient emrServerlessClient = emrServerlessClientFactory.getClient(); DataSourceMetadata dataSourceMetadata = this.dataSourceService.getRawDataSourceMetadata(dispatchQueryRequest.getDatasource()); - dataSourceUserAuthorizationHelper.authorizeDataSource(dataSourceMetadata); + this.dataSourceService.verifyDataSourceAccess(dataSourceMetadata); AsyncQueryHandler asyncQueryHandler = sessionManager.isEnabled() ? new InteractiveQueryHandler(sessionManager, jobExecutionResponseReader, leaseManager) diff --git a/spark/src/test/java/org/opensearch/sql/spark/asyncquery/AsyncQueryExecutorServiceImplSpecTest.java b/spark/src/test/java/org/opensearch/sql/spark/asyncquery/AsyncQueryExecutorServiceImplSpecTest.java index 33fec89e26..446c22e795 100644 --- a/spark/src/test/java/org/opensearch/sql/spark/asyncquery/AsyncQueryExecutorServiceImplSpecTest.java +++ b/spark/src/test/java/org/opensearch/sql/spark/asyncquery/AsyncQueryExecutorServiceImplSpecTest.java @@ -16,12 +16,10 @@ import static org.opensearch.sql.spark.execution.statestore.StateStore.getStatement; import static org.opensearch.sql.spark.execution.statestore.StateStore.updateStatementState; -import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import java.util.HashMap; import java.util.Map; import java.util.Optional; -import org.apache.commons.lang3.StringUtils; import org.junit.Ignore; import org.junit.Test; import org.junit.jupiter.api.Disabled; @@ -255,13 +253,11 @@ public void datasourceWithBasicAuth() { properties.put("glue.indexstore.opensearch.auth.password", "password"); dataSourceService.createDataSource( - new DataSourceMetadata( - "mybasicauth", - StringUtils.EMPTY, - DataSourceType.S3GLUE, - ImmutableList.of(), - properties, - null)); + new DataSourceMetadata.Builder() + .setName("mybasicauth") + .setConnector(DataSourceType.S3GLUE) + .setProperties(properties) + .build()); LocalEMRSClient emrsClient = new LocalEMRSClient(); EMRServerlessClientFactory emrServerlessClientFactory = () -> emrsClient; AsyncQueryExecutorService asyncQueryExecutorService = @@ -514,21 +510,20 @@ public void submitQueryInInvalidSessionWillCreateNewSession() { @Test public void datasourceNameIncludeUppercase() { dataSourceService.createDataSource( - new DataSourceMetadata( - "TESTS3", - StringUtils.EMPTY, - DataSourceType.S3GLUE, - ImmutableList.of(), - ImmutableMap.of( - "glue.auth.type", - "iam_role", - "glue.auth.role_arn", - "arn:aws:iam::924196221507:role/FlintOpensearchServiceRole", - "glue.indexstore.opensearch.uri", - "http://localhost:9200", - "glue.indexstore.opensearch.auth", - "noauth"), - null)); + new DataSourceMetadata.Builder() + .setName("TESTS3") + .setConnector(DataSourceType.S3GLUE) + .setProperties( + ImmutableMap.of( + "glue.auth.type", + "iam_role", + "glue.auth.role_arn", + "arn:aws:iam::924196221507:role/FlintOpensearchServiceRole", + "glue.indexstore.opensearch.uri", + "http://localhost:9200", + "glue.indexstore.opensearch.auth", + "noauth")) + .build()); LocalEMRSClient emrsClient = new LocalEMRSClient(); EMRServerlessClientFactory emrServerlessClientFactory = () -> emrsClient; diff --git a/spark/src/test/java/org/opensearch/sql/spark/asyncquery/AsyncQueryExecutorServiceSpec.java b/spark/src/test/java/org/opensearch/sql/spark/asyncquery/AsyncQueryExecutorServiceSpec.java index c9b4b6fc88..e176a2b828 100644 --- a/spark/src/test/java/org/opensearch/sql/spark/asyncquery/AsyncQueryExecutorServiceSpec.java +++ b/spark/src/test/java/org/opensearch/sql/spark/asyncquery/AsyncQueryExecutorServiceSpec.java @@ -17,7 +17,6 @@ import com.amazonaws.services.emrserverless.model.JobRun; import com.amazonaws.services.emrserverless.model.JobRunState; import com.google.common.base.Charsets; -import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; import com.google.common.io.Resources; @@ -30,7 +29,6 @@ import lombok.Getter; import lombok.RequiredArgsConstructor; import lombok.SneakyThrows; -import org.apache.commons.lang3.StringUtils; import org.junit.After; import org.junit.Before; import org.opensearch.action.admin.indices.create.CreateIndexRequest; @@ -119,38 +117,36 @@ public void setup() { .get(); dataSourceService = createDataSourceService(); DataSourceMetadata dm = - new DataSourceMetadata( - DATASOURCE, - StringUtils.EMPTY, - DataSourceType.S3GLUE, - ImmutableList.of(), - ImmutableMap.of( - "glue.auth.type", - "iam_role", - "glue.auth.role_arn", - "arn:aws:iam::924196221507:role/FlintOpensearchServiceRole", - "glue.indexstore.opensearch.uri", - "http://localhost:9200", - "glue.indexstore.opensearch.auth", - "noauth"), - null); + new DataSourceMetadata.Builder() + .setName(DATASOURCE) + .setConnector(DataSourceType.S3GLUE) + .setProperties( + ImmutableMap.of( + "glue.auth.type", + "iam_role", + "glue.auth.role_arn", + "arn:aws:iam::924196221507:role/FlintOpensearchServiceRole", + "glue.indexstore.opensearch.uri", + "http://localhost:9200", + "glue.indexstore.opensearch.auth", + "noauth")) + .build(); dataSourceService.createDataSource(dm); DataSourceMetadata otherDm = - new DataSourceMetadata( - DSOTHER, - StringUtils.EMPTY, - DataSourceType.S3GLUE, - ImmutableList.of(), - ImmutableMap.of( - "glue.auth.type", - "iam_role", - "glue.auth.role_arn", - "arn:aws:iam::924196221507:role/FlintOpensearchServiceRole", - "glue.indexstore.opensearch.uri", - "http://localhost:9200", - "glue.indexstore.opensearch.auth", - "noauth"), - null); + new DataSourceMetadata.Builder() + .setName(DSOTHER) + .setConnector(DataSourceType.S3GLUE) + .setProperties( + ImmutableMap.of( + "glue.auth.type", + "iam_role", + "glue.auth.role_arn", + "arn:aws:iam::924196221507:role/FlintOpensearchServiceRole", + "glue.indexstore.opensearch.uri", + "http://localhost:9200", + "glue.indexstore.opensearch.auth", + "noauth")) + .build(); dataSourceService.createDataSource(otherDm); stateStore = new StateStore(client, clusterService); createIndexWithMappings(dm.getResultIndex(), loadResultIndexMappings());