Add missing tags and MV support

Signed-off-by: Vamsi Manohar <[email protected]>

spark/src/main/antlr/FlintSparkSqlExtensions.g4

@@ -17,6 +17,7 @@ singleStatement
 statement
     : skippingIndexStatement
     | coveringIndexStatement
+    | materializedViewStatement
     ;
 
 skippingIndexStatement
@@ -76,6 +77,39 @@ dropCoveringIndexStatement
     : DROP INDEX indexName ON tableName
     ;
 
+materializedViewStatement
+    : createMaterializedViewStatement
+    | showMaterializedViewStatement
+    | describeMaterializedViewStatement
+    | dropMaterializedViewStatement
+    ;
+
+createMaterializedViewStatement
+    : CREATE MATERIALIZED VIEW (IF NOT EXISTS)? mvName=multipartIdentifier
+        AS query=materializedViewQuery
+        (WITH LEFT_PAREN propertyList RIGHT_PAREN)?
+    ;
+
+showMaterializedViewStatement
+    : SHOW MATERIALIZED (VIEW | VIEWS) IN catalogDb=multipartIdentifier
+    ;
+
+describeMaterializedViewStatement
+    : (DESC | DESCRIBE) MATERIALIZED VIEW mvName=multipartIdentifier
+    ;
+
+dropMaterializedViewStatement
+    : DROP MATERIALIZED VIEW mvName=multipartIdentifier
+    ;
+
+/*
+ * Match all remaining tokens in non-greedy way
+ * so WITH clause won't be captured by this rule.
+ */
+materializedViewQuery
+    : .+?
+    ;
+
 indexColTypeList
     : indexColType (COMMA indexColType)*
     ;

spark/src/main/antlr/SparkSqlBase.g4

@@ -154,21 +154,26 @@ COMMA: ',';
 DOT: '.';
 
 
+AS: 'AS';
 CREATE: 'CREATE';
 DESC: 'DESC';
 DESCRIBE: 'DESCRIBE';
 DROP: 'DROP';
 EXISTS: 'EXISTS';
 FALSE: 'FALSE';
 IF: 'IF';
+IN: 'IN';
 INDEX: 'INDEX';
 INDEXES: 'INDEXES';
+MATERIALIZED: 'MATERIALIZED';
 NOT: 'NOT';
 ON: 'ON';
 PARTITION: 'PARTITION';
 REFRESH: 'REFRESH';
 SHOW: 'SHOW';
 TRUE: 'TRUE';
+VIEW: 'VIEW';
+VIEWS: 'VIEWS';
 WITH: 'WITH';
 
 

spark/src/main/java/org/opensearch/sql/spark/dispatcher/SparkQueryDispatcher.java

@@ -39,8 +39,8 @@
 import org.opensearch.sql.spark.client.StartJobRequest;
 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.CreateSessionRequest;
 import org.opensearch.sql.spark.execution.session.Session;
 import org.opensearch.sql.spark.execution.session.SessionId;
