From 83bc3d2cdcea815200ad6e0d2b608cea1a88c2e2 Mon Sep 17 00:00:00 2001
From: Tomoyuki MORITA <moritato@amazon.com>
Date: Thu, 29 Aug 2024 12:54:03 -0700
Subject: [PATCH] Add mvQuery attribute in IndexQueryDetails (#2946)

Signed-off-by: Tomoyuki Morita <moritato@amazon.com>
---
 .../dispatcher/model/IndexQueryDetails.java   |  6 ++++++
 .../sql/spark/utils/SQLQueryUtils.java        | 11 +++++++++++
 .../sql/spark/utils/SQLQueryUtilsTest.java    | 19 +++++++++++++++----
 3 files changed, 32 insertions(+), 4 deletions(-)

diff --git a/async-query-core/src/main/java/org/opensearch/sql/spark/dispatcher/model/IndexQueryDetails.java b/async-query-core/src/main/java/org/opensearch/sql/spark/dispatcher/model/IndexQueryDetails.java
index 5596d1b425..2ca997f6b0 100644
--- a/async-query-core/src/main/java/org/opensearch/sql/spark/dispatcher/model/IndexQueryDetails.java
+++ b/async-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
index 0bb9cb4b85..ff08a8f41e 100644
--- 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,6 +13,7 @@
 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;
@@ -20,6 +21,7 @@
 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
index bf6fe9e5db..fe7777606c 100644
--- 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) {