Skip to content

Commit 83bc3d2

Browse files
authored
Add mvQuery attribute in IndexQueryDetails (#2946)
Signed-off-by: Tomoyuki Morita <moritato@amazon.com>
1 parent 69853fe commit 83bc3d2

3 files changed

Lines changed: 32 additions & 4 deletions

File tree

async-query-core/src/main/java/org/opensearch/sql/spark/dispatcher/model/IndexQueryDetails.java

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@ public class IndexQueryDetails {
3131
// materialized view special case where
3232
// table name and mv name are combined.
3333
private String mvName;
34+
private String mvQuery;
3435
private FlintIndexType indexType;
3536

3637
private IndexQueryDetails() {}
@@ -73,6 +74,11 @@ public IndexQueryDetailsBuilder mvName(String mvName) {
7374
return this;
7475
}
7576

77+
public IndexQueryDetailsBuilder mvQuery(String mvQuery) {
78+
indexQueryDetails.mvQuery = mvQuery;
79+
return this;
80+
}
81+
7682
public IndexQueryDetailsBuilder indexType(FlintIndexType indexType) {
7783
indexQueryDetails.indexType = indexType;
7884
return this;

async-query-core/src/main/java/org/opensearch/sql/spark/utils/SQLQueryUtils.java

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,13 +13,15 @@
1313
import lombok.Getter;
1414
import lombok.experimental.UtilityClass;
1515
import org.antlr.v4.runtime.CommonTokenStream;
16+
import org.antlr.v4.runtime.misc.Interval;
1617
import org.antlr.v4.runtime.tree.ParseTree;
1718
import org.opensearch.sql.common.antlr.CaseInsensitiveCharStream;
1819
import org.opensearch.sql.common.antlr.SyntaxAnalysisErrorListener;
1920
import org.opensearch.sql.common.antlr.SyntaxCheckException;
2021
import org.opensearch.sql.spark.antlr.parser.FlintSparkSqlExtensionsBaseVisitor;
2122
import org.opensearch.sql.spark.antlr.parser.FlintSparkSqlExtensionsLexer;
2223
import org.opensearch.sql.spark.antlr.parser.FlintSparkSqlExtensionsParser;
24+
import org.opensearch.sql.spark.antlr.parser.FlintSparkSqlExtensionsParser.MaterializedViewQueryContext;
2325
import org.opensearch.sql.spark.antlr.parser.SqlBaseLexer;
2426
import org.opensearch.sql.spark.antlr.parser.SqlBaseParser;
2527
import org.opensearch.sql.spark.antlr.parser.SqlBaseParser.IdentifierReferenceContext;
@@ -353,6 +355,15 @@ public Void visitAlterMaterializedViewStatement(
353355
return super.visitAlterMaterializedViewStatement(ctx);
354356
}
355357

358+
@Override
359+
public Void visitMaterializedViewQuery(MaterializedViewQueryContext ctx) {
360+
int a = ctx.start.getStartIndex();
361+
int b = ctx.stop.getStopIndex();
362+
String query = ctx.start.getInputStream().getText(new Interval(a, b));
363+
indexQueryDetailsBuilder.mvQuery(query);
364+
return super.visitMaterializedViewQuery(ctx);
365+
}
366+
356367
private String propertyKey(FlintSparkSqlExtensionsParser.PropertyKeyContext key) {
357368
if (key.STRING() != null) {
358369
return key.STRING().getText();

async-query-core/src/test/java/org/opensearch/sql/spark/utils/SQLQueryUtilsTest.java

Lines changed: 15 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -181,11 +181,22 @@ void testExtractionFromFlintCoveringIndexQueries() {
181181
}
182182
}
183183

184+
@Test
185+
void testExtractionFromCreateMVQuery() {
186+
String mvQuery = "select * from my_glue.default.logs";
187+
String query = "CREATE MATERIALIZED VIEW mv_1 AS " + mvQuery + " WITH (auto_refresh = true)";
188+
189+
assertTrue(SQLQueryUtils.isFlintExtensionQuery(query));
190+
IndexQueryDetails indexQueryDetails = SQLQueryUtils.extractIndexDetails(query);
191+
assertNull(indexQueryDetails.getIndexName());
192+
assertNull(indexQueryDetails.getFullyQualifiedTableName());
193+
assertEquals(mvQuery, indexQueryDetails.getMvQuery());
194+
assertEquals("mv_1", indexQueryDetails.getMvName());
195+
}
196+
184197
@Test
185198
void testExtractionFromFlintMVQuery() {
186199
String[] mvQueries = {
187-
"CREATE MATERIALIZED VIEW mv_1 AS query=select * from my_glue.default.logs WITH"
188-
+ " (auto_refresh = true)",
189200
"DROP MATERIALIZED VIEW mv_1",
190201
"VACUUM MATERIALIZED VIEW mv_1",
191202
"ALTER MATERIALIZED VIEW mv_1 WITH (auto_refresh = false)",
@@ -200,6 +211,7 @@ void testExtractionFromFlintMVQuery() {
200211

201212
assertNull(indexQueryDetails.getIndexName());
202213
assertNull(fullyQualifiedTableName);
214+
assertNull(indexQueryDetails.getMvQuery());
203215
assertEquals("mv_1", indexQueryDetails.getMvName());
204216
}
205217
}
@@ -428,8 +440,7 @@ public static IndexQuery index() {
428440
}
429441

430442
public static IndexQuery mv() {
431-
return new IndexQuery(
432-
"CREATE MATERIALIZED VIEW mv_1 AS query=select * from my_glue.default.logs");
443+
return new IndexQuery("CREATE MATERIALIZED VIEW mv_1 AS select * from my_glue.default.logs");
433444
}
434445

435446
public IndexQuery withProperty(String key, String value) {

0 commit comments

Comments
 (0)