@@ -63,9 +63,8 @@ public class SparkQueryDispatcher {
 
   public static final String INDEX_TAG_KEY = "index";
   public static final String DATASOURCE_TAG_KEY = "datasource";
-  public static final String SCHEMA_TAG_KEY = "schema";
-  public static final String TABLE_TAG_KEY = "table";
   public static final String CLUSTER_NAME_TAG_KEY = "cluster";
+  public static final String JOB_TYPE_TAG_KEY = "job_type";
 
   private EMRServerlessClient emrServerlessClient;
 
@@ -190,6 +189,8 @@ private DispatchQueryResponse handleSQLQuery(DispatchQueryRequest dispatchQueryR
     if (SQLQueryUtils.isIndexQuery(dispatchQueryRequest.getQuery())) {
       IndexDetails indexDetails =
           SQLQueryUtils.extractIndexDetails(dispatchQueryRequest.getQuery());
+      fillMissingDetails(dispatchQueryRequest, indexDetails);
+
       if (indexDetails.isDropIndex()) {
         return handleDropIndexQuery(dispatchQueryRequest, indexDetails);
       } else {
@@ -200,17 +201,29 @@ private DispatchQueryResponse handleSQLQuery(DispatchQueryRequest dispatchQueryR
     }
   }
 
+  // TODO: Revisit this logic.
+  // Currently, Spark if datasource is not provided in query.
+  // 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, IndexDetails indexDetails) {
+    if (indexDetails.getFullyQualifiedTableName() != null
+        && indexDetails.getFullyQualifiedTableName().getDatasourceName() == null) {
+      indexDetails
+          .getFullyQualifiedTableName()
+          .setDatasourceName(dispatchQueryRequest.getDatasource());
+    }
+  }
+
   private DispatchQueryResponse handleIndexQuery(
       DispatchQueryRequest dispatchQueryRequest, IndexDetails indexDetails) {
-    FullyQualifiedTableName fullyQualifiedTableName = indexDetails.getFullyQualifiedTableName();
     DataSourceMetadata dataSourceMetadata =
         this.dataSourceService.getRawDataSourceMetadata(dispatchQueryRequest.getDatasource());
     dataSourceUserAuthorizationHelper.authorizeDataSource(dataSourceMetadata);
     String jobName = dispatchQueryRequest.getClusterName() + ":" + "index-query";
     Map<String, String> tags = getDefaultTagsForJobSubmission(dispatchQueryRequest);
-    tags.put(INDEX_TAG_KEY, indexDetails.getIndexName());
-    tags.put(TABLE_TAG_KEY, fullyQualifiedTableName.getTableName());
-    tags.put(SCHEMA_TAG_KEY, fullyQualifiedTableName.getSchemaName());
+    tags.put(INDEX_TAG_KEY, indexDetails.openSearchIndexName());
+    tags.put(JOB_TYPE_TAG_KEY, JobType.STREAMING.getText());
     StartJobRequest startJobRequest =
         new StartJobRequest(
             dispatchQueryRequest.getQuery(),
@@ -221,12 +234,12 @@ private DispatchQueryResponse handleIndexQuery(
                 .dataSource(
                     dataSourceService.getRawDataSourceMetadata(
                         dispatchQueryRequest.getDatasource()))
-                .structuredStreaming(indexDetails.getAutoRefresh())
+                .structuredStreaming(indexDetails.isAutoRefresh())
                 .extraParameters(dispatchQueryRequest.getExtraSparkSubmitParams())
                 .build()
                 .toString(),
             tags,
-            indexDetails.getAutoRefresh(),
+            indexDetails.isAutoRefresh(),
             dataSourceMetadata.getResultIndex());
     String jobId = emrServerlessClient.startJobRun(startJobRequest);
     return new DispatchQueryResponse(jobId, false, dataSourceMetadata.getResultIndex(), null);
@@ -251,6 +264,7 @@ private DispatchQueryResponse handleNonIndexQuery(DispatchQueryRequest dispatchQ
         session = createdSession.get();
       } else {
         // create session if not exist
+        tags.put(JOB_TYPE_TAG_KEY, JobType.INTERACTIVE.getText());
         session =
             sessionManager.createSession(
                 new CreateSessionRequest(
@@ -277,6 +291,7 @@ private DispatchQueryResponse handleNonIndexQuery(DispatchQueryRequest dispatchQ
           dataSourceMetadata.getResultIndex(),
           session.getSessionId().getSessionId());
     } else {
+      tags.put(JOB_TYPE_TAG_KEY, JobType.BATCH.getText());
       StartJobRequest startJobRequest =
           new StartJobRequest(
               dispatchQueryRequest.getQuery(),

spark/src/main/java/org/opensearch/sql/spark/dispatcher/model/IndexDetails.java

@@ -5,56 +5,129 @@
 
 package org.opensearch.sql.spark.dispatcher.model;
 
-import lombok.AllArgsConstructor;
-import lombok.Data;
+import com.google.common.base.Preconditions;
 import lombok.EqualsAndHashCode;
-import lombok.NoArgsConstructor;
+import lombok.Getter;
+import org.apache.commons.lang3.StringUtils;
 import org.opensearch.sql.spark.flint.FlintIndexType;
 
 /** Index details in an async query. */
-@Data
-@AllArgsConstructor
-@NoArgsConstructor
+@Getter
 @EqualsAndHashCode
 public class IndexDetails {
+
+  public static final String STRIP_CHARS = "`";
+
   private String indexName;
   private FullyQualifiedTableName fullyQualifiedTableName;
   // by default, auto_refresh = false;
-  private Boolean autoRefresh = false;
+  private boolean autoRefresh;
   private boolean isDropIndex;
+  // materialized view special case where
+  // table name and mv name are combined.
+  private String mvName;
   private FlintIndexType indexType;
 
+  private IndexDetails() {}
+
+  public static IndexDetailsBuilder builder() {
+    return new IndexDetailsBuilder();
+  }
+
+  // Builder class
+  public static class IndexDetailsBuilder {
+    private final IndexDetails indexDetails;
+
+    public IndexDetailsBuilder() {
+      indexDetails = new IndexDetails();
+    }
+
+    public IndexDetailsBuilder indexName(String indexName) {
+      indexDetails.indexName = indexName;
+      return this;
+    }
+
+    public IndexDetailsBuilder fullyQualifiedTableName(FullyQualifiedTableName tableName) {
+      indexDetails.fullyQualifiedTableName = tableName;
+      return this;
+    }
+
+    public IndexDetailsBuilder autoRefresh(Boolean autoRefresh) {
+      indexDetails.autoRefresh = autoRefresh;
+      return this;
+    }
+
+    public IndexDetailsBuilder isDropIndex(boolean isDropIndex) {
+      indexDetails.isDropIndex = isDropIndex;
+      return this;
+    }
+
+    public IndexDetailsBuilder mvName(String mvName) {
+      indexDetails.mvName = mvName;
+      return this;
+    }
+
+    public IndexDetailsBuilder indexType(FlintIndexType indexType) {
+      indexDetails.indexType = indexType;
+      return this;
+    }
+
+    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;
+    }
+  }
+
   public String openSearchIndexName() {
     FullyQualifiedTableName fullyQualifiedTableName = getFullyQualifiedTableName();
-    if (FlintIndexType.SKIPPING.equals(getIndexType())) {
-      String indexName =
-          "flint"
-              + "_"
-              + fullyQualifiedTableName.getDatasourceName()
-              + "_"
-              + fullyQualifiedTableName.getSchemaName()
-              + "_"
-              + fullyQualifiedTableName.getTableName()
-              + "_"
-              + getIndexType().getSuffix();
-      return indexName.toLowerCase();
-    } else if (FlintIndexType.COVERING.equals(getIndexType())) {
-      String indexName =
-          "flint"
-              + "_"
-              + fullyQualifiedTableName.getDatasourceName()
-              + "_"
-              + fullyQualifiedTableName.getSchemaName()
-              + "_"
-              + fullyQualifiedTableName.getTableName()
-              + "_"
-              + getIndexName()
-              + "_"
-              + getIndexType().getSuffix();
-      return indexName.toLowerCase();
-    } else {
-      throw new UnsupportedOperationException(
-          String.format("Unsupported Index Type : %s", getIndexType()));
+    String indexName = StringUtils.EMPTY;
+    switch (getIndexType()) {
+      case COVERING:
+        indexName =
+            "flint"
+                + "_"
+                + StringUtils.strip(fullyQualifiedTableName.getDatasourceName(), STRIP_CHARS)
+                + "_"
+                + StringUtils.strip(fullyQualifiedTableName.getSchemaName(), STRIP_CHARS)
+                + "_"
+                + StringUtils.strip(fullyQualifiedTableName.getTableName(), STRIP_CHARS)
+                + "_"
+                + StringUtils.strip(getIndexName(), STRIP_CHARS)
+                + "_"
+                + getIndexType().getSuffix();
+        break;
+      case SKIPPING:
+        indexName =
+            "flint"
+                + "_"
+                + StringUtils.strip(fullyQualifiedTableName.getDatasourceName(), STRIP_CHARS)
+                + "_"
+                + StringUtils.strip(fullyQualifiedTableName.getSchemaName(), STRIP_CHARS)
+                + "_"
+                + StringUtils.strip(fullyQualifiedTableName.getTableName(), STRIP_CHARS)
+                + "_"
+                + getIndexType().getSuffix();
+        break;
+      case MATERIALIZED_VIEW:
+        indexName = "flint" + "_" + StringUtils.strip(getMvName(), STRIP_CHARS).toLowerCase();
+        break;
     }
+    return indexName.toLowerCase();
   }
 }

spark/src/main/java/org/opensearch/sql/spark/dispatcher/model/JobType.java

@@ -0,0 +1,37 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.spark.dispatcher.model;
+
+public enum JobType {
+  INTERACTIVE("interactive"),
+  STREAMING("streaming"),
+  BATCH("batch");
+
+  private String text;
+
+  JobType(String text) {
+    this.text = text;
+  }
+
+  public String getText() {
+    return this.text;
+  }
+
+  /**
+   * Get JobType from text.
+   *
+   * @param text text.
+   * @return JobType {@link JobType}.
+   */
+  public static JobType fromString(String text) {
+    for (JobType JobType : JobType.values()) {
+      if (JobType.text.equalsIgnoreCase(text)) {
+        return JobType;
+      }
+    }
+    throw new IllegalArgumentException("No JobType with text " + text + " found");
+  }
+}