Add mvQuery attribute in IndexQueryDetails (#2946) (#2951)

(cherry picked from commit 83bc3d2) Signed-off-by: Tomoyuki Morita <moritato@amazon.com> Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
opensearch-project · Aug 30, 2024 · f3d494b · f3d494b
1 parent 8e04f32
commit f3d494b
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 4 deletions.
diff --git a/...query-core/src/main/java/org/opensearch/sql/spark/dispatcher/model/IndexQueryDetails.java b/...query-core/src/main/java/org/opensearch/sql/spark/dispatcher/model/IndexQueryDetails.java
@@ -31,6 +31,7 @@ public class IndexQueryDetails {
   // materialized view special case where
   // table name and mv name are combined.
   private String mvName;
+  private String mvQuery;
   private FlintIndexType indexType;
 
   private IndexQueryDetails() {}
@@ -73,6 +74,11 @@ public IndexQueryDetailsBuilder mvName(String mvName) {
       return this;
     }
 
+    public IndexQueryDetailsBuilder mvQuery(String mvQuery) {
+      indexQueryDetails.mvQuery = mvQuery;
+      return this;
+    }
+
     public IndexQueryDetailsBuilder indexType(FlintIndexType indexType) {
       indexQueryDetails.indexType = indexType;
       return this;

diff --git a/async-query-core/src/main/java/org/opensearch/sql/spark/utils/SQLQueryUtils.java b/async-query-core/src/main/java/org/opensearch/sql/spark/utils/SQLQueryUtils.java
@@ -13,13 +13,15 @@
 import lombok.Getter;
 import lombok.experimental.UtilityClass;
 import org.antlr.v4.runtime.CommonTokenStream;
+import org.antlr.v4.runtime.misc.Interval;
 import org.antlr.v4.runtime.tree.ParseTree;
 import org.opensearch.sql.common.antlr.CaseInsensitiveCharStream;
 import org.opensearch.sql.common.antlr.SyntaxAnalysisErrorListener;
 import org.opensearch.sql.common.antlr.SyntaxCheckException;
 import org.opensearch.sql.spark.antlr.parser.FlintSparkSqlExtensionsBaseVisitor;
 import org.opensearch.sql.spark.antlr.parser.FlintSparkSqlExtensionsLexer;
 import org.opensearch.sql.spark.antlr.parser.FlintSparkSqlExtensionsParser;
+import org.opensearch.sql.spark.antlr.parser.FlintSparkSqlExtensionsParser.MaterializedViewQueryContext;
 import org.opensearch.sql.spark.antlr.parser.SqlBaseLexer;
 import org.opensearch.sql.spark.antlr.parser.SqlBaseParser;
 import org.opensearch.sql.spark.antlr.parser.SqlBaseParser.IdentifierReferenceContext;
@@ -353,6 +355,15 @@ public Void visitAlterMaterializedViewStatement(
       return super.visitAlterMaterializedViewStatement(ctx);
     }
 
+    @Override
+    public Void visitMaterializedViewQuery(MaterializedViewQueryContext ctx) {
+      int a = ctx.start.getStartIndex();
+      int b = ctx.stop.getStopIndex();
+      String query = ctx.start.getInputStream().getText(new Interval(a, b));
+      indexQueryDetailsBuilder.mvQuery(query);
+      return super.visitMaterializedViewQuery(ctx);
+    }
+
     private String propertyKey(FlintSparkSqlExtensionsParser.PropertyKeyContext key) {
       if (key.STRING() != null) {
         return key.STRING().getText();

diff --git a/async-query-core/src/test/java/org/opensearch/sql/spark/utils/SQLQueryUtilsTest.java b/async-query-core/src/test/java/org/opensearch/sql/spark/utils/SQLQueryUtilsTest.java
@@ -181,11 +181,22 @@ void testExtractionFromFlintCoveringIndexQueries() {
     }
   }
 
+  @Test
+  void testExtractionFromCreateMVQuery() {
+    String mvQuery = "select * from my_glue.default.logs";
+    String query = "CREATE MATERIALIZED VIEW mv_1 AS " + mvQuery + " WITH (auto_refresh = true)";
+
+    assertTrue(SQLQueryUtils.isFlintExtensionQuery(query));
+    IndexQueryDetails indexQueryDetails = SQLQueryUtils.extractIndexDetails(query);
+    assertNull(indexQueryDetails.getIndexName());
+    assertNull(indexQueryDetails.getFullyQualifiedTableName());
+    assertEquals(mvQuery, indexQueryDetails.getMvQuery());
+    assertEquals("mv_1", indexQueryDetails.getMvName());
+  }
+
   @Test
   void testExtractionFromFlintMVQuery() {
     String[] mvQueries = {
-      "CREATE MATERIALIZED VIEW mv_1 AS query=select * from my_glue.default.logs WITH"
-          + " (auto_refresh = true)",
       "DROP MATERIALIZED VIEW mv_1",
       "VACUUM MATERIALIZED VIEW mv_1",
       "ALTER MATERIALIZED VIEW mv_1 WITH (auto_refresh = false)",
@@ -200,6 +211,7 @@ void testExtractionFromFlintMVQuery() {
 
       assertNull(indexQueryDetails.getIndexName());
       assertNull(fullyQualifiedTableName);
+      assertNull(indexQueryDetails.getMvQuery());
       assertEquals("mv_1", indexQueryDetails.getMvName());
     }
   }
@@ -428,8 +440,7 @@ public static IndexQuery index() {
     }
 
     public static IndexQuery mv() {
-      return new IndexQuery(
-          "CREATE MATERIALIZED VIEW mv_1 AS query=select * from my_glue.default.logs");
+      return new IndexQuery("CREATE MATERIALIZED VIEW mv_1 AS select * from my_glue.default.logs");
     }
 
     public IndexQuery withProperty(String key, String value) {