diff --git a/spark/src/main/antlr/FlintSparkSqlExtensions.g4 b/spark/src/main/antlr/FlintSparkSqlExtensions.g4 index f48c276e44..c4af2779d1 100644 --- a/spark/src/main/antlr/FlintSparkSqlExtensions.g4 +++ b/spark/src/main/antlr/FlintSparkSqlExtensions.g4 @@ -79,7 +79,6 @@ dropCoveringIndexStatement materializedViewStatement : createMaterializedViewStatement - | refreshMaterializedViewStatement | showMaterializedViewStatement | describeMaterializedViewStatement | dropMaterializedViewStatement @@ -91,10 +90,6 @@ createMaterializedViewStatement (WITH LEFT_PAREN propertyList RIGHT_PAREN)? ; -refreshMaterializedViewStatement - : REFRESH MATERIALIZED VIEW mvName=multipartIdentifier - ; - showMaterializedViewStatement : SHOW MATERIALIZED (VIEW | VIEWS) IN catalogDb=multipartIdentifier ; 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 6ec67709b8..ff7ccf8c08 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 @@ -36,7 +36,10 @@ import org.opensearch.sql.spark.asyncquery.model.SparkSubmitParameters; import org.opensearch.sql.spark.client.EMRServerlessClient; import org.opensearch.sql.spark.client.StartJobRequest; -import org.opensearch.sql.spark.dispatcher.model.*; +import org.opensearch.sql.spark.dispatcher.model.DispatchQueryRequest; +import org.opensearch.sql.spark.dispatcher.model.DispatchQueryResponse; +import org.opensearch.sql.spark.dispatcher.model.IndexDetails; +import org.opensearch.sql.spark.dispatcher.model.JobType; import org.opensearch.sql.spark.execution.session.CreateSessionRequest; import org.opensearch.sql.spark.execution.session.Session; import org.opensearch.sql.spark.execution.session.SessionId; @@ -53,10 +56,11 @@ public class SparkQueryDispatcher { private static final Logger LOG = LogManager.getLogger(); + public static final String INDEX_TAG_KEY = "index"; public static final String DATASOURCE_TAG_KEY = "datasource"; public static final String CLUSTER_NAME_TAG_KEY = "cluster"; - public static final String JOB_TYPE_TAG_KEY = "type"; + public static final String JOB_TYPE_TAG_KEY = "job_type"; private EMRServerlessClient emrServerlessClient; @@ -103,18 +107,15 @@ public String cancelJob(AsyncQueryJobMetadata asyncQueryJobMetadata) { } private DispatchQueryResponse handleSQLQuery(DispatchQueryRequest dispatchQueryRequest) { - if (SQLQueryUtils.isFlintExtensionQuery(dispatchQueryRequest.getQuery())) { - IndexQueryDetails indexQueryDetails = + if (SQLQueryUtils.isIndexQuery(dispatchQueryRequest.getQuery())) { + IndexDetails indexDetails = SQLQueryUtils.extractIndexDetails(dispatchQueryRequest.getQuery()); - fillMissingDetails(dispatchQueryRequest, indexQueryDetails); + fillMissingDetails(dispatchQueryRequest, indexDetails); - // TODO: refactor this code properly. - if (IndexQueryActionType.DROP.equals(indexQueryDetails.getIndexQueryActionType())) { - return handleDropIndexQuery(dispatchQueryRequest, indexQueryDetails); - } else if (IndexQueryActionType.CREATE.equals(indexQueryDetails.getIndexQueryActionType())) { - return handleStreamingQueries(dispatchQueryRequest, indexQueryDetails); + if (indexDetails.isDropIndex()) { + return handleDropIndexQuery(dispatchQueryRequest, indexDetails); } else { - return handleFlintNonStreamingQueries(dispatchQueryRequest, indexQueryDetails); + return handleIndexQuery(dispatchQueryRequest, indexDetails); } } else { return handleNonIndexQuery(dispatchQueryRequest); @@ -126,59 +127,24 @@ private DispatchQueryResponse handleSQLQuery(DispatchQueryRequest dispatchQueryR // Spark Assumes the datasource to be catalog. // This is required to handle drop index case properly when datasource name is not provided. private static void fillMissingDetails( - DispatchQueryRequest dispatchQueryRequest, IndexQueryDetails indexQueryDetails) { - if (indexQueryDetails.getFullyQualifiedTableName() != null - && indexQueryDetails.getFullyQualifiedTableName().getDatasourceName() == null) { - indexQueryDetails + DispatchQueryRequest dispatchQueryRequest, IndexDetails indexDetails) { + if (indexDetails.getFullyQualifiedTableName() != null + && indexDetails.getFullyQualifiedTableName().getDatasourceName() == null) { + indexDetails .getFullyQualifiedTableName() .setDatasourceName(dispatchQueryRequest.getDatasource()); } } - private DispatchQueryResponse handleStreamingQueries( - DispatchQueryRequest dispatchQueryRequest, IndexQueryDetails indexQueryDetails) { - DataSourceMetadata dataSourceMetadata = - this.dataSourceService.getRawDataSourceMetadata(dispatchQueryRequest.getDatasource()); - dataSourceUserAuthorizationHelper.authorizeDataSource(dataSourceMetadata); - String jobName = dispatchQueryRequest.getClusterName() + ":" + "index-query"; - Map tags = getDefaultTagsForJobSubmission(dispatchQueryRequest); - tags.put(INDEX_TAG_KEY, indexQueryDetails.openSearchIndexName()); - if (indexQueryDetails.isAutoRefresh()) { - tags.put(JOB_TYPE_TAG_KEY, JobType.STREAMING.getText()); - } - StartJobRequest startJobRequest = - new StartJobRequest( - dispatchQueryRequest.getQuery(), - jobName, - dispatchQueryRequest.getApplicationId(), - dispatchQueryRequest.getExecutionRoleARN(), - SparkSubmitParameters.Builder.builder() - .dataSource( - dataSourceService.getRawDataSourceMetadata( - dispatchQueryRequest.getDatasource())) - .structuredStreaming(indexQueryDetails.isAutoRefresh()) - .extraParameters(dispatchQueryRequest.getExtraSparkSubmitParams()) - .build() - .toString(), - tags, - indexQueryDetails.isAutoRefresh(), - dataSourceMetadata.getResultIndex()); - String jobId = emrServerlessClient.startJobRun(startJobRequest); - return new DispatchQueryResponse( - AsyncQueryId.newAsyncQueryId(dataSourceMetadata.getName()), - jobId, - false, - dataSourceMetadata.getResultIndex(), - null); - } - - private DispatchQueryResponse handleFlintNonStreamingQueries( - DispatchQueryRequest dispatchQueryRequest, IndexQueryDetails indexQueryDetails) { + private DispatchQueryResponse handleIndexQuery( + DispatchQueryRequest dispatchQueryRequest, IndexDetails indexDetails) { DataSourceMetadata dataSourceMetadata = this.dataSourceService.getRawDataSourceMetadata(dispatchQueryRequest.getDatasource()); dataSourceUserAuthorizationHelper.authorizeDataSource(dataSourceMetadata); String jobName = dispatchQueryRequest.getClusterName() + ":" + "index-query"; Map tags = getDefaultTagsForJobSubmission(dispatchQueryRequest); + tags.put(INDEX_TAG_KEY, indexDetails.openSearchIndexName()); + tags.put(JOB_TYPE_TAG_KEY, JobType.STREAMING.getText()); StartJobRequest startJobRequest = new StartJobRequest( dispatchQueryRequest.getQuery(), @@ -189,11 +155,12 @@ private DispatchQueryResponse handleFlintNonStreamingQueries( .dataSource( dataSourceService.getRawDataSourceMetadata( dispatchQueryRequest.getDatasource())) + .structuredStreaming(indexDetails.isAutoRefresh()) .extraParameters(dispatchQueryRequest.getExtraSparkSubmitParams()) .build() .toString(), tags, - indexQueryDetails.isAutoRefresh(), + indexDetails.isAutoRefresh(), dataSourceMetadata.getResultIndex()); String jobId = emrServerlessClient.startJobRun(startJobRequest); return new DispatchQueryResponse( @@ -275,12 +242,11 @@ private DispatchQueryResponse handleNonIndexQuery(DispatchQueryRequest dispatchQ } private DispatchQueryResponse handleDropIndexQuery( - DispatchQueryRequest dispatchQueryRequest, IndexQueryDetails indexQueryDetails) { + DispatchQueryRequest dispatchQueryRequest, IndexDetails indexDetails) { DataSourceMetadata dataSourceMetadata = this.dataSourceService.getRawDataSourceMetadata(dispatchQueryRequest.getDatasource()); dataSourceUserAuthorizationHelper.authorizeDataSource(dataSourceMetadata); - FlintIndexMetadata indexMetadata = - flintIndexMetadataReader.getFlintIndexMetadata(indexQueryDetails); + FlintIndexMetadata indexMetadata = flintIndexMetadataReader.getFlintIndexMetadata(indexDetails); // if index is created without auto refresh. there is no job to cancel. String status = JobRunState.FAILED.toString(); try { @@ -289,7 +255,7 @@ private DispatchQueryResponse handleDropIndexQuery( dispatchQueryRequest.getApplicationId(), indexMetadata.getJobId()); } } finally { - String indexName = indexQueryDetails.openSearchIndexName(); + String indexName = indexDetails.openSearchIndexName(); try { AcknowledgedResponse response = client.admin().indices().delete(new DeleteIndexRequest().indices(indexName)).get(); diff --git a/spark/src/main/java/org/opensearch/sql/spark/dispatcher/model/IndexQueryDetails.java b/spark/src/main/java/org/opensearch/sql/spark/dispatcher/model/IndexDetails.java similarity index 55% rename from spark/src/main/java/org/opensearch/sql/spark/dispatcher/model/IndexQueryDetails.java rename to spark/src/main/java/org/opensearch/sql/spark/dispatcher/model/IndexDetails.java index 5b4326a10e..42e2905e67 100644 --- a/spark/src/main/java/org/opensearch/sql/spark/dispatcher/model/IndexQueryDetails.java +++ b/spark/src/main/java/org/opensearch/sql/spark/dispatcher/model/IndexDetails.java @@ -5,6 +5,7 @@ package org.opensearch.sql.spark.dispatcher.model; +import com.google.common.base.Preconditions; import lombok.EqualsAndHashCode; import lombok.Getter; import org.apache.commons.lang3.StringUtils; @@ -13,7 +14,7 @@ /** Index details in an async query. */ @Getter @EqualsAndHashCode -public class IndexQueryDetails { +public class IndexDetails { public static final String STRIP_CHARS = "`"; @@ -21,59 +22,75 @@ public class IndexQueryDetails { private FullyQualifiedTableName fullyQualifiedTableName; // by default, auto_refresh = false; private boolean autoRefresh; - private IndexQueryActionType indexQueryActionType; + private boolean isDropIndex; // materialized view special case where // table name and mv name are combined. private String mvName; private FlintIndexType indexType; - private IndexQueryDetails() {} + private IndexDetails() {} - public static IndexQueryDetailsBuilder builder() { - return new IndexQueryDetailsBuilder(); + public static IndexDetailsBuilder builder() { + return new IndexDetailsBuilder(); } // Builder class - public static class IndexQueryDetailsBuilder { - private final IndexQueryDetails indexQueryDetails; + public static class IndexDetailsBuilder { + private final IndexDetails indexDetails; - public IndexQueryDetailsBuilder() { - indexQueryDetails = new IndexQueryDetails(); + public IndexDetailsBuilder() { + indexDetails = new IndexDetails(); } - public IndexQueryDetailsBuilder indexName(String indexName) { - indexQueryDetails.indexName = indexName; + public IndexDetailsBuilder indexName(String indexName) { + indexDetails.indexName = indexName; return this; } - public IndexQueryDetailsBuilder fullyQualifiedTableName(FullyQualifiedTableName tableName) { - indexQueryDetails.fullyQualifiedTableName = tableName; + public IndexDetailsBuilder fullyQualifiedTableName(FullyQualifiedTableName tableName) { + indexDetails.fullyQualifiedTableName = tableName; return this; } - public IndexQueryDetailsBuilder autoRefresh(Boolean autoRefresh) { - indexQueryDetails.autoRefresh = autoRefresh; + public IndexDetailsBuilder autoRefresh(Boolean autoRefresh) { + indexDetails.autoRefresh = autoRefresh; return this; } - public IndexQueryDetailsBuilder indexQueryActionType( - IndexQueryActionType indexQueryActionType) { - indexQueryDetails.indexQueryActionType = indexQueryActionType; + public IndexDetailsBuilder isDropIndex(boolean isDropIndex) { + indexDetails.isDropIndex = isDropIndex; return this; } - public IndexQueryDetailsBuilder mvName(String mvName) { - indexQueryDetails.mvName = mvName; + public IndexDetailsBuilder mvName(String mvName) { + indexDetails.mvName = mvName; return this; } - public IndexQueryDetailsBuilder indexType(FlintIndexType indexType) { - indexQueryDetails.indexType = indexType; + public IndexDetailsBuilder indexType(FlintIndexType indexType) { + indexDetails.indexType = indexType; return this; } - public IndexQueryDetails build() { - return indexQueryDetails; + public IndexDetails build() { + Preconditions.checkNotNull(indexDetails.indexType, "Index Type can't be null"); + switch (indexDetails.indexType) { + case COVERING: + Preconditions.checkNotNull( + indexDetails.indexName, "IndexName can't be null for Covering Index."); + Preconditions.checkNotNull( + indexDetails.fullyQualifiedTableName, "TableName can't be null for Covering Index."); + break; + case SKIPPING: + Preconditions.checkNotNull( + indexDetails.fullyQualifiedTableName, "TableName can't be null for Skipping Index."); + break; + case MATERIALIZED_VIEW: + Preconditions.checkNotNull(indexDetails.mvName, "Materialized view name can't be null"); + break; + } + + return indexDetails; } } diff --git a/spark/src/main/java/org/opensearch/sql/spark/dispatcher/model/IndexQueryActionType.java b/spark/src/main/java/org/opensearch/sql/spark/dispatcher/model/IndexQueryActionType.java deleted file mode 100644 index 2c96511d2a..0000000000 --- a/spark/src/main/java/org/opensearch/sql/spark/dispatcher/model/IndexQueryActionType.java +++ /dev/null @@ -1,15 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.sql.spark.dispatcher.model; - -/** Enum for Index Action in the given query.* */ -public enum IndexQueryActionType { - CREATE, - REFRESH, - DESCRIBE, - SHOW, - DROP -} diff --git a/spark/src/main/java/org/opensearch/sql/spark/execution/session/InteractiveSession.java b/spark/src/main/java/org/opensearch/sql/spark/execution/session/InteractiveSession.java index 956275b04a..a2e7cfe6ee 100644 --- a/spark/src/main/java/org/opensearch/sql/spark/execution/session/InteractiveSession.java +++ b/spark/src/main/java/org/opensearch/sql/spark/execution/session/InteractiveSession.java @@ -34,8 +34,6 @@ public class InteractiveSession implements Session { private static final Logger LOG = LogManager.getLogger(); - public static final String SESSION_ID_TAG_KEY = "sid"; - private final SessionId sessionId; private final StateStore stateStore; private final EMRServerlessClient serverlessClient; @@ -48,7 +46,6 @@ public void open(CreateSessionRequest createSessionRequest) { createSessionRequest .getSparkSubmitParametersBuilder() .sessionExecution(sessionId.getSessionId(), createSessionRequest.getDatasourceName()); - createSessionRequest.getTags().put(SESSION_ID_TAG_KEY, sessionId.getSessionId()); String jobID = serverlessClient.startJobRun(createSessionRequest.getStartJobRequest()); String applicationId = createSessionRequest.getStartJobRequest().getApplicationId(); diff --git a/spark/src/main/java/org/opensearch/sql/spark/flint/FlintIndexMetadataReader.java b/spark/src/main/java/org/opensearch/sql/spark/flint/FlintIndexMetadataReader.java index d4a8e7ddbf..e4a5e92035 100644 --- a/spark/src/main/java/org/opensearch/sql/spark/flint/FlintIndexMetadataReader.java +++ b/spark/src/main/java/org/opensearch/sql/spark/flint/FlintIndexMetadataReader.java @@ -1,6 +1,6 @@ package org.opensearch.sql.spark.flint; -import org.opensearch.sql.spark.dispatcher.model.IndexQueryDetails; +import org.opensearch.sql.spark.dispatcher.model.IndexDetails; /** Interface for FlintIndexMetadataReader */ public interface FlintIndexMetadataReader { @@ -8,8 +8,8 @@ public interface FlintIndexMetadataReader { /** * Given Index details, get the streaming job Id. * - * @param indexQueryDetails indexDetails. + * @param indexDetails indexDetails. * @return FlintIndexMetadata. */ - FlintIndexMetadata getFlintIndexMetadata(IndexQueryDetails indexQueryDetails); + FlintIndexMetadata getFlintIndexMetadata(IndexDetails indexDetails); } diff --git a/spark/src/main/java/org/opensearch/sql/spark/flint/FlintIndexMetadataReaderImpl.java b/spark/src/main/java/org/opensearch/sql/spark/flint/FlintIndexMetadataReaderImpl.java index a16d0b9138..5f712e65cd 100644 --- a/spark/src/main/java/org/opensearch/sql/spark/flint/FlintIndexMetadataReaderImpl.java +++ b/spark/src/main/java/org/opensearch/sql/spark/flint/FlintIndexMetadataReaderImpl.java @@ -5,7 +5,7 @@ import org.opensearch.action.admin.indices.mapping.get.GetMappingsResponse; import org.opensearch.client.Client; import org.opensearch.cluster.metadata.MappingMetadata; -import org.opensearch.sql.spark.dispatcher.model.IndexQueryDetails; +import org.opensearch.sql.spark.dispatcher.model.IndexDetails; /** Implementation of {@link FlintIndexMetadataReader} */ @AllArgsConstructor @@ -14,8 +14,8 @@ public class FlintIndexMetadataReaderImpl implements FlintIndexMetadataReader { private final Client client; @Override - public FlintIndexMetadata getFlintIndexMetadata(IndexQueryDetails indexQueryDetails) { - String indexName = indexQueryDetails.openSearchIndexName(); + public FlintIndexMetadata getFlintIndexMetadata(IndexDetails indexDetails) { + String indexName = indexDetails.openSearchIndexName(); GetMappingsResponse mappingsResponse = client.admin().indices().prepareGetMappings(indexName).get(); try { diff --git a/spark/src/main/java/org/opensearch/sql/spark/utils/SQLQueryUtils.java b/spark/src/main/java/org/opensearch/sql/spark/utils/SQLQueryUtils.java index c1f3f02576..4816f1c2cd 100644 --- a/spark/src/main/java/org/opensearch/sql/spark/utils/SQLQueryUtils.java +++ b/spark/src/main/java/org/opensearch/sql/spark/utils/SQLQueryUtils.java @@ -20,8 +20,7 @@ import org.opensearch.sql.spark.antlr.parser.SqlBaseParser; import org.opensearch.sql.spark.antlr.parser.SqlBaseParserBaseVisitor; import org.opensearch.sql.spark.dispatcher.model.FullyQualifiedTableName; -import org.opensearch.sql.spark.dispatcher.model.IndexQueryActionType; -import org.opensearch.sql.spark.dispatcher.model.IndexQueryDetails; +import org.opensearch.sql.spark.dispatcher.model.IndexDetails; import org.opensearch.sql.spark.flint.FlintIndexType; /** @@ -43,7 +42,7 @@ public static FullyQualifiedTableName extractFullyQualifiedTableName(String sqlQ return sparkSqlTableNameVisitor.getFullyQualifiedTableName(); } - public static IndexQueryDetails extractIndexDetails(String sqlQuery) { + public static IndexDetails extractIndexDetails(String sqlQuery) { FlintSparkSqlExtensionsParser flintSparkSqlExtensionsParser = new FlintSparkSqlExtensionsParser( new CommonTokenStream( @@ -53,10 +52,10 @@ public static IndexQueryDetails extractIndexDetails(String sqlQuery) { flintSparkSqlExtensionsParser.statement(); FlintSQLIndexDetailsVisitor flintSQLIndexDetailsVisitor = new FlintSQLIndexDetailsVisitor(); statementContext.accept(flintSQLIndexDetailsVisitor); - return flintSQLIndexDetailsVisitor.getIndexQueryDetailsBuilder().build(); + return flintSQLIndexDetailsVisitor.getIndexDetailsBuilder().build(); } - public static boolean isFlintExtensionQuery(String sqlQuery) { + public static boolean isIndexQuery(String sqlQuery) { FlintSparkSqlExtensionsParser flintSparkSqlExtensionsParser = new FlintSparkSqlExtensionsParser( new CommonTokenStream( @@ -118,29 +117,29 @@ public Void visitCreateTableHeader(SqlBaseParser.CreateTableHeaderContext ctx) { public static class FlintSQLIndexDetailsVisitor extends FlintSparkSqlExtensionsBaseVisitor { - @Getter private final IndexQueryDetails.IndexQueryDetailsBuilder indexQueryDetailsBuilder; + @Getter private final IndexDetails.IndexDetailsBuilder indexDetailsBuilder; public FlintSQLIndexDetailsVisitor() { - this.indexQueryDetailsBuilder = new IndexQueryDetails.IndexQueryDetailsBuilder(); + this.indexDetailsBuilder = new IndexDetails.IndexDetailsBuilder(); } @Override public Void visitIndexName(FlintSparkSqlExtensionsParser.IndexNameContext ctx) { - indexQueryDetailsBuilder.indexName(ctx.getText()); + indexDetailsBuilder.indexName(ctx.getText()); return super.visitIndexName(ctx); } @Override public Void visitTableName(FlintSparkSqlExtensionsParser.TableNameContext ctx) { - indexQueryDetailsBuilder.fullyQualifiedTableName(new FullyQualifiedTableName(ctx.getText())); + indexDetailsBuilder.fullyQualifiedTableName(new FullyQualifiedTableName(ctx.getText())); return super.visitTableName(ctx); } @Override public Void visitCreateSkippingIndexStatement( FlintSparkSqlExtensionsParser.CreateSkippingIndexStatementContext ctx) { - indexQueryDetailsBuilder.indexQueryActionType(IndexQueryActionType.CREATE); - indexQueryDetailsBuilder.indexType(FlintIndexType.SKIPPING); + indexDetailsBuilder.isDropIndex(false); + indexDetailsBuilder.indexType(FlintIndexType.SKIPPING); visitPropertyList(ctx.propertyList()); return super.visitCreateSkippingIndexStatement(ctx); } @@ -148,8 +147,8 @@ public Void visitCreateSkippingIndexStatement( @Override public Void visitCreateCoveringIndexStatement( FlintSparkSqlExtensionsParser.CreateCoveringIndexStatementContext ctx) { - indexQueryDetailsBuilder.indexQueryActionType(IndexQueryActionType.CREATE); - indexQueryDetailsBuilder.indexType(FlintIndexType.COVERING); + indexDetailsBuilder.isDropIndex(false); + indexDetailsBuilder.indexType(FlintIndexType.COVERING); visitPropertyList(ctx.propertyList()); return super.visitCreateCoveringIndexStatement(ctx); } @@ -157,9 +156,9 @@ public Void visitCreateCoveringIndexStatement( @Override public Void visitCreateMaterializedViewStatement( FlintSparkSqlExtensionsParser.CreateMaterializedViewStatementContext ctx) { - indexQueryDetailsBuilder.indexQueryActionType(IndexQueryActionType.CREATE); - indexQueryDetailsBuilder.indexType(FlintIndexType.MATERIALIZED_VIEW); - indexQueryDetailsBuilder.mvName(ctx.mvName.getText()); + indexDetailsBuilder.isDropIndex(false); + indexDetailsBuilder.indexType(FlintIndexType.MATERIALIZED_VIEW); + indexDetailsBuilder.mvName(ctx.mvName.getText()); visitPropertyList(ctx.propertyList()); return super.visitCreateMaterializedViewStatement(ctx); } @@ -167,94 +166,28 @@ public Void visitCreateMaterializedViewStatement( @Override public Void visitDropCoveringIndexStatement( FlintSparkSqlExtensionsParser.DropCoveringIndexStatementContext ctx) { - indexQueryDetailsBuilder.indexQueryActionType(IndexQueryActionType.DROP); - indexQueryDetailsBuilder.indexType(FlintIndexType.COVERING); + indexDetailsBuilder.isDropIndex(true); + indexDetailsBuilder.indexType(FlintIndexType.COVERING); return super.visitDropCoveringIndexStatement(ctx); } @Override public Void visitDropSkippingIndexStatement( FlintSparkSqlExtensionsParser.DropSkippingIndexStatementContext ctx) { - indexQueryDetailsBuilder.indexQueryActionType(IndexQueryActionType.DROP); - indexQueryDetailsBuilder.indexType(FlintIndexType.SKIPPING); + indexDetailsBuilder.isDropIndex(true); + indexDetailsBuilder.indexType(FlintIndexType.SKIPPING); return super.visitDropSkippingIndexStatement(ctx); } @Override public Void visitDropMaterializedViewStatement( FlintSparkSqlExtensionsParser.DropMaterializedViewStatementContext ctx) { - indexQueryDetailsBuilder.indexQueryActionType(IndexQueryActionType.DROP); - indexQueryDetailsBuilder.indexType(FlintIndexType.MATERIALIZED_VIEW); - indexQueryDetailsBuilder.mvName(ctx.mvName.getText()); + indexDetailsBuilder.isDropIndex(true); + indexDetailsBuilder.indexType(FlintIndexType.MATERIALIZED_VIEW); + indexDetailsBuilder.mvName(ctx.mvName.getText()); return super.visitDropMaterializedViewStatement(ctx); } - @Override - public Void visitDescribeCoveringIndexStatement( - FlintSparkSqlExtensionsParser.DescribeCoveringIndexStatementContext ctx) { - indexQueryDetailsBuilder.indexQueryActionType(IndexQueryActionType.DESCRIBE); - indexQueryDetailsBuilder.indexType(FlintIndexType.COVERING); - return super.visitDescribeCoveringIndexStatement(ctx); - } - - @Override - public Void visitDescribeSkippingIndexStatement( - FlintSparkSqlExtensionsParser.DescribeSkippingIndexStatementContext ctx) { - indexQueryDetailsBuilder.indexQueryActionType(IndexQueryActionType.DESCRIBE); - indexQueryDetailsBuilder.indexType(FlintIndexType.SKIPPING); - return super.visitDescribeSkippingIndexStatement(ctx); - } - - @Override - public Void visitDescribeMaterializedViewStatement( - FlintSparkSqlExtensionsParser.DescribeMaterializedViewStatementContext ctx) { - indexQueryDetailsBuilder.indexQueryActionType(IndexQueryActionType.DESCRIBE); - indexQueryDetailsBuilder.indexType(FlintIndexType.MATERIALIZED_VIEW); - indexQueryDetailsBuilder.mvName(ctx.mvName.getText()); - return super.visitDescribeMaterializedViewStatement(ctx); - } - - @Override - public Void visitShowCoveringIndexStatement( - FlintSparkSqlExtensionsParser.ShowCoveringIndexStatementContext ctx) { - indexQueryDetailsBuilder.indexQueryActionType(IndexQueryActionType.SHOW); - indexQueryDetailsBuilder.indexType(FlintIndexType.COVERING); - return super.visitShowCoveringIndexStatement(ctx); - } - - @Override - public Void visitShowMaterializedViewStatement( - FlintSparkSqlExtensionsParser.ShowMaterializedViewStatementContext ctx) { - indexQueryDetailsBuilder.indexQueryActionType(IndexQueryActionType.SHOW); - indexQueryDetailsBuilder.indexType(FlintIndexType.MATERIALIZED_VIEW); - return super.visitShowMaterializedViewStatement(ctx); - } - - @Override - public Void visitRefreshCoveringIndexStatement( - FlintSparkSqlExtensionsParser.RefreshCoveringIndexStatementContext ctx) { - indexQueryDetailsBuilder.indexQueryActionType(IndexQueryActionType.REFRESH); - indexQueryDetailsBuilder.indexType(FlintIndexType.COVERING); - return super.visitRefreshCoveringIndexStatement(ctx); - } - - @Override - public Void visitRefreshSkippingIndexStatement( - FlintSparkSqlExtensionsParser.RefreshSkippingIndexStatementContext ctx) { - indexQueryDetailsBuilder.indexQueryActionType(IndexQueryActionType.REFRESH); - indexQueryDetailsBuilder.indexType(FlintIndexType.SKIPPING); - return super.visitRefreshSkippingIndexStatement(ctx); - } - - @Override - public Void visitRefreshMaterializedViewStatement( - FlintSparkSqlExtensionsParser.RefreshMaterializedViewStatementContext ctx) { - indexQueryDetailsBuilder.indexQueryActionType(IndexQueryActionType.REFRESH); - indexQueryDetailsBuilder.indexType(FlintIndexType.MATERIALIZED_VIEW); - indexQueryDetailsBuilder.mvName(ctx.mvName.getText()); - return super.visitRefreshMaterializedViewStatement(ctx); - } - @Override public Void visitPropertyList(FlintSparkSqlExtensionsParser.PropertyListContext ctx) { if (ctx != null) { @@ -266,7 +199,7 @@ public Void visitPropertyList(FlintSparkSqlExtensionsParser.PropertyListContext // https://github.com/apache/spark/blob/v3.5.0/sql/api/src/main/scala/org/apache/spark/sql/catalyst/util/SparkParserUtils.scala#L35 to unescape string literal if (propertyKey(property.key).toLowerCase(Locale.ROOT).contains("auto_refresh")) { if (propertyValue(property.value).toLowerCase(Locale.ROOT).contains("true")) { - indexQueryDetailsBuilder.autoRefresh(true); + indexDetailsBuilder.autoRefresh(true); } } }); diff --git a/spark/src/test/java/org/opensearch/sql/spark/dispatcher/SparkQueryDispatcherTest.java b/spark/src/test/java/org/opensearch/sql/spark/dispatcher/SparkQueryDispatcherTest.java index fc8623d51a..a69c6e2b1a 100644 --- a/spark/src/test/java/org/opensearch/sql/spark/dispatcher/SparkQueryDispatcherTest.java +++ b/spark/src/test/java/org/opensearch/sql/spark/dispatcher/SparkQueryDispatcherTest.java @@ -63,7 +63,11 @@ import org.opensearch.sql.spark.asyncquery.model.AsyncQueryJobMetadata; import org.opensearch.sql.spark.client.EMRServerlessClient; import org.opensearch.sql.spark.client.StartJobRequest; -import org.opensearch.sql.spark.dispatcher.model.*; +import org.opensearch.sql.spark.dispatcher.model.DispatchQueryRequest; +import org.opensearch.sql.spark.dispatcher.model.DispatchQueryResponse; +import org.opensearch.sql.spark.dispatcher.model.FullyQualifiedTableName; +import org.opensearch.sql.spark.dispatcher.model.IndexDetails; +import org.opensearch.sql.spark.dispatcher.model.JobType; import org.opensearch.sql.spark.execution.session.Session; import org.opensearch.sql.spark.execution.session.SessionId; import org.opensearch.sql.spark.execution.session.SessionManager; @@ -122,7 +126,7 @@ void testDispatchSelectQuery() { HashMap tags = new HashMap<>(); tags.put("datasource", "my_glue"); tags.put("cluster", TEST_CLUSTER_NAME); - tags.put("type", JobType.BATCH.getText()); + tags.put("job_type", JobType.BATCH.getText()); String query = "select * from my_glue.default.http_logs"; String sparkSubmitParameters = constructExpectedSparkSubmitParameterString( @@ -177,7 +181,7 @@ void testDispatchSelectQueryWithBasicAuthIndexStoreDatasource() { HashMap tags = new HashMap<>(); tags.put("datasource", "my_glue"); tags.put("cluster", TEST_CLUSTER_NAME); - tags.put("type", JobType.BATCH.getText()); + tags.put("job_type", JobType.BATCH.getText()); String query = "select * from my_glue.default.http_logs"; String sparkSubmitParameters = constructExpectedSparkSubmitParameterString( @@ -233,7 +237,7 @@ void testDispatchSelectQueryWithNoAuthIndexStoreDatasource() { HashMap tags = new HashMap<>(); tags.put("datasource", "my_glue"); tags.put("cluster", TEST_CLUSTER_NAME); - tags.put("type", JobType.BATCH.getText()); + tags.put("job_type", JobType.BATCH.getText()); String query = "select * from my_glue.default.http_logs"; String sparkSubmitParameters = constructExpectedSparkSubmitParameterString( @@ -368,7 +372,7 @@ void testDispatchIndexQuery() { tags.put("datasource", "my_glue"); tags.put("index", "flint_my_glue_default_http_logs_elb_and_requesturi_index"); tags.put("cluster", TEST_CLUSTER_NAME); - tags.put("type", JobType.STREAMING.getText()); + tags.put("job_type", JobType.STREAMING.getText()); String query = "CREATE INDEX elb_and_requestUri ON my_glue.default.http_logs(l_orderkey, l_quantity) WITH" + " (auto_refresh = true)"; @@ -426,7 +430,7 @@ void testDispatchWithPPLQuery() { HashMap tags = new HashMap<>(); tags.put("datasource", "my_glue"); tags.put("cluster", TEST_CLUSTER_NAME); - tags.put("type", JobType.BATCH.getText()); + tags.put("job_type", JobType.BATCH.getText()); String query = "source = my_glue.default.http_logs"; String sparkSubmitParameters = constructExpectedSparkSubmitParameterString( @@ -481,7 +485,7 @@ void testDispatchQueryWithoutATableAndDataSourceName() { HashMap tags = new HashMap<>(); tags.put("datasource", "my_glue"); tags.put("cluster", TEST_CLUSTER_NAME); - tags.put("type", JobType.BATCH.getText()); + tags.put("job_type", JobType.BATCH.getText()); String query = "show tables"; String sparkSubmitParameters = constructExpectedSparkSubmitParameterString( @@ -537,7 +541,7 @@ void testDispatchIndexQueryWithoutADatasourceName() { tags.put("datasource", "my_glue"); tags.put("index", "flint_my_glue_default_http_logs_elb_and_requesturi_index"); tags.put("cluster", TEST_CLUSTER_NAME); - tags.put("type", JobType.STREAMING.getText()); + tags.put("job_type", JobType.STREAMING.getText()); String query = "CREATE INDEX elb_and_requestUri ON default.http_logs(l_orderkey, l_quantity) WITH" + " (auto_refresh = true)"; @@ -596,7 +600,7 @@ void testDispatchMaterializedViewQuery() { tags.put("datasource", "my_glue"); tags.put("index", "flint_mv_1"); tags.put("cluster", TEST_CLUSTER_NAME); - tags.put("type", JobType.STREAMING.getText()); + tags.put("job_type", JobType.STREAMING.getText()); String query = "CREATE MATERIALIZED VIEW mv_1 AS query=select * from my_glue.default.logs WITH" + " (auto_refresh = true)"; @@ -649,168 +653,6 @@ void testDispatchMaterializedViewQuery() { verifyNoInteractions(flintIndexMetadataReader); } - @Test - void testDispatchShowMVQuery() { - HashMap tags = new HashMap<>(); - tags.put("datasource", "my_glue"); - tags.put("cluster", TEST_CLUSTER_NAME); - String query = "SHOW MATERIALIZED VIEW IN mys3.default"; - String sparkSubmitParameters = - constructExpectedSparkSubmitParameterString( - "sigv4", - new HashMap<>() { - { - put(FLINT_INDEX_STORE_AWSREGION_KEY, "eu-west-1"); - } - }); - when(emrServerlessClient.startJobRun( - new StartJobRequest( - query, - "TEST_CLUSTER:index-query", - EMRS_APPLICATION_ID, - EMRS_EXECUTION_ROLE, - sparkSubmitParameters, - tags, - false, - any()))) - .thenReturn(EMR_JOB_ID); - DataSourceMetadata dataSourceMetadata = constructMyGlueDataSourceMetadata(); - when(dataSourceService.getRawDataSourceMetadata("my_glue")).thenReturn(dataSourceMetadata); - doNothing().when(dataSourceUserAuthorizationHelper).authorizeDataSource(dataSourceMetadata); - DispatchQueryResponse dispatchQueryResponse = - sparkQueryDispatcher.dispatch( - new DispatchQueryRequest( - EMRS_APPLICATION_ID, - query, - "my_glue", - LangType.SQL, - EMRS_EXECUTION_ROLE, - TEST_CLUSTER_NAME)); - verify(emrServerlessClient, times(1)).startJobRun(startJobRequestArgumentCaptor.capture()); - StartJobRequest expected = - new StartJobRequest( - query, - "TEST_CLUSTER:index-query", - EMRS_APPLICATION_ID, - EMRS_EXECUTION_ROLE, - sparkSubmitParameters, - tags, - false, - null); - Assertions.assertEquals(expected, startJobRequestArgumentCaptor.getValue()); - Assertions.assertEquals(EMR_JOB_ID, dispatchQueryResponse.getJobId()); - Assertions.assertFalse(dispatchQueryResponse.isDropIndexQuery()); - verifyNoInteractions(flintIndexMetadataReader); - } - - @Test - void testRefreshIndexQuery() { - HashMap tags = new HashMap<>(); - tags.put("datasource", "my_glue"); - tags.put("cluster", TEST_CLUSTER_NAME); - String query = "REFRESH SKIPPING INDEX ON my_glue.default.http_logs"; - String sparkSubmitParameters = - constructExpectedSparkSubmitParameterString( - "sigv4", - new HashMap<>() { - { - put(FLINT_INDEX_STORE_AWSREGION_KEY, "eu-west-1"); - } - }); - when(emrServerlessClient.startJobRun( - new StartJobRequest( - query, - "TEST_CLUSTER:index-query", - EMRS_APPLICATION_ID, - EMRS_EXECUTION_ROLE, - sparkSubmitParameters, - tags, - false, - any()))) - .thenReturn(EMR_JOB_ID); - DataSourceMetadata dataSourceMetadata = constructMyGlueDataSourceMetadata(); - when(dataSourceService.getRawDataSourceMetadata("my_glue")).thenReturn(dataSourceMetadata); - doNothing().when(dataSourceUserAuthorizationHelper).authorizeDataSource(dataSourceMetadata); - DispatchQueryResponse dispatchQueryResponse = - sparkQueryDispatcher.dispatch( - new DispatchQueryRequest( - EMRS_APPLICATION_ID, - query, - "my_glue", - LangType.SQL, - EMRS_EXECUTION_ROLE, - TEST_CLUSTER_NAME)); - verify(emrServerlessClient, times(1)).startJobRun(startJobRequestArgumentCaptor.capture()); - StartJobRequest expected = - new StartJobRequest( - query, - "TEST_CLUSTER:index-query", - EMRS_APPLICATION_ID, - EMRS_EXECUTION_ROLE, - sparkSubmitParameters, - tags, - false, - null); - Assertions.assertEquals(expected, startJobRequestArgumentCaptor.getValue()); - Assertions.assertEquals(EMR_JOB_ID, dispatchQueryResponse.getJobId()); - Assertions.assertFalse(dispatchQueryResponse.isDropIndexQuery()); - verifyNoInteractions(flintIndexMetadataReader); - } - - @Test - void testDispatchDescribeIndexQuery() { - HashMap tags = new HashMap<>(); - tags.put("datasource", "my_glue"); - tags.put("cluster", TEST_CLUSTER_NAME); - String query = "DESCRIBE SKIPPING INDEX ON mys3.default.http_logs"; - String sparkSubmitParameters = - constructExpectedSparkSubmitParameterString( - "sigv4", - new HashMap<>() { - { - put(FLINT_INDEX_STORE_AWSREGION_KEY, "eu-west-1"); - } - }); - when(emrServerlessClient.startJobRun( - new StartJobRequest( - query, - "TEST_CLUSTER:index-query", - EMRS_APPLICATION_ID, - EMRS_EXECUTION_ROLE, - sparkSubmitParameters, - tags, - false, - any()))) - .thenReturn(EMR_JOB_ID); - DataSourceMetadata dataSourceMetadata = constructMyGlueDataSourceMetadata(); - when(dataSourceService.getRawDataSourceMetadata("my_glue")).thenReturn(dataSourceMetadata); - doNothing().when(dataSourceUserAuthorizationHelper).authorizeDataSource(dataSourceMetadata); - DispatchQueryResponse dispatchQueryResponse = - sparkQueryDispatcher.dispatch( - new DispatchQueryRequest( - EMRS_APPLICATION_ID, - query, - "my_glue", - LangType.SQL, - EMRS_EXECUTION_ROLE, - TEST_CLUSTER_NAME)); - verify(emrServerlessClient, times(1)).startJobRun(startJobRequestArgumentCaptor.capture()); - StartJobRequest expected = - new StartJobRequest( - query, - "TEST_CLUSTER:index-query", - EMRS_APPLICATION_ID, - EMRS_EXECUTION_ROLE, - sparkSubmitParameters, - tags, - false, - null); - Assertions.assertEquals(expected, startJobRequestArgumentCaptor.getValue()); - Assertions.assertEquals(EMR_JOB_ID, dispatchQueryResponse.getJobId()); - Assertions.assertFalse(dispatchQueryResponse.isDropIndexQuery()); - verifyNoInteractions(flintIndexMetadataReader); - } - @Test void testDispatchWithWrongURI() { when(dataSourceService.getRawDataSourceMetadata("my_glue")) @@ -1061,15 +903,15 @@ void testGetQueryResponseOfDropIndex() { @Test void testDropIndexQuery() throws ExecutionException, InterruptedException { String query = "DROP INDEX size_year ON my_glue.default.http_logs"; - IndexQueryDetails indexQueryDetails = - IndexQueryDetails.builder() + IndexDetails indexDetails = + IndexDetails.builder() .indexName("size_year") .fullyQualifiedTableName(new FullyQualifiedTableName("my_glue.default.http_logs")) .autoRefresh(false) - .indexQueryActionType(IndexQueryActionType.DROP) + .isDropIndex(true) .indexType(FlintIndexType.COVERING) .build(); - when(flintIndexMetadataReader.getFlintIndexMetadata(indexQueryDetails)) + when(flintIndexMetadataReader.getFlintIndexMetadata(indexDetails)) .thenReturn(flintIndexMetadata); when(flintIndexMetadata.getJobId()).thenReturn(EMR_JOB_ID); // auto_refresh == true @@ -1098,7 +940,7 @@ void testDropIndexQuery() throws ExecutionException, InterruptedException { TEST_CLUSTER_NAME)); verify(emrServerlessClient, times(1)).cancelJobRun(EMRS_APPLICATION_ID, EMR_JOB_ID); verify(dataSourceUserAuthorizationHelper, times(1)).authorizeDataSource(dataSourceMetadata); - verify(flintIndexMetadataReader, times(1)).getFlintIndexMetadata(indexQueryDetails); + verify(flintIndexMetadataReader, times(1)).getFlintIndexMetadata(indexDetails); SparkQueryDispatcher.DropIndexResult dropIndexResult = SparkQueryDispatcher.DropIndexResult.fromJobId(dispatchQueryResponse.getJobId()); Assertions.assertEquals(JobRunState.SUCCESS.toString(), dropIndexResult.getStatus()); @@ -1108,14 +950,14 @@ void testDropIndexQuery() throws ExecutionException, InterruptedException { @Test void testDropSkippingIndexQuery() throws ExecutionException, InterruptedException { String query = "DROP SKIPPING INDEX ON my_glue.default.http_logs"; - IndexQueryDetails indexQueryDetails = - IndexQueryDetails.builder() + IndexDetails indexDetails = + IndexDetails.builder() .fullyQualifiedTableName(new FullyQualifiedTableName("my_glue.default.http_logs")) .autoRefresh(false) - .indexQueryActionType(IndexQueryActionType.DROP) + .isDropIndex(true) .indexType(FlintIndexType.SKIPPING) .build(); - when(flintIndexMetadataReader.getFlintIndexMetadata(indexQueryDetails)) + when(flintIndexMetadataReader.getFlintIndexMetadata(indexDetails)) .thenReturn(flintIndexMetadata); when(flintIndexMetadata.getJobId()).thenReturn(EMR_JOB_ID); when(flintIndexMetadata.isAutoRefresh()).thenReturn(true); @@ -1142,7 +984,7 @@ void testDropSkippingIndexQuery() throws ExecutionException, InterruptedExceptio TEST_CLUSTER_NAME)); verify(emrServerlessClient, times(1)).cancelJobRun(EMRS_APPLICATION_ID, EMR_JOB_ID); verify(dataSourceUserAuthorizationHelper, times(1)).authorizeDataSource(dataSourceMetadata); - verify(flintIndexMetadataReader, times(1)).getFlintIndexMetadata(indexQueryDetails); + verify(flintIndexMetadataReader, times(1)).getFlintIndexMetadata(indexDetails); SparkQueryDispatcher.DropIndexResult dropIndexResult = SparkQueryDispatcher.DropIndexResult.fromJobId(dispatchQueryResponse.getJobId()); Assertions.assertEquals(JobRunState.SUCCESS.toString(), dropIndexResult.getStatus()); @@ -1153,14 +995,14 @@ void testDropSkippingIndexQuery() throws ExecutionException, InterruptedExceptio void testDropSkippingIndexQueryAutoRefreshFalse() throws ExecutionException, InterruptedException { String query = "DROP SKIPPING INDEX ON my_glue.default.http_logs"; - IndexQueryDetails indexQueryDetails = - IndexQueryDetails.builder() + IndexDetails indexDetails = + IndexDetails.builder() .fullyQualifiedTableName(new FullyQualifiedTableName("my_glue.default.http_logs")) .autoRefresh(false) - .indexQueryActionType(IndexQueryActionType.DROP) + .isDropIndex(true) .indexType(FlintIndexType.SKIPPING) .build(); - when(flintIndexMetadataReader.getFlintIndexMetadata(indexQueryDetails)) + when(flintIndexMetadataReader.getFlintIndexMetadata(indexDetails)) .thenReturn(flintIndexMetadata); when(flintIndexMetadata.isAutoRefresh()).thenReturn(false); @@ -1181,7 +1023,7 @@ void testDropSkippingIndexQueryAutoRefreshFalse() TEST_CLUSTER_NAME)); verify(emrServerlessClient, times(0)).cancelJobRun(EMRS_APPLICATION_ID, EMR_JOB_ID); verify(dataSourceUserAuthorizationHelper, times(1)).authorizeDataSource(dataSourceMetadata); - verify(flintIndexMetadataReader, times(1)).getFlintIndexMetadata(indexQueryDetails); + verify(flintIndexMetadataReader, times(1)).getFlintIndexMetadata(indexDetails); SparkQueryDispatcher.DropIndexResult dropIndexResult = SparkQueryDispatcher.DropIndexResult.fromJobId(dispatchQueryResponse.getJobId()); Assertions.assertEquals(JobRunState.SUCCESS.toString(), dropIndexResult.getStatus()); @@ -1192,14 +1034,14 @@ void testDropSkippingIndexQueryAutoRefreshFalse() void testDropSkippingIndexQueryDeleteIndexException() throws ExecutionException, InterruptedException { String query = "DROP SKIPPING INDEX ON my_glue.default.http_logs"; - IndexQueryDetails indexQueryDetails = - IndexQueryDetails.builder() + IndexDetails indexDetails = + IndexDetails.builder() .fullyQualifiedTableName(new FullyQualifiedTableName("my_glue.default.http_logs")) .autoRefresh(false) - .indexQueryActionType(IndexQueryActionType.DROP) + .isDropIndex(true) .indexType(FlintIndexType.SKIPPING) .build(); - when(flintIndexMetadataReader.getFlintIndexMetadata(indexQueryDetails)) + when(flintIndexMetadataReader.getFlintIndexMetadata(indexDetails)) .thenReturn(flintIndexMetadata); when(flintIndexMetadata.isAutoRefresh()).thenReturn(false); @@ -1221,7 +1063,7 @@ void testDropSkippingIndexQueryDeleteIndexException() TEST_CLUSTER_NAME)); verify(emrServerlessClient, times(0)).cancelJobRun(EMRS_APPLICATION_ID, EMR_JOB_ID); verify(dataSourceUserAuthorizationHelper, times(1)).authorizeDataSource(dataSourceMetadata); - verify(flintIndexMetadataReader, times(1)).getFlintIndexMetadata(indexQueryDetails); + verify(flintIndexMetadataReader, times(1)).getFlintIndexMetadata(indexDetails); SparkQueryDispatcher.DropIndexResult dropIndexResult = SparkQueryDispatcher.DropIndexResult.fromJobId(dispatchQueryResponse.getJobId()); Assertions.assertEquals(JobRunState.FAILED.toString(), dropIndexResult.getStatus()); @@ -1234,14 +1076,14 @@ void testDropSkippingIndexQueryDeleteIndexException() @Test void testDropMVQuery() throws ExecutionException, InterruptedException { String query = "DROP MATERIALIZED VIEW mv_1"; - IndexQueryDetails indexQueryDetails = - IndexQueryDetails.builder() + IndexDetails indexDetails = + IndexDetails.builder() .mvName("mv_1") - .indexQueryActionType(IndexQueryActionType.DROP) + .isDropIndex(true) .fullyQualifiedTableName(null) .indexType(FlintIndexType.MATERIALIZED_VIEW) .build(); - when(flintIndexMetadataReader.getFlintIndexMetadata(indexQueryDetails)) + when(flintIndexMetadataReader.getFlintIndexMetadata(indexDetails)) .thenReturn(flintIndexMetadata); when(flintIndexMetadata.getJobId()).thenReturn(EMR_JOB_ID); // auto_refresh == true @@ -1270,7 +1112,7 @@ void testDropMVQuery() throws ExecutionException, InterruptedException { TEST_CLUSTER_NAME)); verify(emrServerlessClient, times(1)).cancelJobRun(EMRS_APPLICATION_ID, EMR_JOB_ID); verify(dataSourceUserAuthorizationHelper, times(1)).authorizeDataSource(dataSourceMetadata); - verify(flintIndexMetadataReader, times(1)).getFlintIndexMetadata(indexQueryDetails); + verify(flintIndexMetadataReader, times(1)).getFlintIndexMetadata(indexDetails); SparkQueryDispatcher.DropIndexResult dropIndexResult = SparkQueryDispatcher.DropIndexResult.fromJobId(dispatchQueryResponse.getJobId()); Assertions.assertEquals(JobRunState.SUCCESS.toString(), dropIndexResult.getStatus()); diff --git a/spark/src/test/java/org/opensearch/sql/spark/execution/session/InteractiveSessionTest.java b/spark/src/test/java/org/opensearch/sql/spark/execution/session/InteractiveSessionTest.java index 14ccaf7708..06a8d8c73c 100644 --- a/spark/src/test/java/org/opensearch/sql/spark/execution/session/InteractiveSessionTest.java +++ b/spark/src/test/java/org/opensearch/sql/spark/execution/session/InteractiveSessionTest.java @@ -13,6 +13,7 @@ import com.amazonaws.services.emrserverless.model.CancelJobRunResult; import com.amazonaws.services.emrserverless.model.GetJobRunResult; +import com.google.common.collect.ImmutableMap; import java.util.HashMap; import java.util.Optional; import lombok.RequiredArgsConstructor; @@ -193,7 +194,7 @@ public static CreateSessionRequest createSessionRequest() { "appId", "arn", SparkSubmitParameters.Builder.builder(), - new HashMap<>(), + ImmutableMap.of(), "resultIndex", DS_NAME); } diff --git a/spark/src/test/java/org/opensearch/sql/spark/flint/FlintIndexMetadataReaderImplTest.java b/spark/src/test/java/org/opensearch/sql/spark/flint/FlintIndexMetadataReaderImplTest.java index 4d809c31dc..3cc40e0df5 100644 --- a/spark/src/test/java/org/opensearch/sql/spark/flint/FlintIndexMetadataReaderImplTest.java +++ b/spark/src/test/java/org/opensearch/sql/spark/flint/FlintIndexMetadataReaderImplTest.java @@ -25,8 +25,7 @@ import org.opensearch.core.xcontent.NamedXContentRegistry; import org.opensearch.core.xcontent.XContentParser; import org.opensearch.sql.spark.dispatcher.model.FullyQualifiedTableName; -import org.opensearch.sql.spark.dispatcher.model.IndexQueryActionType; -import org.opensearch.sql.spark.dispatcher.model.IndexQueryDetails; +import org.opensearch.sql.spark.dispatcher.model.IndexDetails; @ExtendWith(MockitoExtension.class) public class FlintIndexMetadataReaderImplTest { @@ -45,10 +44,10 @@ void testGetJobIdFromFlintSkippingIndexMetadata() { FlintIndexMetadataReader flintIndexMetadataReader = new FlintIndexMetadataReaderImpl(client); FlintIndexMetadata indexMetadata = flintIndexMetadataReader.getFlintIndexMetadata( - IndexQueryDetails.builder() + IndexDetails.builder() .fullyQualifiedTableName(new FullyQualifiedTableName("mys3.default.http_logs")) .autoRefresh(false) - .indexQueryActionType(IndexQueryActionType.DROP) + .isDropIndex(true) .indexType(FlintIndexType.SKIPPING) .build()); Assertions.assertEquals("00fdmvv9hp8u0o0q", indexMetadata.getJobId()); @@ -65,11 +64,11 @@ void testGetJobIdFromFlintCoveringIndexMetadata() { FlintIndexMetadataReader flintIndexMetadataReader = new FlintIndexMetadataReaderImpl(client); FlintIndexMetadata indexMetadata = flintIndexMetadataReader.getFlintIndexMetadata( - IndexQueryDetails.builder() + IndexDetails.builder() .indexName("cv1") .fullyQualifiedTableName(new FullyQualifiedTableName("mys3.default.http_logs")) .autoRefresh(false) - .indexQueryActionType(IndexQueryActionType.DROP) + .isDropIndex(true) .indexType(FlintIndexType.COVERING) .build()); Assertions.assertEquals("00fdmvv9hp8u0o0q", indexMetadata.getJobId()); @@ -88,12 +87,12 @@ void testGetJobIDWithNPEException() { IllegalArgumentException.class, () -> flintIndexMetadataReader.getFlintIndexMetadata( - IndexQueryDetails.builder() + IndexDetails.builder() .indexName("cv1") .fullyQualifiedTableName( new FullyQualifiedTableName("mys3.default.http_logs")) .autoRefresh(false) - .indexQueryActionType(IndexQueryActionType.DROP) + .isDropIndex(true) .indexType(FlintIndexType.COVERING) .build())); Assertions.assertEquals("Provided Index doesn't exist", illegalArgumentException.getMessage()); diff --git a/spark/src/test/java/org/opensearch/sql/spark/flint/IndexQueryDetailsTest.java b/spark/src/test/java/org/opensearch/sql/spark/flint/IndexDetailsTest.java similarity index 71% rename from spark/src/test/java/org/opensearch/sql/spark/flint/IndexQueryDetailsTest.java rename to spark/src/test/java/org/opensearch/sql/spark/flint/IndexDetailsTest.java index e725ddc21e..cf6b5f8f2b 100644 --- a/spark/src/test/java/org/opensearch/sql/spark/flint/IndexQueryDetailsTest.java +++ b/spark/src/test/java/org/opensearch/sql/spark/flint/IndexDetailsTest.java @@ -9,19 +9,18 @@ import org.junit.jupiter.api.Test; import org.opensearch.sql.spark.dispatcher.model.FullyQualifiedTableName; -import org.opensearch.sql.spark.dispatcher.model.IndexQueryActionType; -import org.opensearch.sql.spark.dispatcher.model.IndexQueryDetails; +import org.opensearch.sql.spark.dispatcher.model.IndexDetails; -public class IndexQueryDetailsTest { +public class IndexDetailsTest { @Test public void skippingIndexName() { assertEquals( "flint_mys3_default_http_logs_skipping_index", - IndexQueryDetails.builder() + IndexDetails.builder() .indexName("invalid") .fullyQualifiedTableName(new FullyQualifiedTableName("mys3.default.http_logs")) .autoRefresh(false) - .indexQueryActionType(IndexQueryActionType.DROP) + .isDropIndex(true) .indexType(FlintIndexType.SKIPPING) .build() .openSearchIndexName()); diff --git a/spark/src/test/java/org/opensearch/sql/spark/utils/SQLQueryUtilsTest.java b/spark/src/test/java/org/opensearch/sql/spark/utils/SQLQueryUtilsTest.java index c86d7656d6..01759c2bdd 100644 --- a/spark/src/test/java/org/opensearch/sql/spark/utils/SQLQueryUtilsTest.java +++ b/spark/src/test/java/org/opensearch/sql/spark/utils/SQLQueryUtilsTest.java @@ -15,9 +15,7 @@ import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.junit.jupiter.MockitoExtension; import org.opensearch.sql.spark.dispatcher.model.FullyQualifiedTableName; -import org.opensearch.sql.spark.dispatcher.model.IndexQueryActionType; -import org.opensearch.sql.spark.dispatcher.model.IndexQueryDetails; -import org.opensearch.sql.spark.flint.FlintIndexType; +import org.opensearch.sql.spark.dispatcher.model.IndexDetails; @ExtendWith(MockitoExtension.class) public class SQLQueryUtilsTest { @@ -27,13 +25,13 @@ void testExtractionOfTableNameFromSQLQueries() { String sqlQuery = "select * from my_glue.default.http_logs"; FullyQualifiedTableName fullyQualifiedTableName = SQLQueryUtils.extractFullyQualifiedTableName(sqlQuery); - Assertions.assertFalse(SQLQueryUtils.isFlintExtensionQuery(sqlQuery)); + Assertions.assertFalse(SQLQueryUtils.isIndexQuery(sqlQuery)); Assertions.assertEquals("my_glue", fullyQualifiedTableName.getDatasourceName()); Assertions.assertEquals("default", fullyQualifiedTableName.getSchemaName()); Assertions.assertEquals("http_logs", fullyQualifiedTableName.getTableName()); sqlQuery = "select * from my_glue.db.http_logs"; - Assertions.assertFalse(SQLQueryUtils.isFlintExtensionQuery(sqlQuery)); + Assertions.assertFalse(SQLQueryUtils.isIndexQuery(sqlQuery)); fullyQualifiedTableName = SQLQueryUtils.extractFullyQualifiedTableName(sqlQuery); Assertions.assertEquals("my_glue", fullyQualifiedTableName.getDatasourceName()); Assertions.assertEquals("db", fullyQualifiedTableName.getSchemaName()); @@ -41,28 +39,28 @@ void testExtractionOfTableNameFromSQLQueries() { sqlQuery = "select * from my_glue.http_logs"; fullyQualifiedTableName = SQLQueryUtils.extractFullyQualifiedTableName(sqlQuery); - Assertions.assertFalse(SQLQueryUtils.isFlintExtensionQuery(sqlQuery)); + Assertions.assertFalse(SQLQueryUtils.isIndexQuery(sqlQuery)); Assertions.assertEquals("my_glue", fullyQualifiedTableName.getSchemaName()); Assertions.assertNull(fullyQualifiedTableName.getDatasourceName()); Assertions.assertEquals("http_logs", fullyQualifiedTableName.getTableName()); sqlQuery = "select * from http_logs"; fullyQualifiedTableName = SQLQueryUtils.extractFullyQualifiedTableName(sqlQuery); - Assertions.assertFalse(SQLQueryUtils.isFlintExtensionQuery(sqlQuery)); + Assertions.assertFalse(SQLQueryUtils.isIndexQuery(sqlQuery)); Assertions.assertNull(fullyQualifiedTableName.getDatasourceName()); Assertions.assertNull(fullyQualifiedTableName.getSchemaName()); Assertions.assertEquals("http_logs", fullyQualifiedTableName.getTableName()); sqlQuery = "DROP TABLE myS3.default.alb_logs"; fullyQualifiedTableName = SQLQueryUtils.extractFullyQualifiedTableName(sqlQuery); - Assertions.assertFalse(SQLQueryUtils.isFlintExtensionQuery(sqlQuery)); + Assertions.assertFalse(SQLQueryUtils.isIndexQuery(sqlQuery)); Assertions.assertEquals("myS3", fullyQualifiedTableName.getDatasourceName()); Assertions.assertEquals("default", fullyQualifiedTableName.getSchemaName()); Assertions.assertEquals("alb_logs", fullyQualifiedTableName.getTableName()); sqlQuery = "DESCRIBE TABLE myS3.default.alb_logs"; fullyQualifiedTableName = SQLQueryUtils.extractFullyQualifiedTableName(sqlQuery); - Assertions.assertFalse(SQLQueryUtils.isFlintExtensionQuery(sqlQuery)); + Assertions.assertFalse(SQLQueryUtils.isIndexQuery(sqlQuery)); Assertions.assertEquals("myS3", fullyQualifiedTableName.getDatasourceName()); Assertions.assertEquals("default", fullyQualifiedTableName.getSchemaName()); Assertions.assertEquals("alb_logs", fullyQualifiedTableName.getTableName()); @@ -75,7 +73,7 @@ void testExtractionOfTableNameFromSQLQueries() { + "STORED AS file_format\n" + "LOCATION { 's3://bucket/folder/' }"; fullyQualifiedTableName = SQLQueryUtils.extractFullyQualifiedTableName(sqlQuery); - Assertions.assertFalse(SQLQueryUtils.isFlintExtensionQuery(sqlQuery)); + Assertions.assertFalse(SQLQueryUtils.isIndexQuery(sqlQuery)); Assertions.assertEquals("myS3", fullyQualifiedTableName.getDatasourceName()); Assertions.assertEquals("default", fullyQualifiedTableName.getSchemaName()); Assertions.assertEquals("alb_logs", fullyQualifiedTableName.getTableName()); @@ -94,7 +92,7 @@ void testErrorScenarios() { sqlQuery = "DESCRIBE TABLE FROM myS3.default.alb_logs"; fullyQualifiedTableName = SQLQueryUtils.extractFullyQualifiedTableName(sqlQuery); - Assertions.assertFalse(SQLQueryUtils.isFlintExtensionQuery(sqlQuery)); + Assertions.assertFalse(SQLQueryUtils.isIndexQuery(sqlQuery)); Assertions.assertEquals("FROM", fullyQualifiedTableName.getFullyQualifiedName()); Assertions.assertNull(fullyQualifiedTableName.getSchemaName()); Assertions.assertEquals("FROM", fullyQualifiedTableName.getTableName()); @@ -106,12 +104,10 @@ void testExtractionFromFlintIndexQueries() { String createCoveredIndexQuery = "CREATE INDEX elb_and_requestUri ON myS3.default.alb_logs(l_orderkey, l_quantity) WITH" + " (auto_refresh = true)"; - Assertions.assertTrue(SQLQueryUtils.isFlintExtensionQuery(createCoveredIndexQuery)); - IndexQueryDetails indexQueryDetails = - SQLQueryUtils.extractIndexDetails(createCoveredIndexQuery); - FullyQualifiedTableName fullyQualifiedTableName = - indexQueryDetails.getFullyQualifiedTableName(); - Assertions.assertEquals("elb_and_requestUri", indexQueryDetails.getIndexName()); + Assertions.assertTrue(SQLQueryUtils.isIndexQuery(createCoveredIndexQuery)); + IndexDetails indexDetails = SQLQueryUtils.extractIndexDetails(createCoveredIndexQuery); + FullyQualifiedTableName fullyQualifiedTableName = indexDetails.getFullyQualifiedTableName(); + Assertions.assertEquals("elb_and_requestUri", indexDetails.getIndexName()); Assertions.assertEquals("myS3", fullyQualifiedTableName.getDatasourceName()); Assertions.assertEquals("default", fullyQualifiedTableName.getSchemaName()); Assertions.assertEquals("alb_logs", fullyQualifiedTableName.getTableName()); @@ -122,99 +118,12 @@ void testExtractionFromFlintMVQuery() { String createCoveredIndexQuery = "CREATE MATERIALIZED VIEW mv_1 AS query=select * from my_glue.default.logs WITH" + " (auto_refresh = true)"; - Assertions.assertTrue(SQLQueryUtils.isFlintExtensionQuery(createCoveredIndexQuery)); - IndexQueryDetails indexQueryDetails = - SQLQueryUtils.extractIndexDetails(createCoveredIndexQuery); - FullyQualifiedTableName fullyQualifiedTableName = - indexQueryDetails.getFullyQualifiedTableName(); - Assertions.assertNull(indexQueryDetails.getIndexName()); - Assertions.assertNull(fullyQualifiedTableName); - Assertions.assertEquals("mv_1", indexQueryDetails.getMvName()); - } - - @Test - void testDescIndex() { - String descSkippingIndex = "DESC SKIPPING INDEX ON mys3.default.http_logs"; - Assertions.assertTrue(SQLQueryUtils.isFlintExtensionQuery(descSkippingIndex)); - IndexQueryDetails indexDetails = SQLQueryUtils.extractIndexDetails(descSkippingIndex); - FullyQualifiedTableName fullyQualifiedTableName = indexDetails.getFullyQualifiedTableName(); - Assertions.assertNull(indexDetails.getIndexName()); - Assertions.assertNotNull(fullyQualifiedTableName); - Assertions.assertEquals(FlintIndexType.SKIPPING, indexDetails.getIndexType()); - Assertions.assertEquals(IndexQueryActionType.DESCRIBE, indexDetails.getIndexQueryActionType()); - - String descCoveringIndex = "DESC INDEX cv1 ON mys3.default.http_logs"; - Assertions.assertTrue(SQLQueryUtils.isFlintExtensionQuery(descCoveringIndex)); - indexDetails = SQLQueryUtils.extractIndexDetails(descCoveringIndex); - fullyQualifiedTableName = indexDetails.getFullyQualifiedTableName(); - Assertions.assertEquals("cv1", indexDetails.getIndexName()); - Assertions.assertNotNull(fullyQualifiedTableName); - Assertions.assertEquals(FlintIndexType.COVERING, indexDetails.getIndexType()); - Assertions.assertEquals(IndexQueryActionType.DESCRIBE, indexDetails.getIndexQueryActionType()); - - String descMv = "DESC MATERIALIZED VIEW mv1"; - Assertions.assertTrue(SQLQueryUtils.isFlintExtensionQuery(descMv)); - indexDetails = SQLQueryUtils.extractIndexDetails(descMv); - fullyQualifiedTableName = indexDetails.getFullyQualifiedTableName(); - Assertions.assertNull(indexDetails.getIndexName()); - Assertions.assertEquals("mv1", indexDetails.getMvName()); - Assertions.assertNull(fullyQualifiedTableName); - Assertions.assertEquals(FlintIndexType.MATERIALIZED_VIEW, indexDetails.getIndexType()); - Assertions.assertEquals(IndexQueryActionType.DESCRIBE, indexDetails.getIndexQueryActionType()); - } - - @Test - void testShowIndex() { - String showCoveringIndex = " SHOW INDEX ON myS3.default.http_logs"; - Assertions.assertTrue(SQLQueryUtils.isFlintExtensionQuery(showCoveringIndex)); - IndexQueryDetails indexDetails = SQLQueryUtils.extractIndexDetails(showCoveringIndex); + Assertions.assertTrue(SQLQueryUtils.isIndexQuery(createCoveredIndexQuery)); + IndexDetails indexDetails = SQLQueryUtils.extractIndexDetails(createCoveredIndexQuery); FullyQualifiedTableName fullyQualifiedTableName = indexDetails.getFullyQualifiedTableName(); Assertions.assertNull(indexDetails.getIndexName()); - Assertions.assertNull(indexDetails.getMvName()); - Assertions.assertNotNull(fullyQualifiedTableName); - Assertions.assertEquals(FlintIndexType.COVERING, indexDetails.getIndexType()); - Assertions.assertEquals(IndexQueryActionType.SHOW, indexDetails.getIndexQueryActionType()); - - String showMV = "SHOW MATERIALIZED VIEW IN my_glue.default"; - Assertions.assertTrue(SQLQueryUtils.isFlintExtensionQuery(showMV)); - indexDetails = SQLQueryUtils.extractIndexDetails(showMV); - fullyQualifiedTableName = indexDetails.getFullyQualifiedTableName(); - Assertions.assertNull(indexDetails.getIndexName()); - Assertions.assertNull(indexDetails.getMvName()); - Assertions.assertNull(fullyQualifiedTableName); - Assertions.assertEquals(FlintIndexType.MATERIALIZED_VIEW, indexDetails.getIndexType()); - Assertions.assertEquals(IndexQueryActionType.SHOW, indexDetails.getIndexQueryActionType()); - } - - @Test - void testRefreshIndex() { - String refreshSkippingIndex = "REFRESH SKIPPING INDEX ON mys3.default.http_logs"; - Assertions.assertTrue(SQLQueryUtils.isFlintExtensionQuery(refreshSkippingIndex)); - IndexQueryDetails indexDetails = SQLQueryUtils.extractIndexDetails(refreshSkippingIndex); - FullyQualifiedTableName fullyQualifiedTableName = indexDetails.getFullyQualifiedTableName(); - Assertions.assertNull(indexDetails.getIndexName()); - Assertions.assertNotNull(fullyQualifiedTableName); - Assertions.assertEquals(FlintIndexType.SKIPPING, indexDetails.getIndexType()); - Assertions.assertEquals(IndexQueryActionType.REFRESH, indexDetails.getIndexQueryActionType()); - - String refreshCoveringIndex = "REFRESH INDEX cv1 ON mys3.default.http_logs"; - Assertions.assertTrue(SQLQueryUtils.isFlintExtensionQuery(refreshCoveringIndex)); - indexDetails = SQLQueryUtils.extractIndexDetails(refreshCoveringIndex); - fullyQualifiedTableName = indexDetails.getFullyQualifiedTableName(); - Assertions.assertEquals("cv1", indexDetails.getIndexName()); - Assertions.assertNotNull(fullyQualifiedTableName); - Assertions.assertEquals(FlintIndexType.COVERING, indexDetails.getIndexType()); - Assertions.assertEquals(IndexQueryActionType.REFRESH, indexDetails.getIndexQueryActionType()); - - String refreshMV = "REFRESH MATERIALIZED VIEW mv1"; - Assertions.assertTrue(SQLQueryUtils.isFlintExtensionQuery(refreshMV)); - indexDetails = SQLQueryUtils.extractIndexDetails(refreshMV); - fullyQualifiedTableName = indexDetails.getFullyQualifiedTableName(); - Assertions.assertNull(indexDetails.getIndexName()); - Assertions.assertEquals("mv1", indexDetails.getMvName()); Assertions.assertNull(fullyQualifiedTableName); - Assertions.assertEquals(FlintIndexType.MATERIALIZED_VIEW, indexDetails.getIndexType()); - Assertions.assertEquals(IndexQueryActionType.REFRESH, indexDetails.getIndexQueryActionType()); + Assertions.assertEquals("mv_1", indexDetails.getMvName()); } /** https://github.com/opensearch-project/sql/issues/2206 */